KL-Regularized Reinforcement Learning for Generative Modelling is Designed to Mode Collapse
Authors: Anthony GX-Chen
Jatin Prakash
Jeff Guo
Rob Fergus
Rajesh Ranganath
Venue: ICLR 2026 Poster (2026)
Keywords: reinforcement learning
LLM
diversity
KL divergence
Primary area: reinforcement learning
TL;DR: KL-regularized RL often collapses to a single reward mode by design. We analyze why and propose a lightweight reward-rescaling method that preserves quality while improving mode coverage.

Abstract: Classical intuitions cast minimizing reverse KL as "mode seeking" and forward KL as "mass covering". In KL-regularized reinforcement learning, however, the regularizer determines _both_ the target distribution's shape _and_ the divergence being implicitly minimized, making its role more nuanced than simply inducing generic mode-seeking or mass-covering behaviour. Specifically, the target distribution is defined jointly by the reward function, the reference model, the type of regularizer, and the regularization strength. We show that under common settings—such as low regularization strength and equal verifiable rewards—both forward and reverse KL regularization tend to specify target distributions whose mass concentrates on a single high-reward region. Thus, the objective itself _by construction_ induces diversity collapse, regardless of the policy optimization algorithm used.

Building on this perspective, we introduce a simple and scalable modification that rescales rewards to induce target distributions assigning substantial probability across _all_ high-reward regions. This yields a principled objective that maintains high solution quality while achieving broad reward-mode coverage. Empirically, this approach improves post-training diversity and performance for Large Language Models and Chemical Language Models, and is effective with either forward or reverse KL regularization, while using either naively fails.

Published as a conference paper at ICLR 2026
KL-REGULARIZED REINFORCEMENT LEARNING FOR GENERATIVE MODELLING IS DESIGNED TO MODE COLLAPSE
Anthony GX-Chen'
Jatin Prakash'
Jeff Guo² Rob Fergus!
Rajesh Ranganath*
New York University anthony. gx. www.cnce.ennc...cmm
Roole Polytechnique Fédérale de Lausanne (EPFL.)
ABSTRACT
Classical intuitions cast minimizing reverse KL as "mode seeking" and forward KL "mass covering". In KL-regularized reinforcement learning, however, the regularizer determines both the target distribution' 's shape and the divergence being implicitly minimized, making its role more nuanced than simply inducing generic mode-seeking mass-covering behaviour. Specifically, the target distribution is defined jointly by the reward function, the reference model, the type of regularizer, and the regularization strength. We show that under common settings settings-such as low regularization strength and equal verifiable rewards- both forward and reverse KL regularization tend to specify target distributions whose mass concentrates on a single high-reward region. Thus, the objective itself by construction induces diversity collapse, regardless of the policy optimization algorithm used.
Building this perspective, introduce simple and scalable modification that rescales rewards to induce target distributions assigning substantial probability across all high-reward regions. This yields a principled objective that maintains high solution quality while achieving broad reward- rewerd-mode coverage. Empirically, this approach improves post-training diversity and performance for Large Language Models and Chemical Language Models, and is effective with either forward or reverse KI. regularization, while using either naively fails.
INTRODUCTION
Reinforcement Learning (RL) is the predominant method for post-training foundation models (Ouyang et al., 2022), and the primary way to train generative models in settings where the correct solution is not known priori. At its core, this involves solving regularized RL (contextual bandits) problem, where a policy is trained to maximize some external reward, while preserving "closeness" to base policy (as to e.g. preserve coherence). Output diversity of the policy is crucial. In Large Language Models (LLMs), it drives engagement for tasks such as creative writing and free-form conversation. More generally, diversity underlies the generation of discovery of novel mathematical solutions (Romera-Parede (Castro et al., 2025), and novel algorithms and software (Surina et al., 2025; Novikov al., 2025; Aygün et al., 2025). Furthermore, diversity reflects uncertainty over competing hypotheses, property fundamental to scientific discovery (GX-Chen et al., 2025). Finally, diversity plays an important role during training to drive exploration such that the policy can find and converge to better solutions (Cui et al., 2025).
knowledge, enabling the et al., 2024), cognitive science models
Yet, current empirical evidence suggests RL post-training improves quality at the cost of diversity (Kirk et al., 2023; Cui et al., 2025). As a response, a number of recent works set out to treat this ailment, with a variety of approaches including explicit diversity rewards (Li et al., 2025), changing the KL regularization (Wang et al., 2023), selecting diverse data (Lanchantin al., 2025), and count-based exploration bonuses (Song et al., 2025).
In this work, we take a step back to diagnose a more fundamental problem: does the objective being optimized actually have solution that is diverse? We find that with current set-ups, the answer is often "no", even with unlimited compute, high quality data, and perfect optimization. We prove that under very commonly used settings (such weak KI. regularization with varied rewards, any KI.

---

Published as a conference paper at ICLR 2026
regularization if correct answers have the same rewards but vastly different reference policy supports), the globally optimal solution is often by construction unimodal. To accomplish this, we analyze KL-regularized RL through tools from variational inference (VI. Jordan et al. 1999; Ranganath et al. 2014) to find and dissect optimal policies for different choices KI. regularization.
Section 2 provides preliminaries about KL divergences and reminds the reader of the mode-seeking mass-covering behaviour of minimizing reverse/ forward KL at suboptimality. Section studies KL-regularized reward maximization as implicitly minimizing a divergence between the current policy and a larget distribution. Section 4 further analyzes the shape of this target solution. We focus particularly on how this distribution puts mass over high-reward regions- -ic. multimodality in terms of reward modes (Definition 3.5). This allows us to understand even if we perfectly minimize any divergence between the policy and the target distribution, the resulting policy will still be non-diverse if the targer distribution is defined to be unimodal. Finally, Section 5 shows how one can directly construct the target distribution to cover high-reward modes. We specify one such distribution which puts mass over all high-reward regions above certain threshold, and show this requires only small change to current algorithms. Each section is empirically supported with didactic simulations. Finally, apply method out of the box to LLMs and chemical language models and find that it works for complex, realistic scenarios.
The main contributions can be summarized as follows,
1. We analyze the role of reverse/forward KI. regularization in RI as both defining the target distribution and the implicitly minimized divergence between policy and target.
2. We show the shape of the target distribution is determined by the regularizer, regularization strength, and relative reward and reference probability magnitudes. This has implications on how the target distribution puts mass over high-rewart regions.
3. We show with typical hyperparameters, the target distribution is often constructed to put mass over single high-reward region, making diversity collapse a natural consequence of correctly solving the regularized RL problem (as currently defined), regardless of algorithm.
4. We derive conditions required for broad reward mode coverage, and use this insight to construct simple and theoretically principled RL algorithm (two-line pseudocode, Alg. 1) that puts uniform mass over all high-reward regions, without any external diversity signals.
2
THE
(KL) DIVERGENCE
0.8
Stop 0
Stop 70
0.8
Stop 600
Step 400
Reverse Kt
0.075
0.6
0.6
0.6
0.4
Reverse
0.050
0.2
0.2
0.025
0.0
0.0
-5
-5
-2.5
-2.5
-25 0.0
Stop
Stop 200
Stop 600
Step
Step 50
Step 400
Forward KL.
0.6
0.6
Forward KL
0.075
0.675
0.075
0.4
0.4
0.059
0.2
0.2
0.025
0.025
0.025
0.0
0.000
0.000
0.000
-5
-5
-25
0.
2.5
-2.5
0.0
-25 0.0
(a) Restrictive (Gaussian) approximation
(b) Flexible (categorical) approximation
Figure 1: Illustration of how the choice of approximate distribution family affects KL optimization. With a restrictive approximate distribution (e.g. 1D Gaussian with two parameters), KI. exhibits the typical "mode seeking" and "mass covering" characteristics. This intuition does not necessarily hold for flexible distributions (e.g. independent categoricals, foundational models).
The Kullback-Leibler (KL) divergence (Kullback Leibler, 1951) measures the discrepancy between two probability distributions. In machine learning, is commonly used in variational inference (VI), where minimizing the KL divergence enables a tractable variational distribution q to approximate an intractable posterior p (Jordan et al., 1999; Blei et al., 2017). Following Murphy (2012), refer to P (qlpp p)) E., [log q(y) log p(y)) as the reverse KL divergence, and Dx (pllg) = E, [log ((v))()) log (1)) the forward KI. divergence. Reverse KI. is often described "mode
2

---

Published as a conference paper at ICLR 2026
seeking", avoiding mass where p is small (Figure la, top). while forward KL is often described as "mass covering", putting mass anywhere p has mass (Figure la, bottom). These intuitions hold if the variational family is not sufficiently expressive and (Bishop Nasrabadi, 2006; Murphy, 2012). With a flexible family, however, optimizing either KI.s the global optimum can well-approximate complex posterior (Figure 1b).
at best settle
optimum with >> 0 KL
KL-REGULARIZED REWARD MAXIMIZATION
KL-regularized reward maximization aims to (i) maximize the expected value of reward function R y - IR, mapping from samples to a scalar outcome (e.g. improve human preference), while (ii) keeping the policy л close reference distribution Tif (e.g. maintain grammatical coherence). The objective is J(T = [g)(()) [[(y]] B (TTe. Tet). where (-:--) denotes a divergence between the policy and reference distributions. For brevity,
consider the unconditional generation problem where the policy models distribution No (v)) The problem is the same in the case of conditional generation (e.g. question answering), where the objective is simply defined over the conditional commonly
distribution (x) We do not deal with the sequential decision making modelled using Markov Decision Processes (Puterman, 1994)-in this work. Nevertheless, non-sequential setting is widely used when training generative models with RL.
In this section, we consider the solution target distribution KL-regularized reward maximization, i.c the distribution which maximizes the objective. The central question is:
If we perfectly solve the regularized RL problem to its global optimum, what does the solution (policy) distribution look like?
3.1 SOLUTION OF THE REVERSE KL REGULARIZED OBJECTIVE
The most common Kl.rregularized policy gradient objective uses the reverse KI. divergence,
[R(y)] - 3 D z (Tel Teet)
number of previous works have discussed the solution / optimal distribution of this optimization problem (Korbak et al., 2022; Go et al., 2023; Rafailov et al., 2023; Azar et al., 2024; Zhang & Ranganath, 2025), which we note again below (see Appendix B.1 for detailed derivations).
Remark 3.1. The optimal solution to the reverse-KL KL regularized reward maximization problem, arg maxs, Jz( is given by the target distribution n² = =GB,
1 Galy) (v) exp
R(y) B
where < Tre(t)) (is(R())//))))))) ) dy is the normalizing constant.
Remark 3.1 tells us the distribution maximizing Equation 1 is wy = Gg. However, it may not be
immediately obvious how optimizing Equation 1, Ve Ja(e)), moves Ta toward Gg. We analyze this below (details in Appendix B.2, also see c.g. Zhang & Ranganath (2025)).
Remark 3.2. The gradient of Equation current policy no and the target distribution Gg.
a gradient of the reverse KI. divergence between the
Va DxL (To |l GB) x Vo Ja(To).
(3)
Main Takeaway
Maximizing the reverse-KL. regularized RL objective J/ (Equation 1) is equivalent to doing distribution matching by minimizing reverse KI. toward the target distribution Gg (Equation 2).
3.2 SOLUTION OF THE FORWARD KL REGULARIZED OBJECTIVE
Alternatively, the reward can be maximized with forward KL penalty,
Jiwa( (T)) EE (1) [R(y)] - 9 DKL (Meer]
number of recent works have used forward KL regularization. Some are motivated explicitly by the "mass covering" intuition of the forward KL (Wang al., 2023), while others-suue GRPO (Shao et al., 2024; Guo et al., 2025a) -my have incidentally estimated the forward KL, despite meaning to the reverse KI. (Tang & Munos, 2025).
3

---

