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

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.

Technical details

• Policy Gradient Loss Derivation 100s

  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 220s

  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 600s

  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.

Mentioned resources

• Score Centering Paper (Research Paper)

Channel & topics

• Hugging Face
• Reinforcement Learning
• LLMs
• Policy Gradient
• Score Centering
• Off-policy Learning
• Quantization
• Top-K Approximation
• Score Centering Paper