Source-linked AI summary
Do Dynamic Routers Need Memory? HeRo: History-Aware Routing for Efficient LLM Inference
Hongjin Lin, Wentao Wan, Keze Wang
TL;DR
Existing dynamic routers treat depth-wise decisions as local hidden-state predictions despite their path-dependent interactions. HeRo adds linear-attention router memory that aggregates prior routing scores and residual updates, and it consistently outperforms ten baselines across three backbones while retaining dense-model performance at reduced computation.
Problem
Existing routers condition only on the current hidden representation even though preceding routing decisions shape downstream states and routing objectives couple decisions across depth.
Method
HeRo uses linear-attention memory to maintain explicit routing history and conditions each FFN routing decision on that history and the current hidden representation.
Results
HeRo consistently achieved the highest aggregate performance retention among ten baselines across Llama 3.1-8B, Llama 2-7B, and Llama 2-13B.
Takeaways & Limitations
Explicit routing memory improves dynamic routing quality across model families and scales while training only lightweight routers and adapters on a frozen backbone.
Abstract
from arXiv · showhide
Dynamic layer routing reduces the inference cost of Large Language Models (LLMs) by learning to skip layers for individual tokens. Existing methods, however, treat each routing decision as a local operation conditioned solely on the current hidden state which is a formulation that overlooks the sequential, path-dependent nature of routing across depth: earlier decisions shape the representations seen by downstream routers, and the layer-usage objective couples all decisions jointly. We propose History-Aware Routing (HeRo), a dynamic routing framework that resolves this mismatch by introducing a router memory mechanism to maintain an explicit routing state across model depth. The memory is constructed via linear attention, incrementally aggregating preceding routing scores and their induced residual updates into a compact history representation. At each routed layer, the router conditions jointly on this accumulated state and the current hidden representation to select the executed branch. Instantiated for token-wise FFN routing, HeRo trains only lightweight routers and adapters on a frozen backbone, requiring no modification to pretrained parameters. Across Llama 3.1-8B, Llama 2-7B, and Llama 2-13B, HeRo consistently achieves the highest aggregate performance retention among ten baselines. On Llama 3.1-8B, it bypasses 26.87% of model parameters while achieving 100.24% of dense model performance across seven benchmarks, and retains 97.01% while bypassing 38.82% of model parameters under a tighter computation budget. Ablation studies confirm that removing routing history consistently degrades performance, most notably on multistep reasoning and code generation, validating that explicit routing memory enables more accurate and adaptive dynamic routing than solely conditioning on hidden state.
Introduction
Dynamic routing is a path-dependent sequence across model depth, but existing routers condition only on the current hidden representation. HeRo addresses this mismatch with explicit memory of preceding routing decisions and residual updates.
- Static-depth inference applies the full decoder stack uniformly, although different predictions require different amounts of computation.
- Existing dynamic routers make local decisions from the current representation, omitting directly observable routing history and residual transitions.
- Earlier routing choices shape downstream representations, while the layer-usage objective jointly couples routing decisions across depth.
- HeRo uses linear-attention memory to aggregate preceding routing scores and induced residual updates into an explicit routing state.
- At each routed layer, HeRo combines the current hidden representation with accumulated history to select the pretrained FFN or a lightweight adapter.
Related Work
Adaptive depth methods reduce computation by changing how much computation is performed for each input, while dynamic routing selects computations at different scopes. HeRo’s routing framework is presented against this progression toward input-dependent execution.
- Related Work: Static pruning removes layers or components before inference, producing a fixed smaller model for every input.
- Related Work: Adaptive methods retain the full stack and determine how much of it to execute for each input or token.
- Related Work: HeRo’s overview depicts routing across blocks using current representations, remaining depth, prior scores, and residual transitions.
- Related Work: Dynamic routing spans token-, prompt-, and sequence-level decisions, including assignments to blocks or modules under different routing scopes.
Method
HeRo routes tokens through selected FFNs or adapters using a router state that combines current representations with path features and a depth-wise memory of prior routing. Its thresholded gate is trained with language modeling and utilization objectives while the backbone remains frozen.
- Router state and gate: HeRo processes each routed block by encoding the post-attention representation and path features, reading prior history, then combining history and local logits into an execution gate.The path features track routing progress, residual-stream transitions, and preceding routing scores.
- Branch selection: The router selects either the frozen pretrained FFN or a trainable bottleneck adapter, while attention remains dense at routed blocks.The adapter uses down and up projections with intermediate width smaller than the model width.
- Router state and gate: The execution threshold selects a single branch during training and inference, while scaling the selected branch by its gate value preserves gradients to the router.The threshold itself is not differentiated.
- Router memory: The depth memory uses linear attention to read accumulated states before writing the current state, so each gate sees only preceding routed blocks.Memory states are initialized to zero, updated across routed blocks, and shared encoder and router parameters are used across those blocks.
- Training objective: The training objective combines language-modeling loss with a utilization penalty whose squared per-token sum couples gate values across routed blocks.The coefficient αt controls the utilization penalty’s weight relative to language modeling loss.
Experiments
Experiments compare HeRo with pruning and routing baselines across matched budgets, then test coefficient, component, and memory-dimension sensitivity. HeRo achieves the strongest aggregate retention across backbones, while ablations show that path state, memory, and history conditioning jointly support routing quality.
- Matched-budget comparisons: HeRo achieves the highest aggregate retention among evaluated parameter-skipping methods across Llama-3.1-8B, Llama-2-7B, and Llama-2-13B.On Llama-3.1-8B at the 25% target budget, it leads competing methods on every benchmark.
- Matched-budget comparisons: 26.87% realized parameter skipping matches dense-model aggregate performance on Llama-3.1-8B at the 25% target budget.HeRo leads all competing pruning and routing methods on every benchmark in this setting.
- Matched-budget comparisons: 97.01% dense-performance retention exceeds SkipGPT-RT’s 81.15% at a similar realized skipping rate under the 40% target budget.HeRo leads on six of the seven benchmarks, and the margin over the strongest dynamic baseline grows under the tighter budget.
- Sensitivity analyses: Increasing α_t raises parameter skipping monotonically from 13.19% to 26.87%, while Retain varies by only 1.73 points.The authors select α_t = 10^-3 because it reaches the target budget with only a modest aggregate-performance reduction.
- Ablation studies: Every component ablation lowers aggregate performance, indicating that state construction and history conditioning both contribute to routing quality.Removing the auxiliary path encoder slightly increases parameter skipping, whereas other ablations reduce both Retain and parameter skipping.
- Ablation studies: Retain and average accuracy vary by only 0.12 and 0.05 points across tested memory dimensions.Dimension 256 gives the highest Retain and average accuracy, while dimension 64 gives the highest parameter skipping rate but the lowest Retain.
Conclusion
HeRo introduces explicit routing memory based on linear attention, aggregating prior routing scores and induced residual updates so routers condition on both current hidden states and accumulated routing information. Across three Llama backbones, it outperforms ten baselines while retaining strong performance at substantial parameter-skipping rates.
- Method: HeRo uses linear-attention router memory to aggregate preceding routing scores and induced residual updates into an explicit routing state.Each router conditions its decision on the current hidden state and routing information accumulated across preceding layers.
- Results: HeRo consistently outperforms all 10 baseline methods across Llama 3.1-8B, Llama 2-7B, and Llama 2-13B.The method trains lightweight routers and adapters while keeping the backbone frozen.
- Results: 26.87% parameter bypassing reaches 100.24% of dense performance across seven benchmarks, while 38.82% bypassing retains 97.01%.Removing the routing history head causes performance losses across six benchmarks, especially HumanEval and GSM8K.