Accelerating RL for LLM Reasoning with Optimal Advantage Regression

强化学习（RL）已成为微调大型语言模型（LLMs）以提升复杂推理能力的强大工具。然而，最先进的策略优化方法通常因每次提示需要多次生成以及依赖于当前策略的批评网络或优势估计，而面临高计算开销和内存消耗的问题。在本文中，我们提出了 $A$*-PO，一种新颖的两阶段策略优化框架，该框架直接逼近最优优势函数，从而实现高效训练用于推理任务的 LLMs。在第一阶段，我们通过从参考策略进行离线采样来估计最优值函数 $V$*，从而消除了昂贵的在线值估计需求。在第二阶段，我们使用简单的最小二乘回归损失进行基于策略的更新，并且每次提示仅需一次生成。理论上，我们建立了性能保证，并证明了可以在无需复杂探索策略的情况下优化 KL 正则化的 RL 目标。实证上，$A$*-PO 在广泛的数学推理基准测试中表现出具有竞争力的性能，同时与 PPO、GRPO 和 REBEL 相比，训练时间缩短了多达 2 倍，峰值内存使用减少了超过 30%。$A$*-PO 的实现代码可在 https://github.com/ZhaolinGao/A-PO 找到。

该论文试图解决强化学习中训练大型语言模型（LLMs）进行复杂推理时面临的高计算开销和内存消耗问题。这是一个已知但尚未完全解决的问题，尤其是在需要多次生成和依赖批评网络或优势估计的情况下。

论文提出了一种名为A*-PO的两阶段策略优化框架。第一阶段通过离线采样从参考策略中估计最优值函数V*，从而避免了昂贵的在线价值估计；第二阶段使用单次生成的最小二乘回归损失进行策略更新。这种方法直接逼近最优优势函数，减少了对复杂探索策略的需求，并显著降低了计算和内存成本。

A*-PO在数学推理基准测试中表现出与现有方法相当的性能，同时将训练时间缩短了高达2倍，峰值内存使用减少了超过30%。实验设计涵盖了多种数学推理任务，验证了方法的有效性。此外，代码已开源（https://github.com/ZhaolinGao/A-PO），便于后续研究者复现和改进。未来可以进一步研究如何将A*-PO扩展到其他类型的任务或更复杂的模型架构中。

最近的相关研究包括PPO（Proximal Policy Optimization）、GRPO（Gradient Regularized Policy Optimization）和REBEL等方法，这些方法都在尝试通过不同的方式优化RL中的策略更新过程。例如，PPO通过剪裁目标函数来稳定训练，而GRPO引入梯度正则化以提高泛化能力。此外，还有研究如《Reward is Enough》探讨了奖励信号在RL中的核心作用，以及《Fine-Tuning Language Models from Human Preferences》利用人类反馈微调语言模型。