Technical Summary: UNISCALE

1. Problem Statement

In real-world deployments of Large Language Models (LLMs), balancing inference quality against computational cost is a critical challenge. Existing approaches typically address this trade-off along two independent dimensions:

• Model Routing: Selecting models of different scales from a heterogeneous pool. While this covers a broad quality–cost spectrum, it operates at a coarse granularity, inducing discrete jumps in performance and cost.
• Test-Time Scaling (TTS): Adjusting inference-time compute (e.g., via Best-of-N, Beam Search, or Process Reward Models) within a fixed model. While this offers fine-grained control, it is fundamentally constrained by the intrinsic capacity of the base model and often exhibits diminishing returns as compute increases.

The decoupled design of these mechanisms limits adaptability in dynamic environments where query distributions, user objectives, and model availability shift. Furthermore, treating them separately restricts the ability to exploit synergies between model selection and inference strategy.

2. Methodology

Unified Inference Scaling (UIS)

The authors introduce Unified Inference Scaling (UIS), a paradigm that unifies model routing and TTS into a single optimization space. An inference execution is parameterized by a joint configuration $(M, QP, CP, BS)$, where:

• $M$: The selected base model.
• $QP$: Question Parallelism (number of subtrees explored).
• $CP$: Candidate Parallelism (number of parallel samples per step).
• $BS$: Beam Size (number of validated nodes retained).

This formulation creates a rich design space where TTS can smooth the discrete gaps between model scales, and routing to larger models can mitigate the diminishing returns of aggressive TTS on smaller models.

The UNISCALE Framework

To solve the adaptive UIS problem, the authors propose UNISCALE, an online framework modeled as a contextual multi-armed bandit problem.

• Online Learning via LinUCB: The system treats configuration selection as an online optimization task. It uses a Transformer-based encoder to map queries and action configurations (model + TTS strategy) into a shared semantic latent space. The Linear Upper Confidence Bound (LinUCB) algorithm is employed to learn inference policies, balancing exploitation (high predicted reward) and exploration (uncertainty in high-dimensional spaces).
• Path-Aware Early Exiting: To address the "straggler effect" in TTS workflows, UNISCALE dynamically terminates low-potential inference paths. By leveraging verifier scores as indicators of final answer correctness, the system calculates the theoretical maximum potential of ongoing paths. If a path cannot surpass the current best result even under perfect future performance, it is terminated immediately, significantly reducing computational cost without compromising final selection quality.
• Dense Verification Feedback: Instead of relying on sparse binary correctness signals, the framework incorporates native verifier scores from TTS. This provides a denser supervision signal, allowing the agent to perceive nuances in logical rigor and guide the selection of configurations with superior performance.
• UIS Cost Model: To ensure consistent optimization, the framework utilizes equivalent FLOPs (eFLOPs) to map both computational and memory overheads into a unified metric. This accounts for hardware arithmetic intensity, bridging the gap between compute-bound and memory-bound operations.

The composite reward function combines inference quality (correctness and verifier score) and cost, allowing the system to adapt to varying user priorities (e.g., Cost-Sensitive vs. Quality-Priority).

3. Key Contributions

1. Unified Paradigm: The proposal of Unified Inference Scaling (UIS), which integrates model routing and TTS into a single decision space, overcoming the limitations of coarse-grained routing and capacity-constrained TTS.
2. Adaptive Framework (UNISCALE): An online algorithm based on contextual bandits (LinUCB) that learns adaptive inference policies in non-stationary environments, utilizing semantic representations for efficient cross-action knowledge transfer.
3. Efficiency Mechanisms:
   • Path-Aware Early Exiting: A dynamic pruning mechanism that reduces redundant computation in TTS workflows.
   • Dense Verification Feedback: A method to densify reward signals using verifier scores, improving policy convergence.
   • Unified Cost Modeling: An eFLOPs-based cost model that accurately reflects hardware-level resource consumption.
4. Empirical Validation: Extensive experiments demonstrating that UNISCALE consistently outperforms baselines (including Random, Greedy, MLP, k-NN, Thompson Sampling, and NeuralUCB) across TTS-only, Routing-only, and full UIS scenarios.

4. Experimental Results

The authors evaluated UNISCALE on 210 instances from AIME'24, AIME'25, and MATH-500 datasets using a Qwen3 model series and a Skywork PRM verifier.

• Performance: UNISCALE achieved the highest rewards and lowest cumulative regret across both Cost-Sensitive and Quality-Priority modes. In the full UIS space, it delivered a fine-grained quality–cost frontier superior to all baselines.
• Ablation Studies:
  • Semantic Representations: Removing action semantic features (using one-hot vectors) resulted in higher cumulative regret, confirming that semantic mapping accelerates policy optimization in high-dimensional spaces.
  • Early Exiting: This mechanism reduced computational load by ~87% and memory access by ~78% with negligible impact on accuracy (<1pp).
  • Dense Feedback: Using verifier scores instead of binary correctness improved accuracy by 5.62pp in Quality-Priority mode and reduced cost by 51.3%.
• Robustness: UNISCALE demonstrated strong adaptability to environmental drifts (e.g., model removal/addition and reward mode shifts), rapidly recalibrating policies to maintain optimal trade-offs where baselines like k-NN struggled.

5. Significance and Claims

The paper claims that UNISCALE provides a unified theoretical foundation and decision space for large-scale AI inference orchestration. Its significance lies in:

• Breaking Technical Silos: It integrates model routing and TTS, shifting inference optimization from heuristic rules to data-driven automated policies.
• Green AI: By minimizing redundant computation through early exiting and unified cost modeling, it directly addresses energy consumption challenges in LLM deployment.
• Democratization: The framework offers fine-grained control over quality–cost trade-offs, enabling resource-constrained edge devices to intelligently leverage local capabilities or cloud upgrades, potentially bridging the digital divide.
• Privacy: The abstract semantic decision space is compatible with federated orchestration, allowing global policy evolution through anonymized parameter updates without accessing raw query data.

The authors note that current performance is modulated by the precision of Process Reward Models (PRMs) and identify the development of more generalizable verification mechanisms as a direction for future work.