Lookahead Routing for LLMs
Free while signed in. Answers cite the passages they came from.
Lookahead is a response-aware LLM routing framework that predicts latent representations of potential model outputs to enable more informed routing decisions without full inference. - **Core limitation of query-only routing**: Traditional routers base decisions solely on input queries, missing critical information about actual response quality and semantic intent that emerges during generation. This leads to suboptimal routing on complex or ambiguous queries. - **Dual implementation architecture**: Sequence-level variant uses causal language models (CLM) that concatenate query with model identifier (MID) tokens, extracting hidden states at MID positions as response representations. Token-level variant uses masked language models (MLM) that jointly reconstruct all candidate responses via repeated MID token blocks, aggregating information through [CLS] token attention. - **Curriculum masking strategy**: MLM variant progressively masks from response end to start, increasing masking ratio linearly to 100% over first 40% of training. This smooth transition from partial to full masking enables robust representations and better generalization than uniform random masking. - **Joint training objective**: Combines routing loss (binary cross-entropy on model selection) with response reconstruction loss (next-token prediction for CLM, masked token recovery for MLM). Auxiliary response modeling improves sample efficiency by 6.3x and captures richer semantic information via higher mutual information with oracle responses. - **Performance**: Achieves 7.7% average normalized score gain over SOTA RouterDC across 7 benchmarks (AlpacaEval-2, Arena-Hard, MT-Bench, GSM8K, MATH, HumanEval, MBPP). MLM variant excels on open-ended instruction-following tasks where joint semantic-space encoding enables fine-grained cross-model comparisons. Routes nearly 100% of code queries to specialized Qwen2.5-Coder model, demonstrating strong specialization awareness.
Core limitation of query-only routing: Traditional routers base decisions solely on input queries, missing critical information about actual response quality and semantic intent that emerges during generation. This leads to suboptimal routing on complex or ambiguous queries.
Dual implementation architecture: Sequence-level variant uses causal language models (CLM) that concatenate query with model identifier (MID) tokens, extracting hidden states at MID positions as response representations. Token-level variant uses masked language models (MLM) that jointly reconstruct all candidate responses via repeated MID token blocks, aggregating information through [CLS] token attention.
Curriculum masking strategy: MLM variant progressively masks from response end to start, increasing masking ratio linearly to 100% over the first 40% of training. This smooth transition from partial to full masking enables robust representations and better generalization than uniform random masking.
Joint training objective: Combines routing loss (binary cross-entropy on model selection) with response reconstruction loss (next-token prediction for CLM, masked token recovery for MLM). Auxiliary response modeling improves sample efficiency by 6.3x and captures richer semantic information via higher mutual information with oracle responses.
Performance: Achieves 7.7% average normalized score gain over SOTA RouterDC across 7 benchmarks (AlpacaEval-2, Arena-Hard, MT-Bench, GSM8K, MATH, HumanEval, MBPP). MLM variant excels on open-ended instruction-following tasks where joint semantic-space encoding enables fine-grained cross-model comparisons. Routes nearly 100% of code queries to the specialized Qwen2.5-Coder model, demonstrating strong specialization awareness.
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 this https URL.
Get next week’s papers.
The same picks and the same summaries, in your inbox. Free, and 176 issues deep.
Subscribe on Substack