Mamba 跨架构蒸馏新配方两阶段初始化提效Hacker News AIT2·基座大模型

Mamba 跨架构蒸馏新配方两阶段初始化提效

原标题：Attention to Mamba: A Recipe for Cross-Architecture Distillation

71AI 研判分

研究提出一种从 Transformer 蒸馏到 Mamba 类架构的两阶段方法：先将知识蒸馏到线性化 Attention，再迁移至 Mamba，并配合原则性初始化，从而避免朴素蒸馏的性能损失。该方法还给出教师模型已存在与需从头训练两种场景下的算力最优蒸馏配方，降低大规模蒸馏风险。

关联知识实体

术语状态空间模型术语线性化注意力术语跨架构蒸馏术语Mamba

报道正文

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.

† Concordia University and MILA Research Institute
‡ Johns Hopkins University and Flatiron Institute
** Work done while at Apple

### Understanding Input Selectivity in Mamba

July 4, 2025research area Methods and Algorithms, research area Speech and Natural Language Processing conference ICML

State-Space Models (SSMs), and particularly Mamba, have recently emerged as a promising alternative to Transformers. Mamba introduces input selectivity to its SSM layer (S6) and incorporates convolution and gating into its block definition. While these modifications do improve Mamba’s performance over its SSM predecessors, it remains largely unclear how Mamba leverages the additional functionalities provided by input selectivity, and how these…

关联知识实体

报道正文

Related readings and updates.