Training a 125M-parameter Model for On-Device Piano Autocomplete

A 125M-parameter transformer model can now autocomplete piano performances in real time on mobile hardware, achieving speeds of approximately 108 notes per second on an iPhone 15. The project, implemented in an app called RollTab, demonstrates that high-quality musical continuation is achievable with a relatively small model by focusing on optimized MIDI representation, aggressive data cleaning, and post-training preference optimization.

Optimized MIDI Representation for Real-Time Inference

The primary technical challenge in MIDI modeling is converting continuous musical events into a discrete sequence that a transformer can predict without sacrificing inference speed or musical coherence.

Moving Beyond Note-On/Note-Off

Traditional MIDI representations often use separate NOTE_ON and NOTE_OFF tokens. However, small real-time models frequently suffer from "drift," where they forget to emit a note-off token, leading to hanging notes. While grammar-masked token streams (e.g., [NOTE_ON, PITCH, VELOCITY]) ensure syntactic validity, they require multiple transformer passes per musical note, which slows down generation.

The Single-Token Note Representation

To maximize throughput, the final model uses a representation where the transformer advances the music by one complete note at a time. Each note is represented as a combination of five categorical fields:

Event Type: (e.g., NOTE, PAD, BOS, EOS)
Pitch: 128 MIDI pitches
Delta Onset: Time since the previous note onset (quantized to 24 steps per quarter note)
Duration: Note length (quantized)
Velocity: Note intensity

Instead of a flat token stream, each field has its own embedding. The final note token is the sum of these embeddings. The model employs separate output heads for each field, with a small nested decoder allowing later fields to condition on earlier predicted fields within the same note. This architecture allows the expensive transformer backbone to run only once per note.

Data Engineering and Augmentation

Model performance was driven more by data quality than quantity. The developer found that scaling the dataset to five times its original size actually degraded performance, highlighting the importance of aggressive cleaning.

Dataset Pipeline

The training set consisted of several hundred thousand MIDI files (roughly 300 million note events), primarily focusing on public domain classical music. The cleaning pipeline included:

Filtering for piano-focused material and removing pathological multi-track mixtures.
Deduplication using fingerprints that ignore global transposition and uniform tempo changes.
Grouping alternate versions of the same composition into the same data split to prevent leakage.

Handling the Sustain Pedal

To simplify the modeling problem, sustain pedal events were removed. Instead, sustain is baked into the note duration during preprocessing: if a key is released while the sustain pedal is down, the note's duration is extended to the pedal-up time.

Augmentation for Live Input

Because live human input is imperfect, the model was trained with augmentations to ensure robustness against timing and velocity errors:

Global transposition and uniform tempo scaling.
Duration and velocity jitter.
Dropped prompt notes.

Training and Optimization Strategies

The model is a decoder-only transformer featuring RMSNorm, rotary positional embeddings (RoPE), and SwiGLU/MLP blocks.

Scheduled Sampling

To bridge the gap between training (where the model sees ground-truth pitches) and inference (where it sees its own predictions), the developer implemented scheduled sampling. By gradually increasing the probability of feeding the model its own predicted pitch (up to 50%), the rollout quality improved, even though validation loss increased.

Direct Preference Optimization (DPO)

DPO was the most significant factor in improving the reliability of continuations. The developer used Gemini 3.5 Flash to perform pairwise evaluations of generated continuations, scoring them on two criteria: how well the output followed the prompt (continuation score) and its general musical quality (sounds-good score).

Using a "consensus" dataset—where the evaluator agreed consistently—and a $\beta$ value between 0.01 and 0.03, the model's preference rate jumped to 69.05% over the base pretrained model.

On-Device Deployment

The model was exported to Core ML and quantized to INT8 for iOS deployment. To handle sessions longer than the 512-note training context, the app maintains the most recent 384 notes and rebuilds the KV cache when the limit is reached.

Community Insights and Critiques

While the project was praised for its technical execution and on-device performance, it sparked a debate among musicians and AI researchers regarding the nature of musical improvisation.

"The results strike me as comparable or worse than you could get with a Markov model... you need to either set up a pipeline to decompose music into harmonic sequences and melodic sequences, or develop a better dataset." — @rajivayyangar

"I can't imagine that anyone actually wants to learn about harmony, about voicing, and voice leading... They want to just press some keys and declare that they made what the computer generated." — @bubblegumcrisis

Other users suggested expanding the model to support multi-part accompaniment (e.g., baroque style) or integrating it as a VST/Max 4 Live device for professional music production.