Back to AI Research

AI Research

TRACE-ROUTER: Task-Consistent and Adaptive Online R... | AI Research

Key Takeaways

  • TRACE-Router: Task-Consistent and Adaptive Online Routing for Agentic AI Modern enterprise AI often relies on a mix of large, powerful models and smaller, fa...
  • Routing to select large language models (LLMs) with different cost-quality trade-offs has become a fundamental deployment feature of enterprise AI.
  • Existing routers, primarily make independent routing decisions for each LLM call.
  • However, agentic applications execute as long-horizon workflows whose quality is determined only by a delayed, task-level outcome.
  • This mismatch prevents per-call routers from correctly attributing feedback to individual routing decisions.
Paper AbstractExpand

Routing to select large language models (LLMs) with different cost-quality trade-offs has become a fundamental deployment feature of enterprise AI. Existing routers, primarily make independent routing decisions for each LLM call. However, agentic applications execute as long-horizon workflows whose quality is determined only by a delayed, task-level outcome. This mismatch prevents per-call routers from correctly attributing feedback to individual routing decisions. Towards mitigating this, we present TRACE-Router, a task-level routing framework that aligns routing with the unit of supervision. TRACE-Router assigns each task to a model once at admission using a contextual bandit, pins all subsequent LLM calls to the selected backend, and updates its policy using the task's terminal reward, jointly accounting for accuracy and latency. By leveraging delayed task feedback, TRACE-Router learns routing policies that adapt to the workload while avoiding explicit task-complexity estimation. Across three agentic benchmarks, TRACE-Router consistently improves the accuracy-latency trade-off, achieving non-dominated Pareto frontier points. On tau2-Bench, it outperforms latency-matched interpolation between individual models by 7-8 accuracy points, while on Terminal-Bench it achieves 7.1 higher accuracy points than the strongest single model baseline with 36% lower latency.

TRACE-Router: Task-Consistent and Adaptive Online Routing for Agentic AI
Modern enterprise AI often relies on a mix of large, powerful models and smaller, faster ones. While routing requests to the right model is common, existing methods typically treat every individual prompt as an independent decision. This approach fails for agentic AI, where a single "task" involves a long, complex workflow of multiple model calls, tool uses, and environment interactions. TRACE-Router introduces a new framework that shifts the focus from routing individual requests to routing entire task workflows, ensuring that a single model handles a task from start to finish to improve both accuracy and efficiency.

Aligning Routing with Task Completion

The core innovation of TRACE-Router is "task-consistent" routing. Instead of letting a router switch models mid-task—which can break the continuity of the agent's work and make it impossible to know which model was responsible for a success or failure—the system assigns a model at the very beginning of a task. This assignment is "pinned" to the task using a persistent identifier. By keeping the same model for the entire duration of the workflow, the system can accurately attribute the final outcome (the terminal reward) to the specific model chosen, allowing the router to learn from the results of completed tasks.

Adaptive Learning Through Context

TRACE-Router does not rely on static, pre-trained rules. Instead, it uses a contextual bandit approach to learn the best model for different types of tasks in real-time. When a new task arrives, the system identifies its "context"—a coarse category based on information available at admission—and selects a model based on past performance for that specific category. This allows the system to adapt to the workload as it runs, automatically favoring faster models for simple tasks and more capable models for challenging ones, without needing explicit, complex estimations of task difficulty.

Balancing Accuracy and Latency

The framework allows developers to define a preference parameter that balances the trade-off between accuracy and latency. By adjusting this parameter, TRACE-Router can generate a range of policies that sit on the "Pareto frontier," meaning it achieves the best possible accuracy for any given latency budget. In testing across agentic benchmarks, this approach significantly outperformed standard methods. For example, on one benchmark, it achieved 7.1 higher accuracy points than the strongest single model while simultaneously reducing latency by 36%.

Key Takeaways

TRACE-Router demonstrates that treating agentic workflows as first-class objects—rather than collections of isolated requests—is essential for effective model routing. By using delayed, task-level feedback to update its policies online, the system avoids the need for extensive offline training and remains flexible enough to handle evolving workloads. This method ensures that the model selection process is not just a series of local decisions, but a strategic, long-term optimization that respects the nature of complex, multi-step AI agents.

Comments (0)

No comments yet

Be the first to share your thoughts!