TLDR: A new mid-training algorithm, Reasoning as Action Abstractions (RA3), improves large language models (LLMs) by teaching them high-level reasoning skills. It formalizes how mid-training should prune action spaces and accelerate reinforcement learning (RL) convergence. RA3 iteratively discovers hidden reasoning steps from expert data and uses them to fine-tune LLMs, leading to better performance and faster learning in code generation tasks compared to standard methods.
Large language models (LLMs) have shown incredible potential when combined with reinforcement learning (RL), but truly unlocking this power requires a crucial intermediate step known as mid-training. This phase is vital for identifying a concise set of useful actions and enabling LLMs to select among them quickly using online RL.
A recent research paper, “Learning to Reason as Action Abstractions with Scalable Mid-Training RL,” delves into this challenge, presenting a theoretical framework for how mid-training influences post-training performance. The authors, Shenao Zhang, Donghan Yu, Yihao Feng, Bowen Jin, Zhaoran Wang, John Peebles, and Zirui Wang, highlight that effective mid-training should define an action subspace that minimizes both the error from simplifying the action space and the subsequent RL error during planning.
The Core Idea: Action Abstractions
The paper identifies two main factors determining mid-training effectiveness: pruning efficiency and its impact on RL convergence. Pruning efficiency relates to how well the system can narrow down the vast array of possible actions to only the most useful ones. The analysis suggests that mid-training is most effective when the decision space is compact and the “effective horizon” (the number of steps needed to achieve a goal) is short. This points to the importance of operating in the space of “action abstractions” rather than primitive, low-level actions.
Think of it like this: instead of telling a robot every single joint movement to pick up a cup (primitive actions), you tell it “pick up cup” (an action abstraction). This high-level instruction encapsulates many primitive actions, making the planning process much simpler and more efficient. For LLMs, these abstractions could be high-level reasoning steps or “skills” that are transferable across different tasks.
Introducing RA3: Reasoning as Action Abstractions
Building on these insights, the researchers propose a scalable mid-training algorithm called Reasoning as Action Abstractions (RA3). This algorithm works by iteratively discovering “temporally-consistent latent structures” through RL, followed by fine-tuning the model on the data generated with these new structures. Essentially, RA3 learns to identify the hidden “thoughts” or “intentions” behind expert demonstrations, which are not explicitly available in the training data.
RA3 uses a sequential variational lower bound, optimized through an Expectation-Maximization (EM) procedure. In the “E-step,” an RL process is used to infer latent sequences (the action abstractions) that best explain the expert’s primitive actions. In the “M-step,” the model is fine-tuned on the expert data, now augmented with these inferred latent abstractions. The algorithm introduces two types of latents: <act> for temporally consistent actions (meaning the current high-level intention continues) and <think> for initiating a new rationale. A penalty mechanism ensures that new reasoning steps are only triggered when they significantly improve the log-likelihood of expert actions, thus controlling computational cost.
Experimental Validation in Code Generation
The effectiveness of RA3 was demonstrated through experiments on code generation tasks, using various base models like Qwen and Llama, ranging from 1B to 8B parameters. The mid-training corpus consisted of 3.5 million code snippets, totaling 1 billion tokens. The results were compelling: RA3 consistently improved the average performance on benchmarks like HumanEval and MBPP by 8 and 4 points, respectively, over the base model and the next-token prediction baseline. Furthermore, RA3 achieved faster convergence and higher asymptotic performance in RL with Verifiable Reward (RLVR) on more challenging benchmarks like HumanEval+, MBPP+, LiveCodeBench, and Codeforces.
The experiments showed that fine-tuning on data bootstrapped with action abstractions significantly reduced cross-entropy loss, indicating more efficient learning. The model learned to extract transferable skills, such as “dummy head creation” or “BFS” (Breadth-First Search) in code, which enhanced its generalizability. An ablation study also confirmed the importance of the KL penalty in regulating the frequency of reasoning, striking a balance between generating useful abstractions and avoiding redundant “thinking” steps.
Also Read:
- Guiding Small Language Models to Think: A New Approach to Reasoning Distillation
- Unlocking New Abilities: How Reinforcement Learning Helps Language Models Compose Skills
Conclusion
This research provides a foundational understanding of how mid-training shapes post-training RL for LLMs. By formalizing the benefits of learning temporal action abstractions, the paper offers a principled approach to designing more effective mid-training algorithms. RA3, as a practical implementation, showcases how uncovering these high-level reasoning trajectories can lead to stronger initial policies and significantly accelerate the convergence and performance of subsequent reinforcement learning. For more details, you can read the full research paper here.
