Attention to Mamba: A Recipe for Cross-Architecture Distillation

• URL: http://arxiv.org/abs/2604.14191v1
• Date: Wed, 01 Apr 2026 09:23:08 GMT
• Title: Attention to Mamba: A Recipe for Cross-Architecture Distillation
• Abstract summary: State Space Models (SSMs) such as Mamba have become a popular alternative to Transformer models.<n>We propose a principled two-stage approach: first, we distill knowledge from a traditional Transformer into a linearized version of Attention, using an adaptation of the kernel trick.<n>Overall, the distilled Mamba model is able to preserve the original Pythia-1B Transformer performance in downstream tasks, maintaining a perplexity of 14.11 close to the teacher's 13.86.
• Score: 15.145728951213997
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: State Space Models (SSMs) such as Mamba have become a popular alternative to Transformer models, due to their reduced memory consumption and higher throughput at generation compared to their Attention-based counterparts. On the other hand, the community has built up a considerable body of knowledge on how to train Transformers, and many pretrained Transformer models are readily available. To facilitate the adoption of SSMs while leveraging existing pretrained Transformers, we aim to identify an effective recipe to distill an Attention-based model into a Mamba-like architecture. In prior work on cross-architecture distillation, however, it has been shown that a naïve distillation procedure from Transformers to Mamba fails to preserve the original teacher performance, a limitation often overcome with hybrid solutions combining Attention and SSM blocks. The key argument from our work is that, by equipping Mamba with a principled initialization, we can recover an overall better recipe for cross-architectural distillation. To this end, we propose a principled two-stage approach: first, we distill knowledge from a traditional Transformer into a linearized version of Attention, using an adaptation of the kernel trick. Then, we distill the linearized version into an adapted Mamba model that does not use any Attention block. Overall, the distilled Mamba model is able to preserve the original Pythia-1B Transformer performance in downstream tasks, maintaining a perplexity of 14.11 close to the teacher's 13.86. To show the efficacy of our recipe, we conduct thorough ablations at 1B scale with 10B tokens varying sequence mixer architecture, scaling analysis on model sizes and total distillation tokens, and a sensitivity analysis on tokens allocation between stages.

Related papers

• Apriel-H1: Towards Efficient Enterprise Reasoning Models [6.630534140883356]
  Apriel-H1 family of hybrid LLMs combine transformer attention and SSM sequence mixers for efficient reasoning at 15B model size.<n>We release multiple post-distillation variants of Apriel-H1-15B-Thinker with different SSM-to-MHA ratios and analyse how reasoning performance degrades as more Mamba layers replace MHA.
  arXiv  Detail & Related papers  (2025-11-04T15:17:43Z)
• Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regression [90.93281146423378]
  Mamba is an efficient Transformer alternative with linear complexity for long-sequence modeling.<n>Recent empirical works demonstrate Mamba's in-context learning (ICL) competitive with Transformers.<n>This paper studies the training dynamics of Mamba on the linear regression ICL task.
  arXiv  Detail & Related papers  (2025-09-28T09:48:49Z)
• Diffusion Transformer-to-Mamba Distillation for High-Resolution Image Generation [65.46359545280546]
  This paper introduces diffusion transformer-to-mamba distillation (T2MD) to form an efficient training pipeline.<n>We establish a diffusion self-attention and Mamba hybrid model that simultaneously achieves efficiency and global dependencies.<n>Experiments demonstrate that our training path leads to low overhead but high-quality text-to-image generation.
  arXiv  Detail & Related papers  (2025-06-23T18:01:19Z)
• TransMamba: Fast Universal Architecture Adaption from Transformers to Mamba [66.80624029365448]
  We propose a cross-architecture knowledge transfer paradigm, TransMamba, that facilitates the reuse of Transformer pre-trained knowledge.<n>We propose a two-stage framework to accelerate the training of Mamba-based models, ensuring their effectiveness across both uni-modal and multi-modal tasks.
  arXiv  Detail & Related papers  (2025-02-21T01:22:01Z)
• The Mamba in the Llama: Distilling and Accelerating Hybrid Models [76.64055251296548]
  We show how to distill large Transformers into linear RNNs by reusing the linear projection weights from attention layers with academic GPU resources.<n>The resulting hybrid model achieves performance comparable to the original Transformer in chat benchmarks.<n>We also introduce a hardware-aware speculative decoding algorithm that accelerates the inference speed of Mamba and hybrid models.
  arXiv  Detail & Related papers  (2024-08-27T17:56:11Z)
• Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models [92.36510016591782]
  We present a method that is able to distill a pretrained Transformer architecture into alternative architectures such as state space models (SSMs)<n>Our method, called MOHAWK, is able to distill a Mamba-2 variant based on the Phi-1.5 architecture using only 3B tokens and a hybrid version (Hybrid Phi-Mamba) using 5B tokens.<n>Despite using less than 1% of the training data typically used to train models from scratch, Phi-Mamba boasts substantially stronger performance compared to all past open-source non-Transformer models.
  arXiv  Detail & Related papers  (2024-08-19T17:48:11Z)
• Demystify Mamba in Vision: A Linear Attention Perspective [72.93213667713493]
  Mamba is an effective state space model with linear computation complexity.<n>We show that Mamba shares surprising similarities with linear attention Transformer.<n>We propose a Mamba-Inspired Linear Attention (MILA) model by incorporating the merits of these two key designs into linear attention.
  arXiv  Detail & Related papers  (2024-05-26T15:31:09Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.