Show HN: I trained a 125M model to autocomplete piano on-device
This Show HN details the journey of building an on-device AI model capable of real-time piano autocomplete, akin to GitHub Copilot for music. The author shares a deep dive into the technical challenges, from MIDI tokenization to data cleaning and model training, and how they deployed a 125M-parameter transformer on an iPhone 15. It's a prime example of solving a complex, niche problem with ingenuity and iterative development, resonating with the HN audience's love for practical technical deep dives.
The Lowdown
Simon Edwardsson embarked on a year-long project to create an AI that could autocomplete piano performances in real-time on a mobile device. The goal was to build a 'GitHub Copilot for piano,' where playing a few notes on a MIDI keyboard would prompt the AI to continue the melody.
Key aspects and findings from this extensive project include:
- MIDI Representation Evolution: Initial attempts with
NOTE_ON/NOTE_OFFtokens led to musical drift, while grammar-based token streams were slow. The final, most effective representation encodes an entire musical note with its attributes (pitch,delta_onset,duration,velocity) into a single transformer step, significantly improving real-time performance. - Sustain Pedal Handling: Instead of modeling explicit pedal events, the system bakes sustain into note duration during preprocessing, extending notes that would sound longer due to the pedal.
- Data Curation: The author found that aggressive data cleaning and selection, focusing on piano-centric material and deduplication, was far more impactful than simply scaling up the dataset size.
- Training Objective: The model was trained with cross-entropy over five separate output heads for
event_type,pitch,delta,duration, andvelocity, allowing for detailed tracking of learning. - Augmentation for Real-World Input: To handle imperfect human input, the training data was augmented with global transposition, tempo scaling, duration/velocity jitter, and dropped prompt notes.
- Model Architecture: A standard decoder-only transformer with RMSNorm, rotary positional embeddings, and SwiGLU/MLP blocks was used. The 'large' model weighed in at 125M parameters, running at ~108 notes/second on an iPhone 15.
- Scheduled Sampling: Introducing scheduled sampling during training, where the model sometimes saw its own predictions instead of ground truth, surprisingly improved musical continuations despite increasing validation loss.
- Evaluation Innovation: Manual listening proved tedious and unreliable. Automatic metrics were useful for catching failures but not for selecting the best model. The breakthrough was using Gemini 3.5 Flash for pairwise preference evaluation, specifically asking which continuation was 'better' based on how well it followed the prompt, rather than just musical quality in isolation.
- Direct Preference Optimization (DPO): This post-training step significantly improved the model's reliability, making it much more likely to produce musically coherent and relevant continuations based on the preference dataset gathered with Gemini.
- Deployment: The PyTorch model was exported to Core ML, quantized to INT8 for on-device performance, and included a context management system to handle sessions longer than the 512-note training context.
Despite ongoing challenges like occasional looping and difficulty with very short prompts (akin to GPT-2 for piano), the author expresses satisfaction with the project's outcome, highlighting the enjoyment of collaborative music-making with the AI.