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

AIGC动态5小时前发布 almosthuman2014

49 0 0

文章摘要

苹果近期提出了一种将 Transformer 模型转化为 Mamba 架构的新方法，旨在解决前者计算成本高且后者性能较弱的问题。传统 Transformer 虽强但成本随长度平方级增长，而 Mamba 虽线性高效却性能不足。直接蒸馏会导致性能崩塌，因此研究采用了两步走策略。首先构建中间形态，利用 Hedgehog 方法将 Softmax Attention 转为线性 Attention，并通过蒸馏对齐输出。其次将对齐后的线性注意力嵌入 Mamba 结构，映射核心计算方式并添加归一化，最后微调恢复模型能力。整套方法的关键不在某一个技巧，而在这条路径本身：先让两种模型在表达方式上对齐，再做结构转换，最后通过训练把能力恢复出来。

实验结果显示，该方法显著降低了性能损失。在只用 10B token 训练 1 模型的情况下，困惑度降至 14.11，接近教师模型的 13.86，远超传统蒸馏基线。他们用一个大约 10 训练出来的 1 模型做实验，最后得到的 Mamba 模型能够保留原始 Pythia-1B Transformer 在下游任务中的性能，其困惑度保持在 14.11，接近老师模型的 13.86。下游任务测试表明，新模型保留了相当一部分推理能力和语义结构，在多个基准测试中基本全面超过基线。这证明损失可以被大幅追回来，且效果有方法论支撑而非单纯调参。

消融实验进一步验证了方法论的有效性。架构上门控机制至关重要，训练策略上两阶段数据分配应偏向第二阶段。架构消融实验表明，让 Mamba 好用的关键，不是简单堆结构，而是门控机制。此外，随着数据规模增加，性能稳定上升，证明路线可规模化。这意味着过去积累的大量 Transformer 模型有机会通过流程转制为更高效形态，为开源生态及自研模型提供降本重构的可能性。过去几年积累的大量 Transformer 模型，并不需要全部推倒重来，而是有机会通过一套流程，被转制为更高效的形态。