Transformer可以改装成Mamba了：苹果把推理成本直接打成线性

来源机器之心 本文约2500字建议阅读5分钟一种从Transformer到Mamba的新型跨架构蒸馏方法。最近苹果又整了个活儿很工程、也挺关键把又贵又强的 Transformer改造成又便宜又差不多强的 Mamba。而且性能基本没怎么掉。听起来炸不炸有点像水电煤级别的升级。为啥要折腾这个很简单Transformer 确实猛这十年基本就是靠它打天下。但问题也很现实它越长越贵而且是平方级变贵。短文本还好一旦上到长上下文比如代码、agent、多轮推理那成本就不是有点高而是直接肉疼。于是大家就开始找替代方案。比如线性 attention、RWKV还有这两年很火的 Mamba。这些模型的思路都很统一别再平方爆炸了改成线性。好处也很直接更快、更省显存、推理更丝滑。但问题来了 它们不够强。尤其一旦规模上去和 Transformer 还是有差距。于是就卡住了一个经典局面要性能用 Transformer但贵 要便宜用 Mamba但弱一点那有没有可能「既要又要」能不能不重训把 Transformer 的能力直接搬到 Mamba 上直接蒸馏性能肯定得崩。Transformer 像那种随时翻资料的学霸 Mamba 更像全靠记忆的选手。你突然让一个翻书型选手闭卷考试还不给过渡—— 那基本就是灾难现场。苹果没有硬刚而是搞了一个「两步走」先造一个「中间形态」让 Transformer 先变成一个更简单、更接近 Mamba 的版本。再从这个中间版本转成 Mamba。用来把 Transformer 里原本计算很贵的 Softmax Attention换成一个更省算力的 Mamba 模块。中间加了一个过渡形态先把 Attention 变成一种学出来的线性 AttentionHedgehog再结合Mamba最终得到一个混合模块叫HedgeMamba。第一步把原始的 softmax attention 换成一种线性 attention同时尽量不丢性能。问题在于传统线性注意力一直有明显短板和 softmax attention 的表现差距很大。为了解决这个差距他们依据Mercer 定理用了 Hedgehog 方法用一个小型神经网络MLP去学习一种特征映射让线性 attention 尽量模仿原来的行为。再通过余弦相似度cosine similarity蒸馏让这个新结构在输出上对齐 Transformer。这一步做完相当于得到一个「更便宜但还挺像 Transformer」的中间模型。第二步是把这个已经对齐好的线性 attention进一步嵌入到 Mamba 结构里。他们做了一件很关键的事情把 attention 里的核心计算方式对应映射到 Mamba 的内部参数上让 Mamba 在初始化的时候行为就已经接近前一步的模型而不是从零开始学。这一步本质是在做结构对齐。不过还有一个问题需要处理。原始 attention 自带一种归一化机制而线性版本没有所以他们额外加了一步归一化让输出形式更接近原来的 attention同时又不破坏计算效率。完成这些之后才进入真正的训练阶段。他们对整个模型做微调用标准的交叉熵训练并重新启用 Mamba 原本的能力比如卷积和门控 Gate 机制。这一步的作用是让模型不只是模仿而是用自己的方式把能力重新学出来。整套方法的关键不在某一个技巧而在这条路径本身先让两种模型在「表达方式」上对齐再做结构转换最后通过训练把能力恢复出来。也正是因为这样分步处理才避免了直接蒸馏时常见的性能崩塌问题。效果到底咋样性能几乎没掉但成本逻辑已经变了。论文里最关键的一张表把三类1B模型摆在一起对比Transformer教师模型Pythia、传统蒸馏基线 Hedgehog 以及他们的方法 HedgeMamba 。在只用10B token大约是教师训练数据的 2.7% 的情况下结果非常直接——教师模型的困惑度是 13.86基线方法掉到 14.89而HedgeMamba把这个指标拉回到 14.11已经贴得很近了。他们用一个大约10B token训练出来的1B模型做实验最后得到的 Mamba 模型能够保留原始 Pythia-1B Transformer 在下游任务中的性能其困惑度perplexity保持在 14.11接近老师模型的 13.86。这件事的含义其实挺重的。过去大家默认一个前提只要你把 Transformer 换成另一种架构性能就会明显掉一截。但这篇论文给出的答案是这个损失可以被大幅追回来。而且不只是语言建模指标好看在Arc、PIQA、BoolQ、RACE、LogiQA 这些下游任务上HedgeMamba基本全面超过基线同时整体表现已经逼近教师模型。这说明它保留下来的不只是表面的概率分布而是相当一部分推理能力和语义结构。更关键的是这种效果不是调出来的而是有方法论支撑的。他们尝试过最直接的做法——从 Transformer 一步蒸馏到 Mamba结果是 PPL 直接炸到 100 以上几乎不可用。换句话说两阶段蒸馏在这里不是优化而是绕不过去的结构性条件。后面的消融和分析其实是在解释这条路径为什么成立。比如架构上真正起作用的不是简单叠模块而是门控机制——也就是让模型学会该记什么、不该记什么训练策略上两阶段的数据分配也不是平均最优而是明显偏向第二阶段说明中间表示只是过渡真正的能力是在后半段完成迁移再看数据规模从1B 到 10B token性能是稳定往上走的没有出现不收敛或反复震荡的情况这一点很重要因为它说明这条路线是可以规模化放大的。架构消融实验表明让 Mamba 好用的关键不是简单堆结构而是门控机制。蒸馏的两步S1 和 S2数据到底该怎么分配才最有效两阶段蒸馏是必要的而且最优策略是轻S1 重S2。蒸馏过程中 token 数量训练数据量对效果的影响。似乎只要给足够多的蒸馏数据Mamba 可以逼近 Transformer 的性能。把这些信息合在一起这篇工作的价值就不在「又做了一个更好的模型」而在于它提供了一种新的工程可能性。过去几年积累的大量 Transformer 模型并不需要全部推倒重来而是有机会通过一套流程被「转制」为更高效的形态。如果这件事能稳定复现那整个开源模型生态、甚至很多公司的自研模型都有机会被整体降本重构。参考链接https://arxiv.org/abs/2604.14191编辑文婧关于我们数据派THU作为数据科学类公众号背靠清华大学大数据研究中心分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识努力建设数据人才聚集平台、打造中国大数据最强集团军。新浪微博数据派THU微信视频号数据派THU今日头条数据派THU