Running 125M Parameter Transformer On-Device for Real-Time MIDI Completion

Musicians live and die by latency. When you press a key on a MIDI keyboard, you expect to hear the sound instantly. If you introduce a machine-learning model to accompany your playing, the timing constraints become brutal. A cloud-based API call introduces a round-trip latency of 50 to 150 milliseconds. In musical terms, a 100-millisecond delay at 120 beats per minute is a sixteenth note. It makes real-time collaboration impossible.

To solve this, we must run the model locally. A local-first architecture keeps the data on the machine, bypassing the network stack. But running a generative model on consumer hardware is difficult. While specialized techniques exist for running 70B models on consumer GPUs, real-time audio requires a different approach. The solution lies in choosing the right model size and optimizing the inference engine for CPU caches. A 125-million parameter transformer strikes the balance. It is small enough to fit into CPU cache memory, yet large enough to capture the patterns of classical and jazz piano.

Memory Bandwidth as the Real Bottleneck

When running machine-learning models on consumer hardware, memory bandwidth is the primary bottleneck. Generative music models process tokens sequentially. During inference, the engine must read all model weights from memory for every single token generated.

A 7-billion parameter model requires 7 gigabytes of memory transfer per token, even at 8-bit quantization. On standard desktop hardware with a memory bandwidth of 50 gigabytes per second, the theoretical maximum speed is around 7 tokens per second. This is too slow for complex musical structures that require high-density token generation.

A 125M parameter model changes the physics. At 16-bit precision (FP16), the model weights occupy 250 megabytes. When quantized to 4-bit precision (INT4), the size drops to roughly 62.5 megabytes. Modern consumer CPUs feature L3 caches ranging from 32 to 96 megabytes. By shrinking the model to fit within these cache limits, we reduce reliance on slow system RAM. The CPU runs the matrix multiplications directly from its fastest cache levels, pushing inference speeds past 100 tokens per second.

Representing Music as Tokens

Unlike text, music is multi-dimensional. Notes happen simultaneously and have varying durations. To feed this into a decoder-only transformer, we use a time-shift tokenization scheme.

We convert MIDI events into a sequence of discrete tokens from a vocabulary of 388 tokens:

• Note-On Tokens: 128 tokens, one for each MIDI pitch.
• Note-Off Tokens: 128 tokens to signal the release of a key.
• Time-Shift Tokens: 100 tokens representing time increments from 10 milliseconds to 1 second.
• Velocity Tokens: 32 tokens representing quantized strike force values.

A C-major triad played quickly looks like this in token form:

Note-On:60
Note-On:64
Note-On:67
Time-Shift:500
Note-Off:60
Note-Off:64
Note-Off:67

This sparse representation allows the machine-learning model to process complex performances without wasting compute on silent intervals.

Dataset and Model Training

To train a model of this size specifically for piano music, we used the GiantMIDI-Piano dataset, which contains over 30 million MIDI notes transcribed from classical piano recordings. We supplemented this with the MAESTRO dataset to provide high-resolution velocity and timing data.

The training process ran for 40 epochs on a cluster of four NVIDIA A100 GPUs, taking roughly 48 hours. We used a context window of 2048 tokens. This length allows the model to remember the theme of a piece even after several minutes of improvisation.

Training a model on MIDI events is different from training on text. In text, a model learns semantic grammar. In MIDI, the model must learn temporal grammar. If the model forgets to emit a Note-Off token, a note will ring out indefinitely, ruining the performance. This is a classic example of the structural failures and things AI still gets wrong when handling precise logical sequences. To prevent this, we added a penalty to the loss function during training for orphan Note-On tokens that remained unresolved for more than 5 seconds.

The Transformer Architecture Details

The core of our 125M parameter model is a decoder-only transformer. It uses 12 layers, 12 attention heads, and an embedding dimension of 768. The feed-forward network dimension is set to 3072.

Instead of absolute positional encodings, we use Rotary Position Embeddings (RoPE). Absolute encodings struggle when a musical piece extends beyond the training context window. RoPE allows the model to generalize to longer sequences by encoding relative distances between tokens.

We use FlashAttention-2 in our training pipeline to speed up attention computation, but for local inference, we focus on optimizing the standard attention math for single-batch execution. Since this is an interactive tool, the batch size during inference is always exactly 1.

The On-Device Inference Engine

To avoid the runtime overhead of Python, the inference engine is written in pure C++. We use ggml, a tensor library designed for consumer hardware. The engine compiles to a single binary with zero external dependencies.

During initialization, we load the model weights directly into a single contiguous memory block using memory-mapped files (mmap). This allows the operating system to load model weights on demand and share them across processes.

struct inference_context {
    ggml_context* ctx_w;
    ggml_context* ctx_b;
    std::vector<float> kv_cache;
    int max_context_length;
};

Static memory allocation is critical. We allocate the Key-Value (KV) cache upfront to prevent runtime allocation delays. The KV cache stores the computed key and value vectors for previous tokens, saving us from recalculating them at each step. For a 12-layer model with 768 embedding dimensions and a context window of 2048 tokens, the KV cache size is calculated as:

KV_Cache_Size = 2 * Layers * Heads * Head_Dim * Context_Length * Float_Size

Using our model configurations:

2 * 12 * 12 * 64 * 2048 * 4 bytes, which equals 150.99 megabytes.

By pre-allocating this 150MB block, we eliminate memory fragmentation and garbage collection pauses during live performance.

