Saved in:
| Main Authors: | , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2512.21852 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866918394877444096 |
|---|---|
| author | Shah, Vedant Obando-Ceron, Johan Jain, Vineet Bartoldson, Brian Kailkhura, Bhavya Mittal, Sarthak Berseth, Glen Castro, Pablo Samuel Bengio, Yoshua Malkin, Nikolay Jain, Moksh Venkatraman, Siddarth Courville, Aaron |
| author_facet | Shah, Vedant Obando-Ceron, Johan Jain, Vineet Bartoldson, Brian Kailkhura, Bhavya Mittal, Sarthak Berseth, Glen Castro, Pablo Samuel Bengio, Yoshua Malkin, Nikolay Jain, Moksh Venkatraman, Siddarth Courville, Aaron |
| contents | The reasoning performance of large language models (LLMs) can be substantially improved by training them with reinforcement learning (RL). The RL objective for LLM training involves a regularization term, which is the reverse Kullback-Leibler (KL) divergence between the trained policy and the reference policy. Since computing the KL divergence exactly is intractable, various estimators are used in practice to estimate it from on-policy samples. Despite its wide adoption, including in several open-source libraries, there is no systematic study analyzing the numerous ways of incorporating KL estimators in the objective and their effect on the downstream performance of RL-trained models. Recent works show that prevailing practices for incorporating KL regularization do not provide correct gradients for stated objectives, creating a discrepancy between the objective and its implementation. In this paper, we further analyze these practices and study the gradients of several estimators configurations, revealing how design choices shape gradient bias. We substantiate these findings with empirical observations by RL fine-tuning \texttt{Qwen2.5-7B}, \texttt{Llama-3.1-8B-Instruct} and \texttt{Qwen3-4B-Instruct-2507} with different configurations and evaluating their performance on both in- and out-of-distribution tasks. Through our analysis, we observe that, in on-policy settings: (1) estimator configurations with biased gradients can result in training instabilities; and (2) using estimator configurations resulting in unbiased gradients leads to better performance on in-domain as well as out-of-domain tasks. We also investigate the performance resulting from different KL configurations in off-policy settings and observe that KL regularization can help stabilize off-policy RL training resulting from asynchronous setups. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2512_21852 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | A Comedy of Estimators: On KL Regularization in RL Training of LLMs Shah, Vedant Obando-Ceron, Johan Jain, Vineet Bartoldson, Brian Kailkhura, Bhavya Mittal, Sarthak Berseth, Glen Castro, Pablo Samuel Bengio, Yoshua Malkin, Nikolay Jain, Moksh Venkatraman, Siddarth Courville, Aaron Machine Learning Artificial Intelligence The reasoning performance of large language models (LLMs) can be substantially improved by training them with reinforcement learning (RL). The RL objective for LLM training involves a regularization term, which is the reverse Kullback-Leibler (KL) divergence between the trained policy and the reference policy. Since computing the KL divergence exactly is intractable, various estimators are used in practice to estimate it from on-policy samples. Despite its wide adoption, including in several open-source libraries, there is no systematic study analyzing the numerous ways of incorporating KL estimators in the objective and their effect on the downstream performance of RL-trained models. Recent works show that prevailing practices for incorporating KL regularization do not provide correct gradients for stated objectives, creating a discrepancy between the objective and its implementation. In this paper, we further analyze these practices and study the gradients of several estimators configurations, revealing how design choices shape gradient bias. We substantiate these findings with empirical observations by RL fine-tuning \texttt{Qwen2.5-7B}, \texttt{Llama-3.1-8B-Instruct} and \texttt{Qwen3-4B-Instruct-2507} with different configurations and evaluating their performance on both in- and out-of-distribution tasks. Through our analysis, we observe that, in on-policy settings: (1) estimator configurations with biased gradients can result in training instabilities; and (2) using estimator configurations resulting in unbiased gradients leads to better performance on in-domain as well as out-of-domain tasks. We also investigate the performance resulting from different KL configurations in off-policy settings and observe that KL regularization can help stabilize off-policy RL training resulting from asynchronous setups. |
| title | A Comedy of Estimators: On KL Regularization in RL Training of LLMs |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2512.21852 |