Adaptive Test-Time Compute Allocation for Reasoning LLMs via Constrained Policy Optimization
作者: Zhiyuan Zhai, Bingcong Li, Bingnan Xiao, Ming Li, Xin Wang
分类: cs.LG
发布日期: 2026-04-16
💡 一句话要点
提出基于约束策略优化的自适应推理计算分配方法，提升LLM在有限预算下的性能。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 计算资源分配 约束优化 拉格朗日松弛 自适应推理
📋 核心要点
- 现有LLM推理时，对所有输入采用统一的计算资源分配策略，忽略了不同输入对计算资源的需求差异。
- 论文提出Solve-then-Learn框架，首先通过拉格朗日松弛求解最优计算分配策略，然后训练分类器预测该策略。
- 实验结果表明，该方法在MATH和GSM8K数据集上，使用多个LLM时，均优于统一分配和启发式分配策略。
📝 摘要（中文）
本文提出了一种测试时计算缩放方法，通过重复采样、搜索或扩展推理，在推理过程中分配额外的计算资源，以提高大型语言模型的性能。该方法将问题形式化为约束优化问题（在平均计算预算约束下最大化预期准确率），并使用两阶段的Solve-then-Learn流程解决。在Solve阶段，拉格朗日松弛将全局约束分解为每个实例的子问题，每个子问题都允许一个闭式最优动作，该动作根据成本对准确率进行定价。证明了诱导成本在对偶变量中是单调的，从而可以通过二分搜索实现精确的预算目标。在Learn阶段，训练一个轻量级分类器，从廉价的输入特征中预测oracle动作，从而分摊实时部署的分配规则。理论证明了学习策略的任务级遗憾受其模仿误差乘以最坏情况下的每个实例差距的限制，从而将约束推理简化为监督分类。在MATH和GSM8K上，使用三个LLM（DeepSeek-V3、GPT-4o-mini、Qwen2.5-7B）进行的实验表明，该方法始终优于统一和启发式分配基线，在匹配的预算约束下，在MATH上实现了高达12.8%的相对准确率提升，同时以超过91%的模仿准确率紧密跟踪拉格朗日oracle上限。
🔬 方法详解
问题定义：现有的大型语言模型在推理时，通常采用固定的计算资源分配策略，即对所有输入样本都分配相同的计算量。然而，不同输入样本的难度各异，简单的样本可能不需要过多的计算即可正确回答，而复杂的样本则需要更多的计算资源。因此，如何根据输入样本的难度自适应地分配计算资源，以在有限的计算预算下最大化模型的整体性能，是一个亟待解决的问题。
核心思路：论文的核心思路是将自适应计算资源分配问题建模为一个约束优化问题，目标是在平均计算预算的约束下，最大化模型的预期准确率。为了解决这个约束优化问题，论文采用了Solve-then-Learn的框架。首先，通过拉格朗日松弛将全局约束分解为每个实例的子问题，并求解每个子问题的最优动作（即计算资源分配策略）。然后，训练一个轻量级的分类器，用于预测每个实例的最优动作，从而实现计算资源分配策略的泛化。
技术框架：该方法主要包含两个阶段：Solve阶段和Learn阶段。在Solve阶段，首先使用拉格朗日松弛将全局约束分解为每个实例的子问题。然后，对于每个子问题，求解其闭式最优解，得到每个实例的最优计算资源分配策略。在Learn阶段，训练一个轻量级的分类器，以输入样本的特征作为输入，预测Solve阶段得到的最优计算资源分配策略。该分类器可以使用任何监督学习算法进行训练。
关键创新：该方法最重要的技术创新点在于将自适应计算资源分配问题建模为一个约束优化问题，并采用Solve-then-Learn的框架进行求解。与传统的启发式方法相比，该方法能够更加精确地控制计算资源的分配，从而在有限的计算预算下获得更高的模型性能。此外，该方法还证明了学习策略的任务级遗憾受其模仿误差乘以最坏情况下的每个实例差距的限制，从而将约束推理简化为监督分类。
关键设计：在Solve阶段，论文证明了诱导成本在对偶变量中是单调的，从而可以通过二分搜索实现精确的预算目标。在Learn阶段，论文使用一个轻量级的分类器来预测oracle动作，从而分摊实时部署的分配规则。分类器的具体结构和参数设置可以根据具体的应用场景进行调整。损失函数采用交叉熵损失函数，优化算法采用Adam优化器。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在MATH和GSM8K数据集上，使用DeepSeek-V3、GPT-4o-mini和Qwen2.5-7B三个LLM时，均优于统一分配和启发式分配策略。在MATH数据集上，该方法在匹配的预算约束下，实现了高达12.8%的相对准确率提升，同时以超过91%的模仿准确率紧密跟踪拉格朗日oracle上限。
🎯 应用场景
该研究成果可应用于各种需要进行推理的大型语言模型应用场景，例如问答系统、数学问题求解、代码生成等。通过自适应地分配计算资源，可以在有限的计算预算下提高模型的性能，降低推理成本，并提升用户体验。该方法具有广泛的应用前景，可以促进大型语言模型在资源受限环境下的部署和应用。
📄 摘要（原文）
Test-time compute scaling, the practice of spending extra computation during inference via repeated sampling, search, or extended reasoning, has become a powerful lever for improving large language model performance. Yet deploying these techniques under finite inference budgets requires a decision that current systems largely ignore: which inputs deserve more compute, and which can be answered cheaply? We formalize this as a constrained optimization problem (maximize expected accuracy subject to an average compute budget) and solve it with a two-stage Solve-then-Learn pipeline. In the solve stage, Lagrangian relaxation decomposes the global constraint into per-instance sub-problems, each admitting a closed-form oracle action that optimally prices accuracy against cost. We prove that the induced cost is monotone in the dual variable, enabling exact budget targeting via binary search. In the learn stage, a lightweight classifier is trained to predict oracle actions from cheap input features, amortizing the allocation rule for real-time deployment. We establish that the task-level regret of the learned policy is bounded by its imitation error times the worst-case per-instance gap, yielding a clean reduction from constrained inference to supervised classification. Experiments on MATH and GSM8K with three LLMs (DeepSeek-V3, GPT-4o-mini, Qwen2.5-7B) show that our method consistently outperforms uniform and heuristic allocation baselines, achieving up to 12.8% relative accuracy improvement on MATH under matched budget constraints, while closely tracking the Lagrangian oracle upper bound with over 91% imitation accuracy.