IR3DE: A Linear Router for Large Language Models

Abstract: Foundational LLMs demonstrate proficiency on a wide range of general tasks, and achieve remarkable results on various specialized tasks via domain-expert LLMs. With the ever-growing list of available LLMs, inference routers are being proposed to select the most appropriate LLM for each prompt. However, existing routing methods either optimize cost across weak-to-strong generalist LLMs or require substantial training to support domain-expertise routing. In this paper, we propose IR3DE, a Ridge Regression-based Router for Domain Experts that provides cheap and fast routing decisions for each prompt. We evaluate IR3DE in two Causal Language Modeling (CLM) settings where the tasks are next-token prediction for all domains, and one reasoning setting where each domain has its own distinct reasoning task. Despite being a linear router, IR3DE achieves performance comparable to the other baselines in both CLM settings, and surpassing them in the reasoning setting, with a normalized performance of 98.4%. Moreover, IR3DE enables the addition or removal of new domain experts without requiring the router to be retrained from scratch, allowing a dynamic set of LLMs to be served with minimal disruption to the router itself. Our code is available at: github.com/gensyn-ai/IR3DE.

• The paper introduces a closed-form ridge-regression router over frozen token embeddings that supports decentralized training, low-cost inference, and incremental expert updates without retraining.
• IR3DE matches kNN routing with a 100.0 normalized average score in one CLM setting and achieves the best reasoning result at 98.4, compared with 97.6 for kNN.
• Entropy-filtered voting improves routing on complex reasoning tasks, but the method trails kNN in CLMlarge at 95.3 versus 97.9, highlighting its limited semantic expressiveness.

IR3DE is a lightweight, linear inference router for selecting among domain-expert LLMs. Its central claim is that a ridge-regression classifier over token embeddings—requiring no additional LLM and no centralized training data—matches or exceeds existing LM-based routers across causal language modeling (CLM) and reasoning benchmarks (2606.06098).

Motivation and positioning

Existing LLM routers fall into two camps: cost–quality routers that select between weak and strong generalists based on prompt difficulty (e.g., RouteLLM, GraphRouter), and accuracy-oriented expert routers such as MoDEM and PolyRouter that classify prompts by domain. The latter require either a fine-tuned LLM (MoDEM uses DeBERTa v3 trained on the union of all domain datasets) or an embedding model plus stored training embeddings (PolyRouter's 1NN/kNN routers). Both designs raise practical concerns: they add nontrivial routing cost, require collecting domain datasets in one location (a privacy issue in decentralized settings), and cannot easily accommodate new experts without retraining. IR3DE is designed to remove all three constraints.

Method

IR3DE has two components. The Token Router (TR) applies a pre-trained embedding layer $\mathcal{E}$ to the tokens of an input $x$ and multiplies by a weight matrix $W \in \mathbb{R}^{h \times C}$, producing per-token softmax probabilities over $C$ domains. $W$ is obtained in closed form via regularized least squares:

$$W^* = (\mathcal{E}(\mathcal{T}(X))^\top \mathcal{E}(\mathcal{T}(X)) + \lambda I_h)^{-1} \mathcal{E}(\mathcal{T}(X)) Y,$$

where $Y$ contains one-hot domain labels per token. Because the Gram matrix $A$ and cross-correlation $B$ can be accumulated batch-wise, each domain's statistics can be computed asynchronously on separate nodes; adding a new expert only requires updating these sufficient statistics rather than retraining from scratch. The tokenizer/embedding pair is arbitrary and independent of those used by the experts. Inversion involves only a small ($\sim$1k × 1k) matrix, performed once.

The Sample Route Selector (SRS) computes Shannon entropy for each token's softmax distribution, retains the $x$0 tokens with smallest entropy, and routes the prompt by majority vote among their argmax predictions. Two variants are also evaluated: IR3DE-all (all tokens vote, capped at 1024) and IR3DE-avg (argmax over the mean token embedding). The entropy filter is motivated by the observation that common tokens appearing across domains receive near-uniform, high-entropy predictions from the ridge solution; allowing them to vote injects noise into the majority decision. Empirically, routing accuracy as a function of $x$1 follows an inverted-U shape in all three settings: too few tokens give a shallow signal, too many introduce noise, and intermediate pools of confident tokens perform best.

Experimental setup

Three settings are evaluated on a single H100 GPU:

• CLM: five experts (coding, math, physics, history/events, philosophy/thinking) finetuned from a 115M Llama3 base pre-trained on OpenWebText, using M2D2 domains; metric is perplexity.
• CLMlarge: four experts (math/OpenWebMath, biology/peS2o, legal/Pile of Law, dialogue/UltraChat) finetuned from a 1B Llama3 base.
• Reasoning: Llama3-3B domain experts from MergeBench evaluated on HumanEval (pass@1), GSM8k, M_ARC, and IFEval.

Results are normalized relative to each domain expert's own-domain performance (higher is better; values above 100 are possible due to sampling temperature 0.7 when routing accuracy approaches 100%). Baselines include domain experts alone, expert averaging, random routing, MoDEM-small (44M) and MoDEM-large (304M), and BERT-embedding-based 1NN/kNN routers. Notably, MoDEM-large is larger than the CLM experts themselves, making it impractical for deployment in that setting—an implicit criticism of LM-based routers at scale.

Results

Setting	Best baseline (avg)	IR3DE (avg)	Notes
CLM	kNN / IR3DE-all: 100.0	98.2	IR3DE-all ties kNN at 100.0
CLMlarge	kNN: 97.9	95.3	kNN leads
Reasoning	kNN: 97.6	98.4	IR3DE best overall

In the CLM setting, IR3DE-all achieves an average normalized score of 100.0, matching kNN and meaning it performs, on average, as well as each expert does within its own domain—a strong result given the router's linearity. In the Reasoning setting, IR3DE attains the best average (98.4 vs. 97.6 for kNN) and best or second-best performance in every individual domain, with several scores exceeding 100 (e.g., 100.6 on instruction following). The entropy-based SRS outperforms both variants here, supporting the paper's claim that precision matters more on complex reasoning tasks. In CLMlarge, however, IR3DE trails kNN (95.3 vs. 97.9), indicating the linear router does not dominate uniformly; the advantage is specific to settings where domain signals are more separable.

Limitations and open questions

The authors concede that IR3DE's linearity makes it less expressive than LM-based routers, so it may underperform on queries requiring richer semantic understanding or complex decision boundaries—the CLMlarge gap versus kNN is consistent with this. Three open directions are identified: extending the formulation to kernel ridge regression to capture non-linear structure while retaining closed-form updates; evaluating whether domain-relevance routing suffices for multi-step reasoning tasks where relevance alone may be inadequate; and incorporating system-level costs (compute, latency, memory) into the routing objective for resource-constrained deployment. Additionally, the reported gains depend on normalized metrics and generation randomness for scores above 100, and hyperparameters ($x$2, learning rates for baselines) were selected post hoc on test performance, which should be kept in mind when comparing numbers.

Conclusion

IR3DE demonstrates that a closed-form ridge-regression router over frozen token embeddings, combined with entropy-filtered majority voting, is competitive with—and in the reasoning setting superior to—LM-based expert routers, while being cheaper at inference, compatible with decentralized data collection, and incrementally updatable as experts join or leave. Its main trade-off is reduced expressiveness relative to learned classifiers, which the CLMlarge results make concrete.