1억 2,500만 파라미터 피아노 자동완성 온디바이스 모델 공개

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) Your browser does not support the audio tag. Final Fantasy VI, Terra's Theme (16-note prompt) Your browser does not support the audio tag. Für Elise (16-note prompt) Your browser does not support the audio tag. 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 pitch 2 : 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 vocabulary 3 , 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 m