一句话总结

复旦大学、字节跳动 Seed 与香港中文大学的研究人员提出 FlashMorph，通过将层选择形式化为预算约束下的子集优化，利用可变形模型与线性化正则化捕捉全局层间依赖，将 Transformer 变形为混合注意力模型，在降低选择成本的同时优于启发式策略。

核心贡献

• 将混合层选择形式化为预算约束下的子集优化问题以捕捉层间依赖，而非依赖孤立的逐层评分。
• FlashMorph 在每个全注意力层上构建带有线性注意力分支的可变形模型，随后利用合成长上下文检索数据上的线性化正则化，在全注意力预算下联合优化逐层门控。
• 在 Qwen3 主干和多种线性注意力变体上的评估表明，FlashMorph 保持长上下文检索与推理性能，同时相较于启发式和逐层评分方法显著降低层选择成本。

引言

Transformer 的 softmax 注意力随序列长度呈二次方增长，使得长上下文推理因不断增大的 KV 缓存而代价高昂。仅在部分层保留全注意力而其余层使用线性注意力的混合模型可以缓解该问题，但将预训练 Transformer 转换为高效的混合模型需要精确选择哪些层保留全注意力。先前的选择方法依赖启发式策略，如均匀放置、昂贵的搜索或孤立的逐层评分，这些方法忽略了被联合线性化时层之间的交互，且往往产生巨大的选择开销。作者提出 FlashMorph，该方法将混合层选择视为预算约束下的联合优化。它构建一个可变形模型，每个注意力层可通过可学习的门控在全注意力和线性模式之间连续变化，随后在合成检索数据上对齐门控以捕捉层间互补与冗余，最后离散化为最终的混合架构，从而改善质量—效率权衡并大幅降低选择成本。

方法

作者提出 FlashMorph，一种将标准 Transformer 模型转换为混合注意力架构的高效且可扩展的层选择方法。总体框架将混合层选择形式化为预算约束下的子集优化问题，超越了固定的分配模式和孤立的逐层评分。

流程始于可变形层的构建。为给每个全注意力层配备训练好的线性注意力替代，作者从原始预训练 Transformer 蒸馏出一个全线性模型。在此阶段，原始全注意力模型的参数被冻结，仅训练线性注意力分支。它们使用逐层隐藏状态对齐损失优化这些分支： $\mathcal{L}_{\text{hidden}} = \frac{1}{L} \sum_{l=1}^{L} \left\| \mathbf{H}_{\text{lin}}^{(l)} - \mathbf{H}_{\text{full}}^{(l)} \right\|_{2}^{2}$Lhidden​=L1​∑l=1L​​Hlin(l)​−Hfull(l)​​22​ 由此获得训练好的全线性学生模型，为原始架构中的每一层提供线性注意力替代。

随后，作者通过联合优化与线性化正则化进行层选择。为避免穷举评估所有可能的混合配置带来的不可行性，他们将离散子集选择问题松弛为基于逐层门控的连续优化问题。对每一层 $l$l，标量门控 $\alpha^{(l)}$α(l) 在全注意力分支和线性注意力分支之间插值： $\mathbf{H}_{\text{mix}}^{(l)} = \alpha^{(l)} \mathbf{H}_{\text{full}}^{(l)} + (1 - \alpha^{(l)}) \mathbf{H}_{\text{lin}}^{(l)}$Hmix(l)​=α(l)Hfull(l)​+(1−α(l))Hlin(l)​ 门控值在保持全注意力主干和已训练线性注意力分支均冻结的情况下被优化。为保持原始模型的行为，作者在答案 token 位置计算对齐损失： $\mathcal{L}_{\text{align}} = \frac{1}{L |\mathcal{T}(x)|} \sum_{l=1}^{L} \sum_{t \in \mathcal{T}(x)} \left\| \mathbf{H}_{\text{mix}, t}^{(l)} - \mathbf{H}_{\text{full}, t}^{(l)} \right\|_{2}^{2}$Lalign​=L∣T(x)∣1​∑l=1L​∑t∈T(x)​​Hmix,t(l)​−Hfull,t(l)​​22​ 为鼓励效率，引入线性化正则化项以惩罚对全注意力的依赖： $\mathcal{L}_{\text{reg}} = \sum_{l=1}^{L} \alpha^{(l)}$Lreg​=∑l=1L​α(l) 最终优化目标结合这些项： $\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{align}} + \lambda \mathcal{L}_{\text{reg}}$Ltotal​=Lalign​+λLreg​ 该优化在合成长上下文检索数据上进行，其中随机生成的通行密钥被置于长文档的不同深度。这提供了一个定向信号，用于衡量用线性注意力替换全注意力是否破坏了长距离信息获取。

优化后，学得的门控值指示了每层保留全注意力的相对必要性。给定全注意力层的目标预算 $K$K，作者通过选择门控值最大的层来离散化混合层： $\mathcal{I}_{\text{full}}^{\text{Hybrid}} = \operatorname{TopK} \left( \{ \alpha^{(l)} \}_{l=1}^{L}, K \right)$IfullHybrid​=TopK({α(l)}l=1L​,K) 其余层使用其训练好的线性注意力替代实例化。

