UniScale: 通过模型路由与测试时扩展的在线联合优化实现自适应统一推理扩展

# 自适应统一推理缩放：基于模型路由与测试时缩放在线联合优化

来源：https://arxiv.org/html/2605.30898

Xingyu Wang，Mingze Kong，Zhubo Shi，Yuqian Hou，Hong Xu，Zhongxiang Dai，Minchen Yu，Qingjiang Shi

###### 摘要

在大语言模型（LLM）的实际部署中，平衡推理质量与计算成本已成为核心挑战。现有方法主要沿着两个相对独立的维度解决这一权衡：模型路由（在不同规模的模型之间切换以匹配请求复杂度）和测试时缩放（TTS，在固定模型内调整推理计算量以实现细粒度控制）。然而，这种解耦设计存在固有局限性。模型路由由于模型规模稀疏性而带来粗粒度、离散的性能变化，而单模型TTS常遭遇能力上限，随着计算量增加收益递减。此外，将这两种机制分开处理限制了在动态推理环境中的适应性。为克服这些局限，我们提出**统一推理缩放（UIS）**，将模型路由与TTS统一到单一优化空间。基于此公式，我们提出**UniScale**在线框架，将自适应UIS建模为上下文多臂老虎机问题，并通过LinUCB学习推理策略。该框架融合了效率感知学习和成本建模，以确保在高维动作空间上的稳定可扩展优化。评估表明，UniScale有效利用了UIS空间中的协同效应，在多样化的动态推理场景中实现了细粒度且持续更优的质量-成本权衡。

机器学习，ICML

参见图注

图1：UIS空间下的精度-成本权衡。方块和虚线分别代表路由节点和单模型TTS轨迹。通过跨两个维度联合优化，UIS实现了富有表现力的质量-成本前沿。

## 1 引言

近年来，大语言模型在复杂推理、决策及多步骤问题求解等各类任务中展现出卓越的成功（Guo et al., 2025；Jaech et al., 2024；Team, 2025）。实际LLM应用覆盖了从易到难的任务难度谱系，满足这些需求通常需要相当大的模型容量。然而，更高的能力通常带来更大的计算成本——要么通过调用更大的模型（更高的每token延迟和内存需求），要么通过增加推理计算量（例如，更长的解码或额外的测试时策略）。这些增加的成本会提高响应延迟和资源消耗，对交互式和大规模在线服务影响尤为显著。因此，实际LLM部署必须在推理质量和计算成本之间取得权衡。

为平衡质量与成本，现有方法通常探索两个相对独立的维度，如图1所示。**模型路由方法**（Feng et al., 2025）在不同规模或能力的模型之间进行选择，覆盖较宽的质量-成本谱系；但操作粒度粗，因为切换模型会导致精度和成本发生离散变化。相比之下，**测试时缩放（TTS）方法**（Snell et al., 2025）在运行时动态调整给定模型的推理过程，提供细粒度控制，但根本上受限于模型的内在容量。此外，这两类方法通常孤立设计和配置，限制了它们在动态推理环境中协同运作的能力。

为在实践中实现最优质量-成本权衡，理想方法应满足三个要求。首先，应提供跨越不同规模模型的广阔优化空间，以处理具有不同难度和计算需求的查询（即**广泛覆盖**）。其次，应提供对推理过程的细粒度控制，以实现精确的质量-成本权衡（即**细粒度**）。最后，由于在线部署面临查询分布、用户目标和模型可用性的变化（即**环境漂移**），该方法应能持续调整其推理策略（即**在线自适应性**）。

为实现这些要求，我们提出**统一推理缩放（UIS）**，一种将模型路由和TTS视为**单一统一推理时决策空间**而非独立旋钮的推理范式。在UIS下，推理由联合指定基础模型及其关联TTS策略的配置参数化（见第2.1节）。如图1所示，由此产生的UIS配置集合构成了一个丰富的设计空间，其中路由和TTS相互交互：TTS可缩小离散模型规模间的性能差距，而路由到更大模型则可在需要时缓解较小模型上过度激进的TTS收益递减。

