# Hugging Face Journal Club: Score Centering Stabilizes Off-policy Reinforcement Learning

## Executive summary

This talk provides an in-depth derivation of Score Centering, a mathematical correction designed to stabilize Off-policy Reinforcement Learning (RL) for Large Language Models (LLMs). The core problem addressed is the instability that arises when there is a mismatch between the policy used for training ($P_ heta$) and the policy used for inference/sampling ($Q$). The solution involves introducing a correction term ($ ilde{S}$) that eliminates the 'drift' in the policy gradient loss, ensuring the gradient only depends on the covariance between the reward and the score. To make this correction computationally feasible, the speakers detail a Top-K approximation method.

## Key takeaways

- The Problem of Inference Mismatch: RL learning for LLMs becomes unstable when the training policy ($P_ heta$) does not match the sampling/inference policy ($Q$), leading to a 'drift' in the policy gradient loss. The mismatch is defined as $P_ heta eq Q$.
- Score Centering Correction ($ ilde{S}$): The correction term $ ilde{S}$ is derived using importance sampling to relate expectations between $P_ heta$ and $Q$. This term stabilizes the expected gradient, ensuring the loss only depends on the covariance of the reward ($R$) and the score ($S_Y$), thereby eliminating the drift.
- Computational Efficiency via Top-K Approximation: Calculating the full correction term requires summing over the entire vocabulary, which is computationally expensive. The speakers propose a Top-K approximation, limiting the calculation to the top K heads and tails of the vocabulary while maintaining high accuracy.

## Technical details

- Policy Gradient Loss Derivation: The policy gradient loss is defined as the expectation of samples from $P_ heta$ for the reward $R(Y)$. The gradient of this loss with respect to $ heta$ is shown to be proportional to $ ext{E}_{Y ext{ from } P_ heta} [R(Y) imes abla_ heta ext{log } P_ heta(Y)]$.
- The Drift Mechanism: The key insight is that the policy gradient contains a 'drift' when $P_ heta eq Q$. The correction term $ ilde{S}$ is designed to eliminate this drift, ensuring that the expected gradient is equal to $ ext{Cov}_Q(R, S_Y)$.
- Top-K Approximation Details: The Top-K approximation for $ ilde{S}$ involves defining the vocabulary mass using the top K heads and tails. The approximation sets $ ilde{S} ext{ approx } ext{Sum}_{ ext{Top K}} (Q_V imes S_Y) + ho imes ext{Sum}_{ ext{Tails}} (P_V imes S_Y)$, where $ ho$ is a calculated scaling factor.

## Practical implications

- Stabilizes RL training for LLMs, allowing for reliable off-policy learning even when the training and inference policies diverge.
- The method is shown to be compatible with other advanced techniques like Truncated Importance Sampling (TIS), Masked Importance Sampling (MIS), and Expert Replay.
- The technique can be implemented efficiently using Top-K approximations, mitigating the computational cost of full vocabulary calculations.

## Topics

Reinforcement Learning, LLMs, Policy Gradient, Score Centering, Off-policy Learning, Quantization, Top-K Approximation, Score Centering Paper

Source: https://www.youtube.com/watch?v=mEtRe_KKYm0