Source-linked AI summary
IR3DE: A Linear Router for Large Language Models
Eros Fanì, Oğuzhan Ersoy
TL;DR
Existing LLM routers either trade cost against capability among generalists or require substantial data and model infrastructure for domain-expert routing. IR3DE uses a ridge-regression token router with a sample selector to make cheap, adaptable expert choices. It matches baselines in CLM settings and surpasses them in reasoning, reaching 98.4 normalized performance, while supporting expert-set changes without full retraining.
Problem
Existing routing methods focus on generalist cost-performance trade-offs or require additional models and centralized domain datasets for expert routing.
Method
IR3DE combines a linear ridge-regression token router with sample selection based on token predictions, including entropy filtering and majority voting.
Results
IR3DE is comparable to baselines in both CLM settings and surpasses them in reasoning, achieving a normalized average performance of 98.4.
Takeaways & Limitations
IR3DE enables lightweight routing and adding or removing domain experts without retraining the router from scratch.
Takeaways & Limitations
IR3DE is less expressive than stronger LLM-based routers and may be less effective for queries requiring richer semantic understanding or complex decision boundaries.
Abstract
from arXiv · showhide
Foundational Large Language Models (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.
1. Introduction
IR3DE addresses expert-model routing by using a lightweight linear router that selects domain experts without the centralized training requirements of prior expert routers. It evaluates this design across CLM and reasoning settings, with entropy-based routing performing best on complex reasoning tasks.
- Motivation: Existing routing methods mainly optimize cost-performance trade-offs among generalist models or require additional models and domain datasets for expert routing.The latter requirement can conflict with privacy or data-centralization constraints.
- Contribution: IR3DE routes each query to a suitable domain expert using a ridge-regression token router and sample route selector.The linear design is intended to reduce routing cost and support dynamic expert sets.
- Contribution: IR3DE is cheaper and faster than LM-based baselines because its linear construction avoids using a language model as the router.Its computation uses token embeddings and a one-time inversion of a small matrix.
- Contribution: The router supports adding or removing domain experts without retraining the entire router from scratch.This enables a dynamically changing set of served experts.
- Evaluation: IR3DE is comparable to baselines in two CLM settings and surpasses them in a reasoning setting.The entropy-based sample route selector performs best on complex reasoning tasks, where precision is more important.
2. Related Work
Prior LLM routing work has emphasized either cost-quality selection among generalists or accuracy-oriented selection among specialists. Specialist routers often depend on additional language or embedding models and centralized domain data, motivating lighter alternatives.
- Generalist routing: Early routing methods estimate prompt difficulty to choose among fixed pools of generalist models with different cost-quality profiles.Cascading approaches extend this idea by using multiple models sequentially.
- Expert routing: Accuracy-oriented routers assess a prompt’s domain and forward it to the most relevant specialist model.Other approaches select multiple experts or use hierarchical multi-hop routing.
- Limitations of prior expert routers: MoDEM trains a DeBERTa v3 router on the union of domain datasets used to train the experts.This creates challenges when data cannot be centralized because of privacy, communication, computational, or memory constraints.
- Limitations of prior expert routers: Existing expert-routing methods use additional language models to classify queries or generate token embeddings.The 1NN-router instead uses sentence-transformer embeddings to retrieve a nearest training query.
3. Method
IR3DE combines a token-level ridge-regression router with a sample route selector that aggregates token predictions to choose an expert. Its entropy-based selector retains the most confident tokens before majority voting, while alternative variants use all tokens or averaged embeddings.
- Overview: IR3DE selects the expert expected to maximize a task-specific performance metric for each input text.The target metric can be perplexity for causal language modeling or pass@1 for coding.
- Token Router: The Token Router maps token embeddings through linear weights to produce a softmax probability vector for every token.Any tokenizer and embedding layer pair is allowed when the embedding layer was trained with that tokenizer.
- Sample Route Selector: The Sample Route Selector computes token entropies, retains the min(k,T) lowest-entropy tokens, and routes the input by majority voting over their predicted experts.This filters uncertain token predictions before aggregating the remaining votes.
- Token Router: Ridge regression estimates the router weights from stacked token embeddings and one-hot domain labels using a closed-form solution.The formulation uses Tikhonov regularization and can collect statistics batch-wise.
- Variants: IR3DE-all uses all tokens for majority voting, whereas IR3DE-avg averages token embeddings before selecting the domain.The all-token variant can introduce voting noise, while averaging compresses individual token signals.
4. Experiments
The experiments compare IR3DE with routing baselines across two causal-language-modeling settings and one reasoning setting. IR3DE is competitive in CLM and achieves the strongest average result in Reasoning, while its entropy-token selector performs best with an intermediate pool of confident tokens.
- Setup: Experiments cover CLM, CLMlarge, and Reasoning settings, using next-token prediction for the first two and distinct domain reasoning tasks for the third.The evaluation uses domain-specific expert models and normalized metrics relative to each expert’s performance in its own domain.
- Baselines: IR3DE is compared with domain-expert, expert-average, random, 1NN, kNN, and MoDEM routing baselines.MoDEM-small and MoDEM-large use DeBERTa v3 routers, while kNN uses embedding similarity to select a domain expert.
- CLM results: 100.0 average normalized performance is achieved by both kNN router and IR3DE-all in CLM, matching experts evaluated in their own domains.IR3DE also surpasses all listed baselines in the Coding, Math, and Physics domains in this setting.
- Reasoning results: 98.4 normalized average performance makes IR3DE the best method in Reasoning, ahead of kNN router at 97.6.IR3DE is best or second-best across all individual reasoning domains; scores above 100 can occur because generation uses temperature 0.7.
- Sample Route Selector: Intermediate top-k token pools provide the strongest routing accuracy across CLM, CLMlarge, and Reasoning.Too few tokens provide shallow signal, while too many admit uncertain predictions that add noise; IR3DE therefore filters for confident tokens.
5. Conclusion
IR3DE targets lightweight domain-expert routing as LLM pools become larger and more diverse. Its linear design is competitive across settings, but reduced expressiveness limits performance on semantically rich queries and complex decision boundaries.
- IR3DE addresses routing methods that are costly or require centralized domain datasets by offering a lightweight alternative for expert selection.The conclusion contrasts existing language-model- or embedding-based routers with IR3DE’s reported efficiency and adaptability.
- IR3DE achieves comparable performance across settings, reaches 98.4% normalized performance in Reasoning, and matches in-domain experts on average in CLM.The conclusion identifies IR3DE as a cheap and fast alternative for inference routing.
- Limitations and future work: IR3DE may be less effective on queries requiring richer semantic understanding or complex decision boundaries because linear routing is less expressive than stronger LLM-based routers.Suggested extensions include kernel ridge regression, evaluation on more complex reasoning tasks, and objectives incorporating computation, latency, and memory.
Appendix
The appendix figures vary MoDEM’s learning rate and report final accuracy for small and large routers in CLM and Reasoning settings.
- Figure 4 reports MoDEM final accuracy across learning rates for small and large routers in the CLM setting.The left and right panels correspond to the small and large routers, respectively.
- Figure 5 reports MoDEM final accuracy across learning rates for small and large routers in the Reasoning setting.The passage identifies the small router on the left and the large router on the right.
A. Details on MoDEM router training
MoDEM routers are trained on merged, balanced domain datasets, with training durations varying by setting and router size. Across tested learning rates, the largest rate is selected for final evaluation.
- MoDEM training uses a merged dataset across domains, with 1750 samples fixed per domain to avoid class imbalance.The balancing procedure applies to the corresponding CLM and Reasoning settings.
- 100 epochs are used for both Reasoning routers and the small CLM router, while the large CLM router is trained for 10 epochs after validation accuracy plateaus.The large CLM router required fewer epochs than the other configurations.
- lr = 0.1 is selected in every setting because it provides the best test routing accuracy among the tested learning rates.The tested values range from 10^-5 through 10^-1.