One of the persistent inefficiencies in deploying large language models is that the transformer architecture treats every token identically, spending the same compute on a syntactic function word as on a rare named entity that carries most of the semantic weight. Three techniques have emerged to address this, each attacking a different resource axis: Mixture-of-Experts (MoE) sparsifies the feed-forward network, Mixture-of-Depths (MoD) makes block execution token-adaptive, and KV-cache quantization compresses attention memory. The problem is that all three have been developed and tuned in isolation, despite making decisions that are, as this paper argues, fundamentally coupled.
TriRoute (arXiv:2607.06601) proposes a single lightweight controller that makes all three decisions jointly, per token per layer, under a unified budget constraint. The central claim is that a shared learned policy can Pareto-dominate the best independently-tuned combination of the same three mechanisms at matched inference FLOPs and memory.
Why Joint Routing Matters
The motivating intuition is clean and worth taking seriously. Consider a rare entity token: MoD might correctly identify that its FFN transformation is low-entropy and can be skipped, but that same token likely needs full attention resolution to bind to coreferents, and its key/value representation should be stored at high precision so future queries can retrieve it faithfully. A function word like "the" sits at the opposite extreme on all three axes simultaneously. The right allocation is therefore a vector of coupled choices, not a scalar budget per token.
This coupling argument is not merely intuitive. If you tune MoD, MoE, and KV-quantization independently, each mechanism optimises its own proxy objective without accounting for the complementary slack or demand created by the others. The result is a suboptimal allocation that may, for instance, skip the FFN for a token that simultaneously gets low cache precision, compounding the information loss.
Methodology
The TriRoute controller is a two-layer MLP of width 128, shared across all three routing axes, adding less than 1% FLOPs and 0.3% parameters to the backbone. At each layer, for each token, it emits a coordinated policy over:
- Attention mode: skip, local window, or full causal attention
- Expert selection: top-2 from 8 experts, where one expert is a null expert that recovers MoD-style block skipping
- KV-cache bit-width: a categorical choice governing how the token's key/value is stored for future retrieval
Training these three decisions end-to-end requires handling heterogeneous gradient estimation. The attention mode and bit-width decisions are categorical, so the authors use Gumbel-Softmax with straight-through estimation. Expert selection uses the standard load-balanced top-k gating familiar from MoE literature. The cost model is differentiable in the router logits, meaning the budget constraint directly shapes the policy rather than acting as a post-hoc filter.
A notable practical contribution is the identification and mitigation of what the authors call a cross-axis routing-collapse cascade. In naive joint training, the router can collapse to degenerate solutions where one axis dominates and the others become trivial. The fix involves per-axis normalisation and a coupling-aware balancing loss that penalises correlated collapse across axes. This is the kind of training instability that would likely be discovered empirically by anyone attempting to replicate the naive version, so making it explicit and providing a solution is genuinely useful.
Budget control is handled via an online Lagrangian scheme: a single constraint converts the average compute-and-memory cost into a tunable knob, allowing the same architecture to be deployed at different operating points without retraining.
Experiments use decoder-only models from 160M to 1.3B parameters, matched to Pythia configurations for comparability, trained at compute-optimal token counts (roughly 20 tokens per parameter following Hoffmann et al.). The MoE variants use 8 experts with top-2 routing, sized so that active FFN FLOPs equal the dense baseline.
Results and Interpretability
The headline result is that TriRoute traces a Pareto frontier that dominates the best independently-tuned combination of MoD, MoE, and KV-quantization at matched inference FLOPs and memory across all three model scales. The gains are not marginal: the paper reports matching dense model quality at roughly half the inference cost.
Tail-case robustness is a secondary but important finding. Pure perplexity optimisation tends to erode performance on rare entities, code, and arithmetic, because these token types are statistically underrepresented and a perplexity-minimising router may cheaply route them to save compute. TriRoute's joint policy, by contrast, learns to spend attention resolution and cache bits on precisely these tokens, preserving tail performance that single-axis methods sacrifice.
The post-hoc interpretability analysis is one of the more satisfying parts of the paper. The controller learns structure that aligns with linguistic intuition: sentence-initial tokens, rare subwords, and named entities receive high attention resolution and high-precision cache storage; function words and predictable continuations are cheaply routed. This is structure that no single-axis method can express, since MoD can only skip or not skip, and KV-quantization has no mechanism to condition on token informativeness jointly with FFN routing.
Limitations and Open Questions
Several questions are worth raising for readers considering this work carefully. First, the experiments top out at 1.3B parameters. The scaling behaviour of joint routing is genuinely uncertain: at larger scales, MoE routing already faces load-balancing challenges, and adding two additional axes of joint optimisation may amplify instability in ways that per-axis normalisation does not fully address. The collapse cascade fix is demonstrated at small scale, but the authors do not provide evidence that it remains sufficient at 7B or 70B.
Second, the compute-optimal training regime (20 tokens per parameter) is a reasonable choice for controlled comparison, but production models are increasingly trained well beyond compute-optimal token counts. Whether the routing policy learned under one token budget transfers cleanly to inference on different distribution shifts is not addressed.
Third, the Lagrangian budget controller is elegant, but the paper does not discuss the sensitivity of the learned policy to the choice of Lagrange multiplier schedule or initialisation. Joint training with three heterogeneous gradient estimators is not trivial to reproduce, and more detail on hyperparameter sensitivity would strengthen the practical contribution.
Finally, the claim of Pareto dominance over independently-tuned baselines depends on how carefully those baselines were tuned. The paper matches configurations to Pythia for comparability, which is reasonable, but practitioners with heavily optimised single-axis deployments may find the gap narrower than reported.
These caveats aside, the core contribution is sound. The coupling argument is well-motivated, the collapse cascade is a real problem the authors identify and solve, and the interpretability findings provide genuine insight into what a jointly-trained router learns. For researchers working on efficient inference, the unified formulation and the reference controller are worth engaging with directly.
Full paper: arXiv:2607.06601