Published as a conference paper at ICLR 2026
Remark 3.3. Assume optimization with 3 O. with finite rewards Rmax < 00, and there exist solution(s) where Rly) = Rmur. (e)r(1) > 0. The optimal solution to the forward- KL regularized reward maximization problem. arg maX,T Iuud is given by the distribution:
Gyod(y) B ( R(y)"
A> max R(y),
where unique A exists for each 33 such that Gyvd valid probability distribution. Notably, Equation is completely different distribution family from the reverse KL case (Equation 2). and does not have a simple closed form unnormalized solution. It is also worth noting that if higher- rewarding regions exist outside of Tet's support, Glwd can place nonzero mass on regions where Tyr(y) = O and R(y) Raxx with no preference among y's within this region in terms of density, See Appendix B.3 for more details.
Remark 3.4. Assume we are optimizing within the support of Tech the gradient of Equation is not forward KL gradient,
N DED (h * Vo J[wa(To).
(6)
for any target distribution h that is defined independently ne. and arbitrary reward functions R. Proof. Appendix B.4.
Therefore, while Equation 4 can still be a good objective to optimize, it does not necessarily inherit exactly the same properties and intuitions as a "forward KI. gradient".
What, then, is the gradient of the forward KI. (G m ? It in fact reduces to doing maximum likelihood (e.g. supervised fine-tuning) on trajectories sampled from the target G33 (Remark B.2), which is intractable for generic targets. However, this provides one perspective algorithms such as STaR (Zelikman et al., 2022) and RAFT (Dong et al., 2023; Xiong et al., 2025) that filter high-reward trajectories for maximum likelihood. One interpret filtering as rejection sampling approximate a target distribution (which put high mass over high-reward regions), when reward is bounded and we know GB up normalization. More generally, other methods that approximately sample from G3 to minimize divergence include Naesseth et al. (2020); Khalifa et al. (2020).
Main Takeaway
Maximizing the forward-KI KI. regularized objective Jiwa (Equation 4) does yield forward-KI. gradient, so its behaviour cannot be naively equated to forward-KL optimization.
To summarize: the regularized RI. objective implicitly minimizes a divergence between the current policy Ta and target distribution G. Different choices of regularizer lead to different target distributions G. Importantly, the regularizing divergence (TTT, Tea need not be the same type of divergence as the one implicitly minimized, D( G),, as is the case for forward- KL regularization.
3.3
BOTH KL REGULARIZATION CAN HAVE MULTIMODAL SOLUTION DISTRIBUTIONS
Reverse KL
Reward Function
# 0.01
8 0.1
8 0.15
= 0.5
B 0.75
024
0.3
E32
E35
6.3
0.44
C315
Forward KL
Reference Policy
A 0.01
8 0.1
0.15
A = 0.5
B 0.75
8/6
301
635
6330
5.0%
Figure 2: Final policy distribution (100-dim categorical) from optimizing a reverse/forward KL regularized reward maximization objective, given the same reward function, reference policy, across a range of regularization strengths (3). Both KLs can lead to multimodal solution distributions.
4

---

Published as a conference paper at ICLR 2026
We briefly note that the target distributions for both the reverse (Equation 2) and forward (Equation 5) KL regularization can be multimodal. To ground the discussion, we first define a common-sense notion of "multimodal" in terms of reward modes, which use for the rest of the paper.
Definition 3.5. A target distribution G' for regularized reward maximization is "reward multimodal"
given a high-reward threshold G(y'); and for any Above-threshol samples R(y) and R(y') 2 G(v) ~ G(y').
for any two samples y. y'. R(y) < T T R(y') implies G(y) <
Informally, this means all high-reward samples have high probability, and sampling from G samples approximately equally from all high-reward regions.
We show in didactic example in Figure 2, where given the same reward function containing two high-reward modes, and a reference policy with support over the first half of the action space, optimizing the reverse and forward KL objectives lead wide variety of solutions that depend on the regularization coefficient 3. Both KLs have settings of / that induce reward multimodal target distributions. We analyze the properties of the target distribution the subsequent section, and return to the Figure example in detail in Section 4.3.
ANALYSIS OF KL REGULARIZER OPTIMAL DISTRIBUTION
We have seen in Section 3.3 that both KL-regularized RL objectives have reward multimodal solutions, and in Section 2 that optimizing either KL divergence to global optimum will give us policies that well-approximate the (multimodal) solution. However, the shape of the target distribution depends on the reward, reference distribution, and regularization strength. This raises the central question:
Is the globally optimal solution we commonly define in Kl-regularized RI. actually reward multimodal (Definition 3.5)?
The central tool we use in this section is probability ratio between two samples under a distribution. Intuitively, we want (i) high-reward samples to be much more probable than low-reward samples, and (ii) similarly high-reward samples to have similar high probabilities. Unless otherwise stated, focus our analysis on the solution of the reverse-KI. regularized objective (Equation 2), both for its clean form and because it is the most common way KI.rregularized RI. is formulated.
Proposition 4.1. The (log) probability ratio between any two samples, Yı, y2, under the optimal solution distribution for reverse-Ki KI. regularized RI, G, can be written closed form,
G3(y1) log G3(32) (32)
= log (11 + (B(y)) - -R(y)))
Proof. Because normalization constant cancel out in ratios. See Appendix B.6.
gives us a generic and closed-form way of analyzing how likely one sample is relative to another in the optimal solution, using only Tee and the reward function R, for any reverse-Ke regularized reward maximization objective. This gives us number of consequential insights.
Proposition 4.1
4.1 WITH EQUAL SUPPORTS, SMALL REWARD CHANGE DRIVES LARGE PROBABILITY CHANGE Remark 4.2. For any two samples y1 and y2, (nf(l)) And (22), their probability ratio is:
Ga(U1) G)(
R(v))= R(y2)
exp
(8)
B
Rewards
Mere
(8=1e
Me. (8=0.5)
We first consider the continuous reward func- tion setting where samples have small differ- ences in rewards. This common for set- tings such as alignment (reward model) and drug discovery (e.g. binding affinity). If two samples have the same probability under the reference distribution Tet ("equal support"),, the difference in their final log probabilities is simply the difference in their rewards, scaled by 1/3. Smaller exaggerates the difference
10
1.000
0.512
0.422
0.000 000
0
Actin Indix
Action Indox
Action Index
Action Index
Figure 3: With equal Tret linear difference in rewards
A R) lead to exponential difference in probabilities
5

---

Published as a conference paper at ICLR 2026
between relative probabilities. Note linear difference in rewards result in exponential difference in probabilities: for a 0.1 difference in rewards, and a commonly used 3 = 1e-3, the higher reward sample is pushed to be 2.6 x 1043 times more likely in the solution distribution. This issue is identically present for entropy-only regularization (see Fig. 7 for effect of / on relative probabilities). This suggests for commonly used hyperparameters, the solution is highly concentrated around the max reward mode(s).
We in Figure a didactic experiment that verifies this theory. At low regularization strength (3), the optimized policy TT- mode collapses onto the highest reward action. At high 3, policy achieves better (still not perfect) coverage over high-rewart answers, the cost of having more mass on low reward actions (more details and results in Appendix C.I).
4.2 WITH EQUAL REWARDS, SOLUTION NEVER PREFERS LOWER-SUPPORT SAMPLES
We now consider the case where the correct solutions all have equal reward. This is a standard set-up for RL with verifiable reward (e.g- math), where correct answer is given a reward I, and incorrect answers given
Rewards
Tret prob
fy prob
15x
Remark 4.3. For any two samples with the same reward, R(y)) = R(y2). their probability ratio is:
0.30
Gz(U1) G3(12)
Tny(U1)
0.02
0.000
0.063
(9)
Tre (32)
Action Index
Action index
Action Index
Figure 4: With equal rewards, RL does not change answers' relative probs.
In words, the correct answers' probability ratio in the optimal solution is simply their probability ratio in the reference distribution Tee- This ratio is independent of the regularization strength 3. In other words, by construction, Kl.rregularized RI. with equal rewards never promotes a low-support answer." Figure demonstrates this point empirically: the final policy never favours the equally correct low-support mode (additional results in Appendix C. This is not an issue with exploration; we will see in the subsequent section that with a small change in reward one can optimize for a distribution that equally weights even prefers the loppr-support solution.
Main Takeaway
RL with any KL-regularization does not increase the relative probability of lower-support samples to high-support ones, long their rewards the same. Lowering the KL regularization strength 3 has no effect on up- ppightig luw-support samples in the optimal solution.
We additionally corroborate Appendix C.4 that in practice for LLMs, the shape of the reference dis- tribution and reward function does result in highly skewed target distribution G3, per Proposition 4.1, Remark 4.2, and Remark 4.3.
4.3 FOR UNEQUAL REWARDS and SUPPORTS, REGULARIZATION STRENGTH DETERMINES MODE COVERAGE
When two trajectories have different rewards and different probabilities under the reference policy, unique setting of 3 will induce the two to have the same probability in the target distribution.
Remark 4.4. Two samples have the same probability the target distribution if.
R(y2) - R(y1) = 3( log TTy(yn) - log Twr())2))
(10)
This condition allows to predict, given only the reward and reference policy, when two samples will have the same probabilities the solution to the RL problem. As an example, we know Figure that the two high-reward modes have rewards 0.75 and 1.0, and reference policy probabilities of log (y1) x -4.05 and log Tier x -.95, respectively. This allows us to predict the setting of 3 which will "flip" the target distribution's preference from the high-support mode to the low-support mode to be (1 0.75) (--.5 + 5.95) x 0.132. Indeed, we see Figure 2 for the reverse KL case, the preference between the modes switch
move from 3 0.15 is the true role of the regularization coefficient 3: it is knob that decides between picking higher rewarding, low-support solutions, lower rewarding, high-support solutions.
0.10. This
This observation for both and forward-KR KT. regularized RI..

---

