复制不变，蒸馏所变：线性 ViT 的初始化之道

从今天觉醒,技术赋予每一个人数字生命

---

复制不变，蒸馏所变：线性 ViT 的初始化之道

① 技术背景：当线性注意力遇上预训练权重

视觉 Transformer 的战场早已不是"能不能用"，而是"怎么用得起"。Softmax 注意力的 $O(N^2)$ 复杂度让高分辨率图像成为显存杀手，线性注意力（Linear Attention）用核函数近似把复杂度压到 $O(N)$，代价是精度和训练成本——它必须从零开始预训练，而大多数开源基础模型都是 Softmax 版本。

一个自然的想法是：既然 Softmax ViT 的权重已经学好了"看哪里"，能不能直接迁移给线性 ViT？近期关于注意力迁移的研究确实表明，注意力是 Softmax ViT 之间最有效的可迁移组件。但当我们把这条经验搬到 Softmax→线性的跨算子场景时，事情反了过来。这篇工作给出了一个清晰的结论：注意力权重是算子相关的，复制几乎无用；MLP 权重是算子无关的，直接复制就能带走大部分收益；剩下那部分差距，靠蒸馏注意力的路由行为来补。

对在校学生和转行者来说，这个结论的价值不在论文本身，而在于它示范了一种工程思维：面对"能不能复用"的问题，先拆解组件，判断哪些与实现耦合、哪些与实现解耦，再分别选择复制或蒸馏。 这套方法论可以写进作品集，也能在面试中讲清楚。

② 主流方案盘点：三条迁移路径

路径一：直接复制注意力权重。 职责是把 Softmax 注意力的 $Q{K}^T$ 参数搬进线性注意力的核函数映射中。代表做法是朴素的 state dict 拷贝。问题在于，Softmax 注意力的权重编码的是"指数归一化后的相对重要性"，而线性注意力用的是特征映射 $\varphi (q)\varphi (k{)}^T$，两者的参数空间语义不同。实验显示，这种复制有时甚至比随机初始化更差——它不是"没帮上忙"，而是"帮了倒忙"。

路径二：直接复制 MLP 权重。 MLP 承担的是逐 token 的特征变换，与 token 之间如何路由无关，因此是算子无关的。代码上就是遍历模块名，匹配 mlp 前缀做加载：

def load_mlp_weights(linear_vit, softmax_vit):
    src = softmax_vit.state_dict()
    tgt = linear_vit.state_dict()
    for name in tgt:
        if "mlp" in name and name in src:
            tgt[name] = src[name]
    linear_vit.load_state_dict(tgt)

这一条路径贡献了预训练权重的大部分收益，是整条迁移链的地基。

路径三：蒸馏注意力路由行为。 注意力权重不能复制，但注意力的"输出行为"可以学。用一个设计过的损失，让线性 ViT 的注意力输出逼近 Softmax ViT 的注意力输出，相当于把"路由决策"从教师迁移到学生。关键在损失设计：只对齐注意力输出而非 logits，避免把 Softmax 的归一化特性硬塞进线性算子。

③ 对比与优劣

维度	复制注意力	复制 MLP	蒸馏注意力
迁移对象	算子相关参数	算子无关参数	行为/输出
实现成本	极低	极低	中（需教师前向、损失调参）
单独效果	差，可能劣于随机	好，带走大部分收益	中，需配合 MLP 复制
主要风险	参数语义错配	几乎无	损失权重敏感
适用场景	不推荐单独使用	任何跨算子迁移	追求追平/超越时

三者不是互斥选项，而是有先后顺序的组合：先复制 MLP 打底，再蒸馏注意力补齐。论文的实验也表明，这个组合能让线性 ViT 缩小差距甚至反超 Softmax 版本，且在多种线性变体、模型尺寸和数据集上一致成立。

④ 选型建议

场景一：课程作业/复现实验。 只做 MLP 复制即可。改动量小，能直观看到"迁移有效"的曲线，适合作为理解预训练权重复用的第一个实验。

场景二：想把线性 ViT 用于自己的小数据集。 先复制 MLP，再冻结大部分层做轻量微调。若精度仍不达标，再加入注意力蒸馏，注意损失权重要从小开始调。

场景三：追求与 Softmax 版本持平或反超。 完整走"复制 MLP + 蒸馏注意力"两阶段，并预留教师模型的前向开销。面试常被追问的点是：为什么不能直接复制注意力？答案要落到"参数语义与算子绑定"上，而不是笼统说"不兼容"。

⑤ 未来展望

这条路线揭示了一个更普遍的问题：跨架构迁移时，哪些组件是"表示"、哪些是"机制"？ MLP 承载表示，可以复制；注意力承载机制，需要蒸馏。这个二分法是否能推广到 Mamba、RWKV 等更多非 Softmax 算子，目前仍是开放问题。另一个未解点是蒸馏损失的通用设计——当前依赖人工调参，若能自动化，跨算子迁移的门槛会进一步降低。对学习者而言，这正是值得跟进的方向：先把"复制什么、蒸馏什么"的判断力练出来，再谈具体算子。