From Attention to Mamba: A Practical Recipe for Cross-Architecture Distillation

Why Cross-Architecture Distillation?

State space models (SSMs) like Mamba are emerging as strong alternatives to Transformers. Compared with attention-based models, Mamba offers lower memory usage and higher throughput during generation. However, the community has accumulated extensive experience in training Transformers, and a large number of pretrained models are available. How can we leverage these existing Transformer resources to accelerate the adoption of Mamba-like architectures? Cross-architecture distillation offers a path: transferring knowledge from Attention models into the Mamba architecture.

State space models (SSMs) are a class of models that describe sequences recursively through hidden states, with the continuous-time form (h'(t)=Ah(t)+Bx(t)). Mamba discretizes this and introduces an input-dependent selection mechanism (selective SSM), allowing the state transition matrix to vary with the input, thereby achieving a degree of content awareness while maintaining linear complexity. Compared with Transformer self-attention, SSMs need only maintain a fixed-size state during autoregressive generation, without caching all historical KVs, so memory usage does not grow with sequence length, and the throughput advantage on long sequences is pronounced. However, SSMs' parallel training techniques, initialization, and stability tuning are far less mature than those of Transformers, which constitutes the practical motivation for cross-architecture distillation.

Distillation Recipe: Key Design Choices

The paper aims to establish an effective recipe for distilling attention-based models into Mamba-like architectures. Although the original text does not elaborate on specific technical details, the problem setup implies that distillation must address two core challenges: representation alignment caused by architectural differences, and how to preserve the knowledge acquired by the Transformer.

Addressing Architectural Differences

Attention and SSMs differ fundamentally in their computational paradigms: the former dynamically aggregates information through attention weights, while the latter implicitly models sequences through state transitions. During distillation, appropriate loss functions and intermediate-layer alignment strategies must be designed so that Mamba can mimic the behavior of the Transformer.

Knowledge distillation typically refers to a student model minimizing KL divergence on the teacher's soft labels (logits), with optional matching of intermediate-layer features. The difficulty of cross-architecture distillation lies in the fact that the teacher and student have mismatched hidden-space dimensions, layer counts, and computation graphs. Common practices therefore include: using learnable projection layers to map student hidden states to the teacher's dimension before applying alignment losses; performing temperature-scaled soft-target matching at the output distribution level; and introducing indirect supervision via attention maps or state transition patterns. Together, these strategies provide richer supervisory signals than hard labels, alleviating the alignment difficulties caused by architectural heterogeneity.

Leveraging Pretrained Weights

The recipe emphasizes leveraging existing pretrained Transformers, which means the distillation process can stand on the shoulders of giants. Through carefully designed distillation objectives, the Mamba student model can inherit the teacher model's capabilities while retaining its own advantages in inference efficiency.

Potential Impact and Outlook

If this recipe proves effective, it will significantly lower the barrier to adopting SSMs. Developers will not need to train Mamba models from scratch; instead, they can perform knowledge transfer based on the mature Transformer ecosystem. This is especially valuable for scenarios requiring long-sequence processing and high-throughput generation, such as real-time dialogue systems and document-level translation.

However, whether the distilled Mamba can fully reproduce the complex reasoning capabilities of Transformers still requires further experimental validation. The generality of cross-architecture distillation also remains to be extended to more model families.

Conclusion

Cross-architecture distillation provides a practical solution for migrating from Transformer to Mamba. Through a systematic distillation recipe, we can potentially retain Transformer knowledge while enjoying the efficiency gains brought by SSMs. The value of this work lies in bridging two active research communities and promoting the adoption of more efficient model architectures.

One-line Summary

Cross-architecture distillation transfers pretrained Transformer knowledge to Mamba, balancing efficiency and knowledge reuse.

Paragraph Summary

The article discusses the motivation and recipe design for distilling Attention models into Mamba-like state space models. Mamba offers lower memory usage and higher throughput during generation, but its training experience and pretrained resources lag far behind Transformers. Cross-architecture distillation leverages existing pretrained Transformers to let a student Mamba inherit the teacher's capabilities while retaining its inference efficiency advantages. The recipe must address two major challenges: representation alignment caused by architectural differences, and how to preserve the knowledge acquired by the Transformer, involving loss function design and intermediate-layer alignment strategies. If validated, it will significantly lower the barrier to SSM adoption, especially for long-sequence, high-throughput scenarios, but whether the distilled Mamba can reproduce complex reasoning capabilities still awaits experimental verification.