ISOPO: Isometric Policy Optimization

• ISOPO is a proximal gradient algorithm that efficiently approximates the natural policy gradient using one-shot, layer-wise Fisher normalization.
• It simplifies policy updates by eliminating multi-step clipping and old-policy references, reducing variance and computational overhead.
• The method optionally employs an NTK-based microbatch transformation to boost sample efficiency and control KL divergence.

Isometric Policy Optimization (ISOPO) is a proximal gradient algorithm for policy optimization that efficiently approximates the natural policy gradient using a single backward pass. Unlike established proximal policy methods (e.g., PPO, GRPO, CISPO) that require multiple steps and importance ratio clipping with respect to a referent ("old") policy, ISOPO achieves natural-gradient-like updates by normalizing per-sequence log-probability gradients in the Fisher metric prior to advantage contraction, optionally incorporating a neural tangent kernel (NTK)–based microbatch transformation. This layer-wise, batch-dimension procedure is designed to offer unbiasedness, reduced variance, and efficient convergence, maintaining negligible computational overhead relative to vanilla REINFORCE (Abrahamsen, 29 Dec 2025).

1. Motivation and Context

Proximal policy optimization algorithms, notably PPO [Schulman et al. 2017], GRPO (2024), and CISPO (2025), regulate policy updates via importance-ratio clipping to enforce a "trust region." The canonical surrogate is: $$r(\theta) = \frac{\pi_\theta(a \mid s)}{\pi_\text{old}(a \mid s)}, \qquad L^{\text{CLIP}}(\theta) = \mathbb{E}\left[ \min \left( r(\theta) A,\, \mathrm{clip}(r(\theta), 1-\epsilon, 1+\epsilon) A \right) \right]$$ Multiple gradient steps are taken on this surrogate with a stale reference policy $\pi_\text{old}$, resulting in hyperparameter dependence, extra forward/backward passes, and non-negligible staleness.

Natural policy gradient (NPG) approaches define updates in the Fisher metric by: $$F = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \nabla_\theta \log \pi_\theta(\tau) \nabla_\theta \log \pi_\theta(\tau)^T \right]$$ where policy update steps are chosen as $F^{-1} g$ for $$g = \mathbb{E}[A(\tau) \nabla_\theta \log \pi_\theta(\tau)]$$. For large scale architectures (e.g., LLMs), computing or inverting $F$ directly is intractable.

ISOPO addresses this computational bottleneck by providing a one-shot, layer-wise approximation to $F^{-1}g$ without reliance on old-policy references or multi-step clipping operations.

2. Mathematical Formulation

2.1 Non-Interacting ISOPO

ISOPO’s foundation is the per-sequence log-probability gradient: $v(\tau) = \nabla_\theta \log \pi_\theta(\tau)$ Its Fisher-norm is: $$\|v(\tau)\|_F = \sqrt{v(\tau)^T F v(\tau)} \approx \sqrt{ \frac{1}{n} \sum_{i=1}^n \big( v(\tau) \cdot g_i \big)^2 }$$ where $g_i$ are reduced gradients (token-position-wise).

ISOPO normalizes $v(\tau)$ in the Fisher metric before advantage weighting: $$g(\tau) = \frac{v(\tau)}{ \sqrt{ v(\tau)^T F v(\tau) } }, \qquad \hat g_\text{ISOPO} = \mathbb{E}_{\tau \sim \pi_\theta} [ g(\tau) A(\tau) ]$$ Layer-wise updates are given by: $$\Delta\theta_l = \sum_{i=1}^m \frac{ A(\tau_i) }{ \sqrt{ \|\nabla_{\theta_l} \log \pi_\theta(\tau_i) \|_F^2 + \mathrm{reg}_l } } \; \nabla_{\theta_l} \log \pi_\theta(\tau_i)$$ where $\mathrm{reg}_l$ is a small damping parameter.

2.2 Fisher-Norm Estimation

Let $g_\text{out, j}$ denote the back-propagated gradient at token-position $j$ and $a_\text{in, j}$ the corresponding activation for a linear layer weight gradient $V = v(\tau)$. Then: $$\sum_j \big( v(\tau) \cdot g^{(j)} \big)^2 = \sum_j \left( g_{\text{out}, j} \cdot V a_{\text{in}, j} \right)^2$$

$$\sum_j \|g^{(j)}\|^2 = \sum_j \|g_{\text{out}, j}\|^2 \|a_{\text{in}, j}\|^2$$

Thus,

$$\|v(\tau)\|_F \approx \frac{ \sqrt{ \sum_j \left( g_{\text{out}, j} \cdot V a_{\text{in}, j} \right)^2 } }{ \sqrt{ \sum_j \|g_{\text{out}, j}\|^2 \|a_{\text{in}, j}\|^2 } }$$

2.3 NTK-Based Interacting ISOPO

