SelfBudgeter: Adaptive Token Allocation for Efficient LLM Reasoning

SelfBudgeter通过自适应令牌预算预测和强化学习优化，在MATH数据集上实现74.47%响应长度压缩，同时保持接近原始准确性，显著提升大型推理模型的效率。

Large Language Model, Reinforcement Learning, Reasoning, Efficiency, Pre-training

Zheng Li, Qingxiu Dong, Jingyuan Ma, Di Zhang, Zhifang Sui

National Key Laboratory for Multimedia Information Processing, School of Computer Science, Peking University

Generated by grok-3

Background Problem

近年来，大型推理模型（如O1）在复杂推理任务中表现出色，其成功主要归功于通过强化学习（RL）训练的冗长思维链（Chain of Thought, CoT）过程。然而，这些模型即使面对简单问题也倾向于过度思考，生成过长的推理过程，导致计算资源浪费和用户等待时间延长。例如，QwQ-32B模型在回答简单问题如’2+3’时，生成了比Qwen2.5-72B-Instruct模型多100倍的令牌数。现有方法（如基于提示、集成训练和路由器架构）在控制响应长度方面存在不足，缺乏精确控制、增加额外计算开销或无法自主预测最佳推理长度。SelfBudgeter旨在解决这一问题，通过自适应预测最小令牌预算并生成符合预算的响应，显著减少用户等待时间，同时允许用户预填预算以满足特定需求。

Method

SelfBudgeter是一种自适应推理策略，旨在通过自动预测令牌预算并控制输出长度来减少大型语言模型（LLM）的过度思考。其核心思想是让模型根据问题难度自主估计所需的令牌预算，并在生成响应时严格遵守这一预算（无论是自估计还是用户指定）。具体方法如下：

• 两阶段训练框架：
  • 冷启动阶段：通过微调模型，使其学习在输出响应前以标签形式预测令牌预算。训练数据包括模型自身生成的正确响应及其长度，确保模型熟悉输出格式而不丧失推理能力。
  • 强化学习（RL）阶段：采用GRPO算法优化模型，设计奖励函数（Reward Function）以平衡三个目标：答案正确性、最小化令牌预算、响应长度与预算的一致性。
• 奖励函数设计：包括预算惩罚（Budget Penalty，当预算超过最大可接受值时施加惩罚）、精确预算控制奖励（PreB Reward，通过松弛系数$\alpha$允许响应长度在预算附近波动）和准确性奖励（确保正确答案的最低奖励高于错误答案的最高奖励）。公式如下： $\mathbf{R}(C, F, \ell, b, b_{\max}) = \begin{cases} r_f, & \text{if } F = 0, \\ \mathbf{P_B}(b, b_{\max}) + \text{PreB}(s_{\min}^W, s_{\max}^W, \ell, b, \alpha, b_{\text{best}}^W), & \text{if } F = 1 \text{ and } C = 0, \\ \mathbf{P_B}(b, b_{\max}) + \text{PreB}(s_{\min}^C, s_{\max}^C, \ell, b, \alpha, b_{\text{best}}^C), & \text{if } F = 1 \text{ and } C = 1. \end{cases}$
• 输出格式：模型输出分为两部分：预算值和解决方案，确保预算预测先行，响应生成受限。

批判性思考：虽然方法设计合理，但其依赖于模型对问题难度的准确预估，而论文未充分讨论预估偏差可能导致的性能下降。此外，奖励函数中多个目标的平衡可能在实际应用中难以调优，尤其是当问题难度分布不均时，模型可能倾向于过度压缩或过度扩展响应长度。

Experiment

实验基于DeepSeek-R1-Distill-Qwen-1.5B模型，在GSM8K和MATH数据集上评估SelfBudgeter的性能，比较指标包括准确性（Acc）、平均响应长度（Len）和预算与响应长度的匹配度（Mat）。实验设置如下：

• 训练数据与流程：冷启动阶段使用GSM8K（5,802个正确样本）和s1k（1,000个高难度数学问题）数据集进行全参数微调；RL阶段使用STILL-3-Preview-RL-Data数据集（30K样本，包含MATH、NuminaMathCoT和AIME 1983-2023），训练3个epoch。
• 基线对比：与DeepSeek-R1-Distill-Qwen-1.5B和L1-Max（不同令牌限制设置）进行比较。
• 结果：
  • 在MATH数据集上，SelfBudgeter（s1k初始化，$\alpha=0.2$）实现了74.47%的响应长度压缩（919.27 vs. 3731.25令牌），准确性仅下降2.16%（74.18% vs. 76.34%）。
  • 在GSM8K数据集上，SelfBudgeter（s1k初始化，$\alpha=0.2$）准确性提升3.18%（81.50% vs. 78.32%），响应长度减少62%（662.08 vs. 1737.92令牌）。
  • 匹配度方面，SelfBudgeter显著优于L1-Max，GSM8K和MATH上的匹配率分别高达97.65%和95.82%（L1-Max最高仅65.35%）。
  • 冷启动数据集选择对性能影响显著，s1k初始化模型准确性更高，但响应长度较长；GSM初始化模型响应更简洁，但准确性较低。

分析与批判：实验设置较为全面，涵盖了不同难度的数据集和多种模型配置，结果显示SelfBudgeter在长度压缩和预算遵守方面表现优异。然而，准确性在某些配置下有所下降（如MATH上的GSM初始化模型仅63.46%），表明方法在高难度任务上的稳定性不足。此外，实验未探讨模型在不同规模或领域上的泛化能力，数据集选择对结果的影响也未深入分析，可能存在cherry-picking的风险。预算匹配度的定义（偏差≤50%视为匹配）较为宽松，可能高估了方法的控制精度。

Further Thoughts

SelfBudgeter的令牌预算预测机制为资源受限环境下的推理模型部署提供了新思路，但其对问题难度预估的依赖可能在实际应用中成为瓶颈，特别是在问题分布不均或跨领域任务中。未来研究可以探索结合上下文感知或历史推理数据的动态预算调整策略，以提高预估的鲁棒性。此外，论文中提到的用户预填预算功能启发了我思考人机交互在AI推理中的潜力，或许可以通过引入实时用户反馈机制（如中途调整预算）进一步优化交互体验。另一个值得探索的方向是跨领域资源管理策略，例如借鉴云计算中的动态资源调度算法，可能为令牌分配提供更高效的解决方案。最后，SelfBudgeter在准确性与效率之间的权衡问题让我联想到多目标优化领域的研究，未来可以尝试引入Pareto前沿分析来系统评估不同配置下的性能边界。