PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yunhao Tang, Taco Cohen, David W. Zhang, Michal Valko, Rémi Munos

arXiv:2503.19612v2cs.LGstat.ML

TL;DR

Online and offline language-model RL involve a trade-off: offline methods are cheaper but generally deliver less improvement. AGRO derives losses from generation consistency to use either data setting, and shows compelling improvements over baseline regularized policy gradient in several on- and off-policy settings.

  • Problem

    Offline RL is cheaper to execute than online RL but generally delivers less improvement, leaving a trade-off between the two approaches.

    Source: §1 · p. 1

  • Method

    AGRO derives sample-trainable losses from generation consistency, an optimality condition for regularized policy optimization, to use on-policy and off-policy data.

    Source: §1 · p. 1

  • Results

    AGRO produces compelling performance improvements over baseline regularized policy gradient in a few on-policy and off-policy learning settings.

    Source: §9 · p. 9

  • Takeaways & Limitations

    Generation consistency provides an optimality condition from which alternative losses can naturally use different data sources.

    Source: §9 · p. 9

    Off-policy learning can still become unstable when data generation is taken to the extreme.

    Source: §9 · p. 9

Abstract

from arXiv · show

We introduce a novel reinforcement learning algorithm (AGRO, for Any-Generation Reward Optimization) for fine-tuning large-language models. AGRO leverages the concept of generation consistency, which states that the optimal policy satisfies the notion of consistency across any possible generation of the model. We derive algorithms that find optimal solutions via the sample-based policy gradient and provide theoretical guarantees on their convergence. Our experiments demonstrate the effectiveness of AGRO in both on-policy and off-policy settings, showing improved performance on the mathematical reasoning dataset over baseline algorithms.

1. Introduction

Language-model fine-tuning uses both online and offline RL, but offline methods are cheaper and generally yield less improvement. AGRO uses generation consistency to derive an algorithm for both data settings, with experiments examining performance in mathematical reasoning.

1. Introduction

  • Offline RL is cheaper than online RL but generally delivers less improvement, creating a trade-off between the two approaches.The introduction also notes that the online/offline dichotomy makes it difficult to adapt algorithms directly between settings.Source: §1 · p. 1
  • Generation consistency yields loss functions for searching for the optimal policy that support sample-based gradient estimation.The introduction identifies generation consistency as an optimality condition inherent to regularized policy optimization.Source: §1 · p. 1
  • AGRO is proposed as an RL algorithm for language-model fine-tuning that can leverage both on-policy and off-policy data.Its derivation is based on generation consistency, an optimality condition for regularized policy optimization.Source: §1 · p. 1
  • The experiments examine AGRO's performance in on-policy and off-policy learning, focusing on mathematical reasoning and hyperparameter ablations.The introduction describes the ablations as providing actionable insights for practitioners.Source: §1 · pp. 1–2

2. Reinforcement learning for language model

The section formulates language-model reinforcement learning as optimizing a policy that generates responses to prompts for reward. Its regularized objective balances expected reward against divergence from a reference policy.

2. Reinforcement learning for language model

  • A language model is modeled as a policy that generates a response to a prompt, which is evaluated by a reward function, usually reflecting human preference.The objective is to maximize expected reward; the reward may be supplied by a parametric network or be programmatic.Source: §2 · p. 2
  • The regularized objective maximizes average reward while penalizing KL divergence from a reference policy.The coefficient β controls the trade-off between policy optimization and regularization; the reference policy is usually the supervised-finetuned policy and starting point.Source: §2 · p. 2

3. Generation consistency for regularized RL

At the optimum of regularized RL, the paper identifies a generation-consistency condition: a quantity computed for a prompt and generation depends only on the prompt. It derives on- and off-policy losses whose global minimizer is the optimal policy under stated support conditions.

3. Generation consistency for regularized RL

  • At the optimal policy, the quantity in Theorem 1 computed for a prompt and any generation does not depend on the generation.The theorem states that this quantity is a function of the prompt alone.Source: §3 · p. 2
  • In the tabular off-policy comparison, AGRO converges to the optimal policy while KL-regularized policy gradient does not.The comparison uses data generated from πref and tracks normalized KL(π, π*) during training.Source: Figure 1 · p. 2