基于此公式，我们提出**UniScale**，一种求解自适应UIS问题的在线算法。我们将UIS配置选择建模为在线上下文老虎机问题（Li et al., 2010），该问题自然捕捉了非平稳环境下低延迟决策的特点（见第2.2节）。UniScale采用基于Transformer的编码器提取查询表示，并使用线性上置信界（LinUCB）（Abbasi-Yadkori et al., 2011）算法在线学习推理策略（见第3节），实现对环境变化的持续适应。

实际优化UIS面临高维、异质配置空间和在线推理严格延迟约束等挑战。为确保高效稳定的学习，我们引入三个紧密集成的机制（见第3.3和3.4节）：（1）**路径感知早退**动态识别并终止低潜力推理路径，显著降低计算成本同时保证推理质量，从而优化所有UIS配置的运行时性能；（2）**密集验证反馈**通过纳入来自TTS的原生验证器分数来稠化稀疏的二进制正确性信号，提供更准确的质量评估以引导系统选择性能更优的UIS配置；（3）**UIS成本模型**利用等效FLOPs（eFLOPs）将计算和内存开销映射为统一度量（Sadhukhan et al., 2025），通过提供准确的成本测量确保UIS配置的理论成本与实际成本一致。

大量实验表明，UniScale在多种设置下持续优于基线策略，包括固定模型下的TTS选择（第4.1节）、模型路由（第4.2节）以及完整的统一推理缩放（第4.3节）。我们进一步通过全面的消融研究分析UniScale，阐明每个组件的贡献（第5节）。

## 2 背景与问题设定

### 2.1 统一推理缩放基础

**模型路由**。模型路由旨在根据输入查询动态地从异质模型池中选择最合适的模型。然而，虽然模型路由方法能够覆盖较宽的质量-成本范围，但其操作粒度粗：切换模型通常会在推理质量和计算成本上导致离散且显著的跳跃。

参见图注

图2：TTS方法的统一参数化。不同的TTS策略通过问题并行度（QP）、候选并行度（CP）和波束大小（BS）进行形式化。

**测试时缩放**。测试时缩放技术通过分配额外的推理时计算预算来增强模型推理（Wu et al., 2025）。利用Best-of-N（BoN）（Cobbe et al., 2021）或基于过程奖励模型（PRM）的搜索（例如波束搜索和DVTS），TTS方法通过迭代验证和路径评估来优化推理（详见第3.3节）。如图2所示，我们将这些策略形式化为三维参数空间：**问题并行度**（QP，探索的子树数量），**候选并行度**（CP，每步并行样本数量），和**波束大小**（BS，保留的已验证节点数量）。这种参数化允许对质量-成本权衡进行细粒度控制，但仍从根本上受限于基础模型的内在容量上限。

**统一推理缩放**。为弥合这些技术隔阂，我们将统一推理缩放（UIS）定义为单一统一的推理时决策空间。我们通过联合配置`(M, QP, CP, BS)`来参数化每次推理执行。这一公式通过搜索强度弥合离散模型间隙，并通过模型路由超越搜索平台，构建出富有表现力的质量-成本前沿（见图1）。这一统一空间使得推理策略能够精确适配每次查询独特的逻辑复杂度和资源约束。

### 2.2 通过老虎机实现自适应统一推理缩放

为在环境漂移下实现在线自适应性，我们通过**上下文多臂老虎机**的视角（Li et al., 2010）将UIS配置选择建模为在线优化问题。

**上下文信息**。对于每个到达的查询 `q_t`（在步骤t），系统观测到上下文向量 `x_t`。该向量作为查询复杂度的抽象表示，为智能体估计不同UIS配置的潜在效用提供必要信号。

