Technical Summary: TRACE-Router

Problem Statement

Current Large Language Model (LLM) routing systems predominantly operate at the request level, making independent routing decisions for each individual prompt. While this granularity is suitable for single-turn inference, it is misaligned with agentic workloads, where a task consists of a long-horizon, stateful execution trace involving multiple LLM calls interleaved with tool use, retrieval, and environment transitions.

The core mismatch lies in feedback attribution:

1. Delayed Feedback: In agentic applications, success or failure is determined only by a terminal, task-level outcome (e.g., passing a test suite or achieving a goal), not by the quality of individual intermediate turns.
2. Credit Assignment Failure: Request-level routers cannot correctly attribute this delayed task-level reward to specific intermediate routing decisions. A seemingly optimal local decision (e.g., routing an intermediate step to a smaller, faster model) can lead to global suboptimality if it causes the model to lose accumulated state or fail to complete the complex task.
3. Static Policies: Existing routers often rely on offline-trained classifiers that fix routing policies before deployment, preventing adaptation to workload shifts or correction of mistaken assignments once a task begins.

Methodology: TRACE-Router

The authors propose TRACE-Router, a framework that shifts the unit of routing and learning from the individual request to the entire task trace. The system operates on the principle that the unit of routing must match the unit of feedback.

Core Mechanisms

1. Task-Consistent Routing:

   • Upon task admission, the router selects a single backend model for the entire execution trace.
   • A persistent task identifier is attached to every request within the trace. Subsequent requests reuse the stored backend assignment, ensuring model consistency and preserving model-specific state throughout the execution.
   • This creates a single credit-assignment target: the final task reward is attributed to the one backend that served the whole trace.

2. Context-Conditioned Online Learning:

   • To handle diverse workload regimes (e.g., tasks where latency is critical vs. tasks where accuracy is paramount), the system uses a contextual bandit approach.
   • Tasks are assigned to a coarse routing context (e.g., Easy, Medium, Hard) based on features available at admission (e.g., regex-based complexity classifiers), without requiring additional LLM calls.
   • The system maintains independent bandit policies ($B_c$) for each context, allowing the router to learn distinct accuracy–latency trade-offs specific to different task types.

3. Cold-Start UCB Policy:

   • The routing policy utilizes an Upper Confidence Bound (UCB) algorithm with a cold-start initialization.
   • To address the statistical difficulty of sparse feedback in each context, the system performs forced round-robin exploration (k pulls per arm) before transitioning to the UCB selection rule.
   • The authors found that "warm-starting" with priors did not improve accuracy and often biased the system toward slower models; thus, cold-start is the default.

4. Delayed Reward and Scalarization:

   • Feedback is observed only after task termination. The system combines graded accuracy ($a_{t,m}$) and end-to-end latency ($\ell_{t,m}$) into a scalar reward using a preference parameter $\alpha$:
     $$r^{(\alpha)}_{t,m} = (1 - \alpha) a_{t,m} - \alpha \tilde{\ell}_{t,m}$$
   • By sweeping $\alpha$, the system traces a family of deployment policies, allowing operators to tune the trade-off between accuracy and latency.

Key Contributions

1. Task-Level Routing Formulation: The paper formulates model routing as task-consistent online model selection, aligning the routing decision with the unit of feedback (the task trace) to enable principled credit assignment from delayed rewards.
2. Context-Aware Online Adaptation: TRACE-Router introduces a contextual bandit-based router that learns context-specific policies online, adapting to workload regimes without offline training or explicit task-complexity estimation.
3. Improved Pareto Frontiers: The framework consistently learns routing policies that occupy the empirical accuracy–latency Pareto frontier, outperforming both heuristic baselines (e.g., complexity routers) and learned request-level routers.

Experimental Results

The authors evaluated TRACE-Router across three agentic benchmarks: $\tau^2$-Bench (retail and telecom domains), LiveCodeBench, and Terminal-Bench.

• Pareto Frontier Dominance: Across all benchmarks, TRACE-Router achieved non-dominated Pareto frontier points. It consistently outperformed latency-matched random interpolation between models.
  • On $\tau^2$-Bench, it outperformed latency-matched interpolation by 7–8 accuracy points.
  • On Terminal-Bench, it achieved 7.1 higher accuracy points than the strongest single-model baseline while reducing latency by 36%.
• Context Conditioning: While the aggregate gain from context conditioning was modest, it was critical for determining reachable operating points. The context-conditioned variant held the frontier at specific $\alpha$ values where the context-free variant was dominated, particularly in domains with divergent model performance characteristics (e.g., telecom).
• Cold-Start vs. Warm-Start: Experiments confirmed that cold-start initialization outperformed warm-start variants, which tended to reinforce prior biases and select slower models.
• Scalability: In a four-model ablation study, TRACE-Router successfully identified and avoided unreliable models (e.g., those prone to timeouts) that heuristic routers continued to select, demonstrating robustness as the candidate pool size increased.

Significance and Claims

The paper claims that task-consistent commitment is a practical and superior unit for deployment-time model selection in agentic AI. By treating workflows as first-class scheduling objects rather than collections of independent requests, TRACE-Router resolves the credit assignment mismatch inherent in current request-level routers.

The authors emphasize that their approach does not require explicit task-complexity estimation or offline training. Instead, it leverages delayed task-level feedback to adaptively learn routing policies that balance accuracy and latency. The results suggest that adaptive task-level routing yields significantly better accuracy than probabilistic model mixing at the same latency budget, offering a path toward more efficient and reliable agentic AI deployment.