ICL-Router: In-Context Learned Model Representations for LLM Routing

AI summary

Overview

Research area: Machine learning / large language model systems — specifically LLM routing, in-context learning, and model representation learning.

Technical level: Intermediate. The paper assumes familiarity with embeddings, projectors, LLM-based routers, cross-entropy and reconstruction losses, but the core idea is described conceptually well enough for readers with general ML background.

Scope: The paper proposes ICL-Router, a two-stage training method that represents each candidate LLM's capabilities as compact in-context vectors built from query-performance pairs, enabling scalable, retraining-free routing across a pool of models.

What This Paper Is About

No single large language model is best at everything; different models have complementary strengths, so routing each query to the most suitable model in a pool can raise overall accuracy. The problem is that effective routing requires an accurate representation of what each model can and cannot do, and existing methods either retrain the router whenever a new model is added (RouterDC, EmbedLLM) or depend on hand-crafted capability descriptions tied to specific benchmarks (MODEL-SAT). ICL-Router's goal is to learn semantically rich model representations that allow new LLMs to be added in a plug-and-play way without retraining the router.

Key Contributions

1. A new perspective on model representation: The authors characterize a model's capabilities as vector representations of query-performance pairs — that is, compact in-context vectors summarizing how the model answered a set of challenging queries — rather than as fixed embeddings or hand-written instructions.

2. A two-stage routing method (ICL-Router): Stage one (Query Reconstruction Training) co-trains a projector and an LLM-based router so the router can interpret projected query vectors; stage two (ICL Model Routing Training) trains the router to predict, given a model's capability profile, whether that model will answer a new query correctly.

3. Decoupling of profiling from routing for scalability: Because capability profiling is separated from the routing decision, a newly released LLM only needs to be evaluated on the same small query set to build its capability profile, after which it can be routed to immediately — with no router retraining.

4. Empirical validation on 10 benchmarks: ICL-Router reports state-of-the-art routing accuracy in both in-distribution and out-of-distribution settings, plus ablations on embedding model choice, number of in-context exemplars, and the query reconstruction stage.

Main Findings

• In-distribution advantage: On five in-distribution benchmarks (OlympiadBench, BBH, LogicBench, MMLUPro, MBPP), ICL-Router reaches an average accuracy of 76.30%, the highest reported. By comparison, RouterDC scores 72.32, EmbedLLM 74.16, MODEL-SAT 71.71, and Max Expert 73.19. Section 4.4 states this is an improvement of 4.08% over RouterDC, 2.14% over EmbedLLM, and 4.59% over MODEL-SAT; the introduction rounds these to 3.9, 2.2, and 4.6 points.

• Beats the best single model: The strongest single LLM in the pool, DeepSeek-R1-Distill-Qwen-7B, averages 69.10 on in-distribution tasks; ICL-Router exceeds this by 7.2 absolute points in the introduction's framing, and the paper also reports it surpasses the Max Expert baseline by an average of 3.11%.

• Task-specific gains over Max Expert: Improvements come from MMLUPro (+8.69%), BBH (+4.9%), LogicBench (+1%), and MBPP (+1.32%), suggesting the router matches queries to models at a finer granularity than dataset-level model selection.

• Out-of-distribution generalization: On five held-out OOD benchmarks (AIME, HumanEval, KORBench, AGIEval, MMLU-CF), ICL-Router averages 66.47%, exceeding RouterDC by 3.34%, EmbedLLM by 3.65%, and MODEL-SAT by 3.48%. It also beats Max Expert (65.68 average) on KORBench (+2.73%) and AGIEval (+3.28%) while matching it on AIME, MMLU-CF, and HumanEval.

• Scalability to new models: Adding Falcon-H1-7B-Instruct, Gemma3-12B-IT, DeepSeek-R1-Llama-8B, OpenThinker3-7B, and AceReason-Nemotron-1.1-7B raises in-distribution accuracy from 76.3% to 79.9% and OOD accuracy from 66.4% to 69.9%, whereas MODEL-SAT's gains are described as inconsistent and sometimes fluctuating.

• Stronger embedding models help: In-distribution accuracy rises from 74.45% (mxbai-embed-large-v1, 0.3B/1024) through 75.13 (bge-m3, 0.6B/1024), 75.51 (stella-en-1.5B-v5, 1.5B/2048), 76.03 (gte-Qwen2-7B-instruct, 7B/3584) to 76.30 (Qwen3-8B-Embedding, 8B/4096); OOD accuracy follows the same ordering from 64.55% to 66.47%.

• Moderate exemplar counts are best: In-distribution accuracy is 74.94% with 100 in-context exemplars, peaks at 76.30% with 500, and drops to 75.75% with 1000. OOD accuracy is 65.36% (100 exemplars), 66.47% (500), and 65.71% (1000). The paper notes it tested 100, 300, 500, and 1000 exemplars but only reports the 100/500/1000 values in the text.

• Query Reconstruction Training matters: Removing it lowers accuracy from 76.30% to 74.01% in-distribution (a 2.29% drop) and from 66.47% to 64.06% OOD (a 2.41% drop).

• Low routing overhead: Although the router is a 7B-scale model, it generates only a small number of tokens at inference (for example, 8 tokens for 8 candidate models), which the authors argue keeps routing cost well within an acceptable range relative to the routed model's hundreds or thousands of output tokens.

Methodology in Plain English

ICL-Router has three components: an embedding model, a projector, and an LLM-based router. The embedding model is Qwen3-Embedding-8B and is used only to produce query vectors — it is not trained. The projector is a two-layer MLP that maps embeddings into the router's input space. The router is Qwen2.5-7B-Instruct.

