Lookahead Routing for Large Language Models

TL;DR

The paper tackles efficient routing in multi-LLM systems by moving beyond query-only decisions to look ahead at latent response signals. It introduces Lookahead, a dual-task framework that jointly learns response-aware latent features $\tilde{\mathbf{r}}_t$ and model-selection scores $\hat{c}_t$ using either a sequence-level CLM predictor or a token-level MLM predictor. Across seven benchmarks spanning instruction-following, mathematics, and code, Lookahead consistently outperforms state-of-the-art routing baselines, with the MLM variant delivering the strongest gains on open-ended tasks and achieving an average improvement of about $7.7\%$. The approach also yields data-efficient training, as response modeling boosts sample efficiency and produces richer semantic representations than query-only features, enabling scalable, context-aware routing in heterogeneous LLM ensembles.

Abstract

Large language model (LLM) routers improve the efficiency of multi-model systems by directing each query to the most appropriate model while leveraging the diverse strengths of heterogeneous LLMs. Most existing approaches frame routing as a classification problem based solely on the input query. While this reduces overhead by avoiding inference across all models, it overlooks valuable information that could be gleaned from potential outputs and fails to capture implicit intent or contextual nuances that often emerge only during response generation. These limitations can result in suboptimal routing decisions, particularly for complex or ambiguous queries that require deeper semantic understanding. To address this challenge, we propose Lookahead, a routing framework that "foresees" potential model outputs by predicting their latent representations and uses these predictions to guide model selection, thus enabling more informed routing without full inference. Within this framework, we implement two approaches based on causal and masked language models. Empirical evaluations across seven public benchmarks - spanning instruction following, mathematical reasoning, and code generation - show that Lookahead consistently outperforms existing routing baselines, achieving an average performance gain of 7.7% over the state-of-the-art. Our code is available at https://github.com/huangcb01/lookahead-routing.

Figures (7)

• Figure 1: Effect of response-aware routing across benchmarks. Left: Including responses improves classifier-based routing performance. Right:Lookahead outperforms existing routing methods.
• Figure 2: Overview of the Lookahead framework. Middle (Data Collection): For each input prompt $x$, responses $y_{1:T}$ are sampled from $T$ candidate LLMs. A judge model evaluates these responses to assign quality scores $s_{1:T}$. Top (Query-Only Routing): Conventional routers encode $x$ into a query embedding $\mathbf{e}_x$ and select a model based solely on the input. Bottom (Lookahead Routing): Given $x$ and model identifiers $\text{MID}_{1:T}$, the Feature Predictor estimates latent response representations $\tilde{\mathbf{r}}_{1:T}$, which are then used by a classifier to predict final quality scores $\hat{c}_{1:T}$ and select the best model.
• Figure 3: Architectures for response-aware routing in Lookahead. Left: Sequence-level modeling with a causal language model (CLM), where the hidden state at the MID token encodes the response information. Right: Token-level modeling with a masked language model (MLM), where fully masked responses are reconstructed and summarized via attention over MID tokens.
• Figure 4: Results of ablation studies for the Lookahead framework. (a) Performance drops when response modeling (RM) is removed. (b) Comparison of curriculum masking (CM) strategies in the MLM-based predictor. (c) Effectiveness of the joint response prediction.
• Figure 5: (a) Training efficiency and (b) mutual-information analysis for CLM-based Lookahead.