Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ISOPO: Proximal policy gradients without pi-old

Nilin Abrahamsen

TL;DR

ISOPO addresses the high cost of natural policy gradient methods in RL-based fine-tuning of large language models by introducing an isometric gradient update that operates in a single backward pass. It normalizes each sequence's log-probability gradient in the Fisher metric and contracts it with the corresponding advantage, with a layer-wise transformation that avoids extra forward passes. The paper provides a non-interacting ISOPO variant with Fisher-norm based rescaling, efficient Fisher-norm estimation from unreduced gradients, and a generalized rescaling family, plus an interacting ISOPO variant that leverages an empirical neural tangent kernel for layer-wise preconditioning. Empirical results on gsm8k with Qwen-3 0.6B demonstrate competitive validation performance and reduced KL drift, while incurring negligible overhead relative to REINFORCE, suggesting a practical path to near natural-gradient updates for RLHF in LLM fine-tuning.

Abstract

This note introduces Isometric Policy Optimization (ISOPO), an efficient method to approximate the natural policy gradient in a single gradient step. In comparison, existing proximal policy methods such as GRPO or CISPO use multiple gradient steps with variants of importance ratio clipping to approximate a natural gradient step relative to a reference policy. In its simplest form, ISOPO normalizes the log-probability gradient of each sequence in the Fisher metric before contracting with the advantages. Another variant of ISOPO transforms the microbatch advantages based on the neural tangent kernel in each layer. ISOPO applies this transformation layer-wise in a single backward pass and can be implemented with negligible computational overhead compared to vanilla REINFORCE.

ISOPO: Proximal policy gradients without pi-old

TL;DR

ISOPO addresses the high cost of natural policy gradient methods in RL-based fine-tuning of large language models by introducing an isometric gradient update that operates in a single backward pass. It normalizes each sequence's log-probability gradient in the Fisher metric and contracts it with the corresponding advantage, with a layer-wise transformation that avoids extra forward passes. The paper provides a non-interacting ISOPO variant with Fisher-norm based rescaling, efficient Fisher-norm estimation from unreduced gradients, and a generalized rescaling family, plus an interacting ISOPO variant that leverages an empirical neural tangent kernel for layer-wise preconditioning. Empirical results on gsm8k with Qwen-3 0.6B demonstrate competitive validation performance and reduced KL drift, while incurring negligible overhead relative to REINFORCE, suggesting a practical path to near natural-gradient updates for RLHF in LLM fine-tuning.

Abstract

This note introduces Isometric Policy Optimization (ISOPO), an efficient method to approximate the natural policy gradient in a single gradient step. In comparison, existing proximal policy methods such as GRPO or CISPO use multiple gradient steps with variants of importance ratio clipping to approximate a natural gradient step relative to a reference policy. In its simplest form, ISOPO normalizes the log-probability gradient of each sequence in the Fisher metric before contracting with the advantages. Another variant of ISOPO transforms the microbatch advantages based on the neural tangent kernel in each layer. ISOPO applies this transformation layer-wise in a single backward pass and can be implemented with negligible computational overhead compared to vanilla REINFORCE.
Paper Structure (11 sections, 12 equations, 7 figures)

This paper contains 11 sections, 12 equations, 7 figures.

Table of Contents

  1. Introduction
  2. Related work
  3. Refresher on the natural policy gradient
  4. The ISOPO gradient update
  5. Derivation: Sequence-wise bound on the KL contribution
  6. Estimating the Fisher metric
  7. Efficient implementation
  8. Alternative sequence re-scalings
  9. Results
  10. Interacting variant
  11. Code

Figures (7)

  • Figure 1: ISOPO acts on the batch dimension of each layer before the advantage vector is applied to the log-probability gradients. Right: GRPO and related methods acts on the same batch dimension, but by clipping sequences based on a reference policy instead of based on the gradient. Bottom: Many optimizers act on the parameter dimension, including diagonal multiplers (AdamW) and non-diagonal transformations (K-FAC, Shampoo, SOAP, Muon). The natural policy gradient itself is not tractable.
  • Figure 2: The non-interacting variant of ISOPO scales the sequence gradients based on a stochastic estimate of the Fisher norm of the log-probability gradient.
  • Figure 3: Generalized ISOPO re-scaling. The examples derived in sections \ref{['sec:sample-wise']} and \ref{['sec:alternative']} correspond to $p=-1$ ($q=r=0$) and $r=-2$ ($p=q=0$) respectively.
  • Figure 4: Validation curve for non-interacting ISOPO (red) with $p=0$, $q=-1$ and $r=-2$ (See algorithm \ref{['fig:algo-alternatives']}) and no regularization. Baselines are REINFORCE (green) with no clipping and GRPO (blue) which uses PPO clipping. For every 10 steps the best validation score among 3 runs is shown for each algorithm (except only best of 2 for ISOPO).
  • Figure 5: KL-drift and validation score for non-interacting ISOPO vs GRPO and REINFORCE. Multiple runs are shown for each algorithm and the convex hulls are shaded for visual clarity. ISOPO is shown with two different settings: with Fisher normalization $p=-1$, $q=r=0$ (red) and with sequence-wise Euclidian normalization $q=-1, p=r=0$ (yellow). The $p=-1$ setting decreases the KL drift whereas $q=1$ does not.
  • ...and 2 more figures