从Attention到Mamba：跨架构蒸馏的实用配方

为什么需要跨架构蒸馏？

状态空间模型（SSM）如Mamba，正成为Transformer的有力替代品。相比基于注意力的模型，Mamba在生成时内存占用更低、吞吐量更高。但社区对如何训练Transformer积累了丰富经验，且有大量预训练模型可用。如何利用这些现成的Transformer资源，加速Mamba类架构的落地？跨架构蒸馏提供了一条路径：将Attention模型的知识迁移到Mamba架构中。

状态空间模型（SSM） 是一类用隐状态递归描述序列的模型，其连续时间形式为 (h'(t)=Ah(t)+Bx(t))。Mamba 将其离散化并引入输入依赖的选择机制（selective SSM），使状态转移矩阵随输入变化，从而在保持线性复杂度的同时具备一定的内容感知能力。与 Transformer 的自注意力相比，SSM 在自回归生成时只需维护固定大小的状态，无需缓存全部历史 KV，因此显存占用不随序列长度增长，长序列吞吐优势明显。但 SSM 的并行训练技巧、初始化与稳定性调优远不如 Transformer 成熟，这构成了跨架构蒸馏的现实动机。

蒸馏配方：关键设计选择

论文旨在确定一套有效配方，将基于注意力的模型蒸馏为Mamba类架构。虽然原文未展开具体技术细节，但从问题设定可以推断，蒸馏需要解决两个核心挑战：架构差异带来的表示对齐问题，以及如何保留Transformer中习得的知识。

架构差异的应对

Attention和SSM在计算范式上存在本质区别：前者通过注意力权重动态聚合信息，后者通过状态转移隐式建模序列。蒸馏过程中，需要设计合适的损失函数和中间层对齐策略，使Mamba能够模仿Transformer的行为。

知识蒸馏通常指学生模型在教师输出的软标签（logits）上最小化 KL 散度，并可选地匹配中间层特征。跨架构蒸馏的难点在于教师与学生的隐空间维度、层数和计算图均不匹配，因此常见做法包括：用可学习的投影层将学生隐状态映射到教师维度后再做对齐损失；在输出分布层面做温度缩放的软目标匹配；以及引入注意力图或状态转移模式的间接监督。这些策略共同的作用是提供比硬标签更丰富的监督信号，缓解架构异构带来的对齐困难。

利用预训练权重

配方强调利用现有预训练Transformer，这意味着蒸馏过程可以站在巨人的肩膀上。通过精心设计的蒸馏目标，Mamba学生模型能够继承教师模型的能力，同时保持自身在推理效率上的优势。

潜在影响与展望

如果这套配方被验证有效，将大幅降低采用SSM的门槛。开发者无需从头训练Mamba模型，而是可以基于成熟的Transformer生态进行知识迁移。这对于需要长序列处理、高吞吐生成的场景尤其有价值，比如实时对话系统、文档级翻译等。

不过，蒸馏后的Mamba能否完全复现Transformer的复杂推理能力，仍需进一步实验验证。跨架构蒸馏的通用性也有待拓展到更多模型家族。

总结

跨架构蒸馏为Transformer到Mamba的迁移提供了实用方案。通过系统化的蒸馏配方，我们有望在保留Transformer知识的同时，享受SSM带来的效率提升。这项工作的价值在于 bridging 两个活跃的研究社区，推动更高效模型架构的落地。