Published as a conference paper at ICLR 2026
5 DIRECTLY OPTIMIZING A REWARD MULTIMODAL TARGET
Having identified the various failure cases the Kl-regulariizd RI. objective (Section 4), and the role of regularization in balancing reward differences (Section 4.3), now turn to the question:
Can we construct an objective that, when optimized, naturally gives rise to reward multimodal target distribution?
Indeed, Remark 4.4 already provides the equality condition required to achieve this. We derive simple procedure which will ensure we are optimizing for solution that puts equal probabilities on all high-quality samples (per Defniiion 3.5), using the augmented reward function,
R(y))
if R(y) T
R(y)
(11)
R(=) 3 (log mer(== - log Teer(y)) if R(y) T,
where T S. maxy B(y) is some threshold for "goodness", and is fixed "anchor" sample chosen from the set of high-qualitg samples. We can pick it to be = = arg max, TTr(y) where R(v) 2 т. Because we are choosing the "anchor" be from a high-reward mode, we colloquially refer this as "mode anchoring". and the method as Mode Anchored Reward Augmentation (MARA). See Algorithm for pseudocode with minimal changes (an alternative that augments reward and ner outlined in Algorithm 2, which is equivalent to Alg.1 when using reverse KI. regularization).
Algorithm I Mode Anchored Reward Augmentation (MARA), within sampled batch. Changes from a standard RL algorithm are blue.
1: Given: initial policy Te, reference distribution Tief, reward function i, regularization coefficient
B, threshold of good answers 7 R. S max, R(y), and trajectory batch {31 ne-
2: Pick anchor trajectory: z = arg maXy, (r)( s.t. R(y) 2
3: for each y; in batch do
4: if R(y)) 2 then
5: Augment: F; = R(z) + 3 log )el(() - log (vr(
6: else
7: Keep same: F; = R(v:)
8: end if
9: end for
Optimize policy parameters 0 using augmented rewards (Fr)
Vanilla RL
MARA
Reward Function
8 0.01
1.0
$ 0.01
1=10
Reverse
0.050
0.04
0.066
KL
0.02
0.008
0.006
0.0
0,000
0.00
0.000
0.000
Reference Policy
0.04
Forward
0,050
0.166
0.02
KL
0.008
0.00
Figure 5: MARA stays close to the reference policy in low-reward areas, and puts high, uniform mass over all high-reward areas.
Intuitively, the augmented reward function constructs a new target distribution with uniform high density over regions of high reward, and stays close to the reference Tee in regions of low reward (see Remark B.3 for an analysis of the shape of the MARA target distribution). We see the Figure 5 that vanilla KL-regularized RL result in policy that heavily favours the left (on-support) mode, regardless of the choice of 3 or KL. On the other hand, using MARA results in solutions that put equal high all high quality samples, for both KLs. Note that in cases where the reward function range known, one can directly set threshold r as a constant. If not, one can set 7 on a per-batch basis by e.g. taking upper percentile of sampled rewards (as do below for non-verifiable I. I.LM tasks).
7

---

Published as a conference paper at ICLR 2026
6 EMPIRICAL VALIDATIONS
We evaluate MARA as a drop-in method in a variety of testnirainin tasks. While our theory has mainly been about the final optimal solution RL achieves, empirically investigate whether training, even if stopped early, can still benefit from a more diverse global optimum. To do this, we evaluate MARA in (i) verifiable LLM task with multiple answers, (ii) non-verifialle task with reward models, and (iii) chemical language model task for drug discovery, where mode collapse is detrimental.
6.1 VERIFIABLE 1-2 TASK FOR LLM
We train an LM (Qwen2.. 5-3B) to generate uniform random integers that are either 1 or 2. It gets reward of 1.0 for correct (producing "1" or "2" XML), and 0.0 otherwise (details Appendix C.2). Most runs are able to optimize the reward well and achieve a reward of approximately 1 (Figure 6a, right). Figure 6a (left) shows the number correctly formatted I's the LM generates over the course of training. We see that for naive KL regularization (grey), across a range of 3's and seeds, all but one run collapse into generating only a single answer as result of RL, and most collapse into generating I's, which has higher likelihood under the base policy. MARA (blue), on the other hand, is able to preserve the diversity in the correct answers, with many runs learning to generate I's and 2's with near uniform probability, while still learning to generate with the correct format (Figure 6a, middle). Further, the Pareto front of model checkpoints at different points in training shows that for both reverse and forward KL regularization, MARA is able match vanilla training in terms of correctness, while exceeding vanilla training in terms of generation diversity.
One Counts (binned)
Valid Answers Diversity
Rewards
500
0.6
0.8
Estrpy
0.6
300
2000
300
400
500
100
200
300
400
50c
100
200
300
400
500
Stop
Stop
Step
Modi Archoored. KL. coul - 2001
Mode Archores, KL coul 0.02
Varilla, KL. cowl 0.001
Varilla, KL. coel 0.02
Mode Archoredd KL coof 0.01
Mode Archoreo, coof 0.03
Varila, coof 0.01
Varilla, KL coof 0.03
(a) Valid answer entropy and rewards
Reverse KL Regularized
Forward KL Regularized
1.0
1.0
Rewwaris
Rewards
Anchored KL-0.01
Anchored, X1-0.22
Mode Anchored. KL-0.0
0.8
Mode Anchored, KL=0.1
0.7
Vanila KL-0.501
Vanila, KL-0.01
0.7
0.6
Vanila,
Vanill, KL-0.1
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.4
0.5
Valid Resporse Entropy
Valid Response Entropy
(b) Pareto front of reward (quality) and entropy (diversity)
Figure 6: Performance on verifiable task with multiple solutions, against both reverse & forward KL.,
MARA (blue) compared against baseline GRPO (grey) at different fts (KL coef).
6.2 CREATIVE QUESTION ANSWERING FOR CHAT LLM
We test MARA in a non-verifiable alignment task. We train Qwen3-1. on subset of WildChat text (Zhao et al., 2024), using parametric reward model (Skywork-Rewd rd-v2-Qwen 3-BB). We evaluate the model on a curated test set (Zhang et al., 2025) and report both the training reward (In dist Reward), and test set reward from different reward model (Out dist Reward, using
2). We also report diversity metrics terms of n-grams (Ngrams), semantic embeddings (Semantic Div), and "distinct functional classes" (Mean Distinct), See Appendix C.3 for more details. Here, MARA used as a drop-in replacement in an RLOO style algorithm (Ranganath, 2017; Kool et al., 2020; Ahmadian al., 2024).
8

---

Published as a conference paper at ICLR 2026
As the reward function range is not known beforehand, percentile reward of each batch (we ablate this in Appendix Table 5 to find that lowering the percentile decreases performance slightly. but nevertheless remains competitive). We compare against regular RL training, and a number of rewarding the unlikely (Un1 i iikel, He et al. 2025), best-of-N training (BoN Trai training, Tang et al. 2025), and weight ensembling (Ensemblle, Dang et al. 2025). We see MARA out-performs all baselines in terms of out-of-distribttion rewards, and all-but-one diversity metries (Table 1).
per-batch basis the 90th
baselines including entropy regularization (Ent ropy),
In-dist.
Out-dist. Reward (t)
Ngrams EAD (t)
Semantic Div (t)
Mean
Model
Reward (1)
Distinet (t)
Base Model
10,94
1.166 #0.06
0.413 #0005
0.220 =0.009
4.01 =0.254
GRPO
14.80
1.317 10.102
0.497 #0.014
0.193 =0.009
3,96
=0.249
RLOO
15.56
1,280 10.100
0,514
10.014
0.192
=0.008
3.88
=0.243
Entropy
1.44
0.786 +0073
0.267
#0.009
0.228
=0.009
3.45
=0.193
Unlikely
10.04
1,381 #0.114
0.532
#0.015
0.191
=0.008
4.24 =0.239
BoN Training
16.88
0.596 10.055
0.541 10.010
0.162
0.008
2.29
10.173
Ensemble
1.143 10.086
0.438
10.014
0.211 00010
4.19
..02669
MARA (rev)
15.42
1,451 +0.103
0.543
=0.014
0.186
=0.008
4.14
-0.233
MARA (fwd)
15.33
1.604 113113
0.568 00.02
0.193 =0.009
4.62 00.258
Table Performance on non-verifiable creative task. Mean + bootstrap SEM.
6.3 DRUG DISCOVERY WITH CHEMICAL LANGUAGE MODELS
Finally, we apply MARA to distinctively different domain where diversity and quality is crucial: drug discovery. Chemical language models (CLMs) have seen success in discovering molecules in clinical trials. We adapt two realistic reward functions from Guo et al. (2025b): SYNTH and ALL-AMIDE that jointly reward binding potency and synthesizability. The core CLM optimization problem is also a regularized RL problem: maximize expected reward, while staying close pretrained "prior" model to ensure chemical validity. Unlike the traditional RL setting. CLMs are evaluated based on their ability to generate unique molecules (Yield) given fixed number of reward function evaluations (which are expensive simulations and/or experiments), making diversity an essential quality for any performant CLMs. The REINVENT algorithm (Olivecrona al., 2017; Guo & Schwaller, 2024b) is a state-of-the-an RL-based method on standard benchmarks (Gao et al., 2022). It optimizes the following objective,
[() log (y) ( (log merr o R(w))
y no-
(12)
which is equivalent to KL--regularized reward maximization. We apply MARA drop-in replace- ment to its rewards. Additional evaluation details are in Appendix C.5.
Table 2 shows MARA consistently results in higher average Wield (number of unique high- reward molecules discovered), and lower 051 00 (efficiency in finding high-reward molecules, measured by reward function calls). The screening level (Screen) is the reward threshold above which we accept discovered molecules. Setting MARA's equal to the screening level always results in the highest yield, consistent with MARA's target distribution having uniform density in areas where R(y) > т. Going further, we also assess "global" diversity (which MARA does not explicitly optimize for) in terms of IntDiv Dv and #Circles. Both define more macroscopic differences based on molecular sub-st sub-structures. We find MARA is competitive with the baseline here. Overall, we see MARA further boosts REINVENT's optimization efficiency, while maintaining diversity.
7 CONCLUSION
In this work, we provide an in-depth understanding of the KL-regularized RL objective, particularly in terms of its diversity. We summarize the main take-aways below.
Studying the divergence being implicitly minimized between the policy and target distribu- tion, D G), is more meaningful for understanding optimization behaviour than studying the divergence of the regularizer, D (To Teet), alone.

---

Published as a conference paper at ICLR 2026
Screen
Algorithm
Yield (1)
OB100 ())
IntDiv1 (t)
Circles (1)
0.80
REINVENT
6569 186
1042 + 66
0.766 + 0.011
67 + 3
MARA (7=0.80)
6834 78
1015 55
0.761 0.009
59 =
MARA (7=0.85)
6584 + 231
1042 + 66
0.761 + 0.008
72 + 6
0.85
REINVENT
1614 + 407
4114 109
0.701 + 0.018
7 1
MARA (7=0.80)
1796 210
3654 = 272
0716:00011
6 1
MARA (7=0.85)
2196 394
4010 + 297
= 0
7
SYNTH task
Screen
Algorithm
Yield (1)
OB100 (L)
IntDiv1 (1)
Circles (t)
0.80
REINVENT
5433 184
1427 +
0.768 + 0.012
35 +
MARA (7=0.80)
5635 + 249
1407 + 123
0.766 = 0.008
36 + 3
MARA (T-0.85)
5502 309
1426 +
0.769 + 0.006
+
0.85
REINVENT
1098 + 88
4360 +: 257
0.721 + 0.016
8 + 1
MARA (r=0.80)
1235 130
3943 + 303
0.733 + 0.009
+
MARA (7=0.85)
1438 = 126
4230 + 401
0.725 = 0.008
8 + 1
(b) ALL-AMIDE task
Table Results for different tasks and screening levels (Screen, higher meaning more strict) for two challenging drug discovery tasks. Error bars (#) denote standard deviation over 5 independent seeds. Bold indicates if the performance is statistically significantly better than the alternative method for that screening level (one-sided student's t-test, p 0.05).
The regularizer, B, and reward function together define the target distribution which ma optimizes towards. Forward and reverse KL regularizers define different target distributions.
For common hyperparameters and reward functions used in practice, the target distribution is often defined be unimodal. Thus, perfectly solving the regularized RL objective yields non-diverse optimal policy distributions.
This diversity loss can be fixed by instead defining multimodal target distributions, such as through dynamically augmenting the reward via MARA.
There are a number of exciting future directions to improve MARA. For one, MARA requires setting (the reward threshold above which the target distribution puts uniform probability mass). The choice of T is obvious in some settings (Section 6.1 & 6.3), but a hyperparameter in settings with unbounded reward functions and unknown priori thresholds (Section 6.2). We found that setting batch-specific- helps, but better approaches may be possible. It should be noted that setting an overly high harmless: if no samples meet the threshold, the MARA mechanism does not kick in, and learning reduces to standard RL.
Further, MARA shapes the target distribution to have multimodality, but does not guarantce faster convergence to the target. Having general algorithms to more efficiently reach the target, or to guarantee distributional properties (e.g. multimodality)
aub-optimality. would be of general importance and complement MARA's target distribution. Further, MARA introduces one such target distribution which places uniform mass at regions where R(y) 2 т. This may not be an optimal choice for all tasks, and considering alternative mode-preserving target distributions can be interesting future work. All all, we emphasize that regularized RI.ass as commonly used for generative model training- is inherently distribution matching problem and should be viewed such. Rather than relying on intuitions (c.g. about regularizers), we should directly specify distributions with properties wish to have as the target of policy optimization.
ACKNOWLEDGEMENTS
This work is supported by ONR MURI 4N00014-22-1-2773,, ONR #N00014-21-1-2758, and the National Science Foundation under NSF Award 1922658. This work partly supported by
10

---

Published as a conference paper at ICLR 2026
the NIH/NHLBI Award R01HL148248, NSF Award 1922658 NRT-HDR: FUTURE Foundations, Translation, and Responsibility for Data Science, NSF CAREER Award 2145542. ONR N00014-23- 1-2634, NIH RO1CA296388, NSF 2404476, Optum, and Apple. Additional support by Fellowship from the Columbia Center of AI Technology. This work was also supported by ITTP with grant funded by the MSIT of the Republic of Korea in connection with the Global Al Frontier Lab International Collaborative Research. A.GXC. is supported by the Natural Sciences and Engineering Research Council of Canada (NSERC), PGSD3-559278- 202222221.. Cette recherche été financée par le Conseil de recherches en sciences naturelles et en génie du Canada (CRSNG), PGSD3-559278 2022. J.G. is supported by the Natural Sciences and Engineering Research Council of Canada (NSERC), PGSD-521 528389. We are particularly thankful for fruitful discussions with Aram-Alexane Alexandre Pooladian and Mark Goldstein the early stages of this project that helped shape the theoretical foundations without which this project would not be possible. We additionally thank Anirudh Buvanesh and Manya Wadhwa for helpful discussion and feedback.
provided
REPRODUCIBILITY STATEMENT
We use open-source, publicly available libraries for all experimental code. Didactic experiments constructed in PyTorch (Paszke al., 2019). Reinforcement learning on LLM training is done using the nano bano-aha-ae -oomett (Kazemnejad et al., 2025) and verl (https: // / iiggttuuu thub. com/ vol rolcengine) /vervrer github repos. Chemical language model experiments use the official saturn github repo (Guo & Schwaller, 2024b). We provide detailed experimental information in Appendix C. Pseudo-code is provided in Algorithm and Algorithm
REFERENCES
Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from self- generated mistakes. In The twelfth international conference on learning representations, 2024.
Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzich Fadace, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. Back to basics: Revisiting reinforce style optimization for learning from human feedback in lims. arXiv preprint arXiv:2402. 14740, 2024.
Amr Alhossary, Stephanus Daniel Handoko, Yuguang Mu, and Chee- Keong Kwoh. Fast, accurate, and reliable molecular docking with quickvina 2. Bioinformatics, 31 31(13):2214-2211 2015.
Eser Aygün, Anastasiya Belyaeva, Gheorghe Comanici. Mare Coram, Hao Cui, Jake Garrison, Renee Johnston Anton Kast, Cory Y McLean, Peter Norgaard, Zahra Shamsi, et al. An ai system to help scientists write expert-level empirical software. arXiv preprint 2xXx:22509,06003, 2025.
Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandricllo. A general theoretical paradigm understand learning from human preferences. In International Conference on Artificial Intelligence and Statistics, pp. 4447-4455. PMLR, 2024.
G Richard Bickerton, Gaia V Paolini, Jérémy Besnard, Sorel Muresan, and Andrew I. Hopkins. Quantifying the chemical beauty of drugs. Nature chemistry, 4(2):90 98, 2012.
Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, volume 4. Springer, 2006.
Esben Jannik Bjerrum. Smiles enumeration as data augmentation for neural network modeling of molecules. arXiv preprint arXiv: 770.007076 2017.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians.
Journal of the American statistical Association. 112(518):559-87) 2017.
Dean G Brown and Jonas Bostrom. Analysis of past and present synthetic methodologies medicinal chemistry: where have all the new reactions gone? miniperspective. Journal of medicinal chemistry, 59(10):4443-455 4458, 2016.
11

---

Published as a conference paper at ICLR 2026
Pablo Samuel Castro, Nenad Tomasev, Ankit Anand, Navodita Sharma, Rishika Mohanta, Apama Dev, Kuba Perlin, Siddhant Jain, Kyle Levin, Noémi Éltetō, et al. Discovering symbolic cognitive models from human and animal behavior. bioRxiv, pp. 2025-0222222220 2025.
Binghong Chen, Chengtao Li, Hanjun Dai, and Le Song. Retro*: learning retrosynthetic planning with neural guided a search. International conference on machine learning. pp. 1608-1616. PMLR, 2020.
Daixuan Cheng. Shaohan Huang, Xuekai Zhu, Bo Dai, Wayne Xin Zhao, Zhenliang Zhang, and Furu Wei. Reasoning with exploration: An entropy perspective. arXiv preprint arXiv:2306.14755 2025.
John Joon Young Chung, Vishakh Padmakumar, Melissa Roemmele, Yuqian Sun, and Max Kreminski. Modifying large language model post-training for diverse creative writing. arXiv preprint arXiv:2503.17126 2025.
Ganqu Cui, Yuchen Zhang, Jiacheng Chen, Lifan Yuan, Zhi Wang, Yuxin Zuo, Haozhan Li, Yuchen Fan, Huayu Chen, Weize Chen, et al. The entropy mechanism of reinforcement learning for reasoning language models. arXiv preprint arXiv:2505 22617, 2025.
Xingyu Dang, Christina Back, Kaiyue Wen, Zico Kolter, and Aditi Raghunathan. Weight ensembling improves reasoning in language models. arXiv preprint urXiv:2504.10477 2025.
Hanze Dong, Wei Xiong, Decpanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang. Kashun Shum, and Tong Zhang. Raft: Reward ranked finetuning for generative foundation model alignment. arXiv preprint 0rXiv:2304.06767, 2023.
Yuanqi Du, Arian R Jamasb, Jeff Guo, Tianfan Fu, Charles Harris, Yingheng Wang, Chenru Duan, Pietro Liò, Philippe Schwaller, and Tom Blundell. Machine learning- aiidd generative molecular design. Nature Machine Intelligence, pp. 1-16, 2024.
Wenhao Gao, Tianfan Fu, Jimeng Sun, and Connor Coley Sample efficiency matters: benchmark for practical molecular optimization. Advances in neural information processing systems, 35: 21342-21357, 2022.
I. M. Gelfand and S. V. Fomin. Calculus of Variations. Selected Russian Publications in the Mathematical Sciences. Prentice-Hall Englewood Cliffs, N.J., 1963. Revised English edition.
Dongyoung Go, Tomasz Korbak, Germán Kruszewski, Jos Rozen, Nahyeon Ryu, and Mare Dymet- man. Aligning language models with preferences through f-divergence minimization. arXiv preprint ar arXiv:2302.082/55 2023.
Bean-Bastien Grill, Florent Altché, Yunhao Tang, Thomas Hubert, Michal Valko, Ioannis Antonoglou, and Rémi Munos. Monte-carlo tree search as regularized policy optimization. In International PMLR, 2020.
Conference on Machine Learning, pp.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al. Deepseek-rl incentivizes reasoning in llms through reinforce- ment learning. Nature, (45(8081):633366) 2025a.
Jeff Guo and Philippe Schwaller. Augmented memory: Sample-efficient generative molecular design with reinforcement learning. JACS Au, 2024a.
Jeff Guo and Philippe Schwaller. Saturn: Sample-efficient generative molecular design using memory manipulation. arXiv preprint arXXiv2405.17006 2024b.
Jeff Guo, Víctor Sabanza- Gil, Zlatko Jončev, Jeremy S Luterbacher, and Philippe Schwaller. Gen- erative molecular design with steerable and granular synthesizability control. arXiv preprint arXiv:2505.088744 2025b.
Anthony GX-Chen, Dongyan Lin, Mandana Samiei, Doina Precup, Blake Aaron Richards, Rob Fergus, and Kenneth Marino. Language agents mirror human causal reasoning biases. how can we help them think like scientists? In Second Conference Language Modeling, 2025.
12

---

Published as a conference paper at ICLR 2026
Tuomas Haarmoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning, pp. 1352-1361. PMLR. 2017.
Jubayer Ibn Hamid, Ifdita Hasan Omey, Ellen Xu, Chelsea Finn, and Dorsa Sadigh. Polychromic objectives for reinforcement learning. arXiv preprint arXiv:2509.25444 2025.
Andre He, Daniel Fried, and Sean Welleck. Rewarding the unlikely: Lifting grpo beyond distribution sharpening. arXiv preprint arXiv:2506.0255 02355, 2025.
Aspen K Hopkins, Alex Renda, and Michael Carbin, Can LLMs generate random numbers? evaluat- ing LLM sampling in controlled domains. In ICMI. 2023 Workshop: Sampling and Optimization in Discrete Space, 2023. URL https: htt::///pppreiiiieeemmm wwtpereeiewweeeeeeeeeeeeeeeeeeee net dorumiid-Vhhhn3 VhhhhhK9
Edward J Hu, Moksh Jain, Eric Elmoznino. Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, and Nikolay Malkin. Amortizing intractable inference in large language models. arXiv preprint arXiv:23 2023.
Audrey Huang, Adam Block, Dylan J Foster, Dhruv Rohatgi, Cyril Zhang, Max Simchowitz, Jor- dan T Ash, and Akshay Krishnamurthy, Seff-improvement in language models: The sharpening mechanism. arXiv preprint arXiv: 2024.
Shiyu Huang, Hang Su, Jun Zhu, and Ting Chen. Svqn: Sequential variational soft q-learning networks. In International Conference on Learning Representations, 2019.
Mete Ismayilzada, Antonio Laverghetta Jr, Simone Luchini, Reet Patel, Antoine Bosselut, Lonneke van der Plas, and Roger Beaty. Creative preference optimization. arXiv preprint 144444 2025.
Yash Jhaveri, Harley Wiltzer, Patrick Shafto, Mare G Bellemare, and David Meger. Conver- gence theorems for entropy-regularized and distributional reinforcement learning. arXiv preprint SoXiv:2510008526, 2025.
Michael Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183-233, 1999.
Amirhossein Kazemnejad, Milad Aghajohari, Alessandro Sordoni, Aaron Courville, and Siva Reddy, Nano aha! moment: Single file "rl for 11m" library. https: htt/://tttuu..omMMMii//ii/ com/McGill nano-aha aha-moment, 2025. GitHub repository.
Peter W Kenny. Hydrogen-bond donors in drug design. Journal of medicinal chemistry, 65(21): 14261-14275, 2022.
Muhammad Khalifa, Hady Elsahar, and Mare Dymetman. A distributional approach to controlled text generation. arXiv preprint arXiv:2012. 1165, 2020.
Robert Kirk, Ishita Mediratta, Christoforos Nalmpantis, Jelena Luketina, Eric Hambro, Edward Grefenstette, and Roberta Raileanu. Understanding the effects of rlhf on 1lm generalisation and diversity. arXiv preprint urXiv:2310.06455 2023.
Wouter Kool, Herke van Hoof, and Max Welling. Estimating gradients for discrete random variables by sampling without replacement. arXiv preprini urXiv:2002.06043, 2020.
Tomasz Korbak, Hady Elsahar, Germán Kruszewski, and Marc Dymetman. On reinforcement learning and distribution matching for fine-tuning language models with no catastrophic forgetting. Advances in Neural Information Processing Systems, 35: 15:16203-16200 2022.
Solomon Kullback and Richard A Leibler. On information and sufficiency. The annals of mathemati- cal statistics, 22(1):79-88) 1951.
Oh Joon Kwon, Daiki E Matsunaga, and Kee-Eung Kim. Gdpo: Learning to directly align language models with diversity using gflownets. arXiv preprint arXiv:2410. /5096, 2024.
13

---

Published as a conference paper at ICLR 2026
Jack Lanchantin, Angelica Chen, Shehzaad Dhuliawala, Ping Yu, Jason Weston, Sainbayar Sukhbaatar, and Ili Kulikov. Diverse preference optimization. arXiv preprint 1eXi::5501.18011 2025.
Tianjian Li, Yiming Zhang, Ping Yu, Swarnadeep Saha, Daniel Khashabi, Jason Weston, Jack Lan- chantin, and Tianlu Wang. Jointly reinforcing diversity and quality in language model generations. arXiv preprint vrXiv:2509.02534 2025.
Siyang Liu, Sahand Sabour, Yinhe Zheng, Pei Ke, Xiaoyan Zhu, and Minlie Huang. Rethinking and refining the distinct metric. arXiv preprint arXiv:2202.13587, 2022.
Mark F Mabanglo, Keith S Wong, Marim M Barghash, Elisa Leung, Stephanie HW Chuang, Afshan Ardalan, Emily M Majaesic, Cassandra J Wong, Shen Zhang. Henk Lang, et al. Potent clpp agonists with anticancer properties bind with improved structural complementarity and alter the mitochondrial n-terminome. Structure, 3 (1()):885-2000 2023.
Krzysztof Maziarz, Austin Tripp. Guoqing Liu, Megan Stanley, Shufang Xie, Piotr Gaiński, Philipp Seidl, and Marwin HS Segler. Re-evaluating retrosynthesis algorithms with syntheseus. Faraday Discussions, 256:568-586, 2025.
Kevin P Murphy. Machine learning: probabilistic perspective. MIT press, 2012.
Christian Naesseth, Fredrik Lindsten, and David Blei. Markovian score climbing: Variational inference with kl (pll q). Advances in Neural Information Processing Systems, 1.1115999-55510 2020.
Alexander Novikov, Ngån Vū, Marvin Eisenberger. Emilien Dupont, Po-Sen Huang, Adam Zsolt Wag- ner, Sergey Shirobokov, Borislav Kozlovskii, Francisco JR Ruiz, Abbas Mehrabian, al. Alphae- volve: A coding agent for scientific and algorithmic discovery. arXiv preprint arXiv:2506. 13131, 2025.
Marcus Olivecrona, Thomas Blaschke, Ola Engkvist, and Hongming Chen. Molecular de-novo design through deep reinforcement learning. Journal of cheminformaties, 9:1-14, 2017.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback, Advances neural information processing systems, 35:27730- 27744, 2022.
Laura O' Mahony, Leo Grinsztajn, Hailey Schoelkopf, and Stella Biderman. Attributing mode collapse in the fine-tuning of large language models. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, volume 2, 2024.
Sven Michael Papidocha, Andreas Burger, Varinia Bernales, and Alán Aspuru-Guzik, The elephant in the lab: Synthesizabilityy in generative molel-molecule design. ChemRxiv, 2025.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, perbrerrform deep learning library. Advances in neural information processing systems, 32, 2019.
Daniil Polykovskiy, Alexander Zhebrak, Benjamin Sanchez-Lengeling, Sergey Golovanov, Oktai Tatanov, Stanislav Belyaev, Rauf Kurbanov, Aleksey Artamonov, Vladimir Aladinskiy, Mark Veselov, et al. Molecular sets (moses): benchmarking platform for molecular generation models. Frontiers in pharmacology, 11: 11:565644, 2020.
Martin L Puterman. Markov decision processes: Discrete stochastic dynamic programming, 1994.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning, Stefano Ermon, and Chelsea Finn. Direct Preference Optimization: Your Language Model is Secretly a Re- ward Model. Advances in Neural Information Processing Systems, 1-53728-537411 De- nips cc httpp/pperffllepppppeee/eeepp//e htw.0006556677744400002066-e--a-
cember 2023. URL
/ash/ html.
14

---

Published as a conference paper at ICLR 2026
Rajesh Ranganath. Black Box variational inference: Scalable, generic Bayesian computation and its applications. PhD thesis, Princeton University, 2017.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Artificial intelligence and statistics, pp. 814 822. PMLR, 2014.
Christian P Robert, George Casella, and George Casella. Monte Carlo statistical methods, volume 2. Springer, 1999.
Bernardino Romera-Paredee Mohammadamin Barekatain, Alexander Novikov, Matej Balog, M Pawan Kumar, Emilien Dupont, Francisco JR Ruiz, Jordan S Ellenberg. Pengming Wang. Omar Fawzi, et al. Mathematical discoveries from program search with large language models. Nature, (25(7995):468-444 2024.
Mikolaj Sacha, Mikotaj Blaz, Piotr Byrski, Pawel Dabrowski-Tumankk Mikolaj Chrominski, Rafal Loska, Pawel SooddarzyykPruszynnee and Stanislaw Jastrzehski. Molecule edit graph atten- tion network: modeling chemical reactions sequences of graph edits. Journal of Chemical Information and Modeling, (10n)ng.1111133333330 61(7):3273-3284, 2021.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseckmath: Pushing the limits of mathemat- ical reasoning open language models. arXiv preprint arXiv:2402.03330, 2024.
Alexander Shypula, Shuo Li, Botong Zhang, Vishakh Padmakumar, Kayo Yin, and Osbert Bastani. Evaluating the diversity and quality of Ilm generated content. arXiv preprint arXiv:2504.1252.. 2025.
Raghav Singhal, Zachary Horvitz, Ryan Teehan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath. A general framework for inference-time scaling and steering of diffusion models. arXiv preprint arXiv:2501.06845 2025.
Marta Skreta, Tara Akhound-Sadeghe Viktor Ohanesian, Roberto Bondesan, Alán Sspuru-Guzit Arnaud Doucet, Rob Brekelmans, Alexander Tong, and Kirill Neklyudov. Feynman kac correctors in diffusion: Annealing, guidance, and product of experts. arXiv preprint arXiv:2503.02819, 2025.
Yuda Song, Julia Kempe, and Remi Munos. Outcome-basee exploration for Ilm reasoning. arXiv preprint arXiv:2509.069411 2025.
Megan Stanley and Marwin Segler. Fake until you make it? generative de novo design and virtual screening of synthesizable molecules. Current Opinion in Structural Biology, 82:102655 2023,
Anja Surina, Amin Mansouri, Lars Quaedvlieg, Amal Seddas, Maryna Viazovska, Emmanuel Abbe, and Caglar Gulcehre. Algorithm discovery with 11ms: Evolutionary search meets reinforcement learning. arXiv preprint arXiv:2504.05 2025.
Shidi Tang, Ji Ding, Xiangyu Zhu, Zheng Wang, Haitao Zhao, and Jiansheng Wu. Vina-gpu 2.1: towards further optimizing docking speed and precision autodock vina and its derivatives. IEEE/ACM Transactions on Computational Biology and Bioinformatics, 2024.
Yunhao Tang and Rémi Munos. On a few pitfalls kl divergence gradient estimation for arXiv preprint 2.Xx:2506009477 2025.
Yunhao Tang, Kunhao Zheng, Gabriel Synnaeve, and Rémi Munos. Optimizing language models for inference time objectives using reinforcement learning. arXiv preprint arXhv:2503.19595 2025.
Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry P Vetrov. Generative flow networks reguopy-re rl. In International Conference Artificial Intelligence and Statistics, pp. 4213 1221. PMLR, 2024.
Oleg Trott and Arthur J Olson. Autodock vina: improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading. Journal of computational chemistry, 31(2):455-661 2010.
15

---

Published as a conference paper at ICLR 2026
Chaoqi Wang. Yibo Jiang, Chenghao Yang, Han Liu, and Yuxin Chen. Beyond reverse kl: Gen- eralizing direct preference optimization with diverse divergence constraints. arXiv preprint 16244,, 2023.
Shenzhi Wang, Le Yu, Chang Gao, Chujie Zheng, Shixuan Liu, Rui Lu, Kai Dang, Xionghui Chen, Jianxin Yang, Zhenru Zhang, et al. Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforeement learning for Ilm reasoning. arXiv preprint arXiv:2506.01939, 2025.
David Weininger. Smiles, a chemical language and information system. 1. introduction to methodol- ogy and encoding rules. Journal of chemical information and computer sciences, 28(1):31-36 1988.
Peter West and Christopher Potts. Base models beat aligned models at randomness and creativity. arXiv preprint arXiv:2505.00047, 2025.
Luhuan Wu, Brian Trippe. Christian Naesseth, David Blei, and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36:31372-3111 372231143, 2023.
Yutong Xie, Ziqiao Xu, Jiaqi Ma, and Qiaozhu Mei. How much space has been explored? measur- ing the chemical space covered by databases and machine-generater molecules. In Proc. 11th International Conference on Learning Representations, 2023.
Wei Xiong, Jiarui Yao, Yuhui Xu, Bo Pang, Lei Wang, Doyen Sahoo, Junnan Li, Nan Jiang, Tong Zhang, Caiming Xiong, and Hanze Dong. A minimalist approach to Ilm reasoning: from rejection sampling to reinforce. arXiv preprint arXiv:2504.11444 2025.
Wanqiao Xu, Shi Dong, Dilip Arumugam, and Benjamin Van Roy. Shattering the agent-enviromment interface for fine-tuning inclusive language models. arXiv preprint arXiv:2305. 1455, 2023.
Chenghao Yang and Ari Holtzman. How alignment shrinks the generative horizon. arXiv preprint arXiv:2506. 17777,, 2025.
Jian Yao, Ran Cheng, Xingyu Wu, Jibin Wu, and Kay Chen Tan. Diversity policy optimization for large language model reasoning. arXiv preprint arXiv:2505.2 233.3333,, 2025.
Longfei Yun, Chenyang An, Zilong Wang, Letian Peng, and Jingbo Shang. The price of format: Diversity collapse in llms. arXiv preprint arXiv:2505. 18949, 2025.
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems, 15:15576-15488,, 2022.
Lily H Zhang and Rajesh Ranganath. Preference learning made easy: Everything should be understood through win rate. arXiv preprini arXiv:2502. 10505, 2025.
Yiming Zhang, Harshita Diddee, Susan Holm, Hanchen Liu, Xinyue Liu, Vinay Samuel, Barry Wang, and Daphne Ippolito. Noveltybench: Evaluating language models for humanlike diversity. arXiv preprint 2rXiv:2504.05228, 2025.
Rosie Zhao, Alexandru Meterez, Sham Kakade, Cengiz Pehlevan, Samy Jelassi, and Eran Malach. Echo chamber: RI post-training amplifies behaviors learned in pretraining. arXiv preprint arXiv:2504.07912. 2025.
Wenting Zhao, Xiang Ren, Jack Hessel, Claire Cardie, Yejin Choi, and Yuntian Deng. Wildchat: Im chatGPT interaction logs in the wild. In The Twelfth International Conference on Learning Representations, 2024. URL https: /ppp//peeeeeeeeeeeeeeeeeeeeeeeee net
Brian Ziebart, Andrew L. Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pp. 1433-1438. Chicago, II., USA, 2008
16

---

Published as a conference paper at ICLR 2026
A RELATED WORK
Entropy collapse in RI. There is a growing line of empirical works observing RL. training collapses
the diversity in generation output of the resulting post-trained policy (Kirk et al., 2023; Huang et al., 2024; O' Mahony et al., 2024; Cui et al., 2025; Yang & Holtzman, 2025; Yun et al., 2025; Shypula et al., 2025; West & Potts, 2025; Zhao al., 2025; Dang et al., 2025; Song al., 2025), such as in formats (Zhao et al., 2025), random generation and creativity (West & Potts, 2025), as well as exploration and reasoning (Cui al., 2025; Dang al., 2025; Song et al., 2025). The observations have mostly been empirical.
A few attempts have been made to theoretically understand entropy collapse. Cui al. (2025) analyzes what per-step policy gradient (approximately) does to the entropy of a tabular softmax policy, and finds that entropy decreases if there is strong a strong positive correlation between the action probabilities and corresponding advantage values. Dang et al. (2025) analyzes special case of multi-an arm bandits with K equally good arms and bad arm, and finds that the optimal probabilities correspond to the re- re-normalized reference probabilities of just the good arms. We note this is a special case of our Remark 4.3.
Generally speaking, entropy preservation via regularization is often referred to "max entropy RI.", with numerous seminal and ongoing works (Ziebart et al., 2008; Haarnoja al., 2017; Huang et al., 2019). Our analysis contributes to further understanding of this setting by shedding light on the multimodality of the target distribution when doing entropy regularization.
Training for diversity A number of attempts have been made to empirically address diversity collapse. Wang et al. (2023) generalizes the DPO objective (Rafailov et al., 2023) from reverse- regularized to more general class of f-divergence regularizers, with the key motivation being that reverse-KL can be mode-sceking, therefore reduce diversity. We argue in this work that the full story is more nuanced and is better analyzed through the target distribution. Cui et al. (2025) proposes to directly regularize the update high-covariance tokens. Cheng et al. (2025) incorporates an entropy term in the advantage to encourage better reasoning. Wang et al. (2025) show that focusing gradient updates on a minority of high-entropy "forking" tokens can improve reasoning. He et al. (2025) proposes a rank-based "unlikeliness" rewarrd where more likely samples (under current policy) receives a larger multiplicative penalty to the reward. Similarly, Yao et al. (2025) uses token entropy to encourage diversity. Song et al. (2025) proposes count-based exploration bonus that more highly rewards less frequently seen outcomes (in previous samples), and Hamid et al. (2025) proposes a similar batch-wise reward. Xu et al. (2023) proposes more inclusive approach to preference alignment. Dang et al. (2025) found that combining weights of earlier and later checkpoints can improve pass@k k performance-one specific measure of diversity.
A number of works attempt to directly optimize for diversity. This relies on the existence of additional information that tells us if two samples are different and by how much. In this vein, diverse DPO Lanchantin et al. (2025) and variants (Chung et al., 2025; Ismayilzada et al., 2025) encourage diversity in preference learning by selecting diverse positives/negatties. Similarly related is Li et al. (2025), which an external model evaluate diversity (via semantic classifier) and use the diversity metric to modify the reward. Hamid et al. (2025) proposes to optimize a batch-level objective that is modified by diversity function. We do not require an external model evaluate diversity.
In the unregularized setting, Jhaveri et al. (2025) optimizes for an unregularized policy with specific distributional properties. More distantly, GFlowNets also provide diversity-seeking policies that specifically designed to sample proportionally to reward, albeit they use different algorithms than the KL-regularized policy gradient which is the most commonly used algorithm for LM post-training (Hu et al., 2023; Kwon et al., 2024; Tiapkin et al., 2024).
B MATHEMATICAL DERIVATIONS
B.1 TARGET DISTRIBUTION OF REVERSE-KE KI. REWARD MAXIMIZATION
Proof of Remark 3.1 We provide proof for the maximizer of the generalized reverse-KL and entropy regularized reward maximization objective,
Ja.n (TO) E [R(y)] 8 Dx ( (DKI(Fol|aaa (To]( + (TT).
(13)
17

---

Published as a conference paper at ICLR 2026
The solution arg max J3.n has the un-normalized form,
R(y)
GBm(()) & go.n(y) = Tree(H) utr exp 13+1)
(14)
Proof.
(T,n(TT) Bro(s) R(y) 3 log no(y)
log (v)) 77 log To (o(
(15)
R(y)
-B 17) (x)
log (v))
log कeat
(16)
B +11
++
R(y)
-(B + .1) (xx(s)
log no (v) - log (y) 379 exp 3+
(17)
= - + 17) (x)
log (v) - log (v) + (B + 17) log SB.1
(18)
= (B + 11 (TalG + (B + v)) log Ca,
(19)
where Sa.n = J n (y) dy. One can see that the above is maximized when DEL (Tell which occurs when कe Gp.
Intuitively, one can see the entropy regularizer as playing the role of "tempering" the reference distribution net (larger 7) drives Net become more uniform), while both and 17 lower the reward's effect on the target distribution. For the KL-only case (17 = 0), the target distribution becomes,
R(vy)
Gsly) x mnrrly) exp
(20)
which is the stated result in Remark 3.1. In the ontr case (B 0), the solution is,
Gv) exp R(y)
(21)
All in all, both coefficients play a role in parameterizing the shape of the optimal distribution for the regularized RL problem.
B.2 GRADIENT OF REVERSE- KL REWARD MAXIMIZATION
Proof of Remark 3.2 From Appendix B.I, we have the identity,
Jo([o) = DKL. (xoll))) log Ç
(22)
We can easily show that the gradient is,
Ve ( V DKz Ve log
(23)
Vo DK (Tol(G))
(24)
In other words, they are the same up to constant -3.
VoJa(To) = -3 V. Du (Tol)
(25)
B.3 TARGET DISTRIBUTION OF FORWARD-KL REWARD MAXIMIZATION
We are interested in finding the distribution mo Grod which maximizes.
(ooa(mo) = (x),() R(y)) - |3 D e. (Tree mo)
(26)
Note we can simplify the expression to only terms that depend on πT.
arg max JFud (TT) = arg max E.., [[(()]] - 3 Dx (Treff TT)
(27)
To
/
arg max We
(T( na (v) R(y)) + 3 mTr(y) log (v) dy + const.
(28)
18

---

Published as a conference paper at ICLR 2026
Remark B.1. Assuming the reward is finite and has maximum value Rmar- If any on-support answer(s) have R(y)= = Rmmcs E supp(Trer), the optimal distribution maximizing J will put zero mass outside of supp(Tmn)).
Proof. Let M be the set of on-support, max reward answers: M if R(y) = Rmax, and M S supp(Tee)) We can generically write any distribution that puts non-zero mass outside of supp( (Te) as,
Jiwa(T) C+ S/L Rmax + 8 mer(() log (Vy) dy + /
R(y) dy
(29)
ydlsup((Far)
where C' captures the contribution to the objective from the remaining y E supp y $ M. Note the forward KL penalty for y & supp(Trer) is zero. We show we can always construct an alternative distribution, It, with mass only inside of supp(Tred) and has strictly higher JIwa. We write,
/ Jima(=)) = C+ J (T(D)) + +a(v))) Rmax BTer(y)) log (v) + alv)) dy.
(30)
where a is function that redistributes the mass outside of supp(mer) across M; axy) > 0. (ytsupp(mm))) a(y) dy = SM a(y) dy.
First, we note the reward contribution do not decrease from π (left hand side) to T' (right hand side),
SM Rmax dy + f
()ssupp(mm)) B(v) dy S [ Rma) dy + L aly) dy, (31)
since (ygtsup(== TT (Vy) dy JM avy) dy and R(y)= Rmax- Second, note the (simplified) KL contribution is strictly larger in x' (right hand side),
S log T(y) dy < <fx Teer(y) log (T(() + aly)) dy
(32)
since log is strictly increasing function, and auy) > O. Therefore, we have established that, (twd(T) Jfwd (
(33)
That is, there always exists more optimal solution with support solely inside of supp(Trer)-
Proof of Remark 3.3 Assume the reward R is finite and some samples from within supp(Feet) has R(y) Rmax- We optimize with 3 > 0 over the restricted feasible set II, va II. where TVy)
almost everywhere on supp(Ter) to avoid dividing by zeros.
We write the maximization objective subject to constraints J [T( dy 1, T(y) 2 0 for all y.
C)(T; X1 = /
+ f log (v) dy X T(V) dy 1) + / (())=(t) dy, (34)
= (()))(()) + XT( + 3 (y) log π (0) + (()) dy - x,
(35)
where at the optimal solution, μ(y) 2 0 and ((y) π (y) = O. We take the Gateaux derivative in any perturbation direction (lv), dy O, m()) +e((y)
d [ X] de - + ex; Al
(36)
c=
Defining log 0 = 0 per convention. We first solve,
de L, + [p:
de (n(8) + =e(()) R(y) + ((=()) ep(y))
d
+ Tur log (T()) + Ep(v))
+ *())( +v(()) dy,
(37)
= =f (()) + xp(x)) +
T (v) (v))
T(y) + Evp p(v))(v)dy
(38)
-(((0 ((O) (R(v) + +
+ mlv) dy म T + xp(l)
(39)
19

---

Published as a conference paper at ICLR 2026
- - + ex; epp; de
f p(w) [R(v) + + 3 Tret(V)
=(w
+ m(v)) dy.
(40)
==0
Define the functional derivative to be,
5 f R(y) + + Tve((U) (y)
+ muy)
(41)
To find the optimum x² which gives 4/as L.((T es; A| 0 for all yo, the fundamental lemma of the
calculus of variations (Gelfand & Fomin (1963), Lemma 1) tells us it would imply 6/ōx [T; A] = Solving for this.
R(y) + + 1 (y)
+ puv) = 0.
(42)
T (1)
3TTe(())
ा (y) - R(y) (x())
(43)
(TT
x (v) A - (R(()) +u(y))
define = -
(44)
Per assumption that some reward samples are within supp(Trer) Remark B. states n² (1) 0 for all & supp(Trer). We can thus ignore the (y) regions.
Further observe (her > 0
implies (y) O, thus p(y) (per π (v))(y)) 0). The optimal distribution is therefore,
BTver(y)
Gima(y) A R(y)*
A Rmax
(45)
where A is the unique solution to J ( R(y)) dy To see this solution exists, observe
as A -> Rmax. Grwd at this point goes to infinity. On the other hand, as A -> co, all Giwa (1) - 0. By continuity, some A exists between Rrmux and co which satisfy normalization
Note Grill et al. (2020), Appendix B.3 arrives at similar solution for the setting of discrete action spaces (i.e. Ty is vector).
When does Grwd have mass outside of supp mrer Interestingly, when regularizing with the
forward KL, there are cases where the optimal distribution Gfwd puts probability density on regions outside of the support of Tref. First, note that when mer(y)) = 0, the KL penalty is zero. We can use this to solve for a simplified version of Equation 42,
R(y) + A p(y) =0
(46)
+ A = R(y) + m(y).
-
(47)
This implics a few possible scenarios for regions where Ter((y) = O,
If R(v) < A, then p(v) > 0, implying (vy) 0 to respect μ(y)T(y) = 0,
If R(y) = A, then p(y) 0, meaning Tイy) can be positive, R(y) A is impossible, as yu() 2 0.
Denote Rin = max R(y) as the on-support max reward, and Amu mapppsupp( Ryy)
as the off-support max reward. Per Remark B. I. Grwd will never leave the support of nee long as Rin 2 Roorr We therefore consider the case where better samples can be found outside of supp(Tret, Frt),, Rout Rin
Denote an integral over Tet as,
( dy.
( supp(Taa) c R(y)
(48)
Now consider the off-support set with constant max rewards: M' = {v & supp(Tre) R(y) = Root ). Recall this set has higher reward than anything within the support of meets Irout > -i
Z (Tegou m ) < 1, A < 12 violates impossibility of R(y) A above, while A I implies no mass can be placed off support, without normalization on-support (Z(A) < 1). Thus, the only valid
solution is A = Rout with the leftover 1 Z [ probability mass allocated to M'. On the other
hand, if Z ( 2 1, it implics some A 2 Rosi exists which normalizes the on-support distribution
and mass is placed off m support.
20

---

Published as a conference paper at ICLR 2026
B.4 GRADIENT OF FORWARD- KL REGULARIZED REWARD MAXIMIZATION
Proof of Remark 3.4 We want to know if optimizing the forward-KL regularized RL objective within the support of ree is equivalent optimizing a forward KL divergence. In other words, are interested in whether the following gradient,
Vg (Tw = Ve B (v) [R(v]] 3 DE (Trer) [Te
(49)
is a gradient of forward KL between Th and some target distribution h that is independent of wz. We prove by contradiction. Suppose h exists, follows that the functional derivative of these two objectives must be equivalent up to proportionality,
5 Ôπ
8 (Tw x DDxz
_ 5n
(50)
where both are subject to constraint f T TVV) dy 1.
We have established from Equation 41 that the functional derivative of Jiwa subject to constraint (t)) dy = 1 is.
[ A1 R(y) + 3 Terryy) (1)
8
+ A.
(51)
δπ
To find the functional derivative of the forward-Ki KL, we first write down the forward KL objective subject constraint,
CK| X1 = (KL((llTT) xx( [ dy dy 1),
(52)
= = hly) log hly) - h(y) log (v) dy + / XT(() dy X',
(53)
J XT(y) h(y) log (v) dy |/ / dy x)
(54)
where the right-hand bracket is independent of The Gateaux derivative is,
- CK|T +p, = d
ds / ((((()) + - h(v) log (I(z) + =x(())) dy h(y)p(v)
(55)
da
= f Xv(u)= T(v) + =s(ly) dy.
(56)
d
f p(v) X [x
h(y) πγy)
ds Cx(T +ep
dy
(57)
==0
The functional derivative of the forward KL with respect the right- hand term is therefore,
h(y) (
[k|m,N] =
(58)
8π
Assuming the functional derivative of the two objectives are proportional to each other, we can solve for the target distribution h(()).
8
8
Zn Lx TT, X') x
L A).
(59)
8m
= X h(y)
R(y) + TTe((V) T(V)
X1
for some constant a, (60)
7²y)
+ h(y) ( ad aR(y) a (1)) ()
(61)
Observe one cannot write hyy) independently of ( other than in trivial cases (e.g. if R(y) is constant such that const R(v) 0). Thus, for general reward functions R, optimizing the forward-KL. does not produce a forward KL gradient toward any distribution that can be expressed independently of no-
21

---

Published as a conference paper at ICLR 2026
B.5 GRADIENT OF THE FORWARD KL
Remark B.2. The gradient of the forward KL divergence between policy mo and target G is,
Ve DxL (GollTe) = -EGa Vo log (b)]
(62)
Proof.
Ve DeL (Gol(ae) = Vobag log Ga(y) - log to(y)] = EGG Vo log Gz(y) log (b()]
(63)
(64)
= -Boo Vo log (b))
(65)
We see that optimizing the forward KL gradient amounts to doing maximum likelihood supervised fine-tuning trajectories sampled from the target distribution G3. is also mentioned in some previous works (Agarwal et al., 2024). This is generally intractable as requires sampling from Gg. Nevertheless, estimating expectation under a distribution known only up to normalization (i.e. [. is well-studied Monte-Carlo methods (Robert et al., 1999), and it is conceivable that a number of methods there would prove helpful here. For instance, number of works in diffusion generative modelling proposes to use particle-based samplers, such as sequential Monte Carlo, to approximate the reward-tiil -ttteed distribution (Wu al., 2023; Skreta al., 2025; Singhal et al., 2025).
B.6 PROBABILITY RATIO UNDER OPTIMAL TARGET DISTRIBUTION
Proof of Proposition 4.1 distribution is given by,
For any two samples, Y1 and y/, their probability ratio under the target
Gp(u)) go(v1) S ( ga g3(2))
(66)
(3(22)
g g
which only require the unnormalized likelihood as the normalization constant cancel out. Expanding the terms, we can write the log likeilhood ratio in closed form,
Ga(y1) = log (11) exp log Gr(U2)
R(y1) B
R(yz))
log Tred (V2) exp
(67)
B
= log Tee((y))
R(y)) - R(vz))
(68)
Teet (U2)
1
B.7 TARGET DISTRIBUTION AFTER REWARD AUGMENTATION
Remark B.3. Optimizing the reverse-KL regularized RL objective with the augmented reward function R yields the following target distribution, which puts uniformly high mass over all samples above reward threshold R(y) Z T,
Terly) exp
Riv) R(y) <T,
3
Golu)
(69)
TTe(z) exp
[(i=)
if R(y) 2T.
Proof. We have established already in Appendix B.1 that the target distribution of reward maximiza- tion with reverse KL regularization is,
R(y)
G2(y) x mer(y) exp
(70)
Plug in the augmented reward function.
R(y)) R(x)) + (1ogg )T)((2 - log Te)(
if R(y) T, if R(y) 2 T,
R(y)
(71)
22

---

Published as a conference paper at ICLR 2026
which gives us the augmented target distribution,
R(u)
Go(y) & Ter(y) exp (
(72)
In the R(y) case, R(y) R(y), and there is no change to the (unnormalized) likelihood. In the R(y) 2 case,
R(y) = log Tref (v) - R(y),
log (y) exp
(73)
log (rr(
1
RR(z) + ((log Ter(2) - log mor(()))
(74)
B
R(2)
+ log Tee( (3) + log TTr((2 log (Te(())
(75)
R(2)
+ log (2)
(76)
Therefore we see in the R(y) 2 case we have,
R(y), = (T)(()) exp
R(²) B
(er(y) exp
(77)
This formally shows the target will have uniformly high density proportional to TTe((2) exp (R())/)) for all samples if their original reward R(y) is above threshold т. we pick z be likely under The- e.g. 2 = arg max, Tee (y), we can also see these samples will have the highest probabilities in the target distribution.
B.8 GRADIENT OF REWARD-AUGMENTED OPTIMIZATION
We also note the MARA gradient estimator for an "above threshold" sample Yi (i.e. R(x) > T), when using reverse-Ki KL regularization, can be equivalently constructed as using the anchor sample's 's reference policy probability (z)(2) in lieu of the actual reference probability mer(y) when constructing the KL. gradient estimator. To see this precisely, we know the gradient of the expected reward to be,
V Em
R[v]] En. [R(y) V log [o())]
(78)
and gradient of the reverse- KI. regularizer to be,
v DKL (Tel(Tee) = Exo [log (e( log Teet (v)) log m)()
(79)
= Denote the reward-augmented objective as (me) = R(y) - 3D ( ( where R(y))
R(z)+ + 3 log TTe((z) log ())( and is the "anchor". The gradient estimator of K, for an "above threshold" sample, yr, R(y)) 2 T, can be written as,
TT(()) Treer (yi)
K, (R(y.) log
Vo log to
(80)
Trer( (2)
To (Vi)
[R(z) log
B log Teet (y:)
Vo log ()()))
(81)
Tref
[[(=) 3 log Trer(=)
V log (v)))
(82)
Intuitively, as the anchor is chosen to have high Tiet i.e. (e)((2) > (3:( this can be interpreted as selectively reducing the KL regularization for nigh-rewarding samples. Mechanistically, this also suggest an alternative implementation which produces the same gradient when using reverse-Ke KL regularization (Algorithm 2).
23

---

Published as a conference paper at ICLR 2026
Algorithm 2 Mode Anchored Reward Augmentation, alternative implementation. The gradient of this algorithm equivalent Algorithm when using reverse-Ki regularization.
1: Given: initial policy Te, reference distribution Tief, reward function A. regularization coefficient
B, threshold of good answers т E R, 7 S maxy R(y), and trajectory batch {ys [W me-
2: Pick anchor trajectory: 2 = arg maxy, (1), s.t. R(y,) 2T
3: for each yı in batch do
4: if R(y)) 2 then
5: Augment: reward F; R(z), reference prob pt Ter((2)
6: else
7: Keep same: reward F; R( reference prob p: (er(
8: end if
9:
end for
10: Optimize policy parameters using augmented rewards (Fr and augmented reference policy
probabilities {p:)
C ADDITIONAL EXPERIMENTAL DETAILS
C.I DIDACTIC EXPERIMENTS
We construct didactic experiment vector of size 100 (akin to output space with 100 tokens). We initialize categorical distribution over this output space whose logits are all O's (i.e. uniform distribution over all tokens). Given some reward function and reference distribution defined over this space, we optimize this categorical distribution with the KL-regularized policy gradient for 3000 gradient steps in PyTorch with Adam optimizer, with learning rate 5e-3 and batch size 32.
G Ggly2)
10
0.005
0,01
$= 0.1
1025
1.0
10aa
10
10
0.0
0.2
0.4
R(ys)
R(y2)
Figure 7: Effect of reward difference (A R, x-axis) and reverse-KI. regularization strength (hue) the relative probabilities between two samples in the optimal policy distribution (y-axis)
C.2 THE 1-2 TASK
We ask the LM to generate a uniform random integer that is either or 2 (Hopkins et al., 2023). as illustrated Figure We run for range of KL coefficients (3) and multiple random seeds. Figure 10 shows the training run for just vanilla RL, without MARA.
C.3 CREATIVE QUESTION ANSWERING TASK
We detail the training settings in Table 3, and evaluation settings in Table 4. We follow the evaluation procedures outlined both Kirk al. (2023) and Zhang al. (2025). The specific evaluation metries are defined as follows.
In Dist Rewardh rd: training reward, on training set, using training reward model
Out Dist Reward: evaluation reward on held-out set, using evaluation reward model
Ngram EAD: Expectation-adjusee Distinct N-gram, proposed Liu et al. (2022). We
follow (Kirk et al., 2023) and average EAD for 1,...
24

---

Uniformly randomly generate integer that is either 2. Respond strictly in this format: <think>Your internal or Published as a conference paper at ICLR 2026
Reverse
0.001
0.01
8 0.25
Reward Function
0.4
0.23
Forward KL
Reference Policy
8
0.001
0.01
8 0.25
0.1
0.32
3.2
332
2.
0.25
0.683
Figure Final policy distribution after KL-regularized RL, with equal rewards for all correct answers. Lup-support (yet equally correct) answers never preferred over high-support answers.
Prompt
Example Generation
Let me decide randomly. <hink>>/thinkkaasseeer </answer><<
2</answer>
Figure 9: The 1-2 task to test output distribution of LMs.
One Counts (binned)
Two Counts (binned)
Valid Answers Diversity
500
400
400
300
Count
300
Entropr
0.4
200
2.00
100
100
200
Stop
Stop
Step
Reverse KL Regularization
Forward KL Regularization
Figure 10: Training outcomes using vanilla RL. (Left, Middle) Policy's empirical distribution over valid answers for runs that reached high rewards (counts binned over 8 consecutive training batches), across a range regularization coefficients (3). Right Diversity of the valid answers over the course of training, measured as the entropy of the Bernoulli distribution over answers of I ''s and 2's.
25

---

Published as a conference paper at ICLR 2026
Hyperparameter
Value
Actor Model
Qwen3-1.
Reward Model
Beyoork-Reward. -wwrddv---e-------- Qwen3-
Training Dataset
Wildchat 10k English
Train Batch Size
128
Mini-Batch Size
64
Max Prompt Length
512
Max Response Length
2048
Learning Rate
x 10-6
Entropy Coefficient
Rollout n (per prompt)
Gradient Checkpointing Enabled
Epochs
Table 3: Creative QA Training Setting
Hyperparameter
Value
Evaluation Reward Model
Skywork-Reward Rewa-d--aaardeeeeeeeee-ee Gemmaa- -emma-2-2755 -2-27B-v0 Nove tyRnchh curated
Dataset
Num Generations / Prompt
10
Max Tokens
4000
Temperature
1.0
Enable Thinking (qwen)
False
Table Creative QA Evaluation Setting
Semantic Div: semantic embedding diversity measured by averaged cosine distance, using embedding model all -MiniMM-L----
Mean Distinct: Estimates a notion of "# of distinct concepts", as introduced in Zhang et al. (2025).
We run additional baselines for the effect of the batch level threshold to set in Table 5.
C.4 EVIDENCE FOR UNIMODAL TARGET DISTRIBUTIONS LMS
We show additional evidence in existing LLMs settings, the shape of the reference distribution and reward function leads to highly skewed target distributions.
First, we draw 8192 samples from Owen2 5-3B using prompt for the 1-2 task (Appendix C.2). We filter for correct answers (R(y) = 1), leaving 2944 samples (35.9% correct). We see in Figure lla the distribution of "1" and "2", with 1 being over-represented, pointing a skew in the base reference distribution favouring "1", despite the model being prompted to uniformly randomly generate an integer.
Out-dist
EAD (t)
Semantic
Distinct
Model
Reward (1)
Div (t)
(f)
MARA (rev, 0.90)
1.451 10.103
0.543 10.014
0.186 10.008
4.14 10233
MARA (fwd, 0.90)
1.604 10.113
0.568
10.012
0.193
10.009
4.62 10.258
MARA (rev, 0.75)
1.498 10.117
0.547
20.033
0.183
20.008
4.41 40.262
MARA (fwd, 0.75)
1.325 10.097
0.508 20.014
0.196
10.009
4.07 +0.243
Table 5: Ablation of batch-level threshold for т, set at either 90th percentile (0.90) or 70th percentile
(0.75). Mean 1 bootstrap SEM.
26

---

Published as a conference paper at ICLR 2026
Correct Answers
Ref Log Prob
Rewards
Prob Ratio (8 00001
2000
-2
15
995
Times more Ie
10th
-4
0.00
10th
1000
-6-
10
5
-8
10a
0
0.19.00
D
0.05
124
0
17
34
51
17
34
5
10
15
Answer
Sample
Sample Index
Sample
(a) 1-2 Task correct
(b) Reference probs log (v) of highest reward
(c) Relative (to idx 0) probs of most likely samples under G3
answers distribution
samples in Creative QA task; A log shows range
Figure 11: Evidence for highly skewed target distributions Gp in LLM tasks
For creative QA (Appendix C.3), we draw 1024 samples from Qwen3-1, 7B using a single question in WildChat text. We evaluate the responses using skywork-Rewardd rd-v--Qwen----------------
(reward mean: 12.33, min: -2.125, max: 18.38), and filter for answers above the 95th percentile (R = 16.37). This leaves 52 samples (reward mean: 17.0, min: 16.38, max: 18.38). We see in Figure 1lb that among these high-reward answers, log (er(y) has a difference of up to 7.76, corresponding to a probability ratio of 2 under G3 (Remark 4.3).
We then estimate the probability ratio taking into account both the reference probability and the reward (Proposition 4.1). We the same 1024 samples and compute their unnormalized likelihood GR, Fee((y) exp( R3)) with 3 = 0.01. We take the top 16 most likely samples and calculate their relative probabilities with respect the 16th most likely sample, such that the lowest probability ratio is 1 x. We observe in Figure 1lc that amongst the sampled responses, the most probable sample (R 18.4, log (y) 5.911) is 2 10a9 times more likely under the target distribution G3 than the 16th most likely sample (R = 17.3, log (y) = -7.09), despite having only slightly higher rewards and reference probabilities.
C.5 DRUG DISCOVERY
Chemical language models (CLMs) that generate molecules in string-based formats, e.g., a SMILES string (Weininger, 1988), have been experimentally validated with numerous generated molecules in clinical trials (Du et al., 2024). Recently, the field has focused on addressing "synthesizability" i... can generated molecules actually be synthesized in the lab? (Stanley & Segler, 2023: Papidocha et al., 2025). Accordingly, we adapt two reward functions from Guo et al. (2025b): SYNTH and SYNTH AMIDE that jointly reward binding potency and synthesizability. REINVENT (Olive- crona et al., 2017) is a state-of-the-ae The recent Saturn CLM (Guo Schwaller, 2024b) notably improves optimization efficiency by using data augmentation (Bjerrum, 2017; Guo & Schwaller, 2024a), but continues to use REINVENT's RL algorithm.
RL-based CLM on standard benchmarks (Gao et al., 2022).
In the drug discovery experiments adapted from Guo et al. (2025b), the reward functions are comprised of numerous individual optimization objectives, and defines nulti-parameter optimization task. Concretely, these objectives are:
1. Minimize the molecular docking score using Vina2-iin2GGPP (Trott & Olson, 2010; Al- hossary et al., 2015; Tang et al., 2024). Docking simulates binding of molecules to target protein and predicts a crude binding affinity value. Docking was performed against the ATP-dependent Clp protease proteolytic subunit (CIpP) Mabanglo et al. (2023).
2. Maximize the quantitative estimate of drug-likeness. likeness (QED) (Bickerton al., 2012), which is itself comprised of various physico-chemin properties, e.g., molecular weight. Maximizing QED can prevent generated molecules from being too large and lipophilic.
3. Constrain the number of hydrogen-bont donors (HBDs): HBDs This can improve absorption, Distribution, metabolism, and excretion (ADME) properties (Kenny, 2022) of the generated molecules.
4. Satisfy the "Synthesizability" constraint. Synthesizability is quantified by using a ret- rosynthesis model on each generated molecule. Retrosynthesis models predict a plausible
27

---

Published as a conference paper at ICLR 2026
synthesis route to synthesize target molecule using commercially available precursors. The precursors set is from the eMolecules catalogue extracted from Chen et al. (2020). Retrosynthesis models typically start with "single-step" model which predicts precur- sors given a target molecule. Since molecules may require multiple steps to synthesize, "Multi-step Retrosynthesis" commonly couples search algorithm with single-step models to iteratively decompose a target molecule. In this work, we use the MEGAN (Sacha et al., 2021) single-step model with the Retro* (Chen el al., 2020) search algorithm using the Syntheseus (Maziarz et al., 2025) package. Finally, molecule is considered synthesizable if the retrosynthesis model successfully proposes synthesis route.
Both the SYNTH and SONTH-ALL-AMIDE: reward functions are comprised of the above objectives. The only difference is that in the SYNTH-ALI--AMTD. case, a molecule is only considered synthe- sizable if all the chemical reactions involved to synthesize are "amide coupling reactions". Amide couplings are one of the most common reactions performed in the pharmaceutical industry (Brown & Bostrom, 2016), and is generally a robust, widely compatible reaction. Subsequently, the reward function is defined product of each individual component above. Given molecule,
R(x)= DS(x)) QED(a) HBD(a) Synlheseus(z)) [0,
(83)
where "DS" is docking score. The HBD and Syntheseus objectives are binary, i.e., 1 if satisfied and 0 otherwise. QED E [0, 1] and is used as is. The Quick Vina2-GPU docking score reward shaped using reverse sigmoid function following Guo et al. (2025b) and gives higher reward to lower docking scores, as desired.
Our goal in this section is to investigate the potential for MARA to be drop-in replacement for the REINVENT (Olivecrona et al., 2017) RL-based algorithm for molecular design. RE REINVENT is amongst the most performant molecular design algorithms (Gao et al., 2022) and the Saturn model (Guo & Schwaller, 2024b) adapts this algorithm and leverages data augmentation (Bjerrum, 2017; Guo & Schwaller, 2024a) further improve optimization efficiency.
We evaluate all models with a fixed budget of 10,000 reward function evaluations, which is standard in benchmarks. We contrast the algorithms' performance molecular design metrics that measure optimization efficiency and diversity. Yield is the number of unique molecules above a reward threshold. CB100 the number of reward evaluations required generate 100 molecules above the same threshold. IntDiv Didivll (Polykovskiy et al., 2020) and WCirclen es (Xie et al., 2023) are diversity metrics based on molecular sub-structure based features, and measure intra-set similarity and sphere packing, respectively.
Tables 6 and show the optimization results for the SYNTH and SYNTH-ALL-AMIDE reward functions at the 0.80 and 0.85 screening thresholds, respectively. MARA is trained with r = 0.80 in both. In general, MARA matches or outperforms REINVENT particularly for the more challenging ALL-AMIDE reward function. In this environment, MARA can find more high-reward molecules (Yield) and using less reward evaluations (OB10C) than REINVENT.
28

---

Published as a conference paper at ICLR 2026
Table Results at Threshold 0.8 (1 larger better; 4. smaller is better). "SYNTH" and "AMIDE" denote the SYNTH and HALTH---LL-MAA reward functions, respectively.
Task
Algorithm
Sigma
Gen Yield
OB100
IntDiv1 (f)
Circles
(f)
(L)
(f)
SYNTH REINVENT
128
6569 = 186
1042 + 66
0.766 0.011
67 = 3
256
6618 + 93
1080 + 89
0.756 + 0.012
57 + 8
512
6746 + 161
1067 + 74
0.752 + 0.016
55 + 5
MARA
128
6834 = 78
1015 - 55
0.761
59 =
256
6750 = 139
1068 + 50
0.760
60 = 4
512
6793 267
1065 +
0.751 0.015
=
AMIDE REINVENT
128
5433 = 184
1427 + 63
0.768 0.012
35 +
256
5544 = 172
1406 + 59
0.768 0.009
34 = 5
512
5334 = 165
1445 111
0.776 0.008
33 = 4
MARA
128
5635 249
1407 + 123
0.766
0.008
36 3
256
5353 + 114
1393 + 42
0.009 0.769
33 4
512
5377
152
1343 + 77
0.763 0.008
31 3
Table 7: Results Threshold 0.85 (1 larger better; 1. smaller better). "SYNTH" and "AMIDE" denote the SYNTH and SYNTH-ALII-AMTON reward functions, respectively.
Task
Algorithm
Sigma
Gen Yield
OB100
IntDiv1 (f)
Circles
(t)
(L)
(f)
SYNTH REINVENT
128
1614 = 407
4114 109
0.701 0.018
7 1
256
1552 = 242
3940 + 371
0.699 + 0.030
6 1
512
1484 + 45
3717 + 201
1771+00000
6
MARA
128
1796 = 210
3654 272
0.005
6 1
256
1530
+ 126
3957 + 335
0.705
8 1
512
1550
347
4016 = 234
0.689
0.024
6 1
AMIDE REINVENT
128
1098 +:
4360 257
0.721
0.016
8+1
256
1488 = 280
4290 + 141
0
8
512
1054 = 152
4620 1 438
0.739 0.000
810
MARA
128
1235 130
3943 + 303
0.733
0.009
8+1
256
1404 + 261
4079 - 172
0.730
+ 0.010
7
512
1341 + 86
3930 + 400
0.723 + 0.004
7+1
29


Reviews:
[{"review_id": "coh1QtjuCV", "reviewer": "Reviewer_W9ms", "summary": "This paper looks at the behavior of policy gradient optimization with respect to the forward vs. reverse KL-regularized expected reward, where it is demonstrated both theoretically and empirically that (i) either choice of divergence can, in the right circumstances, cover the various modes of the target distribution and (ii) there are circumstances in which neither will appropriately cover the various modes. The authors propose a modification to the reward function, MARA, such that policy gradient optimization with respect to the forward or reverse KL-regularized expected modified reward assigns high probability to all actions that induce a sufficiently high reward.", "questions": "Do you have an explanation for why MARA has lower \"semantic embedding\" diversity than the baseline?\n\nHow sensitive is the behavior of the MARA learned policy to choice of tau? It would be helpful to have a figure or result dedicated to this (I was not able to find it in the paper, please point it out to me if I overlooked it). Also, how did you go about selecting the value of tau in experiments? \n\nIn the REINVENT task, is there a good explanation for why using MARA as a drop-in replacement did not lead to improvement on the two diversity related metrics? Have you compared the number of, e.g., ECFP4-based Tanimoto clusters for compounds sampled from REINVENT vs. REINVENT with MARA objective?\n\nHow sensitive is MARA to the choice of the anchor z?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper does a good job highlighting that common intuitions about the mode-seeking vs. mass-covering nature of different divergence choices does not obviously result in the expected behavior that is often highlighted by 1d examples. Indeed, for expected reward regularized divergence minimization, the authors demonstrate in a convincing way both formally and through experiments how the choice of reference policy, regularization strength, the shape of the reward function are all highly relevant factors in determining whether the learned policy covers the various modes of the reward function. I found the authors solution to this (mode anchored reward augmentation) to be an original one and well motivated by the paper's exposition. The paper's experiments are well constructed and show on a variety of pertinent RL problems that MARA generally improves diversity of the resulting policy.", "weaknesses": "The MARA reward function itself feels a bit unsatisfying in that it requires knowledge about reasonable choices of tau (and therefore some knowledge/belief about the distribution of attainable reward values). I wonder to what extent this makes MARA suitable in settings where rewards can be potentially unbounded in magnitude (e.g., certain Atari games). The need to have a high quality anchor example also feels like a limitation on the problems that MARA may be applicable to, e.g., problems where the reward signal is very sparse, or where high reward examples are rare a priori. Further, it seems that redefining the reward to a constant for all sufficiently high reward examples may not be appropriate for all problems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YHm7lI9HC9", "reviewer": "Reviewer_mZ9a", "summary": "This paper studies KL regularization in RL finetuning of language models from first principles. This paper lays out various facts about the global solutions of the forward and reverse kl regularization + RL finetuning process:\n\n1) RL + reverse KL regularization is equivalent to minimizing reverse kl with its globally optimal solution (known previously)\n2) RL + forward Kl regularization does not optimize any form of forward KL.\n3) In RL + forward KL, the beta coefficient acts as an exponential decay / growth term in the solution. Standard values can essentially make the best solution orders of magnitude more likely than other solutions. This is known but often overlooked.\n4) When rewards are equal, off policy solutions are discouraged. This is also obvious (looking at equation 2), but still overlooked.\n\nFinally, for KL regularization + RL finetuning, they provide a simple heuristic to boost multi-modality and present experimental results verifying their approach.", "questions": "See weaknesses", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "I think the summary^ is a fair assessment of the paper's strengths. It is presents simple results in a very clear and easy to read manner. According to me, the heuristic presented seems to be not that significant in terms of usability in frontier models for example. But I could be wrong on that. For the theoretical results, they are simple in hindsight, but are often overlooked in the literature. Remark 3.3 seems to be the most non trivial and novel result in my opinion.\n\nOverall I enjoyed reading this paper and I think this paper should be accepted.", "weaknesses": "The only weakness I see is a lack of strong baselines. LLM post training is not my primary area of research but surely their have to be many multi-modality inducing algorithms for post training (RL exploration, gflownets, feature  covariance maximization, etc). I do not expect your algorithm to be better than them, but a comparison might still be useful for readers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LrZXlIWBeb", "reviewer": "Reviewer_v1oM", "summary": "This paper studies the effect of KL regularization in the objective function of LLM post-training from the perspective of diversity, e.g., the multi-modals in the learned distributions from the final policy. The authors compare the optimal policy under the reverse and forward KL regularizations under the KL-regularized objective in LLM post-training. Based on the analysis, they proposed to put equal weights on all high-quality samples to encourage the multi-modal target. Experiments are conducted on several LLM post-training settings.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "1. Understanding the learning behaviors, e.g., the multi-modality of the learned distribution, of the KL-regularized objective function in LLM post-training is an important research problem.\n\n2. The usage of illustration improves the readability, although in simple settings.", "weaknesses": "1. **The scope of the paper is questionable**. The title of this paper targets the general RL, but I am surprised that the scope of this paper only focuses on the KL-regularized objective function used in LLM post-training. There are a flurry of prior RL works on pure KL-regularized RL, such as MaxEnt RL, Soft Actor Critic, and natural policy gradient, that focus on the MDP setting of the general RL setting. By contrast, the KL-regularized objective in LLM only assumes a specific contextual bandit setting with iid samples, which is far less representative of the general RL setting. Using such a title is very misleading. Researchers in LLM post-training should respect prior works in the general KL-regularized policy gradient methods if they really want to target a general research problem. Alternatively, it is required to narrow down the scope and title of the paper to the LLM setting without causing confusion, which is just a specific application of general RL in the language scenario. Lots of analyses are very divergent.\n\n2.  **Less convincing theoretical analysis**. \n* The paper relies on the fact that reverse and forward KL have the effect of mass covering or mode seeking, but I do not think they are established in a fundamental way. The authors are suggested to provide sufficient references, either from information theory or statistics; otherwise, the subsequent analysis would be less convincing. Also, some illustrations on experiments are too simple to make a rigorous conclusion. This issue occurs many times across the paper. \n* Is the diversity the sole objective in the policy optimization? In statistics, the reason why people use MLE or equivalent KL divergence is that it ensures statistical efficiency, asymptotic unbiasedness, and consistency, often on iid samples. However, if you manually introduce the sample selection bias by focusing on the diversity to encourage multi-modality, there is definitely some price to pay. However, from the theoretical perspective, I did not find such a deep analysis or conclusion regarding this point. After reading Line 102, does it mean that in LLM the foundation models are often less expressive such that we need to analyze the difference of two KL divergences? A lot of details are unclear to me. \n* This paper makes a lot of strong claims, such as ‘diversity collapse is a natural consequence of solving the RL problem.’ It is very risky to accept these claims without finding rigorous proof or strong empirical evidence or articulating the specific setting of the RL problem they are studying among the whole RL literature. \n* In Eq. 5, there is no closed-form solution under forward KL divergence, where $\\Lambda$ is the negative of the Lagrangian multiplier. However, it does not necessarily mean that the optimizers under the two KL divergences are completely different in the general case (line 163)! \n* The analysis in Section 4 is performed under very specific scenarios. Although sometimes they provide some insights, in general they are very unlikely to hold in the training process. Such an analysis is fragile and less convincing. The result in Proposition 4.1 is trivial. The logic in Section 4 is very hard to follow.\n\n3. **Unclear and informal writing**. I personally think a lot of writing is not formal or rigorous.   \n* What is the definition of ‘solution distribution’? I understand what it is by guess, but the authors did not give a definition or explanation when they used it many times.   \n* What is the approximate inference in Line 53? I do not think it is a scientific word.   \n* In Line 92. What is the definition of two-parameter Gaussian? Is the foundation model as an example of a flexible distribution in line 93?   \n* It is inaccurate to say we aim to maximize a reward function in Line 107. Instead, we are maximizing the average rewards or the expectation of the returns. What is the definition of \"off-support samples in Line 264. There are lots of expressions that are not formal, which is discouraged to me.   \n* There are also many grammar errors, such as ‘This let’ in Line 232 and ‘KL exhibit’ in Line 91. The writing needs substantial improvements.\n\n4. **Unclear empirical validations**. Since the paper focuses on the improvement of diversity, I could not find a detailed explanation about the metrics of diversity. Learning curves in Figure 7 are blurred, and thus it is hard to make a clear conclusion on it. Also, it makes more sense to me if the LLM focuses on the diversity; the efficiency would be reduced to some extent. Sometimes this price is worthy to pay, but I could not find such an analysis in the experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sdh09zURCZ", "reviewer": "Reviewer_K4Yn", "summary": "This paper analyzes the reasons for \"mode collapse\" when using KL-regularized RL for LLM post-training, linking it fundamentally to the shape of the target distribution. For this, the paper revisits earlier-made arguments about the shape of the target distribution that Reverse KL-regularized RL approximates and makes (new?) arguments about the target distribution of Forward KL-regularization. Using insights from this analysis, the authors propose a new MARA algorithm that is claimed to approximate a target distribution with uniform high probability modes for highly-rewarded samples, and sticks to the reference distribution otherwise.", "questions": "- Q1) L31 why do you say RL is the only way to train models where the solution is not known a priory? You already cite Go et al. (2023). Wouldn't this fit this bill?\n-Q2) What is the role of remarks 3.2, 3.3 and 3.4 in the overall story of the paper? I felt that they were made and forgotten.\n-Q3) L174: when you say that ML training on samples from $G_\\beta$ is intractable, you mean for an unbounded reward function, right? You mention STaR, but you could make it more explicit that this algorithm is doing exactly what you are saying is intractable for a particular class of rewards on which we can do rejection sampling. More importantly, note that you don't necessarily need to sample from $G_\\beta$ to approximate this objective. $G_\\beta$ is particular case of target distributions described as \"distributional constraints\" in https://arxiv.org/abs/2012.11635, and the forward KL objective can be approximated using the DPG algorithm.\n- Q4) What is the policy you are using for Fig. 2 and following?\n- Q5) Definition 3.5: just to double check, do you really mean that _all_ high-reward samples must have high probability? OK, but this doesn't match my intuition for multimodality. Maybe you are looking for a different concept?\n- Q6) When you say \"off-support\", you mean $p(y)<\\epsilon$, right? Correct me if I'm wrong, but I think mathematically off-support means $p(y)=0$. You might need to define this other concept of support to use it here.\n- Q7) Many recent papers set $\\beta=0$ (e.g. Dr. GRPO). What happens when your algorithm for this particular case in the scenario of verifiable rewards? It looks to me that it boils down to regular RL, and cannot prevent the loss in diversity.\n- Q8) L361 \"stays close to the reference in regions of low reward\" -> why is that the case?\n- Q9) Following up, can you say something about the target distribution of MARA?\n- Q10) Figure7: the KL parameter in the legend is $\\beta$?\n- Q11) What is In dist and out dist reward?\n- Q12) Could you add a few words about the REINVENT algorithm?\n- Q13) Why does the threshold affects REINVENT?\n- Q14) \"In this work, we provide an in-depth understanding of the KL-regularized RL objective, particularly in terms of its diversity\" -> could you summarize it in a few words?\n\nTypos/style/corrections\n- L107 \"maximize a reward function\" -> \"maximize the expected value of a reward function\"?\n- L149 \"we can regluarized\" -> we can regularize\n- L149 \"maximizaztion\" -> \"maximization\"\n- L208 \"token space\" -> sequence/output space ??\n- Figure 7 caption: \"entrpy\" -> entropy\n- L426 \"adpat\" -> adapt", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- S1) The paper addresses a highly debated topic at the moment on which a lot of noise and confusion exists.\n- S2) The paper is quite clearly written and the explanations are easily followable.\n- S3) The paper introduces a new method based on a theoretically clear argument.\n- S4) The analysis of the forward KL-regularized objective is quite interesting if new, even though I'm not sure it has much impact on the paper's main narrative.", "weaknesses": "- W1) I feel that part of the discussion increases --rather than dispelling--  the confusion. \nIn particular, starting from the first two sentences of the abstract:\n  \"It is commonly believed that optimizing the reverse KL divergence result in “mode\nseeking”, while optimizing forward KL result in “mass covering”, with the latter\nbeing preferred if the goal is to sample from multiple diverse modes. We show—\nmathematically and empirically—that this intuition does not necessarily transfer\nwell to doing reinforcement learning with reverse/forward KL regularization (e.g.\nas commonly used with language models).\"\nI think this is either false or misleading at best. One thing is the kind of regularization used in the loss and another is the divergence implicitly optimized by the training objective (remark 3.2). It is the latter that can be ascribed \"mode-seeking\" or \"mode-covering\" behavior, whereas here the authors are referring to the divergence used in the regularization term. If the goal of the paper is clarifying these aspects of LLM post-training using RL, it could do a much better job at keeping this distinction clear.\n- W2) Furthermore, there are a number of inaccuracies sprinkled across the paper. One of the most important ones is in one of the main messages of the paper, which is that the shape of the target distribution is responsible for the lack of diversity (again, in the abstract, \"Further, we show commonly used settings such as low regularization strength and equal verifiable rewards tend to specify uni-modal target distributions, meaning the optimization objective is by construction non-diverse.\"). This looks like a misconceived conclusion from the analysis in the paper. The illustration in Section 3.3 corresponds to a reward function that assigns non-uniform values to rewarded samples, whereas the illustration in Figure 5 compares \"on-support\" with \"off-support\" sequences (see also Q6 about this terminology). Actually, as noted in Remark 4.3 \"[the] relative probabilities in the solution is simply the relative probabilities in the reference\ndistribution, and do not depend on the KL-regularization strength $\\beta$\". Doesn't then this contradict the sentence in the abstract?\n- W3) While the proposed algorithm is interesting, there is no comparison with other diversity-enhancing baselines in the empirical part such as rewarding the unlikely or pass@k training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)