LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches
作者: Shaokun Zhang, Yifan Zhang, Jian Hu, Yueying Li, Hao Zhang, Binfeng Xu, Jan Kautz, Yi Dong
分类: cs.CL
发布日期: 2026-10-05
备注: 16 pages, 6 figures
🔗 代码/项目: GITHUB
💡 一句话要点
提出LoGRA以解决大规模语言模型强化学习的内存瓶颈问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 强化学习 大规模语言模型 低秩梯度 内存优化 策略同步 KL步长控制 模型训练
📋 核心要点
- 现有的强化学习方法在大规模语言模型的训练中面临内存需求高的问题，限制了其应用范围。
- LoGRA通过低秩梯度草图保留有用的学习信号，结合预测的KL步长控制来优化内存使用和更新稳定性。
- 实验结果表明，LoGRA在推理任务中将平均训练内存减少了45.7%，并成功在资源有限的环境中训练大型模型。
📝 摘要（中文）
强化学习（RL）在提升大规模语言模型（LLMs）能力方面取得了显著进展，但其内存需求仍然是广泛应用的障碍。本文提出LoGRA，一种通过保留低秩梯度草图中的有用学习信号来减少内存的RL后训练方法。这些紧凑的表示支持模型更新和高效的策略同步。为防止过大的更新干扰学习，我们结合梯度压缩与预测的KL步长控制，估计每次更新前的策略变化并相应调整其幅度。在推理任务中，LoGRA将平均训练内存减少了多达45.7%，而不牺牲性能。此外，它还使得在单个八GPU节点上稳定训练27B参数模型超过1100步成为可能，而密集的Adam方法则因内存不足而无法实现，从而使得之前内存不可行的RL训练变得实用。代码可在Molt库中获取。
🔬 方法详解
问题定义：本文旨在解决大规模语言模型在强化学习训练中面临的高内存需求问题。现有方法在处理大模型时常常因内存不足而无法进行有效训练。
核心思路：LoGRA通过使用低秩梯度草图来压缩学习信号，从而减少内存占用。同时，结合预测的KL步长控制，确保更新过程的稳定性，避免过大的更新对学习造成干扰。
技术框架：LoGRA的整体架构包括两个主要模块：低秩梯度草图生成模块和KL步长控制模块。前者负责生成压缩的梯度表示，后者则在每次更新前预测策略变化并调整更新幅度。
关键创新：LoGRA的主要创新在于将低秩梯度草图与KL步长控制相结合，显著降低了内存需求，同时保持了模型性能。这一方法与传统的梯度更新方法相比，能够在内存受限的环境中实现有效的RL训练。
关键设计：在设计中，LoGRA采用了特定的低秩分解技术来生成梯度草图，并通过动态调整KL步长来控制更新幅度，确保训练过程的稳定性和有效性。
🖼️ 关键图片
📊 实验亮点
实验结果显示，LoGRA在推理任务中将平均训练内存减少了45.7%，并成功在单个八GPU节点上稳定训练27B参数模型超过1100步，而传统的密集Adam方法则因内存不足无法实现。这一显著的内存优化使得大规模RL训练变得更加可行。
🎯 应用场景
LoGRA的研究成果在多个领域具有潜在应用价值，尤其是在需要大规模语言模型进行强化学习的场景中，如自然语言处理、对话系统和智能助手等。通过降低内存需求，该方法使得在资源有限的环境中进行大规模模型训练成为可能，推动了AI技术的普及与应用。
📄 摘要（原文）
Reinforcement learning (RL) has greatly advanced the capabilities of large language models (LLMs), but its memory demands remain a barrier to broader adoption. We introduce LoGRA, an approach to RL post-training that reduces memory by retaining useful learning signals in low-rank gradient sketches. These compact representations support both model updates and efficient policy synchronization. To prevent overly large updates from disrupting learning, we complement gradient compression with predicted-KL step control, which estimates policy changes before applying each update and adjusts its magnitude accordingly. Across reasoning tasks, LoGRA reduces average training memory by up to 45.7\% without sacrificing performance. It also enables stable training of a 27B-parameter model for over 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory, making previously memory-infeasible RL training practical. Code is available in the \href{https://github.com/skzhang1/labs-molt/tree/logra/examples/scripts/logra}{Molt library}.