条件深度路由的辅助损失：一次经验研究
摘要
条件深度执行在每个受控层上将部分 token 通过轻量廉价前馈网络（FFN），而其余 token 执行标准完整 FFN。核心难点在于门控训练：门控决策必须在影响语言模型（LM）损失之前经历多个层，导致梯度稀疏且噪声较大。常规做法通过堆叠辅助损失来稳定训练，但不同辅助损失之间的相互作用——尤其是预测性辅助损失与显式评分监督之间的关系——在受控条件下尚未系统比较。我们在 1.575 亿参数的解码器-only 模型下评估了两种门控设计，采用 controller-only 训练、50% 完整路径预算，并在 fineweb-edu 子集上进行 3 次独立实验。MLP 门控（G1）将当前隐藏状态映射到效用得分；JEPA 引导的门控（G3）增加一个动作条件预测器， forecasts 在低维潜空间中执行完整 vs. 廉价 FFN 的结果，与固定目标头对齐。标准配方下，G3 在 3/3 随机种子中实现更早至中期优化优势（平均 LM 降低、阈值触发更快、约 10.3 倍梯度范数降低），且 2 万步端点 LM 距离经验参考仍在 0.005 范围内。关键发现（消融实验 A3）：联合移除 util/rank 可在 3/3 随机种子中改善最佳/平均 LM 及阈值触发速度，两者门控均获益；G3 对 G1 的早至中期优势随之消失。我们追溯到这一现象：基于全后续层均执行完整 FFN 的保守标签（off-policy oracle label）与当前仅路由部分通过完整 FFN 的实际执行方式不符，使得 util/rank 在当前配方下产生负效应。移除 util/rank 还将训练 FLOPs 代理从约 1.53 倍降至约 1.07 倍（V100-32GB 上从 2.87 小时缩短至 1.75 小时，约 39%）。结论局限于所研究的 regimes。
引用
@article{arxiv.2604.17228,
title = {Revisiting Auxiliary Losses for Conditional Depth Routing: An Empirical Study},
author = {Qingwei Lin},
journal= {arXiv preprint arXiv:2604.17228},
year = {2026}
}
备注
23 pages, 4 figures. Preprint. Controlled empirical study with 3-seed runs at 157.5M parameters; includes a negative result on oracle-style utility/rank supervision for conditional depth routing