ORBIT: On-policy Exploration-Exploitation for Controllable Multi-Budget Reasoning
作者: Kun Liang, Clive Bai, Xin Xu, Chenming Tang, Sanwoo Lee, Weijie Liu, Saiyong Yang, Yunfang Wu
分类: cs.LG, cs.AI
发布日期: 2026-01-13
备注: Preprint
💡 一句话要点
ORBIT:面向可控多预算推理的On-policy探索-利用框架
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 多预算推理 强化学习 On-policy蒸馏 思维链 可控推理
📋 核心要点
- 现有方法难以准确估计推理所需的最小计算量,且训练时固定了推理成本与准确率的权衡。
- ORBIT通过多阶段强化学习发现不同计算量下的最优推理策略,再用on-policy蒸馏融合到单一模型。
- 实验表明,ORBIT实现了多模式下的可控推理,并保证了各模式内的推理效率和整体性能。
📝 摘要(中文)
现有的大型推理模型(LRM)通过利用长篇的思维链(CoT)推理实现了强大的性能,但在推理时统一应用过长的推理过程会产生巨大的且通常是不必要的计算成本。为了解决这个问题,之前的工作探索了各种策略来从输入中推断出适当的推理预算。然而,这种方法在最坏的情况下是不可靠的,因为估计所需的最小推理工作量从根本上是困难的,并且它们在训练期间隐式地固定了推理成本和准确性之间的权衡,限制了在不同部署场景下的灵活性。 鉴于这些局限性,我们提出了ORBIT,一个可控的多预算推理框架,具有由输入触发的明确分离的推理模式。 ORBIT采用多阶段强化学习来发现每个工作量下的帕累托最优推理行为,然后进行on-policy蒸馏,将这些行为融合到单个统一模型中。 实验表明,ORBIT实现了(1)多个模式下的可控推理行为,(2)每个模式内具有竞争力的推理密度,以及(3)将这些前沿策略集成到单个统一的学生模型中,同时保持清晰的模式分离和高每模式性能。
🔬 方法详解
问题定义:现有的大型推理模型虽然性能强大,但推理过程计算成本高昂,且难以根据输入自适应调整推理预算。现有方法在估计最小推理工作量时存在困难,并且在训练时就固定了推理成本和准确率之间的权衡,限制了模型在不同部署场景下的灵活性。
核心思路:ORBIT的核心思路是学习多个不同计算预算下的推理策略,并允许模型根据输入选择合适的策略。通过多阶段强化学习,模型可以探索不同计算量下的最优推理行为,然后使用on-policy蒸馏将这些行为融合到一个统一的模型中,从而实现可控的多预算推理。
技术框架:ORBIT框架包含以下几个主要阶段:1) 多阶段强化学习:针对不同的计算预算,使用强化学习训练多个推理策略。每个策略都旨在在给定的计算资源约束下最大化推理准确率。2) 帕累托最优策略发现:在每个计算预算下,通过强化学习找到帕累托最优的推理行为。3) On-policy蒸馏:将多个帕累托最优策略蒸馏到一个统一的学生模型中。学生模型可以根据输入选择合适的推理模式。
关键创新:ORBIT的关键创新在于:1) 可控的多预算推理:模型可以根据输入选择不同的推理模式,从而实现计算成本和准确率之间的灵活权衡。2) 多阶段强化学习:通过多阶段强化学习,模型可以有效地探索不同计算量下的最优推理策略。3) On-policy蒸馏:使用on-policy蒸馏将多个策略融合到一个统一的模型中,避免了离线蒸馏可能导致的性能下降。
关键设计:ORBIT的关键设计包括:1) 强化学习奖励函数:奖励函数需要平衡推理准确率和计算成本。2) On-policy蒸馏损失函数:蒸馏损失函数需要保证学生模型能够学习到教师模型的推理行为,同时保持模式之间的分离性。3) 网络结构:学生模型的网络结构需要能够支持多模式推理,并能够根据输入选择合适的推理模式。
📊 实验亮点
实验结果表明,ORBIT在多个推理任务上取得了显著的性能提升。具体来说,ORBIT实现了可控的推理行为,并保证了每个模式内的推理效率。此外,ORBIT成功地将多个前沿策略集成到一个统一的学生模型中,同时保持了清晰的模式分离和高每模式性能。相较于基线模型,ORBIT在不同计算预算下均取得了更好的性能。
🎯 应用场景
ORBIT适用于需要根据计算资源约束和精度要求进行权衡的各种推理任务,例如问答系统、对话系统和决策支持系统。该框架可以部署在资源受限的设备上,例如移动设备和嵌入式系统,也可以用于降低大型推理模型的部署成本。未来,ORBIT可以扩展到支持更复杂的推理任务和更多的计算预算。
📄 摘要(原文)
Recent Large Reasoning Models (LRMs) achieve strong performance by leveraging long-form Chain-of-Thought (CoT) reasoning, but uniformly applying overlong reasoning at inference time incurs substantial and often unnecessary computational cost. To address this, prior work explores various strategies to infer an appropriate reasoning budget from the input. However, such approaches are unreliable in the worst case, as estimating the minimal required reasoning effort is fundamentally difficult, and they implicitly fix the trade-off between reasoning cost and accuracy during training, limiting flexibility under varying deployment scenarios. Motivated by these limitations, we propose ORBIT, a controllable multi-budget reasoning framework with well-separated reasoning modes triggered by input. ORBIT employs multi-stage reinforcement learning to discover Pareto-optimal reasoning behaviors at each effort, followed by on-policy distillation to fuse these behaviors into a single unified model. Experiments show that ORBIT achieves (1) controllable reasoning behavior over multiple modes, (2) competitive reasoning density within each mode, and (3) integration of these frontier policies into a single unified student model while preserving clear mode separation and high per-mode performance.