Learning to Reason as Action Abstractions with Scalable Mid-Training RL

TL;DR

提出RA3算法,通过中期训练学习动作抽象,提升代码生成性能,平均提升8点。

cs.LG 🔴 高级 2025-09-30 47 次浏览
Shenao Zhang Donghan Yu Yihao Feng Bowen Jin Zhaoran Wang John Peebles Zirui Wang
强化学习 动作抽象 中期训练 大模型 代码生成

核心发现

方法论

本文建立理论框架,分析中期训练如何塑造后续强化学习的动作子空间,提出利用时间变分下界(ELBO)优化潜在结构。通过RL发现时间一致的高层次动作抽象,结合微调提升性能。算法核心包括动作空间剪枝、潜在结构学习与RL优化,具体实现为迭代EM流程,结合变分推断与策略优化。

关键结果

  • 在CodeGen、Qwen和Llama模型上,RA3在HumanEval和MBPP任务中平均提升性能8点和4点,超越单纯的下一词预测基线。实验显示RA3在RLVR中收敛更快,性能更优,特别是在HumanEval+和MBPP+等扩展任务中表现出更高的极限性能。
  • RA3在中期训练中有效学习动作层级,显著减少了后续强化学习的样本需求,提高了样本效率。具体数据表明,RA3在中期训练阶段利用3.5M代码片段,提升了模型的推理和生成能力,缩短了训练时间。
  • 通过潜在变量的时间一致性设计,算法实现了高效的抽象学习,降低了推理成本,增强了模型的迁移能力。对比传统的原子动作训练,RA3在抽象层面表现出更强的泛化和适应性。

研究意义

该研究突破了大模型中动作空间设计的理论瓶颈,为中期训练引入高层次动作抽象提供了坚实基础。其创新的变分ELBO框架和潜在结构学习机制,为未来大规模模型的高效训练和推理提供了新思路。实践中,RA3显著改善了代码生成、推理任务的样本效率和性能,推动了AI在复杂任务中的应用落地。长远来看,该方法有望引领模型自主学习高层次策略,降低训练成本,提升泛化能力,具有重要的学术和工业价值。

技术贡献

论文首次系统性地结合动作抽象与中期训练理论,提出基于时间变分下界的潜在结构学习方法,明确了动作剪枝与RL收敛的关系。算法实现为迭代EM流程,结合RL和微调,显著提升模型的推理能力和样本效率。理论分析提供了动作空间剪枝效率与RL收敛速度的定量关系,为未来动作层级设计提供指导。

AI 总览摘要

本研究提出RA3算法,旨在通过中期训练学习大模型的动作抽象,从而提升代码生成等任务的性能。传统方法多依赖于原子动作,面临状态空间庞大和样本不足的问题。本文建立理论框架,分析中期训练如何塑造动作子空间,提出利用时间变分下界(ELBO)优化潜在结构。该方法通过RL发现时间一致的高层次动作抽象,结合微调显著提升模型性能。在多个模型和任务上,RA3实现了平均8点性能提升,且收敛速度快、效果优越。实验结果验证了动作抽象在提升样本效率和泛化能力中的作用,为未来大模型的自主学习和高效训练提供新思路。该方法不仅在代码生成中表现出色,也为复杂推理任务的中期训练提供了理论基础和实践方案。未来,结合多模态数据和更长序列的潜在结构学习,有望推动AI自主策略的广泛应用,降低训练成本,增强模型的泛化能力。

深度分析

研究背景

近年来,随着大规模预训练模型的发展,强化学习(RL)在提升模型推理和生成能力方面展现出巨大潜力。早期工作如GPT系列、BERT等主要依赖自监督学习,但在复杂推理任务中,单纯的下一词预测存在局限。引入RL后,模型能通过奖励信号优化策略,增强推理能力。中期训练阶段,利用专家示范进行策略微调,逐步引入动作层级结构,提升模型的抽象能力。相关研究如动作层次(Options)和动作抽象(Hierarchical RL)为本工作提供理论基础。尽管如此,如何在大规模模型中高效学习高层次动作抽象,仍是当前挑战。现有方法多依赖人工设计或有限的示例,难以充分利用海量数据。本文突破在于提出基于时间变分的潜在结构学习框架,为中期训练引入高效的动作抽象机制。

核心问题

核心问题在于大模型中动作空间庞大,直接强化学习难以高效收敛,且样本需求巨大。传统方法多在原子动作层面操作,导致探索空间过大,训练成本高昂。中期训练阶段的目标是提取高层次的动作抽象,减少决策复杂度,但缺乏理论指导,导致效果不稳定。如何设计既能有效剪枝动作空间,又能保证后续RL收敛速度的中期训练策略,是当前亟待解决的问题。本文试图通过理论分析,明确动作剪枝的效率与RL收敛的关系,为设计高效的中期训练算法提供依据。

核心创新

主要创新包括:1)提出基于时间变分下界(ELBO)的潜在结构学习框架,有效捕获高层次动作抽象;2)引入时间一致性机制,确保潜在变量的连续性和合理性;3)设计迭代EM流程,将RL与微调结合,实现自监督学习与策略优化的协同提升;4)理论分析动作剪枝效率与RL收敛速度的关系,为动作空间设计提供指导。这些创新使得模型能在中期训练阶段自主学习高层次策略,显著提升样本效率和泛化能力。

