← 返回信息流
技术博客arXiv cs.AI·4 小时前

MARS:并行大模型推理的边际对抗风险可控停止策略

原标题:MARS: Margin-Adversarial Risk-controlled Stopping for Parallel LLM Test-time Scaling

速览

针对并行测试时扩展中因等待完整推理轨迹导致的计算开销问题,研究提出MARS策略。该方法通过探测中间检查点,利用边际对抗规则估计活跃轨迹的答案变化概率,并在领先答案安全时提前停止。实验显示,MARS在保持精度的同时,可节省25-47%的自洽性Token,显著优于现有基线。

AI 深度解读

MARS:并行大语言模型测试时扩展中的边际对抗风险可控停止机制

背景

在大语言模型(LLM)的推理过程中,测试时扩展(Test-time Scaling) 是一种显著提升模型准确率的主流策略。其中,并行测试时扩展 通过采样多个推理轨迹(reasoning traces),并采用多数投票(majority-voting)机制来整合这些轨迹的答案,从而获得更可靠的最终结果。

然而,这种策略存在一个显著的瓶颈:为了获得准确的投票结果,所有的推理轨迹必须运行至完成(run to completion)。这意味着即使某个轨迹在早期已经明显偏离正确路径或变得冗余,系统仍需消耗大量的计算资源将其完整生成。这种“全量运行”的方式导致了巨大的计算开销(computational overhead),限制了并行测试时扩展在资源受限场景下的实际应用效率。

核心内容

针对上述问题,研究团队提出了一种名为 MARS(Margin-Adversarial Risk-controlled Stopping,边际对抗风险可控停止)的新方法。MARS 的核心思想是引入“早期停止”(early-stopping)机制,在推理过程中动态评估哪些轨迹可能改变其答案,并在保证最终投票结果安全的前提下,提前终止那些不再影响最终结果的轨迹。

1. 核心观察:中间检查点的投票演变

研究人员观察到,在推理轨迹的中间检查点(intermediate checkpoints)对部分轨迹进行探测,可以在不干扰生成过程的情况下提取当前的答案状态。随着生成的进行,这些中间答案会形成一个不断演变的聚合投票(evolving aggregate vote)。这一观察为在推理中途判断轨迹的“稳定性”提供了数据基础。

2. MARS 的工作原理

MARS 是一种基于边际(margin)的停止规则,其工作流程如下:

  • 估计轨迹变更概率:MARS 旨在估算当前活跃的轨迹中,哪些轨迹极有可能改变其最终答案。
  • 保守边界下的安全判定:一旦领先者(leader,即当前得票最多的答案)在未来的投票移动中,即使在最保守的边界估计下依然保持安全(即其领先优势足以抵御其他轨迹的翻转),MARS 就会停止那些不再影响结果的轨迹。
  • 分离不确定性来源:MARS 将不确定性分为两个部分处理:
    1. 轨迹级切换概率(Trace-level switch probabilities):学习当前边际(margin)中有多少比例可能被保留。这决定了轨迹改变答案的可能性。
    2. 切换后的落点分布:处理更困难的问题,即如果轨迹切换了答案,它们会落在哪里。MARS 通过一个对抗性边界(adversarial bound)来处理这一部分,该边界是从预热轨迹(warmup traces)中校准得出的。

3. 理论保证与实践模型

  • 理论保证:在拥有真实的切换概率(true switch probabilities)的情况下,MARS 能以高概率保证,早期停止后的答案与使用完整预算(full-budget)运行后的投票结果一致。
  • 实践模型:在实际应用中,研究团队使用一个包含五个特征的逻辑回归模型(five-feature logistic model)来近似 oracle(理想情况)的切换行为,该模型表现良好。

4. 实验结果

在三个推理模型和三个竞赛数学基准测试(competition-math benchmarks)上,MARS 展现了显著的效率提升:

  • 节省计算资源:相比传统的 Self-Consistency 方法,MARS 节省了 25-47% 的 token 消耗。
  • 超越强基线:相比 DeepConf Online(一种已经能够过滤和截断弱轨迹的强大置信度加权基线),MARS 进一步节省了 14-29% 的计算成本。
  • 保持准确率:在大幅节省计算资源的同时,MARS 保持了与相应完整预算基线相当的准确率。

关键要点

  • 问题痛点:并行测试时扩展(如 Self-Consistency)需要所有推理轨迹运行至完成,导致严重的计算浪费。
  • 核心创新:MARS 是一种基于边际对抗的风险可控停止规则,允许在推理中途安全地终止不影响最终投票结果的轨迹。
  • 技术机制
    • 利用中间检查点的投票演变来预测轨迹稳定性。
    • 分离处理“轨迹是否切换”的概率问题和“切换后落在哪里”的分布问题。
    • 使用对抗性边界校准切换后的分布,确保保守估计。
  • 性能优势
    • 相比 Self-Consistency,节省 25-47% 的 token。
    • 相比 DeepConf Online,额外节省 14-29% 的 token。
    • 在节省资源的同时,准确率与完整运行基线持平。
  • 适用场景:特别适用于需要高准确率且计算资源敏感的并行推理场景,如数学推理、复杂逻辑任务等。

意义与影响

MARS 的提出为大规模语言模型的推理效率优化提供了新的思路。它证明了在并行推理中,动态评估和早期停止 是可行的,并且可以在不牺牲准确率的前提下显著降低计算成本。

  1. 降低推理成本:通过减少不必要的完整轨迹生成,MARS 直接降低了 LLM 在测试时的 token 消耗,这对于部署大规模模型、降低 API 调用成本以及提升实时响应速度具有重要意义。
  2. 提升资源利用率:在 GPU/TPU 等昂贵计算资源受限的环境中,MARS 使得更多的推理请求可以在相同的硬件资源下并行处理,提高了集群的整体吞吐量。
  3. 推动测试时扩展的实用化:此前,测试时扩展因计算开销过大而难以在资源受限的场景中广泛应用。MARS 通过引入风险可控的停止机制,降低了这一门槛,使得更复杂的并行推理策略在工业界落地成为可能。
  4. 方法论贡献:MARS 提出的“分离不确定性来源”和“对抗性边界校准”的方法论,为后续研究如何更精细地控制推理过程中的不确定性提供了参考框架。
查看原文 →arxiv.org