Deterministic Continuous Replacement: Fast and Stable Module Replacement in Pretrained Transformers¶

会议: NeurIPS 2025 (ScaleOPT Workshop)
arXiv: 2511.18670
代码: 暂未公开（作者称将在扩展版本中发布）
领域: 模型压缩
关键词: 模块替换, 确定性混合, 梯度方差, 知识蒸馏, Vision Transformer

一句话总结¶

DCR 通过确定性退火权重 α(t) 混合 teacher 和 student 模块输出，消除了随机门控（如 BERT-of-Theseus）带来的梯度方差，在冷启动模块替换场景下实现更快收敛和更强的特征对齐。

研究背景与动机¶

领域现状：随着训练成本攀升，模型适配（model adaptation）成为关键方向。主流做法包括用更小的代理模块替换原始块（压缩）和用高效注意力变体（Linformer/Performer 等 O(n) 复杂度）替换标准自注意力。
现有痛点：在冻结的预训练骨干网络中替换模块时，冷启动的随机初始化模块会产生分布外特征，导致下游层收到异常输入、优化不稳定、梯度更新无效、恢复缓慢。
核心矛盾：知识蒸馏需要昂贵的全 teacher 前向传播且强制刚性特征匹配；BERT-of-Theseus 等随机替换方法使用 Bernoulli 门控 \(z_\ell(t) \sim \text{Bernoulli}(p(t))\)，在 p(t) 处于中间值时引入大量梯度方差。
本文要解决什么？ 核心问题是"如何在冻结骨干中稳定地集成一个随机初始化的新模块"——即模块替换的稳定性问题。
切入角度：将随机门控替换为确定性混合权重，从理论上消除门控导致的梯度方差项。
核心 idea 一句话：用确定性退火的 α(t) 代替随机伯努利门控，理论上消除 gate-induced 梯度方差，实践中实现更快收敛。

方法详解¶

整体框架¶

给定预训练网络 F 的 L 个模块，对替换子集 \(\mathcal{I} \subseteq \{1,...,L\}\) 中的每个模块 ℓ，保留冻结的 teacher 模块 \(T_\ell\) 并训练 student 模块 \(S_\ell(\cdot;\theta_\ell)\)。DCR 在残差分支上做确定性混合：

\[x_{\ell+1}(t) = x_\ell(t) + [\alpha(t) T_\ell(h_\ell(t)) + (1-\alpha(t)) S_\ell(h_\ell(t); \theta_\ell)]\]

其中 \(\alpha(t) \in [0,1]\)：\(\alpha(0) = 1\)（纯 teacher）→ \(\alpha(T) = 0\)（student 完全接管），\(h_\ell = \text{LN}(x_\ell)\)。

关键设计¶

确定性退火门控（Deterministic Gate）:
做什么：用全局确定性权重 α(t) 线性混合 teacher 和 student 输出
核心思路：aggr20 调度——前 10% 训练 \(\alpha: 1.0 \to 0.3\)，10%-20% \(\alpha: 0.3 \to 0.0\)，之后 \(\alpha = 0\) student 独立运行
设计动机：确定性门控完全消除了 gate-induced 梯度方差。对于 Theseus 硬门控 \(z \sim \text{Bernoulli}(p)\)，梯度方差的门控分量为 \(p(1-p)\mathbb{E}\|a\|^2\)；DCR 的确定性 α 使条件方差 \(\text{Var}(\nabla_{\theta_\ell} L | X) = 0\)
Deep Feature Guidance (DFG):
做什么：在替换位置添加辅助 L2 对齐损失
核心思路：\(\mathcal{L}_{\text{DFG}} = \sum_{\ell \in \mathcal{I}} \|S_\ell(h_\ell) - T_\ell(h_\ell)\|_2^2\)，因为 DCR 已经计算了两个分支的输出，所以 DFG 几乎零额外成本
设计动机：与标准蒸馏不同，不需要全 teacher 模型前向传播，仅在替换层做局部对齐。DFG 权重 λ 与 α 同步退火
曲率偏差消除（Curvature Bias）:
做什么：消除随机混合通过非线性层时的曲率偏差
核心思路：Theseus 随机选择在非线性函数 ψ 后产生期望偏差 \(|\mathbb{E}[\psi(Y)] - \psi(\mu)| \leq \frac{M}{2} p(1-p)\|\Delta\|^2\)，而 DCR 的确定性路径使 \(\mathbb{E}[\psi(Y_\alpha)] = \psi(Y_\alpha)\)，不存在混合偏差