During autoregressive generation, we feed the last generated token back into the model. Instead of running the self-attention calculation over the entire sequence of historical tokens, we retrieve the keys and values of the history from the cache and only compute the keys and values for the new token. This reduces the computational complexity of the attention step from O(N^2) to O(N), where N is the sequence length.

Optimizing for Hardware Constraints

To achieve sub-10ms latency, the inference loop must be optimized for the target hardware. We focus on two primary platforms: ARM64 (Apple Silicon) and x86_64 (Intel/AMD).

On ARM64, we utilize NEON assembly instructions to vectorize matrix-vector multiplications. On x86_64, we compile with AVX2 or AVX-512 support. The core operation in transformer inference is the GEMV (General Matrix-Vector multiplication), which occurs when multiplying the input token representation by the weight matrices.

#if defined(__ARM_NEON)
inline float neon_dot_product(const float* a, const float* b, int size) {
    float32x4_t sum_vec = vdupq_n_f32(0.0f);
    for (int i = 0; i < size; i += 4) {
        float32x4_t va = vld1q_f32(a + i);
        float32x4_t vb = vld1q_f32(b + i);
        sum_vec = vmlaq_f32(sum_vec, va, vb);
    }
    return vaddvq_f32(sum_vec);
}
#endif

Thread pinning is another important optimization. Modern CPUs mix performance cores with efficiency cores. If the operating system moves our inference thread from a performance core to an efficiency core mid-performance, latency spikes. We pin the inference thread to a specific physical performance core using platform-specific APIs.

#ifdef __linux__
cpu_set_t cpuset;
CPU_ZERO(&cpuset);
CPU_SET(2, &cpuset);
pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
#endif

Quantization and Musical Fidelity

Quantizing weights from FP16 to INT8 or INT4 reduces memory bandwidth pressure, but it introduces quantization noise. In text generation, minor noise might cause a slightly different word choice. In music, it can destroy the performance.

We tested three quantization schemes: FP16, INT8, and INT4.

• FP16: Complete fidelity. The model captures micro-timings and subtle velocity variations. The memory footprint is 250 megabytes.
• INT8 (RTN - Round-to-Nearest): Minimal degradation. The model still generates coherent chord progressions, but velocity dynamics become slightly compressed. The memory footprint is 125 megabytes.
• INT4 (Group-wise quantization): Noticeable degradation. The model struggles with long-term structure and sometimes repeats chords. It also loses the ability to generate soft notes, defaulting to a flat velocity distribution. The memory footprint is 62.5 megabytes.

For INT8 quantization, we used symmetric per-channel quantization. Each weight tensor is scaled so that the maximum absolute value fits into the range [-127, 127]. The scale factor is stored alongside the quantized weights.

For INT4 quantization, we used block-level quantization (specifically Q4_0 format). In this format, weights are grouped into blocks of 32 values. Each block shares a single 16-bit float scale factor. This grouping reduces the quantization error compared to global tensor-level quantization, but still struggles with the high-dynamic range of MIDI velocity values. Since MIDI velocities range from 0 to 127, a small quantization error in the output layer can shift a note from a soft touch to a loud strike, making the performance sound mechanical.

For live performance, INT8 represents the sweet spot. It cuts the memory bandwidth requirement in half while preserving the musicality of the model.

The Real-Time MIDI Loop

To use the model, we set up a virtual MIDI loop. The architecture consists of three components: the physical MIDI keyboard, the inference engine, and the software synthesizer (DAW).

The input queue captures incoming MIDI events and converts them to tokens. When the musician stops playing for a set threshold (e.g., 200 milliseconds), the inference engine wakes up. It reads the last 512 tokens of performance history from the KV cache, generates the next 32 tokens, and pushes them to the output queue.

To make the generated notes sound human, we apply top-p (nucleus) sampling. We set p = 0.9 and temperature T = 0.85. This prevents the model from choosing highly improbable notes while avoiding repetitive loops.

The generated tokens are converted back to MIDI events and scheduled for playback. Because the inference engine runs in less than 5 milliseconds per step, the transition between human playing and model completion is imperceptible.

Latency Benchmarks

We measured the performance of our 125M parameter model across three different hardware configurations. The benchmark measures the time taken to generate a single token (in milliseconds) with a context window of 512 tokens.

On desktop and laptop hardware, the latency remains well under the 10ms threshold. Even on the Raspberry Pi 4, the INT4 model approaches the limit of real-time usability, though the musical quality degradation makes it less practical.

Handling Time Drift in Live Performance

One challenge with time-shift tokens is drift. If the model generates a sequence of time-shift tokens that do not align precisely with the DAW grid, the accompaniment will sound out of tempo.

To fix this, we implement a MIDI clock synchronization step in the output queue. The engine reads the current tempo (BPM) from the DAW and rounds the generated time-shift tokens to the nearest sixteenth-note grid step. This ensures that even if the model generates a slightly off-beat token, the output remains locked to the project tempo.

int quantize_time_shift(int ms, int bpm) {
    double ms_per_beat = 60000.0 / bpm;
    double sixteenth_note_ms = ms_per_beat / 4.0;
    double steps = std::round(ms / sixteenth_note_ms);
    return static_cast<int>(steps * sixteenth_note_ms);
}

Local-First Music AI

Moving machine-learning models out of the cloud and onto local hardware changes how we build creative tools. By keeping the model small and optimizing the C++ inference engine for CPU caches, we get zero-latency musical accompaniment. The architecture shows that you do not need giant cloud clusters to run useful generative AI. Sometimes, 125 million parameters and a well-written loop are all it takes to keep the music playing.