**统一动作空间**。我们定义离散化动作空间 `\mathcal{A}`，其中每个动作 `a \in \mathcal{A}` 对应一个具体的UIS配置`(M, QP, CP, BS)`。通过将路由和TTS视为联合动作，UniScale能够捕捉独立方法忽视的跨维度依赖关系。

**优化目标**。执行动作`a_t`后，智能体收到奖励`r_t`，该奖励表征为推理质量和计算成本的联合函数。智能体的目标是最小化累积遗憾：
```
R_T = E[ \sum_{t=1}^{T} ( r(q_t, a_t^*) - r(q_t, a_t) ) ],
```
其中`a_t^*`是上下文`x_t`下的最优配置。通过最小化这一遗憾，算法学习到一个自适应策略，能够持续选择适应演化环境的最优UIS配置。

算法1 UniScale: 基于LinUCB的自适应UIS

1: 初始化: `A_0 \leftarrow \lambda I`, `b_0 \leftarrow 0`, `x_{0,a_0} \leftarrow 0`, `r_0 \leftarrow 0`, 动作嵌入 `{s_a}_{a \in \mathcal{A}}`
2: for t = 1 to T do
3:   更新 `A_t \leftarrow A_{t-1} + x_{t-1, a_{t-1}} x_{t-1, a_{t-1}}^{\top}`
4:   更新 `b_t \leftarrow b_{t-1} + r_{t-1} x_{t-1, a_{t-1}}`
5:   更新 `\hat{\theta}_t \leftarrow A_t^{-1} b_t`
6:   观测查询 `q_t` 并提取嵌入 `s_{q_t}`
7:   for all a in \mathcal{A} do
8:       `x_{t,a} \leftarrow concat(s_{q_t}, s_a)`
9:   end for
10:  选择UIS配置 `a_t = \arg\max_{a \in \mathcal{A}} ( \hat{\theta}_t^{\top} x_{t,a} + \alpha \sqrt{x_{t,a}^{\top} A_t^{-1} x_{t,a}} )`
11:  使用配置 `a_t` 执行推理
12:  根据公式4获取奖励 `r_t`
13: end for

## 3 UniScale框架

**概述**。UniScale作为一个在线闭环系统运行（见算法1），旨在导航联合UIS决策空间。在每次迭代t中，系统基于历史特征-奖励对更新其奖励估计器 `\hat{\theta}_t`（第3.1节）。随后，它通过最大化LinUCB采集函数为传入查询选择最优UIS配置`a_t`（第3.2节）。选择后，UniScale在基础模型`M_t`上执行推理过程，利用路径感知早退优化运行时的质量-成本权衡（第3.3节）。最后，系统通过密集验证反馈和UIS成本模型评估执行结果，产生复合奖励`r_t`用于持续策略优化（第3.4节）。

### 3.1 更新奖励估计器

在每次迭代t开始时，UniScale使用上一轮观测到的反馈`(x_{t-1, a_{t-1}}, r_{t-1})`更新奖励估计器。我们将期望奖励形式化为线性关系 `\hat{r}_t = \langle x_{t, a_t}, \theta \rangle`，其中 `\theta` 是跨整个UIS动作空间共享的可学习参数向量。为优化此估计器，系统增量更新格拉姆矩阵 `A_t` 和特征-奖励向量 `b_t`，以得到当前参数估计 `\hat{\theta}_t`，详见算法1（第3-5行）。实践中，我们利用Sherman-Morrison公式通过高效的一阶秩更新来更新 `A_t^{-1}`。这避免了从头重新计算逆矩阵，将每轮计算复杂度从 `O(d^3)` 降低到 `O(d^2)`。这种增量机制确保了奖励估计器以最小计算开销持续适应环境漂移。

**联合语义表示**。为捕捉用户需求与系统能力之间的内在对齐，UniScale通过Transformer编码器将查询和配置映射到共享潜在空间。