损失函数 / 训练策略¶

\[\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}}(\hat{y}, y^\star) + \lambda \mathcal{L}_{\text{DFG}}\]

训练分两阶段：(1) 分类头预热 6 epochs，lr \(1 \times 10^{-3}\)；(2) 全模型训练 50 epochs，lr \(5 \times 10^{-4}\)，AdamW，权重衰减 0.05，梯度裁剪 1.0，batch size 128，label smoothing 0.1，混合精度 BF16。

实验关键数据¶

主实验¶

在 ImageNet 预训练 ViT-Small 上，CIFAR-100 微调后做注意力模块自替换实验。

方法	梯度方差	额外计算	需要特征匹配	收敛速度
知识蒸馏	低	高（全 teacher 前向）	是（刚性）	中
Theseus (随机)	高（门控引入）	低	否	慢
DCR (本文)	低（确定性）	低（仅替换层 teacher）	否（可选 DFG）	快

DCR+DFG 在 epoch 和 wall-clock 两个维度上均最快达到目标精度，最终精度约 78-80%。DCR+DFG 在所有层（Block 0/7/11）的 teacher-student 接口余弦相似度均显著高于随机基线。

消融实验¶

配置	接口对齐质量	收敛速度	说明
DCR + DFG	最高	最快	确定性混合 + 特征引导
DCR only	高	快	纯确定性混合
GUM (Gumbel)	中	中	软随机门控仍有方差
GUM + DFG	中偏高	中	软门控 + 特征引导
BERN (Theseus)	低	慢	硬门控方差最大
Student-only	低	最慢	无渐进替换

关键发现¶

确定性混合使下游层从一开始就收到分布内特征，避免了 GUM/BERN 中深层收敛被 gate-induced 梯度饥饿延迟的问题
DFG 在深层（Block 11）提升最显著，确认了近零成本特征引导与确定性混合的协同效应
即使在非计算饱和的小规模实验中，DCR 也展现了明确的优势

亮点与洞察¶

零额外前向传播的特征对齐：DCR 已经计算了 teacher 和 student 输出用于混合，DFG 利用这个"免费"的信号做对齐，与需要全 teacher 前向的标准蒸馏形成鲜明对比
理论驱动的方法设计：完整的方差分解理论（Prop 1-4）不仅解释了方法为何有效，还精确量化了 DCR 相对于随机方法的优势大小 \(p(1-p)\mathbb{E}\|a\|^2\)
可迁移到异构算子替换：虽然实验在自替换（attention→re-init attention）上验证，但方法和理论都是为异构替换（attention→Linformer/Performer）设计的

局限性 / 可改进方向¶

仅在 ViT-Small + CIFAR-100 的小规模设定下验证，单种子实验，缺乏大规模模型和数据集结果
自替换实验（同构替换）隔离了稳定性问题但未验证异构替换的实际效果
全局 α(t) 调度没有考虑层间差异——自适应的按层调度（基于接口相似度）可能进一步提升效果
未讨论 Batch Norm 架构或 post-norm Transformer 下的行为
作为 Workshop 论文，实验对比不够全面，缺少 Net2Net、CKA matching 等更强基线

评分¶

新颖性: ⭐⭐⭐⭐ 确定性混合消除门控方差是直觉简单但理论扎实的贡献
实验充分度: ⭐⭐ 小规模单种子实验，未涉及异构替换和大模型
写作质量: ⭐⭐⭐⭐ 理论推导清晰，实验局限性自知且诚实
价值: ⭐⭐⭐ 为模块替换提供了理论基础，但需后续工作验证大规模可行性