Reproduce “On a few pitfalls in KL divergence gradient estimation for RL” on OpenResearch

About this paper

Yunhao Tang, Rémi Munos

A paper from Meta FAIR and Meta GenAI identifies widespread implementation errors in KL divergence gradient estimation, a crucial component in reinforcement learning, particularly for large language models. It demonstrates that common methods yield biased gradients or optimize the unintended reverse KL, and provides correct, unbiased estimators, showing their efficacy in tabular and LLM experiments including Llama 3 models.

Key insights

• Auto-differentiating common Monte-Carlo estimates of KL divergence (e.g., vanilla or variance-reduced) as a loss function yields biased gradients, often resulting in no effective regularization or minimizing the unintended reverse KL divergence.
• For sequence generation tasks like those with large language models, token-level KL losses produce only partial gradients, as they fail to account for the auto-regressive dependencies between tokens.
• Correctly estimating the KL divergence gradient requires specific forms derived from the score function, and for sequences, unbiased sequence-level estimators that capture full temporal dependencies are crucial for effective optimization.

Summary by alphaXiv.