KL-Regularized Reinforcement Learning is Designed to Mode Collapse

TL;DR

The paper investigates KL-regularized reinforcement learning as a distribution-matching problem, showing that common settings often yield unimodal solutions and hence diversity collapse. It analyzes the exact optima for reverse KL and forward KL regularization, highlighting how the regularization strength and reward scales shape modality; it also introduces Mode Anchored Reward Augmentation (MARA) to directly optimize for multimodal targets. The authors provide closed-form target distributions for both KL types, derive conditions that govern multimodality, and demonstrate MARA as a drop-in method across language-model and chemical-language-model tasks, achieving higher quality and diversity without extra diversity signals. This work offers a principled framework to preserve diversity in post-training generation and suggests practical directions for designing explicit multimodal objectives in RL settings.

Abstract

It is commonly believed that optimizing the reverse KL divergence results in "mode seeking", while optimizing forward KL results in "mass covering", with the latter being preferred if the goal is to sample from multiple diverse modes. We show -- mathematically and empirically -- that this intuition does not necessarily transfer well to doing reinforcement learning with reverse/forward KL regularization (e.g. as commonly used with language models). Instead, the choice of reverse/forward KL determines the family of optimal target distributions, parameterized by the regularization coefficient. Mode coverage depends primarily on other factors, such as regularization strength, and relative scales between rewards and reference probabilities. Further, we show commonly used settings such as low regularization strength and equal verifiable rewards tend to specify unimodal target distributions, meaning the optimization objective is, by construction, non-diverse. We leverage these insights to construct a simple, scalable, and theoretically justified algorithm. It makes minimal changes to reward magnitudes, yet optimizes for a target distribution which puts high probability over all high-quality sampling modes. In experiments, this simple modification works to post-train both Large Language Models and Chemical Language Models to have higher solution quality and diversity, without any external signals of diversity, and works with both forward and reverse KL when using either naively fails.

Figures (10)

• Figure 1: Illustration of how the choice of approximate distribution family affects KL optimization. With a restrictive approximate distribution (e.g. two-parameter Gaussian), KL exhibits the typical "mode seeking" and "mass covering" characteristics. This intuition does not necessarily hold for flexible distributions (e.g. independent categoricals, foundational models).
• Figure 2: Final policy distribution from optimizing a reverse/forward KL regularized reward maximization objective, given the same reward function, reference policy, across a range of regularization strengths ($\beta$). Note that both KLs can lead to multimodal solution distributions.
• Figure 3: With equal ${\pi_{\text{ref}}}$, linear difference in rewards ($\Delta R$) lead to exponential difference in probabilities
• Figure 4: With equal rewards, RL does not change answers' relative probs.
• Figure 5: MARA stays close to the reference policy in low-reward areas, and puts high, uniform mass over all high-reward areas.

Theorems & Definitions (18)

• Remark 3.1
• Remark 3.2
• Remark 3.3
• Remark 3.4
• proof
• Definition 3.5
• Proposition 4.1
• proof
• Remark 4.2
• Remark 4.3