Training happens in two stages. In the Query Reconstruction Training stage, each query is embedded and projected into a vector, and the router is trained to reconstruct the original query text from that vector alone. This forces the projector to produce vectors whose meaning the router can actually read. The projector is first trained alone for one epoch at a learning rate of 2e-5, then the projector and router are trained jointly for two epochs with the router at 5e-6, with a batch size of 32.

In the ICL Model Routing Training stage, the researchers build a set of 500 challenging queries — 125 each where exactly 1, 2, 3, or 4 of the 8 pooled models answered correctly — so that all-correct or all-wrong queries, which cannot discriminate between models, are excluded. Every candidate LLM is run on this query set, producing for each model a capability profile: a list of (query vector, correct/incorrect) pairs. The projector and router are then jointly trained for five epochs (learning rates 1e-5 and 2e-6) with cross-entropy loss to predict whether a given model will answer a new query correctly given its profile. Training is repeated three times with different random seeds.

At inference, a new query is embedded and projected, paired with each model's capability profile, and the router scores the probability that each model answers correctly; the highest-probability model is selected. Because the profiles are separate from the router's weights, a new model is onboarded simply by evaluating it on the same 500 queries to build its profile — no retraining. Evaluation samples each routed model 10 times with temperature 0.3 and top-p 1.0 and reports average accuracy.

The in-distribution benchmarks follow RouterDC's setup, split 7:3 into training and test sets; AIME, HumanEval, KORBench, AGIEval, and MMLU-CF are held out entirely for OOD testing. Baselines are Random Router, LLM Router (a prompt-based Qwen2.5-7B-Instruct), Max Expert (best model per dataset), RouterDC, EmbedLLM, and MODEL-SAT.

Why This Matters

Impact on research. The paper reframes model representation for routing: instead of learning a fixed embedding per model (which locks the model pool) or writing a capability instruction by hand for a specific benchmark, it derives a model's "profile" from observed behavior on a shared query set and expresses it as vectors the router can consume. This connects the in-context vector literature (Vector-ICL and related work) to routing, and it offers a concrete answer to the field's scalability problem — the need to repeatedly re-evaluate and retrain as new models are released. The ablation showing that query reconstruction training is worth roughly 2.3–2.4 accuracy points also isolates a specific, reusable design ingredient.

Real-world applications.

• Cost- and quality-aware model gateways: An API gateway or orchestration layer that dispatches each incoming request to whichever of its subscribed models is most likely to succeed, rather than sending everything to one model.
• Rapid onboarding of new releases: Teams can swap in or add a newly released open-source model by running it over one fixed query set, without rebuilding the routing system.
• Domain-specific agent pipelines: Settings such as coding assistants, math tutoring, or logical-reasoning tools where the paper's benchmark suite (MBPP and HumanEval for code, OlympiadBench and AIME for math, LogicBench and KORBench for reasoning) mirrors real workloads.
• Budget-constrained deployments: Since the router emits only a few tokens per decision, routing overhead stays small relative to generation, making it viable even where inference budget matters.

Industry relevance. Model pools in production change constantly, and the dominant routing approaches described here (RouterDC, EmbedLLM) require retraining on a fixed pool. A method whose accuracy keeps improving as more models are added — 76.3% to 79.9% in-distribution in this paper — and whose onboarding cost is one evaluation pass is directly relevant to anyone maintaining a multi-model serving stack. The limitation to small-parameter models (up to about 9B in the pool) is the main caveat for enterprise-scale deployment.

Future Directions

1. Scaling to larger models. The authors explicitly note that computational and data-collection constraints meant the pool consisted of small-parameter LLMs; whether the capability-profile approach transfers to much larger models is untested.

2. Benchmarks beyond general evaluation. The paper states its benchmarks do not specifically assess chat or instruction-following ability, leaving open whether in-context capability profiles capture those behaviors.

3. Better query-set design. The 500-query challenging set is a fixed hand-designed heuristic (125 queries each for 1–4 correct models). Automatically selecting maximally discriminative queries — and understanding the drop in accuracy beyond 500 exemplars — is an open question.

4. Cost-aware and sequential routing. This work is in the performance-maximizing branch of routing, not the performance-versus-cost branch; extending in-context capability profiles to jointly optimize accuracy and inference cost is a natural next step.

Target Audience

This paper is most useful to machine learning researchers and engineers working on LLM routing, model selection, and multi-model serving infrastructure; to practitioners building LLM gateways or agent systems that must choose among several models per request; and to readers interested in in-context learning, in-context vectors, and how model capabilities can be encoded as representations rather than text. Readers need comfort with embedding models, projectors, and cross-entropy training objectives, but not deep familiarity with routing literature, since the paper positions its baselines clearly.

Authors’ abstract

Large language models (LLMs) often exhibit complementary strengths. Model routing harnesses these strengths by dynamically directing each query to the most suitable model, given a candidate model pool. However, routing performance relies on accurate model representations, and adding new models typically requires retraining, limiting scalability. To address these challenges, we propose a novel routing method using in-context vectors to represent model capabilities. The method proceeds in two stages. First, queries are embedded and projected into vectors, with a projector and LLM-based router trained to reconstruct the original queries, aligning vector representations with the router's semantic space. Second, each candidate model is profiled on a query set, and the router learns -- based on in-context vectors of query and model performance -- to predict whether each model can correctly answer new queries. Extensive experiments demonstrate that our method achieves state-of-the-art routing performance in both in-distribution and out-of-distribution tasks. Moreover, our method allows for seamless integration of new models without retraining the router. The code is available at https://github.com/lalalamdbf/ICL-Router.