Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Projected Microbatch Accumulation yields reference-free proximal policy updates for reinforcement learning

Nilin Abrahamsen

TL;DR

PROMA introduces a reference-free proximal update for RL-based large-language-model fine-tuning by projecting the accumulated gradient orthogonally to the span of sequence-wise log-probability gradients across microbatches, implemented efficiently with a layer-wise backward pass. The approximate projection runs in $O(kd)$ FLOPS, where $k$ is the microbatch size and $d$ is the parameter count, and reduces overlaps with $\text{grad\_log\_prob}$ to control local KL divergence. In GSM8K with Qwen-3 0.6B, PROMA delivers competitive validation performance compared to PPO-based GRPO, while maintaining higher entropy for longer and lower $KL$ divergence from the initial policy than REINFORCE, along with smoother successive-policy updates. The method provides a scalable, reference-policy-free proximal update mechanism for RL-based LLM fine-tuning, with connections to proximal optimization concepts observed in computational physics.

Abstract

This note introduces Projected Microbatch Accumulation (PROMA), a proximal policy update method for large language model fine-tuning. PROMA accumulates policy gradients across microbatches by projecting out sequence-wise gradient components before microbatch aggregation. The projection is applied layer-wise during the backward pass, enabling efficient implementation without additional forward or backward passes. Empirically, PROMA enforces tighter control of local KL divergence than GRPO, resulting in more stable policy learning. Unlike PPO and GRPO, PROMA achieves proximal updates without inducing entropy collapse and does not rely on a reference policy or likelihood-ratio clipping.

Projected Microbatch Accumulation yields reference-free proximal policy updates for reinforcement learning

TL;DR

PROMA introduces a reference-free proximal update for RL-based large-language-model fine-tuning by projecting the accumulated gradient orthogonally to the span of sequence-wise log-probability gradients across microbatches, implemented efficiently with a layer-wise backward pass. The approximate projection runs in FLOPS, where is the microbatch size and is the parameter count, and reduces overlaps with to control local KL divergence. In GSM8K with Qwen-3 0.6B, PROMA delivers competitive validation performance compared to PPO-based GRPO, while maintaining higher entropy for longer and lower divergence from the initial policy than REINFORCE, along with smoother successive-policy updates. The method provides a scalable, reference-policy-free proximal update mechanism for RL-based LLM fine-tuning, with connections to proximal optimization concepts observed in computational physics.

Abstract

This note introduces Projected Microbatch Accumulation (PROMA), a proximal policy update method for large language model fine-tuning. PROMA accumulates policy gradients across microbatches by projecting out sequence-wise gradient components before microbatch aggregation. The projection is applied layer-wise during the backward pass, enabling efficient implementation without additional forward or backward passes. Empirically, PROMA enforces tighter control of local KL divergence than GRPO, resulting in more stable policy learning. Unlike PPO and GRPO, PROMA achieves proximal updates without inducing entropy collapse and does not rely on a reference policy or likelihood-ratio clipping.
Paper Structure (5 sections, 3 figures)

This paper contains 5 sections, 3 figures.

Table of Contents

  1. Background
  2. Projected Microbatch Accumulation (PROMA)
  3. Results
  4. Discussion
  5. Acknowledgements

Figures (3)

  • Figure 1: $\texttt{accumulate\_PROMA}$ is called during the backward pass of the policy gradient loss and modifies the gradient accumulation. Before adding the microbatch gradient, the partially accumulated gradient is projected to be orthogonal to the gradient of each sequence in the current microbatch ($\texttt{mcb\_seq\_grads}$).
  • Figure 2: Training performance and policy statistics for PROMA, GRPO, and REINFORCE. (a) Validation performance. (b) KL divergence from the initial policy. (c) Policy entropy. (d) KL divergence between the current policy and a lagged reference policy (running mean over 40 steps). Only GRPO (blue) uses PPO clipping of the likelihood ratio.
  • Figure 3: To exactly project away the $k$ sequences in the microbatch ($\texttt{project\_to\_complement\_v1}$) takes $\approx 2k^2d$ FLOPS where $d$ is the number of parameters. The results in this section used the iterative projection $\texttt{project\_to\_complement\_v2}$ which takes $O(kd)$ FLOPS. I did not test if this was faster in practice.