TriRoute: Unified routing jointly optimizes attention, experts, KV-cache

A new paper from Andrii Balashov and Olena Ponomarova introduces TriRoute, a lightweight controller that makes three interdependent decisions per token per layer: (i) whether to skip, use local, or full attention, (ii) which FFN experts to activate (including a null expert that recovers Mixture-of-Depths), and (iii) the bit-width of the KV-cache. The controller is trained end-to-end using a heterogeneous relaxation—Gumbel-Softmax with straight-through estimation for categorical choices and load-balanced top-k gating for experts—under a Lagrangian budget constraint that allows fine-grained control over compute and memory.

The authors identify a key training failure: naively jointly training all three axes leads to a collapse cascade where collapse on one axis propagates to the others. To fix this, they introduce per-axis normalization and a coupling-aware balancing loss. On decoder-only models from 160M to 1.3B parameters at compute-optimal token counts, TriRoute Pareto-dominates the best independently optimized combination of MoD, MoE, and KV quantization at matched inference FLOPs and memory. Moreover, it better preserves accuracy on rare entities, code, and arithmetic—tasks that pure perplexity optimization typically erodes. Post-hoc analysis reveals interpretable patterns: the controller allocates full attention and high-precision cache to sentence-initial positions, rare subwords, and named entities, while cheaply routing function words.