🎹 Training a 125M-parameter Model to Autocomplete Piano

By Simon Edwardsson | August 20, 2026 | simedw@blog

TL;DR: I developed a 125M-parameter transformer capable of real-time piano performance autocompletion, achieving speeds of approximately 108 notes/second on an iPhone 15.

The primary breakthroughs in this project stemmed from three areas:

1. Optimizing the MIDI representation.
2. Lazy data collection Aggressive cleaning of training data.
3. Implementing DPO (Direct Preference Optimization) post-training.

About a year ago, I envisioned a system where I could connect a MIDI keyboard to my phone, play a few bars, and let an AI finish the composition. After 14 different experiments, I've reached a version I'm excited to share.

Note: The video quality is low because my primary phone was dedicated to running the MIDI model.

The resulting application, RollTab, is available for free for anyone with an iPhone/iPad and a MIDI keyboard.

---

🎵 1. Sound Samples

The following examples demonstrate the model's ability to continue a piece based on a short initial prompt:

• Pokémon, Pallet Town: (Started with an 8-note prompt) $\rightarrow$ [Audio Sample]
• Final Fantasy VI, Terra's Theme: (Started with a 16-note prompt) $\rightarrow$ [Audio Sample]
• Für Elise: (Started with a 16-note prompt) $\rightarrow$ [Audio Sample]

---

🛠️ The Technical Challenge: Representing Music

To understand the model, one must first understand the difference between audio formats:

Since my goal was specifically piano continuation, I filtered the data to retain only piano-like material.

The Evolution of Tokenization

I tested several ways to turn MIDI events into discrete tokens for the transformer:

1. Naive Mapping: Using tokens like NOTE_ON_60_80 (pitch 60, velocity 80).
   • Problem: The vocabulary explodes. $128 \text{ pitches} \times 128 \text{ velocities} + 128 \text{ off-tokens} = 16,512$ tokens. This leads to sparse data and poor learning.
2. Grammar-based: Using a sequence like [NOTE_ON, PITCH, VELOCITY].
   • Problem: The model suffered from "drift"—it would forget to close notes or lose track of the state, which is unacceptable for a small, real-time model.
3. Explicit Duration: Using [NOTE, PITCH, VELOCITY, DURATION].
   • Problem: While it solved the drift, it required four autoregressive steps to produce a single note, slowing down generation.

The Final Representation

I eventually settled on a consolidated NOTE event: NOTE(pitch, delta_onset, duration, velocity)

• No TIME_SHIFT tokens: Silence is handled by the delta_onset (the time elapsed since the previous note started).
• Chords: Represented as multiple notes where delta_onset = 0, sorted by pitch.

Example Sequence:

# Play C4, wait 24 steps, then play D4
NOTE(C4, delta=0, duration=12, velocity=80)
NOTE(D4, delta=24, duration=12, velocity=80)

# A C-Major Chord
NOTE(C4, delta=24, duration=24, velocity=80)
NOTE(E4, delta=0, duration=24, velocity=78)
NOTE(G4, delta=0, duration=24, velocity=82)

Model Architecture

Instead of a flat token stream, the model uses a multi-head output. The transformer backbone runs once per note, not once per attribute.

The embedding for a note is calculated as: $\text{note} = E_{\text{type}} + E_{\text{pitch}} + E_{\text{delta}} + E_{\text{duration}} + E_{\text{velocity}}$

A small nested decoder allows later fields (like velocity) to be conditioned on earlier ones (like pitch), but the heavy lifting is done by the transformer in a single pass. This allows the iPhone to hit $\sim 108$ notes/sec.

---

🎹 Handling the Sustain Pedal

The sustain pedal complicates MIDI because notes keep ringing after the key is released. To keep the model simple, I baked sustain into the duration during preprocessing:

• If the pedal is down when a key is released $\rightarrow$ Extend duration to the moment the pedal is lifted.
• If the same pitch is played again while the pedal is down $\rightarrow$ Cut off the previous note at the retrigger point.

This removes the need for the model to predict pedal events explicitly, simplifying the problem to just four variables: pitch, onset, duration, and velocity.

---

📚 Dataset

I focused on public domain classical music. However, the quality of available datasets was inconsistent, so I ended