Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Informed Routing in LLMs: Smarter Token-Level Computation for Faster Inference

Chao Han, Yijuan Liang, Zihao Xuan, Daokuan Wu, Wei Zhang, Xiaoyu Shen

TL;DR

The Lightweight Feature Forecaster (LFF) is proposed, a small predictive module that estimates a unit's output before routing decisions are made that enables a flexible execute-or-approximate policy that preserves model fidelity while drastically reducing computation.

Abstract

The deployment of large language models (LLMs) in real-world applications is increasingly limited by their high inference cost. While recent advances in dynamic token-level computation allocation attempt to improve efficiency by selectively activating model components per token, existing methods rely on greedy routing--a myopic execute-or-skip mechanism that often leads to irreversible information loss and suboptimal token selection. This paper introduces informed routing, a new paradigm that proactively addresses these issues. The key insight is to assess not only a token's immediate importance but also its recoverability, i.e., how well its transformation can be approximated. To this end, we propose the Lightweight Feature Forecaster (LFF), a small predictive module that estimates a unit's output before routing decisions are made. This enables a flexible execute-or-approximate policy that preserves model fidelity while drastically reducing computation. Extensive experiments on both language modeling and reasoning tasks show that informed routing achieves state-of-the-art efficiency-performance trade-offs across multiple sparsity levels. Notably, even without final LoRA fine-tuning, our method matches or surpasses strong baselines that require full fine-tuning, all while reducing training time by over 50%. The code is available at: https://github.com/EIT-NLP/informed-routing

Informed Routing in LLMs: Smarter Token-Level Computation for Faster Inference

TL;DR

The Lightweight Feature Forecaster (LFF) is proposed, a small predictive module that estimates a unit's output before routing decisions are made that enables a flexible execute-or-approximate policy that preserves model fidelity while drastically reducing computation.

Abstract

The deployment of large language models (LLMs) in real-world applications is increasingly limited by their high inference cost. While recent advances in dynamic token-level computation allocation attempt to improve efficiency by selectively activating model components per token, existing methods rely on greedy routing--a myopic execute-or-skip mechanism that often leads to irreversible information loss and suboptimal token selection. This paper introduces informed routing, a new paradigm that proactively addresses these issues. The key insight is to assess not only a token's immediate importance but also its recoverability, i.e., how well its transformation can be approximated. To this end, we propose the Lightweight Feature Forecaster (LFF), a small predictive module that estimates a unit's output before routing decisions are made. This enables a flexible execute-or-approximate policy that preserves model fidelity while drastically reducing computation. Extensive experiments on both language modeling and reasoning tasks show that informed routing achieves state-of-the-art efficiency-performance trade-offs across multiple sparsity levels. Notably, even without final LoRA fine-tuning, our method matches or surpasses strong baselines that require full fine-tuning, all while reducing training time by over 50%. The code is available at: https://github.com/EIT-NLP/informed-routing
Paper Structure (30 sections, 8 equations, 3 figures, 9 tables)

This paper contains 30 sections, 8 equations, 3 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Static Pruning
  4. Dynamic Computation Allocation
  5. Error Compensation
  6. Methodology
  7. Preliminaries and Notation
  8. Lightweight Feature Forecaster
  9. Three-Stage Optimization
  10. Experiments
  11. Experimental Setup
  12. Results
  13. Inconsistency Between Language Modeling and Reasoning
  14. Training Stability and Efficiency Gains of LFF Initialization
  15. Superiority After Fine-tuning and Underlying Mechanisms
  16. ...and 15 more sections

Figures (3)

  • Figure 1: The limitation of greedy routing and the promise of informed routing. Under the same sparity ratio, informed routing (a) reduces the perplexity and (b) increases feature similarity.
  • Figure 2: (a), (b), and (c) present the architectural comparison diagrams of the vanilla LLM, greedy routing, and our proposed informed routing paradigm. (d) illustrates how the LFF initialization is performed.
  • Figure 3: Layer-wise Token Allocation. The hatched area represents the proportion of tokens processed by the efficient LFF branch and the colored areas show the tokens retained for full computation in the Attention (green) and FFN (orange) modules.