UniScale: Adaptive Unified Inference Scaling via Online Joint Optimization of Model Routing and Test-Time Scaling

2026-05-29

AI Summary

This work addresses the fundamental trade-off between inference quality and computational cost in large language model deployment, where existing approaches decouple model routing and test-time scaling, leading to coarse control and poor adaptability. To overcome these limitations, the paper proposes Unified Inference Scaling (UIS), a novel framework that jointly models both decisions as a contextual multi-armed bandit problem for the first time. UIS employs LinUCB for online learning of adaptive policies and incorporates efficiency-aware learning with fine-grained cost modeling to enable stable optimization in high-dimensional action spaces. Experimental results demonstrate that UIS significantly outperforms current methods across diverse dynamic scenarios, achieving higher inference quality at the same computational cost or substantially reducing overhead while maintaining equivalent performance.

Abstract

In real-world deployments of large language models (LLMs), balancing inference quality and computational cost has become a central challenge. Existing approaches tackle this trade-off along two largely independent dimensions: model routing, which switches among models of different scales to match request complexity, and test-time scaling (TTS), which adjusts inference-time compute within a fixed model for fine-grained control. However, this decoupled design introduces inherent limitations. Model routing yields coarse-grained, discrete performance changes due to the sparse set of model scales, while single-model TTS often encounters capacity ceilings and exhibits diminishing returns as compute increases. Moreover, treating the two mechanisms separately restricts adaptability in dynamic inference environments. To overcome these limitations, we introduce Unified Inference Scaling (UIS), which unifies model routing and TTS in a single optimization space. Building on this formulation, we propose UniScale, an online framework that models adaptive UIS as a contextual multi-armed bandit problem and learns inference policies via LinUCB. The framework incorporates efficiency-aware learning and cost modeling to ensure stable and scalable optimization over high-dimensional action spaces. Evaluation shows that UniScale effectively exploits the synergy in the UIS space to deliver a fine-grained and consistently better quality-cost trade-off across diverse, dynamic inference scenarios.

Problem

• inference scaling
• model routing
• test-time scaling
• quality-cost trade-off
• large language models

Innovation

• Unified Inference Scaling
• model routing
• test-time scaling
• online joint optimization
• contextual multi-armed bandit