改进条件深度路由的辅助损失设计，发现传统方法可能适得其反。
Revisiting Auxiliary Losses for Conditional Depth Routing: An Empirical Study
- 用预测性辅助损失与显式评分监督联合训练门控机制
- 移除传统辅助损失后，模型优化更快、性能更稳，梯度噪声降低10.3倍
- 适合关注高效推理与门控训练优化的研究者
条件深度执行在每层仅让部分令牌通过轻量级计算路径，其余走完整路径。核心挑战是门控训练：门控决策需经过多层才影响语言建模损失，导致梯度弱且噪声大。常用辅助损失用于稳定训练，但其相互作用——尤其是预测性辅助与显式评分监督之间——尚未在受控条件下系统比较。本研究在157.5M参数的解码器单向模型上，以控制器仅训练、50%全路径预算、fineweb-edu子集上进行3次随机种子实验，评估两种门控设计：基于MLP的门（G1）直接从隐藏状态映射出效用得分；基于JEPA引导的门（G3）引入动作条件预测器，在低维潜在空间中预测每令牌执行完整或轻量路径的后果，并对齐固定目标头。在标准方案（使用真值效用回归与成对排序监督）下，G3在3个种子中均表现更优（平均语言建模损失更低、阈值达成更快，梯度范数降低约10.3倍），20,000步终点损失接近0.005参考值。关键发现（消融A3）：同时移除效用/排序辅助损失后，两种门控在3个种子中均实现更好/平均的LM性能和更快的阈值命中速度，且G3的优势消失。这归因于一个非策略真值标签——假设后续所有层都执行完整路径，而实际仅部分路径如此——导致效用/排序监督在当前设置下整体为负。移除该损失后，训练浮点运算量代理从约1.53倍降至约1.07倍（V100-32GB上耗时从2.87小时降至1.75小时，节省约39%）。结论限于本研究设定范围。
原文摘要 · Abstract (English)
Conditional depth execution routes a subset of tokens through a lightweight cheap FFN while the remainder execute the standard full FFN at each controlled layer. The central difficulty is gate training: the gate decision must propagate through many layers before it influences the language modeling (LM) loss, so the resulting gradients are weak and noisy. Auxiliary losses are commonly stacked to stabilise training, yet the interactions among them -- particularly between a predictive auxiliary and explicit score supervision -- have not been systematically compared under controlled conditions. We evaluate two gate designs under a 157.5M-parameter decoder-only model with controller-only training, 50% full-path budget, and 3-seed runs on a fineweb-edu subset. The MLP gate (G1) maps the current hidden state to a utility score; the JEPA-guided gate (G3) adds an action-conditional predictor that forecasts, in a low-dimensional latent space, the outcome of executing full vs. cheap per token, aligned against a fixed target head. Under the standard recipe with oracle-style utility regression and pairwise rank supervision (util/rank), G3 improves early-to-mid optimisation over G1 in 3/3 seeds (lower avg LM, faster threshold hits, ~10.3x lower grad norms), with 20k-step endpoint LM within a 0.005 heuristic reference. A key finding (ablation A3): jointly removing util/rank improves best/avg LM and threshold-hit speed in 3/3 seeds for both gates, and the early-to-mid advantage of G3 over G1 disappears. We trace this to an off-policy oracle label that assumes all subsequent layers execute full, whereas gated execution routes only a fraction through full -- making util/rank net-negative under the current recipe. Removing util/rank also cuts the training FLOPs proxy from ~1.53x to ~1.07x full-only (2.87h to 1.75h on a V100-32GB, ~39%). Conclusions are scoped to the studied regime.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。