KL-Regularized Reinforcement Learning is Designed to Mode Collapse (2510.20817v1)

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 LLMs). 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 LLMs and Chemical LLMs 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.

• The paper analyzes KL-regularized reinforcement learning to show that the interplay between regularization strength, reward functions, and the reference policy inherently leads to mode collapse.
• It introduces Mode Anchored Reward Augmentation (MARA), a method that modifies rewards to equalize high-reward samples and enforce multimodal policy distributions.
• Empirical evaluations in LLM tasks, creative question answering, and chemical language models demonstrate MARA's effectiveness in preserving output diversity and improving optimization efficiency.

KL-Regularized Reinforcement Learning and Mode Collapse: A Formal Analysis

Introduction

This paper presents a rigorous theoretical and empirical analysis of KL-regularized reinforcement learning (RL) objectives, with a particular focus on their implications for output diversity in policy optimization. The authors challenge the prevailing intuition that reverse KL regularization is inherently "mode seeking" and forward KL is "mass covering," arguing that these characterizations do not hold in the context of RL with expressive policy classes such as those used in LLMs. Instead, the diversity of the optimal policy is determined by the interplay between the regularization strength, the reward function, and the reference policy, rather than the choice of KL direction. The paper further demonstrates that standard RL setups often induce unimodal solution distributions, leading to mode collapse by design, and introduces a simple, theoretically justified algorithm—Mode Anchored Reward Augmentation (MARA)—to directly optimize for multimodal solutions.

KL Divergence in RL: Theory and Misconceptions

The KL divergence is a central regularizer in RL for policy post-training, typically used to maintain proximity to a reference policy while maximizing external rewards. The paper revisits the classical variational inference intuition: reverse KL ($D_{KL}(q||p)$) is "mode seeking," while forward KL ($D_{KL}(p||q)$) is "mass covering." However, this dichotomy is shown to be an artifact of restrictive variational families (e.g., Gaussians), and does not generalize to flexible policy classes such as categorical distributions over large discrete spaces.

Figure 1: Restrictive (Gaussian) approximation highlights classical mode-seeking/mass-covering behavior, which does not generalize to flexible policy classes.

KL-Regularized Reward Maximization: Solution Distributions

The paper formalizes the solution distributions for both reverse and forward KL-regularized RL objectives. For reverse KL, the optimal policy is:

$$G_\beta(y) = \frac{1}{Z} \pi_{\text{ref}}(y) \exp\left(\frac{R(y)}{\beta}\right)$$

where $Z$ is the normalizing constant, $R(y)$ is the reward, $\pi_{\text{ref}}$ is the reference policy, and $\beta$ is the regularization coefficient. For forward KL, the solution is:

$$G_{\text{fwd}}(y) = \frac{\beta \pi_{\text{ref}}(y)}{\Lambda - R(y)}$$

with $\Lambda$ chosen to normalize the distribution. Notably, both forms can yield multimodal solutions depending on the reward landscape and regularization strength.

Figure 2: Final policy distributions for reverse/forward KL regularization across regularization strengths; both can produce multimodal solutions.

Analysis of Mode Collapse: Probability Ratios and Diversity

A key contribution is the closed-form analysis of probability ratios under the optimal policy. For reverse KL, the log-probability ratio between two samples $y_1$ and $y_2$ is:

$$\log \frac{G_\beta(y_1)}{G_\beta(y_2)} = \log \frac{\pi_{\text{ref}}(y_1)}{\pi_{\text{ref}}(y_2)} + \frac{1}{\beta}(R(y_1) - R(y_2))$$

This reveals several critical properties:

• Equal reference support, reward difference: Small reward differences are exponentially amplified in probability ratios, especially for small $\beta$, leading to strong mode concentration.
• Equal rewards, reference support difference: The optimal policy preserves the reference policy's relative probabilities; RL does not upweight low-support answers even if they are equally correct.

Figure 3: Valid answer entropy and rewards, demonstrating diversity collapse in naive KL-regularized RL.

Figure 4: Effect of reward difference and regularization strength on relative probabilities; small $\beta$ induces extreme mode concentration.

Figure 5: Final policy distribution after KL-regularized RL with equal rewards; low-support correct answers are never preferred.

Direct Optimization for Multimodal Solutions: MARA

To address the inherent mode collapse, the authors propose Mode Anchored Reward Augmentation (MARA), which modifies the reward function such that all high-reward samples above a threshold $\tau$ are assigned augmented rewards to equalize their probabilities in the optimal policy. The algorithm is simple to implement and requires only minor changes to standard RL pipelines.

Figure 6: MARA maintains proximity to the reference policy in low-reward regions and assigns uniform high probability to all high-reward modes.

Empirical Validation

The paper validates MARA across three domains:

• Verifiable LLM Task (1-2 generation): MARA preserves diversity in correct answers, achieving near-uniform output over multiple valid solutions, while naive KL-regularized RL collapses to a single answer.

  Figure 7: Vanilla RL training outcomes; diversity collapses to a single answer across regularization strengths.

• Creative Question Answering: MARA outperforms GRPO and RLOO in out-of-distribution reward and diversity metrics, demonstrating improved generalization and output variety.
• Chemical LLMs for Drug Discovery: MARA increases the yield of unique high-reward molecules and improves optimization efficiency, matching or exceeding baseline diversity metrics.

Implications and Future Directions

The findings have significant implications for RL-based post-training of generative models. The analysis demonstrates that mode collapse is not merely an artifact of imperfect optimization or limited exploration, but a direct consequence of the objective's construction. The choice of regularization strength and reward scaling is critical; naive settings will almost always induce unimodal solutions. MARA provides a principled mechanism to enforce diversity without external diversity signals or complex reward engineering.

Theoretically, the work calls for a re-examination of RL regularization practices, advocating for explicit construction of target distributions with desired properties. Practically, MARA is immediately applicable to a wide range of RL post-training scenarios, including LLM alignment and molecular design.

Conclusion

This paper provides a comprehensive theoretical and empirical analysis of KL-regularized RL objectives, demonstrating that mode collapse is a natural consequence of standard formulations. The proposed MARA algorithm offers a simple, scalable solution for enforcing diversity in policy optimization. Future work may explore more expressive target distributions, improved gradient estimators for MARA, and deeper analysis of forward KL regularization dynamics. The results underscore the necessity of explicit target distribution design in RL for generative models.