自适应并行推理:高效推理扩展的新范式

TL;DR

自适应并行推理 (Adaptive Parallel Reasoning, APR) 是一种全新的范式,使大语言模型 (LLMs) 能够在推理时动态地在并行操作和串行操作之间分配计算资源。通过允许模型决定何时启动独立的推理线程以及如何协调它们,APR 降低了端到端延迟,并防止了与长顺序推理链相关的“上下文退化 (context-rot)”。

顺序推理扩展的问题

顺序推理的扩展与探索量呈线性关系,从而导致三个主要瓶颈:

  1. 上下文退化 (Context-Rot):随着模型生成更多的中间探索路径,从干扰信息中区分有用信息变得越来越困难,从而降低性能。
  2. 延迟:生成时间随推理长度成比例增长,对于复杂任务,有时会导致长达数十分钟甚至数小时的等待时间。
  3. 计算强度:沿着输出序列长度进行扩展会使推理变慢且可靠性降低。

并行推理通过允许模型独立地(不依赖彼此的上下文)且并发地(同时执行)探索多个线程来解决这些问题。

从固定并行到自适应控制

虽然之前已经存在并行推理方法,但它们通常依赖于从模型外部施加的结构。APR 将这种控制权转移到了模型自身。

非自适应方法的局限性

  • 简单的 Fork-and-Join (Self-consistency, BoN):由于轨迹是独立采样的,这些方法通常会产生冗余计算。
  • 基于启发式的结构化搜索 (Tree/Graph/Skeleton of Thoughts, MCTS):这些方法需要预先知道分解策略,而这些策略对于所有问题可能都是未知的。
  • 固定阶段变体 (ParaThinker, GroupThink, Hogwild! Inference):无论问题是否能从中受益,这些方法都会施加特定的并行结构。

APR 的优势

自适应并行推理允许模型根据问题的复杂程度决定并行化程度。这提供了三个关键优势:

  • 无需特定领域的启发式方法:模型通过强化学习 (RL) 学习通用的分解策略,发现诸如同步自我验证之类的涌现模式。
  • 减少冗余:与 Best-of-N (BoN) 不同,模型可以在分支前产生独特的、非重叠的子任务。
  • 动态分配:模型可以选择不对简单问题进行并行化,从而避免在不需要并行化的地方产生开销。

用于自适应并行的推理系统

执行自适应并行分支通常遵循 fork-join 设计:模型将问题拆分为并发的子任务,然后将其合并为一个最终答案。主要的挑战在于在 KV cache 层面聚合结果。

修改引擎的方法 (Multiverse)

像 Multiverse、Parallel-R1 和 NPR 这样的方法通过修改推理引擎,将来自独立线程的非连续 KV cache 内存块进行复制和拼接。

  • 权衡:这避免了第二次 prefill,但引入了系统脆弱性(例如错误的指针或缓存剔除)以及位置编码中的分布偏移,需要大量的训练和修改后的 attention masks 来对齐行为。

与引擎无关的方法 (ThreadWeaver)

ThreadWeaver 将编排移至客户端。客户端拼接来自独立分支的文本输出,然后引擎执行第二次 prefill 以生成最终结论的 KV cache。

  • 权衡:这在 prefill 期间引入了一些计算冗余,但由于 prefill 的成本远低于 decoding,因此仍然保持高效。它避免了对引擎的修改,并适用于标准的 causal attention,使其更容易适配顺序模型。

训练与奖励设计

教导模型使用并行控制流 token 需要演示数据和特定的激励结构。

演示与 SFT

监督微调 (SFT) 用于教授并行控制流的语法。目前关于 SFT 是诱导了新的推理能力,还是仅仅将现有能力对齐到特定的 token 语法,仍存在争论。

奖励机制

仅靠结果准确性通常不足以激励并行化。研究人员探索了各种奖励设计:

  • 仅结构奖励:奖励线程数量或结构的正确使用很容易被模型通过生成无用的线程来“刷分”。
  • 关键路径奖励:为了优化延迟,奖励侧重于关键路径(因果相关的 token 的最长序列)。ThreadWeaver 使用 $1 - L_{\mathrm{critical}} / L_{\mathrm{total}}$ 作为奖励,随着关键路径占总 token 的比例变小,该奖励会增加。
  • 正确性门控:为了防止奖励低效但错误的路径,只有在最终答案正确时才会授予并行化奖励:$R = \mathbf{1}(\text{Correctness}) + \mathbf{1}(\text{Correctness}) \times (\text{parallelization metric})$。

性能与开放性问题

对 APR 的评估因目标而异。Multiverse 和 ThreadWeaver 优先考虑准确性和延迟,而 NPR 专注于真实并行率 (Genuine Parallelism Rate),Parallel-R1 将 APR 作为训练中的探索脚手架,以在 RL 后提升性能。

剩余挑战

  • 推理 vs. 训练:目前尚不清楚推理时的并行化是否能持续提高准确性,或者其主要价值是否在于作为训练时的探索脚手架。
  • 稳定性:如果放宽并行化奖励,模型往往会退回到顺序推理。
  • 硬件感知:未来的训练可以根据可用的计算预算,使并行化决策具备硬件感知能力。
  • 并行深度:当前的结构是扁平的;通过递归语言模型 (RLMs) 探索递归并行化(深度 > 1)是一个潜在的下一步。

Sources

相关

  • Dispatch
  • 项目
  • Dispatch
  • 项目
  • 项目