3.1. Deriving loss functions from generation consistency

  • The on-policy and off-policy losses are based on variance of the regularized reward across generations, with off-policy sampling supplied by a behavior policy μ independent of π.The off-policy setting includes a fixed dataset or replay buffer.Source: §3.1 · p. 3
  • Both losses have the optimal policy as their unique global minimum among policies sharing the reference policy's support, assuming μ belongs to that policy space.The derivation connects zero variance of the regularized reward across generations to the optimal policy.Source: §3.1 · p. 3

4. Gradient of the consistency losses

The consistency-loss gradients connect off-policy learning to the RLHF gradient under on-policy sampling, while the on-policy loss preserves the same optimum with different learning dynamics.

4.1. Gradient of the off-policy loss

  • With on-policy sampling, the off-policy AGRO gradient is equivalent to the RLHF gradient, combining reward maximization with squared-loss divergence minimization against πref.The squared-loss gradient aligns with the KL-regularization gradient when µ = π.Source: §4.1 · p. 3

4.2. Gradient of the on-policy loss

  • The on-policy loss gradient combines pathwise-derivative and likelihood-ratio estimates, with both terms of order O(β^2).The pathwise-derivative estimate is aligned with the gradient of the original objective G.Source: §4.2 · p. 4
  • Both the full on-policy loss gradient and its pathwise-derivative component lead to the same optimal policy π*, but generate different learning dynamics.The pathwise-derivative component is the off-policy loss gradient evaluated with µ = π.Source: §4.2 · p. 4

5. Any-Generation Reward Optimization

AGRO uses generation consistency to optimize a policy from on- or off-policy responses; its off-policy gradient converges to the optimal policy where a naive KL-regularized policy-gradient adaptation generally does not.

5. Any-Generation Reward Optimization

  • AGRO optimizes from prompts paired with responses and rewards generated by a behavior policy off-policy or the current policy on-policy.The algorithm derives from generation consistency and is presented as finding the optimal policy.Source: §5 · p. 4

5.1. The off-policy AGRO algorithm

  • With one response per prompt, minimizing the off-policy loss can use a prompt-dependent value function; with n ≥ 2 responses, an empirical average avoids learning that auxiliary function.The authors note that avoiding the auxiliary value function may be more desirable for implementation simplicity.Source: §5.1 · p. 4
  • Off-policy AGRO converges to π*, while the adapted KL-regularized policy-gradient algorithm generally fails to converge because KL regularization is not correct off-policy.Figure 1 tracks KL(π, π*) over updates.Source: §5.1 · p. 5

5.2. The on-policy AGRO algorithm

  • The on-policy AGRO gradient estimate is unbiased, and its full gradient combines a pathwise-derivative component with a likelihood-ratio component.The pathwise-derivative term aligns with the usual on-policy RL gradient, while the likelihood-ratio term is intended to reduce variance.Source: §5.2 · p. 5

6. Implementation of AGRO at the token level

AGRO’s sequence-level gradients require token-level adaptation for language models, where alternative estimates exploit zero-mean terms and next-token expectations to reduce variance.

  • LLM responses are autoregressive token sequences, and the sequence-level policy probability factors into conditional next-token probabilities.The paper writes the joint probability as a product of token probabilities with the conditioning context suppressed.Source: §6 · p. 6
  • The three token-level estimates are unbiased; removing lower-triangular zero-mean terms reduces variance, and computing local next-token expectations can reduce it further.The next-token distribution is available during the forward pass.Source: §6 · p. 6
  • AGRO variants use a vanilla estimate, remove lower-triangular terms, or compute local expectations; token-level KL terms can be calculated by summing over tokens.These choices correspond to estimates ĝ1, ĝ2, and ĝ3.Source: §6 · p. 6
  • The experiments used vanilla AGRO because variance-reduction gains were limited at small training scale; the impact in expensive large-scale training remains for further ablation.The authors leave the large-scale impact to practitioners to investigate.Source: §6 · p. 6

7. Discussions on related work and extensions

The discussion connects generation consistency to pairwise-preference objectives and regularized value learning, including a token-level formulation suggested for dense rewards.

