mlsanityData to 13 Oct 2026

Revisiting Auxiliary Losses for Conditional Depth Routing: An Empirical Study

arXiv 2604.17228 ↗First posted 19 Apr 2026cs.LGcited by 0

Not on the map. Why?

Abstract (arXiv)

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.

Key facts. No citing papers in the map yet (data through 13 Oct 2026). Builds on 18 of its 25 references that are in the map. First posted on arXiv 19 Apr 2026.

Builds on

18 of 25 references.

References · 18

  1. 2024Qwen2.5 Technical Reportcited by 4,840
  2. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,463
  3. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,044
  4. 2016Categorical Reparameterization with Gumbel-Softmaxcited by 2,726
  5. 2013Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computationcited by 2,056
  6. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,618
  7. 2024The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scalecited by 1,126
  8. 2023Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecturecited by 1,018
  9. 2021GLaM: Efficient Scaling of Language Models with Mixture-of-Expertscited by 896
  10. 2015Scheduled Sampling for Sequence Prediction with Recurrent Neural Networkscited by 786
  11. 2022Mixture-of-Experts with Expert Choice Routingcited by 619
  12. 2016Adaptive Computation Time for Recurrent Neural Networkscited by 491
6 more
  1. 2022ST-MoE: Designing Stable and Transferable Sparse Expert Modelscited by 396
  2. 2022Confident Adaptive Language Modelingcited by 302
  3. 2024LayerSkip: Enabling Early Exit Inference and Self-Speculative Decodingcited by 301
  4. 2024Mixture-of-Depths: Dynamically allocating compute in transformer-based language modelscited by 250
  5. 2019Depth-Adaptive Transformercited by 230
  6. 2024Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Expertscited by 200

Built on by

Not cited yet.