On a few pitfalls in KL divergence gradient estimation for RL

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Tang, Yunhao, Munos, Rémi
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913889244938240
author Tang, Yunhao
Munos, Rémi
author_facet Tang, Yunhao
Munos, Rémi
contents We point out a few pitfalls in implementing gradient estimation for KL divergence in RL training for LLM, as seen in a number of open source projects and papers. The first major pitfall is to differentiate through the KL estimate as loss functions to minimize KL divergence. We show that such implementations are generally incorrect and do not produce the desired KL gradient. Secondly, we show that some implementations do not account for the sequential nature of the estimation problem and produce a partial gradient at best. We demonstrate the impact of such issues with illustrative tabular and LLM experiments, and show the correct way to implement the KL gradient.
format Preprint
id arxiv_https___arxiv_org_abs_2506_09477
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle On a few pitfalls in KL divergence gradient estimation for RL
Tang, Yunhao
Munos, Rémi
Machine Learning
We point out a few pitfalls in implementing gradient estimation for KL divergence in RL training for LLM, as seen in a number of open source projects and papers. The first major pitfall is to differentiate through the KL estimate as loss functions to minimize KL divergence. We show that such implementations are generally incorrect and do not produce the desired KL gradient. Secondly, we show that some implementations do not account for the sequential nature of the estimation problem and produce a partial gradient at best. We demonstrate the impact of such issues with illustrative tabular and LLM experiments, and show the correct way to implement the KL gradient.
title On a few pitfalls in KL divergence gradient estimation for RL
topic Machine Learning
url https://arxiv.org/abs/2506.09477