方法详解

  • �� 设计时间变分ELBO,最大化专家示范数据的潜在结构似然;
  • �� 引入潜在变量z,表示高层次动作抽象,定义其先验分布以保证时间一致性;
  • �� 采用EM算法,E步通过RL优化潜在变量的后验分布,M步通过微调策略实现动作生成;
  • �� 利用RL发现潜在动作的时间一致性,减少采样成本;
  • �� 结合奖励设计,强化潜在抽象的表达能力,提升模型推理能力;
  • �� 通过潜在空间剪枝,有效缩小动作决策空间,加快RL收敛速度。

实验设计

在Python代码生成任务中,选用Qwen-2.5-1.5B、Llama-3.2-1B和Llama-3.1-8B模型,使用3.5M代码片段作为中期训练数据。对比基线包括纯下一词预测和传统微调,评估指标为准确率和损失。采用多轮EM迭代,每轮400次梯度更新,前40次为RL策略梯度。通过AB测试验证RA3在HumanEval、MBPP及扩展任务中的性能提升,观察到平均提升8点和4点,收敛速度明显加快。还分析了潜在结构学习对样本效率的影响,验证了动作抽象的有效性。

结果分析

RA3在多个模型和任务中实现了显著性能提升,平均8点在HumanEval,4点在MBPP,超越传统微调和下一词预测基线。模型在RLVR任务中的收敛速度提升30%以上,极限性能也优于对比方法。潜在结构学习显著减少了样本需求,提高了模型的推理能力。实验还验证了动作抽象在长序列任务中的优势,模型表现出更强的泛化能力和鲁棒性。

应用场景

该方法适用于大规模模型的代码生成、推理和决策任务。通过中期训练学习高层次动作抽象,模型能更高效地进行复杂任务规划,减少样本和计算成本。未来可扩展到多模态任务、长序列推理,推动AI自主策略学习,降低工业应用门槛。

局限与展望

当前方法依赖大量专家示范,数据偏差可能影响潜在结构的学习效果。潜在变量的推断复杂,计算成本较高,特别是在大模型中。对于极端复杂或多模态任务,潜在抽象设计仍需优化。未来需解决潜在空间的泛化能力和推断效率问题。

通俗解读 非专业人士也能看懂

想象你在学习做菜,平时只会做简单的菜,比如炒蛋或煮面。这些都是基本动作,但如果你想做一道复杂的菜,比如火锅或烤鸡,就需要把做菜的步骤拆分成更高层次的“技能”——比如准备调料、腌制、火候控制等。这些高层次技能可以反复用在不同菜肴中,帮你更快做出好菜。本文的算法就像教你学会这些“厨艺技能”,让模型在学习过程中掌握一些“高层次的动作”,而不是每次都从头开始。这样,模型就能更快理解任务、减少试错,最终做出更好的“菜”。通过学习这些“技能”,模型在处理复杂任务时变得更聪明、更高效,就像你学会了厨房里的“秘密武器”。

简单解释 像给14岁少年讲一样

想象你在学校学新技能,比如打篮球。刚开始,你只会投篮或运球这些基本动作,但如果你学会了“传球”、“防守”这些更高层次的技能,你就能更快地玩好比赛。每次比赛都像一场大任务,里面有很多步骤。现在,假设你能把这些步骤变成一些“高层次的策略”,比如“组织进攻”或“防守反击”。这样,你就不用每次都从零开始思考,而是用这些“策略”来指导行动。论文里的算法就像教模型学会这些“策略”,让它在面对复杂问题时,能用更聪明的方法解决。这样,模型不仅能做得更快,还能做得更好,就像你变成了篮球高手!

原文摘要

Large language models excel with reinforcement learning (RL), but fully unlocking this potential requires a mid-training stage. An effective mid-training phase should identify a compact set of useful actions and enable fast selection among them through online RL. We formalize this intuition by presenting the first theoretical result on how mid-training shapes post-training: it characterizes an action subspace that minimizes both the value approximation error from pruning and the RL error during subsequent planning. Our analysis reveals two key determinants of mid-training effectiveness: pruning efficiency, which shapes the prior of the initial RL policy, and its impact on RL convergence, which governs the extent to which that policy can be improved via online interactions. These results suggest that mid-training is most effective when the decision space is compact and the effective horizon is short, highlighting the importance of operating in the space of action abstractions rather than primitive actions. Building on these insights, we propose Reasoning as Action Abstractions (RA3), a scalable mid-training algorithm. Specifically, we derive a sequential variational lower bound and optimize it by iteratively discovering temporally-consistent latent structures via RL, followed by fine-tuning on the bootstrapped data. Experiments on code generation tasks demonstrate the effectiveness of our approach. Across multiple base models, RA3 improves the average performance on HumanEval and MBPP by 8 and 4 points over the base model and the next-token prediction baseline. Furthermore, RA3 achieves faster convergence and higher asymptotic performance in RLVR on HumanEval+, MBPP+, LiveCodeBench, and Codeforces.

cs.LG cs.AI cs.CL stat.ML