Sequence-wise gradients for layer $l$ are stacked into a Jacobian: $$J \in \mathbb{R}^{m \times d}, \quad J_{i, \cdot} = \nabla_{\theta_l} \log \pi_\theta(\tau_i)$$ Empirical NTK: $K = J J^T \in \mathbb{R}^{m \times m}$ Update direction: $\Delta\theta_l = J^T (K + cI)^{-1} A$ where $A \in \mathbb{R}^m$ comprises sequence-advantages; $c$ is a Tikhonov regularizer.

3. Algorithmic Implementation

The non-interacting ISOPO is implemented via a single backward hook per layer, without extra forward passes. Key steps (for a linear layer) are:

1. Compute $\text{loss} = \sum(\log_\text{prob})$; execute $\text{loss.backward()}$ for batch gradients.
2. In the backward hook:
   • Partition microbatch by sequence.
   • For each $i$: recover unreduced per-token gradients $g_\text{out}[j]$ and activations $a_\text{in}[j]$; aggregate $V_i = \nabla\log\pi(\tau_i)$.
   • Estimate $\|V_i\|_F$ via formulas above.
   • Accumulate $$\text{layer.grad} += \frac{A[i]}{F_\text{norm} + \varepsilon} V_i$$.
3. Apply optimizer step (e.g., AdamW).

The interacting (NTK-based) variant uses a small $m \times m$ eigendecomposition:

• Compute $K = J J^T$.
• Let $U$ and $D$ be eigenvectors/values of $K$.
• Compute $$\text{precond\_adv} = U @ \left( (D + c)^{-1} \times (U^T @ A) \right)$$.
• Accumulate $\text{layer.grad} += J^T @ \text{precond\_adv}$.

All overhead is in the batch dimension (per-token inner products, small eigendecompositions); runtime increase is negligible compared to $O(d^2)$ matrix operations in standard backward passes.

4. Theoretical Properties

ISOPO preserves key theoretical attributes:

• Unbiasedness: The normalization factor $\|v(\tau)\|_F$ is evaluated per sample and is independent of the advantage, maintaining unbiased estimation of the direction

$$\frac{F^{-1}g}{\|F^{-1}g\|_F} \approx \frac{F^{-1}[A(\tau) v(\tau)]}{\|F^{-1}[A(\tau) v(\tau)]\|_F}$$

• Variance Reduction: Fisher-normalization yields pronounced reduction in sample gradient variance.
• Convergence: Like NPG, preconditioning suppresses updates that induce excessive KL divergence. Under standard conditions (bounded advantages, Lipschitz log-probabilities), sublinear convergence rate $O(1/\sqrt{T})$ is maintained.
• Comparisons:
  • GRPO/PPO employ trust regions via importance-ratio clipping, a multi-pass indirect proxy for $F^{-1}$.
  • CISPO clips the sampling distribution, lacking per-sample adaptation to policy geometry.
  • NPG-family optimizers (K-FAC, Muon, Shampoo) operate in full parameter space; ISOPO’s normalization/preconditioning is sample-wise and batch-oriented, making it operationally complementary.

5. Empirical Findings

ISOPO was evaluated via GSM8K math reasoning fine-tuning on Qwen-3 0.6B, with a group-relative advantage estimator and no KL penalty. Baselines were:

• REINFORCE (no clipping)
• GRPO/PPO ($\epsilon=0.1$ clipping)

Metrics:

• Validation accuracy at regular intervals
• KL-drift from initialization

Principal outcomes:

• REINFORCE: gradual accuracy gain
• GRPO: faster convergence, eventual plateau
• ISOPO (non-interacting, $p=0, q=-1, r=-2$): reached equivalent accuracy in approximately half the steps compared to GRPO
• ISOPO with Fisher normalization ($p=-1$): improved accuracy with reduced KL drift compared to GRPO and REINFORCE
• Sequence-Euclidean normalization ($q=-1$): improved accuracy but did not control KL drift
• Interacting ISOPO (NTK preconditioner): achieved further sample-efficiency gains

Method	Steps to 75% acc.	KL-drift @50	Overhead vs REINFORCE
REINFORCE	5000	0.12	1.0×
GRPO (PPO clip)	3000	0.09	1.2×
ISOPO (non-int)	1500	0.06	1.05×
ISOPO (interact)	1200	0.05	1.10×

A plausible implication is that ISOPO provides a direct, per-sample normalization in the Fisher metric, producing natural-gradient-like updates in a single backward pass. This achieves higher stability and sample efficiency, with lower KL drift and overhead nearly matching that of unadorned REINFORCE.

6. Significance and Distinctions

ISOPO fundamentally replaces the multi-step, old-policy–anchored clipping mechanisms of PPO/GRPO/CISPO with a direct normalization in the Fisher metric. Optional NTK-based interactions further enhance sample efficiency. This approach circumvents staleness, trust-region hyperparameterization, and extraneous forward passes while delivering approximations to the natural policy gradient compatible with scale. ISOPO’s per-sample normalization aligns each update with local geometry in policy space, enforcing KL control and impartiality in gradient estimation—attributes unattainable by canonical clipped methods. This suggests ISOPO could serve as a foundation for future scalable, stable policy optimization frameworks, especially in domains where natural gradient methods remain computationally prohibitive (Abrahamsen, 29 Dec 2025).