Revisiting Auxiliary Losses for Conditional Depth Routing: An Empirical Study

• URL: http://arxiv.org/abs/2604.17228v1
• Date: Sun, 19 Apr 2026 03:20:40 GMT
• Title: Revisiting Auxiliary Losses for Conditional Depth Routing: An Empirical Study
• Abstract summary: Gate decision must propagate through many layers before it influences the language modeling (LM) loss.<n> 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.<n>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.
• Score: 31.968379218484746
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: 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.

Related papers

• The Dark Room in the Reward Channel: Dense Prediction Rewards Collapse GRPO-Trained LLM Agents -- and What Actually Works [3.040305634723913]
  We show that under group-normalized RL, a potential-based prediction reward drives every run into a degenerate absorbing state.<n>A single-episode ablation ablations the cause -- removing only GRPO's normalization turns the same reward from catastrophic into baseline parity.<n>Our central insight generalizes this: what z-scoring amplifies is a dense signal's within-group variance while all-fail groups dominate.
  arXiv  Detail & Related papers  (2026-07-23T12:50:18Z)
• Adaptive Runge-Kutta Step Control Buys Training Loss, Not Generalization: An Honest Compute-Matched Study of RK-Adam Optimizers [0.0]
  We build a representative Adam variant (Bogacki-Shampine 3(2) RK pair, FSAL reuse, local-error step control)<n>We evaluate it under a strict compute-matched protocol giving every method the same gradient-evaluation budget.
  arXiv  Detail & Related papers  (2026-07-16T03:12:58Z)
• CIGPO: Contextual Information-Gain Policy Optimization for Multi-Turn Evidence-Reading LLM Agents [1.4298580363875282]
  Training multi-turn evidence-reading agents with outcome-only reinforcement learning is unstable because intermediate turns receive little direct credit.<n>In HotpotQA experiments with Qwen2.5-3B-Instruct, GRPO initially improves (standard F1 0.430) but subsequently collapses to 100% format-violating outputs.<n>We propose a variance-textual strategy: by assigning per-turn rewards to intermediate evidence-reading turns, we prevent the group reward distribution from collapsing to a single value.
  arXiv  Detail & Related papers  (2026-06-26T11:50:29Z)
• Hidden-State Privacy Has an Empty Middle [51.56484100374058]
  We prove a complementary Fisher-ball lower bound: every full-rank Gaussian release at $O(1)$ Fisher utility admits a direction whose Mahalanobis signal grows linearly in hidden width.<n>A split-memory transformer trained from scratch reaches $G_mathrmMah in [20, 33]$ at 90M and maintains a $6$--$24times$ advantage over same-budget GPT baselines from 30M to 1B at a fixed-token language loss penalty.
  arXiv  Detail & Related papers  (2026-05-21T20:12:09Z)
• Verifier-Free RL for LLMs via Intrinsic Gradient-Norm Reward [69.99652051809737]
  In this work, we propose Verifier-free Intrinsic Gradient-Norm Reward (VIGOR)<n>VIGOR is a simple reward that uses only the policy model itself.<n>It exhibits cross-domain transfer to code benchmarks when trained only on math data.
  arXiv  Detail & Related papers  (2026-05-11T03:15:37Z)
• When Is Rank-1 Steering Cheap? Geometry, Granularity, and Budgeted Search [61.73232505100331]
  Activation steering offers a lightweight way to control LLMs without retraining, but its effectiveness varies sharply across concepts.<n>We formalize rank-1 steering as a budget-constrained optimization over intervention layer and coefficient.<n>We present textitGRACE, a Granularity- and Representation-Aware Concept Engineering framework that uses activation geometry to diagnose the dominant source of steering difficulty.
  arXiv  Detail & Related papers  (2026-05-09T14:26:49Z)
• REAL: Regression-Aware Reinforcement Learning for LLM-as-a-Judge [83.2858110368572]
  We propose textbfREAL (underlineREgression-underlineAware Reinforcement underlineLThought), a principled RL framework designed to optimize regression rewards.<n>We show that REAL consistently outperforms both regression-aware SFT baselines and standard RL methods.
  arXiv  Detail & Related papers  (2026-03-17T21:19:08Z)
• Overconfident Errors Need Stronger Correction: Asymmetric Confidence Penalties for Reinforcement Learning [17.384089089363382]
  We identify a root cause that existing methods overlook: the uniform penalization of errors.<n>Current approaches treat all incorrect rollouts within a group identically.<n>We propose the Asymmetric Confidence-aware Error Penalty (ACE)
  arXiv  Detail & Related papers  (2026-02-24T22:46:43Z)
• Late-Stage Generalization Collapse in Grokking: Detecting anti-grokking with Weightwatcher [1.6615337656760856]
  emphMemorization in neural networks lacks a precise operational definition and is often inferred from the grokking regime.<n>We identify a previously unreported third phase of grokking in this training regime: emphanti-grokking, a late-stage collapse of generalization.
  arXiv  Detail & Related papers  (2026-02-02T22:09:14Z)
• TARG: Training-Free Adaptive Retrieval Gating for Efficient RAG [46.122203287541005]
  Training-free Adaptive Retrieval Gating (TARG) is a single-shot policy that decides when to retrieve using only a short, no-latency draft from the base model.<n>On NQ-Open, TriviaQA, and PopQA, TARG consistently shifts the accuracy-efficiency frontier.
  arXiv  Detail & Related papers  (2025-11-12T23:09:52Z)
• What Layers When: Learning to Skip Compute in LLMs with Residual Gates [66.23658560048241]
  GateSkip is a residual-stream gating mechanism that enables token-wise layer skipping in decoder-only LMs.<n>Each Attention/MLP branch is equipped with a sigmoid-linear gate that condenses the branch's output before it re-enters the residual stream.
  arXiv  Detail & Related papers  (2025-10-13T16:31:50Z)
• AdaGC: Improving Training Stability for Large Language Model Pretraining [18.163318397205533]
  Large LanguageText Models (LLMs) face increasing loss spikes during scaling.<n>While global clipping mitigates this, traditional approaches mitigate specific variations.<n>We show that AdaGC converges 25% faster than global clipping.
  arXiv  Detail & Related papers  (2025-02-16T08:13:23Z)
• Improved techniques for deterministic l2 robustness [63.34032156196848]
  Training convolutional neural networks (CNNs) with a strict 1-Lipschitz constraint under the $l_2$ norm is useful for adversarial robustness, interpretable gradients and stable training. We introduce a procedure to certify robustness of 1-Lipschitz CNNs by replacing the last linear layer with a 1-hidden layer. We significantly advance the state-of-the-art for standard and provable robust accuracies on CIFAR-10 and CIFAR-100.
  arXiv  Detail & Related papers  (2022-11-15T19:10:12Z)
• GDP: Stabilized Neural Network Pruning via Gates with Differentiable Polarization [84.57695474130273]
  Gate-based or importance-based pruning methods aim to remove channels whose importance is smallest. GDP can be plugged before convolutional layers without bells and whistles, to control the on-and-off of each channel. Experiments conducted over CIFAR-10 and ImageNet datasets show that the proposed GDP achieves the state-of-the-art performance.
  arXiv  Detail & Related papers  (2021-09-06T03:17:10Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.