最后，作者应用 logits 蒸馏和长上下文微调以恢复所选混合注意力模型的质量。混合模型通过 Kullback-Leibler 散度从原始全注意力教师蒸馏： $\mathcal{L}_{\text{KD}} = D_{\text{KL}} (p_{\text{T}}(\mathbf{X}) \parallel p_{\text{H}}(\mathbf{X}))$LKD​=DKL​(pT​(X)∥pH​(X)) 随后，混合模型在标准语言建模目标下于长上下文序列上进行微调： $\mathcal{L}_{\text{FT}} = - \sum \log p_{\text{H}}(x_{t} \mid x_{<t})$LFT​=−∑logpH​(xt​∣x<t​) 至此完成 FlashMorph 的 Transformer 到混合模型的转换流程。

实验

FlashMorph 采用联合优化，以远超先前选择方法的效率选择全注意力层，所需的 selection tokens 少数个数量级，同时在各种混合比例和线性注意力主干上均实现强大的长上下文检索与回想性能。它保持了有竞争力的常识推理能力，并显著提升回想密集型任务性能，尤其是在检索导向的合成监督指导下。选择的混合架构带来显著的推理加速和显存降低，从而在超越全注意力基线的前提下实现可扩展的长序列处理。

FlashMorph 是唯一同时具备非孤立和基于优化性质的混合层选择方法，并在基于搜索的方法中实现了最低的选择成本。相比之下，其他方法要么牺牲联合审视层的能力（非孤立），要么放弃优化，或者需要远超此方法的 token 预算。FlashMorph 同时满足非孤立选择和基于优化选择，而仅需 20M tokens，远低于 KL‑LS（20B）和 PostNAS（50B）。HALO 的 token 成本较低，为 234M，但既非孤立也非基于优化；KL‑LS 和 PostNAS 虽然是基于优化的，但无法做到非孤立或伴随极高的选择成本。

FlashMorph 仅使用 2000 万层选择 token 即实现了强大的 needle-in-a-haystack 检索。原始模型在 64K 以上崩溃，YaRN 无法恢复长上下文性能，而先前的混合方法在长序列上的多针变体中丢失精度。FlashMorph 在单针检索上接近或达到完美，且多针准确率显著提升，这种增益在 1.7B 规模上更为明显。原始 Qwen3 模型在 32K 和 64K 处达到完美检索，但在 128K 和 256K 的所有 NIAH 变体上准确率降至零。应用 YaRN 会使所有上下文长度的性能退化至接近零，始终无法恢复长上下文检索能力。均匀混合注意力在单针上保持高精度，但在 256K 的更困难多针设置上得分很低（例如降至百分之十几）。KL‑LS 和 HALO 相较均匀有所改善，但在 256K 的 NIAH‑Single‑2 和 Single‑3 上仍急剧下降。FlashMorph 在 0.6B 主干上实现近乎完美的 NIAH‑Single‑1，在 1.7B 主干上则实现完美 NIAH‑Single‑1。在多针任务上，FlashMorph 相较于先前的层选择方法提供明显增益，尤其在短至中等的序列长度上。

在 0.6B 主干的零样本常识推理上，FlashMorph 在不同注意力主干间保持竞争力，同时在回想密集型任务上的性能大幅超越其他高效注意力变体。在 Lightning 注意力下，它达到线性方法中最高的回想平均分，所需 selection tokens 远少于其他方案，并大幅缩小了与原始全注意力模型的差距。FlashMorph 搭配 Lightning 注意力的回想密集型平均分为 60.1，优于 HALO（55.9）和均匀线性转换（51.1）。FlashMorph 仅用 20M tokens 进行层选择，而 KL‑LS 需 20B、HALO 需 234M，同时取得与其他高效方法相似的常识平均分（52.7）。在挑战性的 SQuAD 任务上，FlashMorph 将准确率从 HALO 的 34.7 提升至 41.7，显著缩小了与全注意力 Qwen3（44.1）的差距。

FlashMorph 仅用 2000 万 token 和 2.1 GPU 小时完成层选择，大幅降低了先前方法的开销。与 PostNAS 相比，GPU 小时减少超过 1200 倍；与 KL‑LS 相比减少 510 倍；与 HALO 相比减少 7.3 倍，同时所需 token 数量低几个数量级。FlashMorph 使用 20M tokens，使 token 成本比 PostNAS（50B）低 2500 倍，比 KL‑LS（20B）低 1000 倍。GPU 小时从 PostNAS 的 2561.3 降至 FlashMorph 的 2.1，降低 1219.7 倍。FlashMorph 实现了最低的 FLOPs（2.5e17），显著低于 HALO 的 6.5e17 FLOPs。

实验评估了 FlashMorph，一种仅用 2000 万 token 联合优化层分配的混合层选择方法，涵盖长上下文 needle-in-a-haystack 检索、零样本推理、回想密集型任务和计算效率。当原始模型在 64K 以上崩溃，且 YaRN、均匀混合注意力、KL‑LS 和 HALO 等方法性能急剧退化时，FlashMorph 实现了近乎完美的单针准确率和显著的多针提升，同时保持了有竞争力的常识推理，并在如 SQuAD 的回想基准上取得明显改进。关键的是，FlashMorph 以 GPU 小时减少超 1200 倍、token 成本降低数个数量级的方式完成层选择，展示了选择成本与下游性能之间的帕累托最优平衡。