Token-Efficient RL for LLM Reasoning

📄 arXiv: 2504.20834v4 📥 PDF

作者: Alan Lee, Harry Tong

分类: cs.LG, cs.AI

发布日期: 2025-04-29 (更新: 2025-06-11)

备注: Title updated to "Token-Efficient RL for LLM Reasoning" to better reflect algorithmic focus. Revised abstract, intro, and conclusion. Paper shortened and typos fixed

💡 一句话要点

提出Token高效强化学习方法，解决LLM推理中内存和计算资源限制问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 大语言模型 LoRA微调 Token选择 策略梯度

📋 核心要点

  1. 现有LLM推理方法在资源受限场景下，面临着内存占用高、计算成本大的挑战，尤其是在微调时。
  2. 论文提出基于token选择的强化学习策略，仅在少量关键token上进行优化，降低内存需求并稳定训练过程。
  3. 实验表明，该方法在Qwen2-1.5B模型上，显著提升了SVAMP基准的准确率，并在多位数乘法任务中表现出色。

📝 摘要（中文）

本文提出了一系列强化学习（RL）策略，专门为大语言模型（LLM）在严格的内存和计算资源限制下进行推理而设计，特别关注与LoRA微调的兼容性。在早期带有基线减法的策略梯度方法的基础上，我们设计了无Critic的方法，这些方法在输出token的一个小的、信息丰富的子集上运行，以减少内存使用并稳定训练。我们引入了S-GRPO，一种Group Relative Policy Optimization的随机变体，以及T-SPMO，一种用于细粒度信用分配的token级别前缀匹配方法。应用于Qwen2-1.5B，我们的方法将SVAMP基准上的准确率从46%提高到70%以上，并在多位数乘法上表现出强大的性能。令人惊讶的是，LoRA下的全token GRPO未能超过基础模型，这表明选择性的token级别优化可能在低参数训练方案中充当隐式正则化器。

🔬 方法详解

问题定义：现有的大语言模型（LLM）推理方法，在资源受限的环境下，特别是使用LoRA进行微调时，面临着巨大的内存和计算压力。传统的强化学习方法需要处理大量的token，导致训练不稳定且效率低下。因此，如何在有限的计算资源下，有效地训练LLM进行复杂推理，是一个亟待解决的问题。

核心思路：论文的核心思路是只关注对推理过程影响最大的token子集，通过选择性地优化这些token，来降低内存占用和计算复杂度。这种方法类似于一种隐式的正则化，可以避免在低参数训练方案中出现过拟合现象。

技术框架：整体框架基于策略梯度强化学习，但移除了Critic网络，简化了训练流程。主要包含以下几个模块：1) Token选择模块：用于选择信息量最大的token子集；2) 策略优化模块：使用改进的策略梯度算法，如S-GRPO和T-SPMO，对选定的token进行优化；3) 奖励函数设计：根据任务目标，设计合适的奖励函数，引导模型学习正确的推理过程。

关键创新：论文的关键创新在于提出了两种新的token级别的强化学习算法：S-GRPO（Stochastic Group Relative Policy Optimization）和T-SPMO（Token-level Prefix Matching Optimization）。S-GRPO是GRPO的一种随机变体，能够更好地探索策略空间。T-SPMO则通过token级别的前缀匹配，实现更细粒度的信用分配，从而更有效地训练模型。与传统的全token优化方法相比，这两种方法都更加token-efficient。

关键设计：S-GRPO的关键设计在于引入了随机性，使得策略更新更加平滑，避免了训练过程中的震荡。T-SPMO的关键设计在于使用前缀匹配来确定每个token对最终结果的贡献，从而实现更精确的信用分配。此外，论文还特别关注了与LoRA微调的兼容性，确保提出的方法能够有效地利用LoRA的参数效率。

🖼️ 关键图片

fig_0

📊 实验亮点

实验结果表明，在Qwen2-1.5B模型上，使用提出的方法在SVAMP基准测试中，准确率从46%提升到70%以上，取得了显著的性能提升。此外，该方法在多位数乘法任务中也表现出强大的性能。令人惊讶的是，全token GRPO在LoRA微调下未能超过基础模型，这进一步验证了token选择策略的有效性。

🎯 应用场景

该研究成果可广泛应用于资源受限的场景，例如移动设备、嵌入式系统等。通过token高效的强化学习方法，可以在这些平台上部署高性能的LLM推理应用，例如智能助手、自动问答系统等。此外，该方法还可以用于训练更小、更高效的LLM，降低模型部署的成本和能耗。

📄 摘要（原文）

We propose reinforcement learning (RL) strategies tailored for reasoning in large language models (LLMs) under strict memory and compute limits, with a particular focus on compatibility with LoRA fine-tuning. Building on early policy gradient methods with baseline subtraction, we design critic-free methods that operate on a small, informative subset of output tokens to reduce memory usage and stabilize training. We introduce S-GRPO, a stochastic variant of Group Relative Policy Optimization, and T-SPMO, a token-level prefix matching approach for fine-grained credit assignment. Applied to Qwen2-1.5B, our methods raise accuracy on the SVAMP benchmark from 46% to over 70% and show strong performance on multi-digit multiplication. Surprisingly, full-token GRPO under LoRA fails to improve over the base model, suggesting that selective token-level optimization may act as an implicit regularizer in low-parameter training regimes.