Attention to Mamba: A Recipe for Cross-Architecture Distillation

The work develops a way to transfer the capabilities of pretrained Transformer language models into Mamba, a state-space architecture that uses less memory and can generate tokens faster. Directly distilling a Transformer into Mamba has often failed to preserve the teacher’s performance, leading earlier approaches to combine attention and state-space layers. The central claim here is that a carefully designed initialization can make it possible to distill into a Mamba model without any attention blocks.

The proposed recipe has two stages. First, knowledge from a conventional Transformer is distilled into a linearized attention model using an adaptation of the kernel trick. That intermediate model then serves as the teacher for an adapted Mamba student. In experiments, a 1-billion-parameter Mamba distilled from Pythia-1B reaches a perplexity of 14.01, close to the Transformer’s 13.86. Its performance largely transfers to downstream tasks, though it falls further behind on tasks that depend heavily on recalling information from context. The study also tests the effects of different sequence-mixer designs, model sizes, total distillation data, and how tokens are divided between the two stages, offering evidence about when and how the recipe works.