8. Experiments

Experiments on MATH compare on-policy and off-policy AGRO with regularized policy gradient, then examine KL efficiency and regularization. On-policy AGRO leads the reported test-set comparison; off-policy AGRO performs better than the regularized policy-gradient baseline in the offline setting.

8. Experiments

  • The experiments use an 8B Llama 3 model fine-tuned on MATH with answer-match rewards and four generations per prompt for gradient estimates.The on-policy comparison uses β = 0.001.Source: §8 · p. 7, §8.1 · p. 7

8.1. On-policy learning

  • On-policy AGRO makes faster reward progress and requires less compute than alternatives to reach the same training performance.The authors treat epochs as an approximation of total compute because per-iteration costs are similar.Source: Figure 2 · p. 7, §8.1 · p. 7
  • On-policy AGRO achieves a +7% lift in performance compared to πref, while other baselines have a +4% lift.The evaluation curves are reported as fairly correlated with training rewards.Source: Figure 3 · p. 8, §8.1 · p. 8

8.3. Ablation and comparison

  • Off-policy learning is as KL-efficient as on-policy learning, but on-policy updates raise KL more and converge faster per learning step.In the synthetic replay setting, off-policy AGRO and regularized policy gradient show no significant difference.Source: Figure 4 · p. 8, §8.3 · pp. 8–9

8.2. Off-policy and offline learning

  • In offline evaluation, off-policy AGRO beats KL-regularized policy gradient at peak performance and at fixed epochs; weak regularization can cause performance to crash during training.Both methods use β = 0.01 for the reported comparison.Source: §8.2 · p. 8, Figure 5 · p. 9

8.3. Ablation and comparison

  • For squared regularization, β = O(1/T) keeps regularization at a constant scale; per-token updates instead place β roughly at O(1).The authors note this per-token scale may simplify hyperparameter tuning.Source: §8.3 · p. 9
  • Higher β makes updates more KL-efficient and conservative, but slows reward progress and reduces deviation from πref at a fixed update count.The ablation varies β across regularized policy gradient and off-policy AGRO settings.Source: Figure 6 · p. 9, §8.3 · p. 9

9. Discussion and limitation

AGRO uses generation consistency to support regularized policy optimization with either on-policy or off-policy data. The authors report theoretical guarantees and performance improvements in several learning settings, while noting that extreme off-policy data can still make learning unstable.

9. Discussion and limitation

  • AGRO uses generation consistency to optimize regularized policies with both on-policy and off-policy data, with theoretical guarantees and reported improvements over regularized policy gradient in several settings.The conclusion characterizes those results as applying to a few on-policy and off-policy learning settings.Source: §9 · p. 9
  • Extreme off-policy data generation can still destabilize learning, so AGRO does not eliminate off-policy instability.The authors identify further investigation of instability and possible importance-sampling techniques as future work.Source: §9 · p. 9

A. Experimental details

The appendix records experimental settings and provides additional theoretical proofs and discussion. It also describes a low-variance estimate for ∇LRL and notes implementation trade-offs.

A. Experimental details

  • Experiments use Llama 3 8B, batches of 64 prompts, four generations per prompt, and temperature 1 with top-p 1.Training and evaluation use 7,500 MATH training examples and 5,000 test examples; supervised fine-tuning warms up RL training.Source: A. Experimental details · p. 13

B.1. Proof of theoretical results

  • The off-policy proof excerpt uses responses generated by a behavior policy μ and identifies the sampled responses as i.i.d.The provided excerpt begins from the off-policy AGRO gradient.Source: §B.1 · p. 13

B.2. Equivalence between KL divergence and performance under regularized reward

  • The appendix formalizes a connection between KL divergence and performance under regularized reward, citing a special case of O’Donoghue (2022).The proof invokes the stated form of the optimal policy.Source: §B.2 · p. 13

B.3. Variance reduction for estimating ∇LRL(π)

  • The low-variance estimate for ∇LRL is assembled by handling zero terms, regrouped terms, and cross terms with explicit one-step expectations.The authors caution that these estimates can be more sophisticated to implement given practical trade-offs.Source: §B.3 · p. 14
Loading 2503.19612v2…