RL-finetuning LLMs from on- and off-policy data with a single algorithm

The proposed Any-Generation Reward Optimization (AGRO) algorithm introduces a single reinforcement learning framework capable of fine-tuning large language models effectively across both on-policy and off-policy data settings by leveraging the concept of generation consistency. This method is significant because it derives optimal policies through sample-based learning with theoretical guarantees, demonstrating improved performance on mathematical reasoning datasets over existing baseline algorithms.

The gist

AGRO leverages the concept of generation consistency, which states that the optimal policy satisfies a notion of consistency across any possible generation of the model.

Reinforcement Learning Formulation and Consistency Condition

The paper frames language model fine-tuning as a regularized policy optimization problem: maximizing the objective function defined by Equation (1), which aims to maximize expected reward while maintaining proximity to a reference policy through KL regularization. A central theoretical finding is Theorem 1, which establishes generation consistency at optimality: for any generation y, the quantity r(x, y) − β log π∗(yx)πref(yx) does not depend on y; it is a function of the prompt x only. This implies that the optimal policy satisfies a specific relationship where the regularized reward depends only on the prompt.

Derivation of Loss Functions

Based on generation consistency, two primary loss functions are introduced to search for the optimal policy:

1. The on-policy loss, defined as:

L(π) = 1/2 Ex∼ρ [Vy∼π(·x) Rπβ(x, y)].

2. The off-policy loss, defined as:

Lµ(π) = 1/2 Ex∼ρ [Vy∼µ(·x) Rπβ(x, y)].

The paper demonstrates that globally minimizing both L(π) and Lµ(π) leads to the optimal policy.

Gradient Estimates for Policy Optimization

The gradient of the off-policy loss, ∇Lµ(π), is given by Lemma 3:

∇Lµ(π) = −βE [Rπβ(x, y)] − Rπ(x, µ) / ∇ log π(yx).

When the sampling is on-policy (µ = π), this gradient simplifies to an expression equivalent to the RLHF gradient from Equation (1), aligning with the standard on-policy RL objective. The full on-policy AGRO gradient, ∇L(π), is decomposed into two terms: ∇PDL(π) and ∇LRL(π).

Token-Level Implementation and Variance Reduction

For LLM applications, the sequence-level generations are adapted to a token level. Three unbiased stochastic estimates for the regularization term g(π) are presented:

1. gˆ1: The vanilla estimate of g(π).

2. gˆ2: An estimate that ignores lower-triangular terms of gˆ1, which are expected to be zero in expectation, reducing variance compared to gˆ1.

3. gˆ3: An estimate that further reduces variance by computing the conditional expectation over each single next token, utilizing the available information from the forward pass.

Experimental Results and Performance

Experiments on the mathematical reasoning dataset MATH show that AGRO outperforms baseline algorithms in both on-policy and off-policy settings. Specifically, Figure 3 indicates that on-policy ARGO achieves the best performance overall, with a +7% lift in performance compared to πref, while off-policy AGRO seems to obtain better performance overall compared to KL-regularized policy gradient algorithm. The analysis shows that while on-policy algorithms drive up the KL divergence more, they achieve faster progress toward training reward. Strong regularization (large β) makes the algorithm more KL-efficient but can lead to instability during off-policy learning.

Theoretical Guarantees

The paper provides theoretical guarantees through Lemma 5, which states that the KL divergence is related to performance under the regularized policy as: βKL(π, π∗) = G(π∗) − G(π). This connects the policy optimization objective directly to the KL divergence between policies. Furthermore, Proposition 4 confirms that for off-policy AGRO with on-policy sampling, ∇Lb(π) is an unbiased estimate of ∇L(π), ensuring it leads to the optimal policy π∗.

Extensions and Future Directions

The formulation can be extended to pairwise preference data using the Bradley-Terry assumption, which implies that reward differences are related to the log ratio of policies. The generation consistency condition also relates to general consistency for regularized value-based learning. Potential future work includes a more careful investigation into "the source of off-policy instability and how one might leverage complementary techniques such as importance sampling.

As a fastidious researcher, I have analyzed the AGRO framework presented in this paper. The core contribution is establishing generation consistency as an optimality condition for RLHF and deriving a unified algorithm (AGRO) that effectively leverages both on-policy and off-policy data streams without requiring complex auxiliary value functions.

Here are the specific improvements to AI systems that can be made using this paper, categorized by capability:

)

Large Language Model (LLM) Fine-Tuning for Mathematical Reasoning and Complex Instruction Following.

1. A unified fine-tuning pipeline capable of optimizing LLMs using both on-policy human feedback (RLHF/DPO style) and static, pre-existing datasets (like MATH or other reasoning corpora).

2. The ability to adapt the learning algorithm seamlessly between these two data regimes without switching codebases or fundamentally altering the objective function structure.

Specific improvements achievable:

1. Inference on complex mathematical proofs and problem-solving tasks where high accuracy is required, leveraging the improved performance observed on the MATH dataset (+7% lift in evaluation performance).

2. Instruction following for multi-step, chain-of-thought reasoning tasks, as the reward function directly optimizes for correct reasoning paths rather than just final answers.

3. Handling complex preference data (Pairwise Preference) by directly optimizing the policy based on the Bradley-Terry model implied by generation consistency, leading to more robust alignment with human preferences.

)

Improved Sample Efficiency and Training Stability in RLHF/Alignment.

1. Significantly reduced sample complexity during fine-tuning, as AGRO is demonstrated to be more data-efficient than vanilla KL-regularized policy gradient, requiring fewer total training steps (epochs).

2. Enhanced training stability, particularly in off-policy settings where the standard KL regularization approach often leads to performance crashes; AGRO maintains consistent performance across different learning dynamics.

3. Training high-quality alignment models with smaller datasets or fewer total gradient updates, drastically reducing computational costs associated with expensive LLM fine-tuning runs.

4. Deployment of RLHF agents in production environments where data collection is costly, as the algorithm can effectively utilize existing offline interaction data (off-policy) while converging to the optimal policy faster than traditional methods.

)

Robust and Low-Variance Gradient Estimation for Large Language Models.

1. Implementation of highly stable gradient estimators (like those derived from generation consistency) that minimize variance without introducing significant bias, especially when dealing with high-dimensional sequence data (token level).

2. The ability to implement low-variance estimates for complex regularization terms in the token-level policy gradient, which is crucial for scaling AGRO to very long sequences and high-resolution token predictions.

3. Training LLMs on extremely long contexts or sequence lengths (e.g., 4K+ tokens) where standard gradient estimation suffers from high variance, leading to more accurate token-level policy updates.

4. More precise control over the trade-off between reward maximization and KL divergence minimization during training, resulting in policies that are both highly aligned with human preferences and maintain strong generative capabilities.

)

Generalization Across On-Policy and Off-Policy Data Regimes.

1. Creation of a single, general-purpose RL fine-tuning algorithm (AGRO) that is inherently agnostic to whether the data used for policy updates is generated by the current model (on-policy) or an external behavior policy (off-policy).

2. The ability to derive optimal policies by leveraging diverse data sources simultaneously—using on-policy experience for immediate feedback and off-policy data for broad coverage—without needing separate specialized algorithms for each scenario.

3. Developing a single, robust fine-tuning tool that can be used in various industrial settings (e.g., fine-tuning on a small high-quality human preference set on-policy, then leveraging a large, older dataset off-policy).

4. Mitigating the online vs. offline trade-off by ensuring that the theoretical convergence guarantees hold for both scenarios under the AGRO framework, allowing practitioners to choose data sources based purely on cost/availability rather than algorithmic limitations.