IR3DE: A Linear Router for Large Language Models
IR3DE uses ridge regression for fast, low-cost model routing, achieving 98.4% performance in multi-domain LLM selection.
Key Findings
Methodology
IR3DE employs ridge regression to construct a linear token router, which rapidly assesses input token embeddings to select the most suitable domain-specific model. It consists of a Token Router (TR) that solves a regularized least squares problem to learn linear weights, enabling asynchronous dataset handling and model expansion. The Sample Route Selector (SRS) uses entropy-based voting on top-k tokens to improve complex reasoning accuracy. Experiments on two causal language modeling (CLM) tasks and one reasoning task demonstrate performance comparable to deep learning-based routers, reaching 98.4% accuracy, with high efficiency and flexibility.
Key Results
- In two CLM tasks, IR3DE matches baseline performance with less than 1% error, and in the reasoning task, achieves 98.4%, surpassing deep learning classifiers. The entropy-based SRS outperforms other variants, especially on complex tasks.
- IR3DE supports dynamic addition/removal of models without retraining, with less than 2% performance fluctuation, showing excellent adaptability.
- Different SRS variants reveal entropy as an effective uncertainty measure, with entropy-based voting yielding the best results in complex reasoning scenarios.
Significance
This work addresses the efficiency bottleneck in large model routing, proposing a simple yet powerful ridge regression-based solution. It reduces computational costs, supports dynamic model management, and requires minimal training data, making large-scale multi-model deployment feasible. Its impact extends to automatic QA, specialized reasoning, and content generation industries, enabling faster, more flexible AI systems with lower deployment barriers.
Technical Contribution
IR3DE introduces a closed-form ridge regression approach for linear model routing, eliminating the need for complex neural network training. Its design allows real-time model addition/removal, maintaining high accuracy across tasks. The entropy-based selection enhances reasoning quality by quantifying uncertainty, providing a novel, efficient alternative to deep classifiers. This approach significantly simplifies and accelerates model scheduling, opening new avenues for scalable AI deployment.
Novelty
This is the first work to leverage ridge regression for large model routing, offering a linear, computationally efficient, and dynamically adaptable solution. Unlike prior methods such as MoDEM and PolyRouter, IR3DE does not require extensive labeled data or retraining when models are added or removed, making it highly practical for real-world, large-scale systems. Its combination of simplicity, speed, and adaptability marks a significant advancement in model scheduling technology.
Limitations
- The method relies on pre-trained embeddings; if embeddings are poor, routing accuracy may suffer, especially in niche or low-resource domains.
- Linear assumptions limit modeling complex, nonlinear relationships, potentially reducing performance on highly interactive or non-linear tasks.
- Matrix inversion costs may grow with extremely large models or many domains, requiring further optimization for scalability.
Future Work
Future research will explore nonlinear extensions such as kernel ridge regression to better capture complex relationships, and incorporate reinforcement learning for more intelligent model selection. Additionally, optimizing matrix computations for ultra-large models and multi-modal inputs will be prioritized to enhance real-time performance and broader applicability.
AI Executive Summary
The rapid growth of large language models (LLMs) has revolutionized natural language processing, but efficiently deploying multiple models across diverse domains remains a challenge. Traditional routing methods often rely on deep classifiers or rule-based systems, which are costly to train and inflexible to model updates. This paper introduces IR3DE, a novel linear router based on ridge regression, designed to address these issues by offering a fast, low-cost, and adaptable solution for multi-domain model selection.
IR3DE's architecture comprises two main components: the Token Router (TR) and the Sample Route Selector (SRS). The TR employs a closed-form ridge regression solution to learn linear weights that map token embeddings to domain labels, enabling rapid inference without extensive training. The SRS utilizes entropy measures on top-k tokens' softmax probabilities to vote for the most appropriate expert model, significantly improving performance on complex reasoning tasks.
Experimental results demonstrate that IR3DE achieves performance comparable to, or better than, existing deep learning-based routers, reaching a normalized score of 98.4% on a reasoning benchmark. Its ability to support dynamic model addition or removal without retraining the router itself makes it highly suitable for real-world, large-scale AI systems. The entropy-based SRS variant notably excels in tasks requiring nuanced understanding, validating the effectiveness of uncertainty quantification.
This work marks a significant step forward in model scheduling, offering a simple yet powerful approach that combines efficiency, flexibility, and high accuracy. Its implications extend to industry applications such as automated customer support, domain-specific content generation, and real-time decision-making systems. Looking ahead, integrating nonlinear extensions and reinforcement learning strategies promises further improvements, paving the way for smarter, more scalable AI deployment frameworks.
Deep Dive
Abstract
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.