论文

重新审视条件深度路由的辅助损耗:一项实证研究

Revisiting Auxiliary Losses for Conditional Depth Routing: An Empirical Study

模型推理测试时计算扩展

摘要

条件深度执行通过轻量级廉价 FFN 路由词元子集,而其余词元在每个受控层执行标准完整 FFN。核心困难是门训练:门决策必须在影响语言建模(LM)损失之前通过许多层传播,因此产生的梯度很弱且有噪声。辅助损失通常被堆叠起来以稳定训练,但它们之间的相互作用——特别是预测辅助和显式分数监督之间的相互作用——尚未在受控条件下进行系统比较。我们在 1.575 亿参数的纯解码器模型下评估了两种门设计,该模型具有纯控制器训练、50% 全路径预算以及在fineweb-edu 子集上运行的 3 种子。 MLP门(G1)将当前隐藏状态映射到效用分数; JEPA 引导门(G3)添加了一个动作条件预测器,该预测器在低维潜在空间中预测每个词元的完整执行与廉价执行的结果,与固定目标头对齐。在具有预言式效用回归和成对排序监督(util/rank)的标准配方下,G3 在 3/3 种子中比 G1 改进了早期到中期的优化(较低的平均 LM,更快的阈值命中,约 10.3 倍较低的梯度规范),在 0.005 启发式参考内具有 20k 步端点 LM。一个关键发现(消融 A3):联合删除 util/rank 提高了两个门的 3/3 种子中的最佳/平均 LM 和阈值命中速度,并且 G3 相对于 G1 的早中期优势消失了。我们将其追溯到 离策略 预言机标签,该标签假设所有后续层都执行完整,而门控执行仅路由完整的一小部分 - 使 util/rank 在当前配方下净负。删除 util/rank 还将训练 FLOPs 代理从约 1.53 倍减少到约 1.07 倍(在 V100-32GB 上从 2.87 小时减少到 1.75 小时,约 39%)。结论仅限于所研究的制度。