Policy-Guided Stepwise Model Routing for Cost-Effective Reasoning
作者: Wenwen Si, Insup Lee, Osbert Bastani
分类: cs.AI
发布日期: 2026-05-07
💡 一句话要点
提出基于策略引导的逐步模型路由方法,以实现大语言模型推理的高效能与低成本平衡。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 推理优化 思维链 强化学习 模型路由 计算效率 约束决策
📋 核心要点
- 现有路由策略依赖于手工规则,限制了推理性能,或依赖于训练昂贵的大型过程奖励模型,导致部署难度大。
- 将逐步模型路由建模为约束决策问题,利用强化学习训练轻量级控制策略,并引入阈值校准实现动态性能-成本权衡。
- 在多个数学基准测试中,该方法在保持高准确率的同时显著降低了推理成本,优于传统手工路由方案。
📝 摘要(中文)
推理时计算显著提升了大语言模型(LLM)在复杂推理任务中的表现,但往往伴随着高昂的推理成本。一种解决方案是将思维链(CoT)的中间状态路由至不同规模的模型进行处理。然而,现有方法要么依赖于性能受限的手工路由策略,要么需要训练难以在实际应用中部署的大型过程奖励模型(PRM)。本文将逐步模型路由建模为约束决策问题,通过强化学习训练小型控制策略,并结合阈值校准来精细化调节性能与效率的权衡。在GSM8K、MATH500和OmniMath三个数学基准测试中,该方法在开放和闭源模型上均表现优异,在准确率-成本权衡方面显著优于手工策略,并达到了与训练大型过程奖励模型相当的性能水平。
🔬 方法详解
问题定义:论文旨在解决大语言模型在复杂推理任务中,如何通过动态分配不同计算资源(模型规模)来优化推理成本与准确率之间的权衡问题,克服了现有手工规则的僵化与大型奖励模型训练的复杂性。
核心思路:将推理过程中的每一步决策视为一个马尔可夫决策过程(MDP),通过训练一个轻量级的控制策略(Policy)来决定当前步骤是否需要调用更强大的模型,从而在保证推理质量的前提下最小化计算开销。
技术框架:整体架构包含一个轻量级策略网络,该网络根据当前的思维链状态(CoT states)输出路由决策。系统通过强化学习进行策略优化,并引入阈值校准机制,允许用户根据实际需求灵活调整性能与成本的偏好。
关键创新:核心创新在于将路由问题转化为约束决策问题,通过轻量级策略替代了笨重的大型过程奖励模型,实现了推理过程中的自适应路由,显著降低了对计算资源的依赖。
关键设计:采用了基于强化学习的策略训练方法,通过定义包含准确率奖励与成本惩罚的复合目标函数,引导策略网络学习最优路由路径。同时,通过阈值校准技术,在推理阶段动态调整路由触发条件,以适应不同的应用场景需求。
🖼️ 关键图片
📊 实验亮点
实验在GSM8K、MATH500及OmniMath数据集上验证了该方法的有效性。结果表明,该方法在准确率-成本权衡上显著优于手工路由基线,且在无需训练大型过程奖励模型的情况下,达到了与此类复杂方法相当的性能表现,证明了轻量级策略在复杂推理任务中的高效性。
🎯 应用场景
该方法适用于对推理成本敏感且任务难度差异大的场景,如企业级数学求解器、自动化代码生成系统及复杂逻辑推理助手。通过动态路由,系统能在保证高准确率的同时,大幅降低云端推理的算力消耗,具有极高的工业部署价值和成本优化潜力。
📄 摘要(原文)
Inference-time computation has greatly enhanced the performance of large language models (LLMs) on challenging reasoning tasks, but this strategy can incur high inference costs. One solution is to route intermediate chain-of-thought (CoT) states to language models of different sizes; however, existing approaches rely on handcrafted routing strategies that limit performance, or on training large process reward models that may be infeasible in many applications. We formulate stepwise model routing as a constrained decision-making problem, which we solve by training a small control policy using reinforcement learning in conjunction with threshold calibration to tune the performance-efficiency tradeoff. We validate our method on three math benchmarks (GSM8K, MATH500, and OmniMath) on both open and closed models. Our method consistently improves the accuracy-cost tradeoff compared to handcrafted approaches, while achieving a comparable tradeoff to methods that require training large process reward models.