TL;DR: I trained a 125M-parameter transformer to autocomplete piano performances in real time (~108 notes/sec on an iPhone 15). The biggest improvements came from finding the right MIDI representation, cleaning the training data aggressively, and adding DPO post-training.
Almost a year ago, I started tinkering with an idea: connect my MIDI piano to my phone, play something, and have AI autocomplete the song for me. Think GitHub Copilot, but for piano.
It turned out to be a deeper rabbit hole than I expected. Fourteen experiments later, it is finally at a point where I am happy enough with it to write about.
The app, RollTab, is available for free here if you have a MIDI keyboard and an iPhone/iPad. 1
A few sound samples
Each audio starts with a short prompt, followed by the model's continuation.
Pokémon, Pallet Town (8-note prompt)
Final Fantasy VI, Terra's Theme (16-note prompt)
Für Elise (16-note prompt)
What’s in a MIDI File?
A MIDI file is quite different from an MP3 or other audio formats. Rather than storing recorded sound, it stores music as a sequence of events: a key is pressed at a certain pitch and velocity, a key is released, the sustain pedal changes state, and so on. Other events include switching instruments or changing volume.
These events are often organised into multiple tracks. A pop or game MIDI might have melody, chords, bass, drums, strings, and several synth parts. This project is focused on piano continuation, so I mostly kept piano-like material and removed or reduced the rest.
How Do You Tokenize Music?
To train a transformer on these performances, I first needed to turn the MIDI events into a discrete sequence the model could read and predict. The most obvious mapping is to make a token for every MIDI event:
NOTE_ON_60_80 # {pitch}_{velocity}
NOTE_OFF_60 # {pitch}
TIME_SHIFT_12 # {time step}
If you include pitch and velocity directly in a NOTE_ON token, the vocabulary can grow quickly. There are 128 MIDI pitches and 128 velocity values, so the naive combined note-on vocabulary has up to:
128 * 128 + 128 = 16,512
tokens just for note-on and note-off. In practice you would probably bucket velocity, but the basic issue remains: many combinations are rare, and the model has to learn a lot of structure from sparse tokens.
A common improvement is to factor the representation with a grammar:
[NOTE_ON, PITCH, VELOCITY] | [NOTE_OFF, PITCH] | [TIME_SHIFT, DURATION]
Now the output spaces are smaller:
NOTE_ON / NOTE_OFF / TIME_SHIFT
PITCH: 128 values
VELOCITY: ~16
DURATION: ~100
You can enforce the grammar during generation by masking invalid next tokens.
After NOTE_ON, only pitch tokens are valid. After pitch, only velocity tokens
are valid. This guarantees syntactically valid output.
I tried note-on/note-off style representations, but my models tended to drift. They would forget to emit note-off, leave hanging notes, or lose track of active state. That was especially bad for my target: a small model running close to real time on a laptop or phone.
Another representation I tried was closer to:
[NOTE, PITCH, VELOCITY, DURATION] | [TIME_SHIFT, DURATION]
This avoids note-off drift because note duration is explicit. The time shift token advances the playhead when no note is played.
This worked better musically, but it was slow. One musical note took roughly four autoregressive transformer steps. It also burns through the context window quickly.
The final representation
The representation I eventually settled on was:
NOTE(pitch, delta_onset, duration, velocity)
There is no separate TIME_SHIFT event in the final version. Silence is represented by delta_onset on the next note: the time since the previous note onset.
For example:
NOTE(C4, delta=0, duration=12, velocity=80)
NOTE(D4, delta=24, duration=12, velocity=80)
means: play C4, wait 24 time steps before the next note onset, then play D4.
Chords are represented as multiple notes with delta_onset = 0, sorted by
pitch2:
NOTE(C4, delta=24, duration=24, velocity=80)
NOTE(E4, delta=0, duration=24, velocity=78)
NOTE(G4, delta=0, duration=24, velocity=82)
It's also not a flat token stream like:
NOTE, PITCH, DELTA, DURATION, VELOCITY
Instead of spending four transformer passes generating the attributes of a note, the transformer advances the music by one complete note at a time.
In practice, this gets the large model to about 108 notes/second on an iPhone, well above anything a human would need for live playing.
Internally each note has five categorical fields, each with its own vocabulary3, with timing quantized to fixed steps.4
[event_type, pitch_id, delta_id, duration_id, velocity_id]
Each field gets its own embedding. The note token is the sum of all the embeddings:
note =
event_type_embedding[NOTE]
+ pitch_embedding[C4]
+ delta_embedding[12]
+ duration_embedding[24]
+ velocity_embedding[80]
The model then has separate output heads: pitch, delta, duration, and so on.
There is a small nested decoder between the fields, so later fields can condition on earlier predicted fields. But the expensive transformer backbone runs only once per note, not once per field.
Sustain Pedal
As you might know, pressing down the sustain pedal on a piano makes notes play even after you release them. I didn't want to muddy the implementation with adding sustain pedal events. Instead, sustain is baked into note duration during preprocessing.
If the key is released while the sustain pedal is down, the note is extended to the pedal-up time. If the same pitch is played again first, the earlier note is cut off at the retrigger. The result is a note duration that approximates the actual sounding duration.
This loses the explicit pedal gesture, but it makes the modeling problem much simpler: the model only has to predict pitch, onset, duration, and velocity.
Dataset
I searched through a lot of publicly available datasets and collections, focusing mostly on older classical music in the public domain. The quality varied wildly, so I ended up writing quite a few cleaning scripts.
The final dataset contained a few hundred thousand MIDI files, representing roughly 300 million note events.
The final pipeline:
- selected piano-focused material
- removed or reduced pathological multi-track mixtures
- filtered by density and pitch/time coverage
- deduplicated by fingerprints that ignore global transposition and uniform tempo changes
- grouped alternate versions of the same composition into the same split
I tried scaling the dataset to roughly 5x the size, hoping it would improve performance, but the resulting models were worse. Cleaning and selecting the data mattered more than simply adding more of it.
Training
Initially, training is just cross-entropy over the five output heads, summed together:
type_loss
+ pitch_loss
+ delta_loss
+ duration_loss
+ velocity_loss
This makes it easy to track pitch, duration, and velocity accuracy separately, rather than relying on a single aggregate next-token loss.
Still, the training objective has an important limitation: music continuation does not have a single correct answer. A held-out song only gives the model one "correct" next note, even though there are often many continuations that would work musically. Cross-entropy is useful for learning the mechanics of music, but not a great proxy for how good a full continuation sounds.
Augmentation
Augmentation was important because the live input is not a pristine MIDI file. It is me playing piano, badly enough that notes might be slightly early, late, too hard, etc.
In the end I settled on the following augmentations:
- global transposition
- uniform tempo scaling
- duration/velocity jitter
- dropped prompt notes
Model
The architecture is essentially a fairly standard decoder-only transformer: RMSNorm, rotary positional embeddings, causal self-attention, SwiGLU/MLP blocks, and autoregressive generation.
I mainly trained three model sizes:
small: about 33M parameters
medium: about 64M parameters
large: about 125M parameters
The small model was great for quick experiments, but the medium model almost always beat it. The large model performed better, although not by a huge margin.
I am currently trying to get the medium model close to the large model's quality, mostly to reduce footprint and latency in the iOS app.
Scheduled Sampling
My best base model used scheduled sampling between the fields of each note. Normally, during training, the duration and velocity predictions get to see the correct pitch. But at inference time they have to work with whatever pitch the model actually predicted.
So during training I sometimes fed the model its own predicted pitch instead. I started at 0% for the first few epochs, then gradually increased it during training, up to 50% in the best model.
Funnily enough, this increased validation loss but improved the continuations.
Evaluation
At first, evaluation was just me listening.
I generated continuations from held-out songs using prompts of 4-32 notes, then compared model outputs manually. This was slow and annoying and after a while everything sounded like noise.
Four-note prompts were the hardest: there simply was not much musical context to work with. Eight notes worked better, while 16–32 note prompts were substantially more reliable because the model had enough structure to infer what was happening.
Unprompted generation is very much hit or miss, but that isn't the use-case I'm gunning for.
I also wrote a bunch of automatic metrics:
- repeated pitch n-grams
- pitch entropy
- pitch-class entropy
- pitch range
- note density
- long pauses
- chord density
These metrics were useful for catching obvious failures, but they were not good enough to select the best model.
Eventually I used Gemini 3.5 Flash for pairwise evaluation. Asking it to give a single absolute score was inconsistent. Asking instead, "given A and B, which continuation is better?" worked much better, especially when I mirrored every comparison to reduce position bias. 5 This let me build a reasonably large preference dataset, which I then used for DPO.
Initially, Gemini overindexed on how good a continuation sounded in isolation, rather than how well it followed from the prompt. The outputs often sounded better on their own, but felt disconnected from what I had just played.
Better prompting helped, but I eventually split the evaluation into two criteria: a continuation score, measuring how well the output follows from the prompt, and a sounds-good score, measuring its musical quality in isolation. I used the continuation score as the primary signal for DPO.
DPO: Direct Preference Optimization
DPO made the biggest difference after pretraining. It took the model from occasionally producing a good continuation to doing it much more reliably.
For each prompt, I generated multiple continuations and used pairwise evaluation to choose a better and worse one:
prompt -> chosen continuation
prompt -> rejected continuation
DPO trains the model to make the chosen continuation more likely than the rejected one, while keeping it reasonably close to the original model.
After DPO, more than 69% of continuations were preferred over the base model in my pairwise evaluation.
The β value controls how strongly DPO penalizes moving away from the base model. In my sweep, β=0.01 and β=0.03 improved the model, while β=0.10 pushed too hard and made it worse.
I also tried a "consensus" dataset: instead of trusting every noisy preference judgment, I only kept preference pairs where the evaluator agreed consistently. That produced the best result in this sweep.
My gut feeling is that the base model had already learned a reasonable mental model of music, just not what makes a good continuation.
What Did Not Work
A lot did not work:
- Note-on/note-off drifted too much for small real-time models.
- Grammar-masked token streams were valid but slow.
- Broader data made results worse when the data was noisy.
- Bigger models helped, but did not magically solve loops.
- Mirostat reduced repetition but often made outputs incoherent.
- Extra local auxiliary losses made training slower without clear listening wins.
- Absolute scalar Gemini ratings were worse than pairwise judging.
- Validation loss alone missed important differences in rollout quality.
- Born-again networks (retraining a model on its own soft predictions) didn't improve quality here.
Packaging it
I exported the PyTorch model to Core ML and quantized the weights to INT8. The first launch is still annoyingly slow while Apple's runtime optimizes the model for the available hardware.
The model was only trained with contexts of up to 512 notes, but I wanted to support longer sessions. Whenever the context gets close to the limit, I keep the most recent 384 notes, rebuild the context from those, and continue from there. This means rebuilding the KV cache, but the model is fast enough that it hasn't been a major problem.
I used RoPE for positional encoding, so in theory I could do something neater with shifted positions and a ring buffer. Unfortunately Core ML does not expose Q, K, and V directly.
At that point, though, I was mostly just happy that it worked.
Conclusion
This has been a very fun project. There are plenty of interesting papers on music generation, but I deliberately avoided reading too deeply into them at first. I wanted the fun of working through the problem myself, rather than just implementing someone else’s research. Only afterward did I go back and compare my approach with the existing literature. 6
It is still far from perfect. It loops occasionally, short prompts are difficult, and there is plenty I want to improve. Think GPT-2, but for piano.
But I have finally reached the point where I actually enjoy sitting down at the piano, playing a few notes, and seeing what we come up with together.
- The first version took 11 days to get approved. I have a new version pending review that allows you to choose between top-k, top-p, min-p, XTC, top-h, and Mirostat v2 sampling. ↩
- We sort by pitch so that during training we don't get penalized when one song encoded a C-major chord as CEG and another as EGC. ↩
- The exact vocabularies are:
event_type: PAD, BOS, EOS, NOTE, MASK pitch: 0 unused/pad + 128 MIDI pitches delta: 0..48 steps, plus 72, 96, 144, 192 duration: 1..96 steps, plus 144, 192, 288, 384 velocity: 4, 12, 20, ..., 124 - Timing uses 24 steps per quarter note. That gives enough resolution for common straight and triplet subdivisions, including the "almost but not quite on the beat" timing I tend to produce when playing live. This was also chosen after looking at the timing distribution in my training dataset. ↩
- In one test over 200 songs, Gemini gave the same preference 70% of the time after reversing A and B. ↩
- Some recent transformer-based models for symbolic MIDI generation include Aria: Scaling Self-Supervised Representation Learning for Symbolic Piano Performance, Moonbeam: A MIDI Foundation Model Using Both Absolute and Relative Music Attributes, MIDI-GPT, Anticipatory Music Transformer, PianoBART, and MIDI-LLM. ↩