用合理初始化提升Mamba模型,让其从Transformer高效迁移性能。
Attention to Mamba: A Recipe for Cross-Architecture Distillation
- 分两阶段蒸馏:先转线性注意力,再转Mamba架构。
- 蒸馏后Mamba在下游任务中保持14.11的困惑度,接近教师模型的13.86。
- 适合想用Mamba替代Transformer又不想重训练的研究者。
状态空间模型(SSMs)如Mamba因其更低的内存消耗和更高的生成吞吐量,正成为Transformer的有力替代方案。然而,社区对Transformer的训练积累了丰富经验,且大量预训练模型可用。为在利用现有Transformer的同时推动SSMs应用,我们提出一种有效的跨架构蒸馏方法,将注意力模型知识迁移到Mamba类结构中。以往研究显示,直接从Transformer蒸馏到Mamba会显著损失性能,通常需结合注意力与SSM模块的混合方案解决。本文核心观点是:通过为Mamba引入合理的初始化策略,可实现更优的蒸馏效果。为此,我们提出两阶段方法:第一阶段将传统Transformer知识蒸馏至线性注意力模型,采用改进的核技巧;第二阶段将该线性模型进一步蒸馏至不包含任何注意力块的适配版Mamba。最终得到的蒸馏模型在下游任务中保持了接近原Pythia-1B Transformer的性能,困惑度为14.11,仅略高于教师模型的13.86。我们通过100亿词、10亿参数规模的详尽消融实验,验证了方法的有效性,涵盖序列混合器结构变化、模型规模扩展、总蒸馏词数影响及各阶段数据分配敏感性分析。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。