Reinforcement Learning for LLM Reasoning Under Memory Constraints

本文提出了S-GRPO和T-SPMO两种内存高效、无批评者的强化学习方法，结合LoRA微调，在有限硬件资源下显著提升了大型语言模型在数学推理任务上的性能，其中T-SPMO在需要细粒度信用分配的任务上表现尤为突出。

Reinforcement Learning, Large Language Model, Reasoning, Parameter-Efficient Fine-Tuning, Efficiency

Alan Lee, Harry Tong

University of Michigan

Generated by gemini-2.5-flash-preview-04-17

Background Problem

大型语言模型（LLMs）通过强化学习（RL）微调在数学和结构化问题解决方面取得了显著进展。然而，标准的RL方法（如PPO和GRPO）对于计算资源有限的研究人员来说并不实用，原因包括：1）它们通常在整个输出轨迹上计算损失；2）它们通常设计用于全模型微调；3）PPO还需要一个批评网络，进一步增加内存需求。此外，监督微调依赖于高质量、昂贵且难以大规模收集的人工标注思维链数据。因此，核心问题是在有限的计算资源（例如，单个40GB GPU）下，能否使用参数高效的更新方法（如LoRA）通过RL微调来提高LLM的推理性能。

Method

本文提出了两种无批评者（critic-free）的RL方法，专门针对内存和计算受限的环境：

S-GRPO (Stochastic Group Relative Policy Optimization): 这是GRPO的一个轻量级变体。它通过从输出轨迹中采样一部分token来计算损失，而不是使用全部token。采样规则是混合的：前 $\alpha$ 个token总是被包含，而后面的token则以概率 $P$ 随机采样，直到达到最大token数 $k$ 。这在学习稳定性和内存效率之间取得了平衡。它与LoRA结合使用，并且每问题只进行一次策略更新，不使用PPO风格的裁剪目标。
T-SPMO (Token-Specific Prefix Matching Optimization): 这是一种新颖的token级RL方法，实现了细粒度的信用分配。对于每个prompt，生成多组完成（|G|个）。构建一个前缀树来识别所有唯一的 (前缀p, 下一个token v) 对。基于这些对，计算token级的优势 $A(v \mid p) = \mathbb{E}[R \mid p \circ v] - \mathbb{E}[R \mid p]$ ，其中期望通过经验平均近似。策略更新的目标是最大化 $\pi_{\theta}(v \mid p) \cdot A(v \mid p)$ 的总和，并带有LoRA参数的L2正则化。它也与LoRA结合使用，每问题更新一次。T-SPMO还引入了基于回放的重采样机制，可以从成功或失败的完成中采样，并从序列的中间位置重新开始生成，以构建后续token的前缀树，从而学习序列后半部分的策略。
全轨迹GRPO基线: 为了隔离稀疏token采样的影响，作者还实现了原始的GRPO目标（使用所有token计算损失），并通过梯度累积（对每个完成进行反向传播）来适应内存限制，作为基线进行比较。

Experiment

实验使用Qwen2-1.5B模型，通过LoRA在单个分区A100 40GB GPU上进行微调。LoRA模块插入到最后1/3（SVAMP）或1/4（乘法）注意力层的query和value投影中，秩分别为16（SVAMP）和8（乘法）。优化器为AdamW，学习率为1e-4，正则化系数 $\beta = \lambda = 0.01$ ，使用float32精度。评估任务是SVAMP数据集和自定义的3x3位数字乘法任务。奖励基于最终提取整数的精确匹配准确率。S-GRPO和GRPO基线每prompt采样|G|=8个完成，T-SPMO采样|G|=50个完成。训练步数基于早停或验证集峰值性能确定。

结果：

SVAMP: Base Qwen2-1.5B准确率45.0%，GRPO 46.7%，S-GRPO 70.3%，T-SPMO 71.6%。S-GRPO和T-SPMO均显著提高了性能，而GRPO基线提升不明显。
乘法: Base Qwen2-1.5B准确率3.9%，GRPO 4.4%，S-GRPO 22.9%，T-SPMO 70.0%。T-SPMO在该任务上表现出色，远超基线和S-GRPO。S-GRPO虽有提升但提升有限，GRPO基线几乎没有提升。

消融实验：

S-GRPO对 $\alpha$ 和 $k$ 相对鲁棒，但在乘法任务上表现出对特定参数组合的敏感性，表明两种可能的策略：手工选择token或保守地少量采样。增加有效批量大小会降低S-GRPO在两个任务上的性能。
T-SPMO对|G|和回放策略敏感。增加|G|提高了乘法性能。从成功完成中回放对于乘法任务至关重要，而从失败完成中回放则会降低性能。

Further Thoughts

本文的核心贡献在于证明了在资源受限的情况下，通过精心设计的RL算法结合参数高效微调（LoRA）仍然可以有效提升LLM的推理能力。GRPO和S-GRPO在乘法任务上表现不佳，而T-SPMO表现出色，这突出了解释性任务（如数学推理）中细粒度信用分配的重要性。GRPO和S-GRPO基于组内奖励均值和标准差进行优势估计，并将优势应用于采样的（或全部）token上，这种”广播式”的信用分配方式在序列中存在早期错误时，可能会不公平地惩罚后续正确的步骤，导致模型难以学习到完整的、正确的推理链。而T-SPMO通过构建前缀树，计算特定 (前缀, token) 对的优势，即使序列后续出错，模型也能从正确的前缀-token转移中获得正向反馈，从而更好地学习局部正确的推理步骤。这与人类学习复杂任务的过程有相似之处，我们通常会分解任务，并从完成部分子任务中获得反馈。T-SPMO的回放机制，特别是从成功完成中回放，进一步强化了对”好”路径的学习，这可以看作是一种有偏见的经验回放，将学习资源集中在有希望的轨迹上。这对于长序列或需要精确步骤的任务尤其重要。未来的工作可以探索如何将T-SPMO的细粒度信用分配思想应用于更复杂的推理任务，例如结合外部工具使用或更复杂的逻辑推理，以及如何设计更通用的、能够适应不同任务结构的信用分配机制。