TriRoute: Unified Learned Routing for Joint Adaptive Attention, Experts, and KV-Cache Allocation

Language models keep getting larger, but the cost of running them token‑by‑token still limits where they can be deployed. Recent work shows that we can save computation by skipping parts of the network, choosing cheaper experts, or compressing the KV‑cache, yet each of these tricks is usually applied in isolation. TriRoute proposes a single, tiny controller that decides all three things together for every token at every layer, aiming to keep quality high while lowering the average inference budget.

What You Need to Know

The paper introduces TriRoute, a shared routing network that sits alongside each transformer layer. For a given token, the controller outputs three discrete decisions: (1) an attention mode—skip, low‑resolution (e.g., windowed), or full attention; (2) an expert ID from a Mixture‑of‑Experts feed‑forward network; and (3) a KV‑cache bit‑width (e.g., 2‑bit, 4‑bit, 8‑bit). All three outputs are produced by the same small multilayer perceptron that takes the token’s hidden state as input.

During training, TriRoute learns to balance two objectives: the usual language‑modeling loss and a differentiable cost term that estimates FLOPs and memory traffic from the three decisions. The cost term uses analytically derived formulas for attention complexity, expert feed‑forward work, and cache storage, allowing gradients to flow back to the routing network without explicit reinforcement‑learning tricks. A straight‑through estimator handles the discreteness of the choices during back‑propagation.

Experiments are run on a 1.3 B‑parameter decoder‑only model with 32 experts per FFN layer and a standard KV‑cache. TriRoute is compared against baselines that apply only one of the three savings techniques (e.g., static MoE, static MoD, uniform 4‑bit KV‑cache). The results show that TriRoute achieves comparable perplexity to the full‑precision model while reducing average attention FLOPs by ~38 %, expert FLOPs by ~45 %, and KV‑cache memory by ~55 % on the WikiText‑103 benchmark.

Why It Matters

Attention, expert selection, and KV‑cache precision are not independent: a token that needs rich contextual information often benefits from both higher‑resolution attention and a higher‑precision cache, regardless of which expert processes it. Treating these decisions separately can lead to mismatched allocations—spending cache bits on tokens that use sparse attention, or wasting expert capacity on tokens that already skip attention. By jointly optimizing them, TriRoute can route “hard” tokens to the most expensive combination of settings and “easy” tokens to the cheapest, yielding a tighter coupling between model quality and resource use.

This coupling matters for real‑world deployment on devices with tight memory and power budgets. Lower KV‑cache bit‑width directly cuts the memory bandwidth needed for generation, while attention mode changes affect both compute and cache traffic. Expert selection influences the compute density of the feed‑forward sub‑layer. A unified controller lets system designers trade off these three dimensions in a principled way, rather than tuning three separate heuristics.

Key Details

• Controller architecture: Two‑layer MLP with 64 hidden units, ReLU activation, and a softmax over the discrete choice sets for each axis.
• Loss formulation: ℒ = ℒ_LM + λ·Cost, where Cost = α·FLOPs_att + β·FLOPs_ffn + γ·Mem_KV; λ anneals from 0 to 0.1 over training.
• Straight‑through estimator: Enables gradient flow through the discrete argmax by treating the forward pass as a hard selection and the backward pass as the softmax probabilities.
• Ablation: Removing any one of the three outputs raises perplexity by 0.2–0.4 points while saving less than half the resources, confirming the benefit of joint routing.
• Hardware‑aware cost model: Uses measured latencies for attention kernels at different sparsities and for expert matmul, making the cost term reflective of actual accelerator behavior.

What’s Next

Future work could extend TriRoute to encoder‑decoder architectures, incorporate more granular choices such as variable head counts, or learn the cost coefficients directly from target hardware profiles. Scaling the controller to larger numbers of experts or to multimodal models (vision‑language, audio) would test whether the same joint‑routing principle holds when the search space grows substantially. Finally, integrating the routing decisions into a compiler‑level schedule could allow static generation of efficient kernels tailored to each token’s policy.

📌 Source: Arxiv Ml