arXiv cs.CL· Abhinav Moudgil, Ningyuan Huang, Eeshan Gunesh Dhekane, Pau Rodr\'iguez, Luca Zappella, Federico Danieli·· 12 小时前AI 评分31
Attention to Mamba：一种跨架构知识蒸馏方案
Attention to Mamba: A Recipe for Cross-Architecture Distillation
AI 导读
研究者提出两阶段跨架构蒸馏方案，先将 Transformer 知识蒸馏到线性化 Attention，再蒸馏到不含 Attention 模块的改造版 Mamba。
正文 · 原文
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 retains most of the original Pythia-1B Transformer performance, reaching a perplexity of 14.01 close to the teacher's 13.86. This largely carries over to downstream tasks, although some gaps appear on tasks relying more on in-context recall. 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.
| Comments: | v2: fixes an implementation error in v1 that affected the SSM decay gates; all experiments re-run on Mamba-2 with updated results; adds inference-efficiency results |
| Subjects: | Computation and Language (cs.CL); Machine Learning (cs.LG) |
| MSC classes: | 68T07 |
| Cite as: | arXiv:2604.14191 [cs.CL] |
| (or arXiv:2604.14191v2 [cs.CL] for this version) | |
| https://doi.org/10.48550/arXiv.2604.14191 arXiv-issued DOI via DataCite |
Submission history
From: Federico Danieli [view email]
[v1]
Wed, 1 Apr 2026 09:23:08 UTC (14,935 KB)
[v2]
Tue, 3 Nov 2026 22:13:00 UTC (14,637 KB)
来源：arXiv cs.CL · arxiv.org