Better Estimation of the Kullback--Leibler Divergence Between Language Models
Authors: Afra Amini
Tim Vieira
Ryan Cotterell
Venue: NeurIPS 2025 poster ()
Keywords: KL divergence
Monte Carlo estimation
RLHF
post training
alignment
Primary area: deep_learning

Abstract: Estimating the Kullback--Leibler (KL) divergence between language models has many applications, e.g., reinforcement learning from human feedback (RLHF), interpretability, and knowledge distillation. However, computing the exact KL divergence between two arbitrary language models is intractable. Thus, practitioners often resort to sampling-based estimators. While it is easy to fashion a simple Monte Carlo (MC) estimator that provides an unbiased estimate of the KL divergence between language models, this estimator notoriously suffers from high variance and can even result in a negative estimate of the KL divergence, a non-negative quantity. In this paper, we introduce a Rao--Blackwellized estimator that is unbiased and provably has variance less than or equal to that of the standard Monte Carlo estimator. 
In an empirical study on sentiment-controlled fine-tuning, we show that our estimator provides more stable KL estimates and reduces variance substantially. Additionally, we derive an analogous Rao--Blackwellized estimator of the gradient of the KL divergence, which leads to more stable training and produces models that more frequently appear on the Pareto frontier of reward vs. KL compared to the ones trained with the MC estimator of the gradient.

Better Estimation of the Kullback-Leibler Divergence Between Language Models
Afra Amini
Tim Vieira ETH Zürich
Ryan Cotterell
fafra. amini, ryan. cot ethz. ch tim. f. viei com
Abstract
Estimating the Kullback-Leiblee (KL) divergence between language models has many applications, e.g., reinforcement learning from human feedback (RLHF), interpretability, and knowledge distillation. However, computing the exact KL divergence between two arbitrary language models is intractable. Thus, practitioners often resort to sampling-hane based estimators. While it is easy to fashion simple Monte Carlo (MC) estimator that provides unbiased estimate of the KL. divergence between language models, this estimator notoriously suffers from high variance and can even result in negative estimate of the KL divergence, a non-negative quantity. In this paper, we introduce a Bao-Blackwellizer estimator that unbiased and provably has variance less than equal that of the standard fine-tuning,
Monte Carlo estimator. In an empirical study sentiment-controlle we show that our estimator provides more stable KL estimates and reduces variance substantially. Additionally, we derive an analogous Rao Rao-Blackwellinn Blackwellized estimator of the gradient of the KL divergence, which leads more stable training and produces models that more frequently appear on the Pareto frontier of reward vs. KI. compared the ones trained with the MC estimator of the gradient.
http:://ggttubb.co///yyolbb////
Introduction
The Kullback-Leible [KL; 19] divergence is a statistical divergence that quantifies how one probabil- ity distribution differs from another. Measuring the KL divergence between probability distributions is a well-established problem that has been studied extensively in the statistics literature [7, 12, inter alia]. some special cases, e.g., in the case that we wish to measure the KL divergence between two Gaussian measures, the KL divergence has an analytical solution. However, in the general case, exact computation of the KL divergence not analytically tractable or approximable with an efficient algo- rithm [14]. This paper treats the case of computing the KI. divergence between two language models (LMs),, a fundamental task natural language processing with numerous practical applications.
The KI. divergence plays a central role across multiple applications. In reinforcement learning from human feedback [RLHF: 26. 35], it used as a regularization term to constrain the fine-tuned model from drifting too far from a reference model, preserving fluency and preventing reward over- optimization. In interpretability research, KL divergence quantifies how specific prompt shifts the model distribution by comparing the model's 's distributions before and after controlled interventions [9, 27, 36]. As an evaluation metric, KL divergence used assess how well language models approximate target distributions [4, 37]. In knowledge distillation, KL divergence is minimized to align student model with a teacher model [1].
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

The above applications demonstrate that measuring the KL divergence between two language models useful and widespread. However, in the case of neural language models, it is far from straight- forward. It is easy to see why: Given an alphabet of symbols E and two language models p and q. distributions over x,, the KL divergence is given by the following expression:
ply) log ply) ((y)
KI. (1
Σ VEE
Recalling that Σn a countably infinite set, we cannot expect, in general, to compute Eq. (1) exactly in finite time without additional assumptions. 3 While in some very special cases, e.g. where p and q are deterministic finite-state automata, there exist efficient algorithms, [4, 20, 22], we should not expect such an algorithm to exist in the case where p and q are neural language models, c.g., those based on the transformer [25, 28, 39J. Thus, most researchers turn to approximation, with Monte Carlo estimation being the most widely used method.
The Monte Carlo (MC) estimator for KL divergence (Eq. (I)) involves sampling M strings p(Y(m))
y'i),.. ... Y(M) iid. p and then averaging log q( Even though this estimator is unbiased, it often
exhibits high variance, which means the approximation can be noisy and unreliable." 4 More patholog- ically, the naive MC estimator can result in negative estimates of KL, which may be undesirable in practice. 5 To address these issues, practitioners adopt alternative techniques to ensure non- non-negatin negativity. For example, Schulman [33] proposes an unbiased, non-negative KL estimator that is widely used in practice [13, 40). However, the proposed method, in its original form, does theoretically yield an estimator with lower variance, and, as we show empirically, can exhibit enormous variance.
In this paper, we derive improved estimator of KL using Rao- Mho-BBaacwellizatiinn [RB; 3, 11, 30], a well-established variance-reduction technique from statistics. 6 This results in an estimator that is provably unbiased and has variance that is always less than or equal to that of the standard Monte Carlo estimator, while requiring no additional computational overhead. As a point of comparison to RB estimator, we also provide a comprehensive formal analysis of various existing methods for estimating KL divergence, examining their bias and variance.
We empirically validate our theoretical findings using the sentiment-controoled generation task [29] as a testbed. Specifically, we measure the KI. divergence between a GPT-2 model [28] before and after fine-tuning, where the fine-tuning objective is to steer the model toward generating positive movie re- views. Our experimental results confirm that our proposed estimator significantly reduces the variance of the Monte Carlo estimator, yielding the most stable and reliable estimates among all methods stud- ied, In contrast, alternative estimators from the literature fail to achieve meaningful variance reduction, and in some cases, lead to unbounded variance. We further examine how using our derived estimator in the fine-tuning loop of RLHF impacts the downstream performance. Our results suggest that using our Rao Blackwellized estimator reduces the instability across different RLHF runs. We further look
"Throughout this paper log denotes the natural logarithm function; thus, KL divergence is measure in nats rather than bits. We also note that terms of the form p(y) log v(v) in Eq. (1) where p(y) = 0 can correctly be taken a(v)
to equal zero because limp-ut plogp/q O.
In conditional tasks like dialogue generation, language models are prompted with an input æ E E', inducing a conditional distribution p(- x). KL divergences are typically averaged over a set of prompts. For simplicity, we omit æ in notation and write p(y); all of our results extend straightforwardly to the conditional case.
In general, computing the KL divergence between two arbitrary LMs exactly undecidable. To see why, assume that each of the two language models is probabilistic context- free grammar. In this case, deciding whether their KL. divergence is zero is undecidable, as it follows directly from the undecidability of testing equivalence between two unweighted context-free grammars [155]. In the more restrictive case of probabilistic finite-state language models, it is PSPACE-hard. Importantly, however, the intractablity of exact computation does not imply that approximation is intractable. In practice, one can often ohtain good Monte Carlo estimates of the KL divergence, provided that its variance is coll-controlled. We study very practical methods for improving variance in this paper.
"Monte Carlo estimation formally requires that the underlying random variable have finite variance; if the variance is unbounded, the estimator no longer converges in the limit.
'Since the KL divergence is non- non-negative by definition, a negative estimate can be problematic when KL is used as part of a loss function, as may destabilize the learning dynamics.
'Despite its simplicity, our proposed estimator is absent from existing literature and open-source RLHF libraries [13, 17, 34, 40J. highlighting gap believe is worth addressing.

---

at the Pareto frontier of average rewards achieved by the model vs. its KL divergence with the refer- ence model. We observe that models fine-tuned using the RB estimator appear significantly more often on the Pareto frontier of reward vs. KI. compared to the models fine-tuned with the MC estimator.
2 Preliminaries
2.1 Language Models
Let Σ be an alphabet, a finite, non-empty set of symbols. A string a finite sequence of symbols from Σ. Let Σ" denote the set of all such strings. language model p is distribution over Σ' The prefix probability function p of prefix a Σ is
p(x)
Σ p(ay).
(2)
VEE-
which is the cumulative probability of all strings in the language that have æ as their prefix. We denote the conditional prefix probability as ply .)
play)
)a) where we additionally define
P[BOS| y) def p(y) Mvy) Language models can be factored into the product of distributions using the
chain rule probability, i.c., for any string y = y1 YN Σ we can write
p = ( y) II
plun ycn))
(3)
==1
def where Ven y1 yn-1 and EOS & a distinguished end-of-string symbol. Let Σ _ Σu {EOS). In Eq. (3), each pl enn ) can fruitfully be viewed as distribution over Σ. Despite the overloading of the notation, whether p(- yen) refers a prefix probability, function which takes argument from Σ", or distribution over E will always be clear from context. Throughout this paper, to represent the string-valued random variable sampled from When taking M i.i.d. samples from we use y (m) to denote the mh sample.
Y
2.2 Monte Carlo KL Estimation
A simple way to estimate the KL divergence is with the Monte Carlo estimator (MC) defined
M
pf (
M
Puc = M m=1
Σ
log
Σ f(y(m)), m==
(4)
q( Y (m))
M
where (1) .... Y (M) iid p and J(Y) del
Throughout the paper, assume that the KL divergence is finite, i.e., KL(P q) < co. straightforward to show pue is unbiased, i.e.. [[use]] (p || g) and the variance of this estimator is Var]umc] = उ Var[f(Y)]. In App. A, we discuss the Horvitz-Thompson estimator, another unbiased KL estimator.
log p(Y) q(Y)'
Note that while the exact KL value is always non-negative, S(Y) may be positive or negative. Consequently, the Monte Carlo estimate Mse may also be negative. This happens because the estimate is based on limited number of samples, and some sample draws can lead to negative values. This can be problematic during RLHF, which depends on the KL divergence being non-negative.
2.3 Control Variate Monte Carlo Estimation
general approach to reduce estimator variance through control variates [32, $8.2]. For KL
divergence between language models, this technique popularized by Schulman [33] and is widely used in RLHF libraries [13, 17, 34]. Formally, a control variate any function g: Σ - for which
G def E. [g
Y)]
can be efficiently computed.? We define the control variate Monte Carlo estimator as
M
Mcy
1
Σ (m)) + (m(Y(m))
(m)) G).
(5)
M
m=l
"One could also consider control variates of the form g: E" Rd for [10].
3

---

where E IR is calibration parameter that must be chosen priori. The proposition below characterizes the variance of Pey as function of a, which will tell use how to choose optimally. Proposition Consider the control variate MC estimator Pey defined in Eq. (5), and assume that E[g(Y) < oc. Then Pev is unbiased estimator, and its variance is given by
Var [f] + o² Var[gl + 2a Cov [f.g.
Varlucv]
(6)
M
Proof. See App. B.
Assume 0 < Var[g] < 0o. It is straightforward to show that a² Cov [J, gl/ Var is the value that minimizes the variance. If plug a* in Eq. (6), and simplify, see that
Varluev] M
Varlf[ Corr[f,
which directly translates to reducing the variance of the MC estimator. The magnitude of the correlation between and g determines the degree of variance reduction. Note that the value of a" may be estimated from a pilot sample when it cannot be computed analytically."
KL Estimation with a Control Variate. A specific control variate for KL estimation was proposed by Schulman [33], who defined g(Y) the KL divergence with this control variate
x(Y) Substituting this into Eq. (5), the MC estimator of p(Y)
1 M
M
P (m)
(((((m))
-1)
Nev
log
+a.
p(Y(m))
(8)
Remarks. This proposal has some notable properties. First,
1, meaning G' is P(Y)
known in advance. Second, Cov [f = is zero if and only if if equal to q (Prop. 8), meaning that when the two distributions not equal, and is chosen suitably, are guaranteed strictly reduce variance. Schulman [33] proposes setting a = 1; the benefit this suboptimal choice that the resulting estimator is always non-negative (Prop. 7). However, setting a 1 will increase variance when a < 1 (Remark 10). Indeed, our experiments ($5.1) confirm that a = 1 is a poor choice in practice- it is better to estimate a to ensure that the control variate is correctly calibrated.
Rao-Blackwellized Monte Carlo
In this article, we propose the application of another classical technique to reduce the variance of the Monte Carlo estimation of the KL divergence-RRo-BBleece waclllielllliizttiiio [RB; 5, 11]. Despite its standing in the statistics literature, a Rao-Blackwellizer Monte Carlo estimator has yet to gain traction in the context of RLHF [13, 17, 34, 40].
We define the Rao- Mho-Blackwellized Monte Carlo estimator PRu as follows:
[Y]
PABB
def
1
Σ
KL
Y
(m) <n |
(())))
M
The key benefit of this estimator is that we analytically compute the expectation over the 1hth symbol in each string rather than relying on the single sampled value at that position.
nho-Blaakwellizzttion Background. The starting point of noo-Blackwellizzation is the following
inequality involving the conditional variance: Var [E[ T]] S Var[u], where p. is an unbiased
estimator of [[] and T' is a statistic for which we can explicitly compute B8/ | T'. This
"Note that the control variate method can be straightforwardly extended support multiple control variates. Notee that for this to hold, we need have q(y) = 0 whenever p(y) = 0. which is different from the support condition we assumed for KL(p 4) < oc.
Noot that a = is the only value of a that guarantees nonnegativity (see proof of Prop. 7).
"Prop. 9 provides the exact conditions
required for variance reduction.
Note that 7' is function of {Y (m) M We have suppressed this function's 's arguments in our notation for improved readability.

---

technique is often referred to as Rao- Blo--laacceeelzzation because the inequality is associated with the -lackBlll theorem [21]. 13
Notation. Before moving forward, we introduce some convenient notation. Let Y denotes the EOS- EOS-padded version of y Additionally, we extend the definition of the prefix probabilities and o to strings containing padding symbols, i.e., y & D", with the following additional case:
Ply y) def { EOS}, and Gly y) def ( = EOS}.
Understanding Our Rao- Bro-Blackwellized Estimator. We now present proof that our Rao- Blackwellized estimator MRn is unbiased and indeed does result in a variance reduction. One might wonder why this requires more than a straightforward application of the Rao-Blackwell theorem. The reason is that JAи does not arise directly from the standard formulation of Rao Blackwellization. In- stead, we apply Rao R alalztellizttion separately each summand, where each summand corresponds to an estimator over strings of particular length. In general, Rao-Blackweele -ee-Bllllllliiingggg the summands pointwise does not guarantee a reduction in the variance of their sum, since the summands may be correlated. However, in the specific case of Maa, we can prove that the overall variance is reduced, despite Rao- Blackwellizing the summands independently, as we state the following theorem.
Suppose the MC estimator Paac has finite variance, i.e., Var [ar] co. Then the following properties regarding PRo hold:
Theorem 2.
Elpaol = kl.(p ()
(unbiasedness)
(ii) Var[HRo] S
[arcu (variance reduction)
Proof. See App. C for detailed proof. However, provide the proof sketch below for the reader to quickly understand the structure of the argument, which broken down into three steps.
(1) wise Estimation. We begin by Rao-Blackwellizing the wise Monte Carlo estimator
for any n > 0,
(m)
(m)
P
PMC
M
log
(m)
(m)
(10)
Intuitively, Mec measures the average KL of the nh symbol. The next step is to define
T, (Y = Y <n and apply Rao-Blackwellization to each ple as follows:
Pk
def
E (
Tn)
(11a)
((r)
M
p (m)
(mL)
E
log
<n
Ta
(11b)
M
(m)
q( (m) n
(m)
<
M
KL(P(
(m)))
(11c)
M
m==1
Now, it is clear from the Rao- nho-Blaakwellization theorem that pi is unbiased and it provides variance reduction (i.e., Varlulel I Var (uMe)) for all n > 0. Inuitively, the source of the variance reduction in the Mžs estimator is that we compute the expectation over the nh symbol exactly rather relying on the sampled symbol at that position.
(N) del E PaR Intuitively, (2) Truncated Estimation. Next, we define (N) del Σ Pae and PRa
these estimators target the KL divergence for symbols up to maximum length of N. In
Lemma 11, we prove that E = El and Var MRn (N)
(N) for all > < VVr AMC
using the law of total expectation and Jensen's 's inequality, where the latter is used in manner analougsly to have it is used in the original Rao- Noo--Baackellizationn theorem.
1 Note that in the Rao Blackwell theorem, we get the stronger result that when is a sufficient statistic, Var[E[H []] is optimal, i.e., it is the mininal-variane unbiased estimator. However, we can perform Rao Blackwellization even when I' is not sufficient and are still guaranteed that the variance is worse [31].
5

---

(3) Complete Estimation. Now, we consider the complete estimation. First, we observe that (N)
the limit of the truncated estimators converge PMe lim PMC (Lemma 12), and analogously, jRB = yim Muu estimator of the KI. divergence with variance less than equal to that of juac
( Thereby, we are able to show that PRa is an unbiased
Remarks. Notably, Pas is guaranteed to be non-negative, property desired by some when de- signing estimators for the KL divergence between two language models, as the KL divergence itself is always non-negative (ef. remarks in $2.3). In the case of our Rao-Blackwellizeer estimator, non- nonegativity follows from the fact that each step-wise estimator computes the exact KI. divergence between the two next-symbol distributions, conditioned on the sampled context y Since all terms in Eq. (1 Ic) are non-negative, jsg remains non-negative as well.
Complexity Analysis. Although computing (ga might seem more expensive than AMe, the overall runtime is dominated by the cost of forward passes. Because forward pass already involves producing the full distribution over Σ at each position n, the additional O (MM[E) work required by Junn is negligible compared to the M forward passes needed for both jlev and juac-
4 Estimating the Gradient
KL estimation essential in many applications, especially in fine-tuning large language models. In reinforcement learning from human feedback (RLHF), for example, the objective includes KL regularization term to balance reward maximization with staying close to a reference model. Since the language model is differentiable function of parameters 0 optimized via gradient descent, this setup requires computing the gradient of the KL divergence with respect to 0:
po (Y)
G deff (pKL(pe || q)
E log
Va lo (Y)
(12)
((Y)
Therefore, the Monte Carlo estimator of this gradient is
M
pol ( (m)) (((m))
SM
Σ
log
Vo log po (m)))
(13)
Now, we derive the Rao Blackwellized Monte Carlo estimator of the gradient. We start with restating Theorem 2.2 in Malagutti al. [24], which will prove useful.
Theorem 3 (Malagutti et al. [24]; Theorem 2.2). Let and q be language models over Σ. The following equality holds
KL 9) FO)KL 9) 1[ v)),
(14)
NESR
where we treat p( y) and q( y) as probability distributions over Σ'.
We refer the reader to Malagutti et al. [24] for the proof. Next, derive the Rao-Blackwellized estimator of the gradient, we take the gradient of the local KL as stated in the following theorem.
Theorem 4. Let pe and q be two language models over E and pe the prefix probability function of Po- Then, the following equality holds
(Y y) ( y) (Vo Volooger (Y v) + Ve Volog
VOKL(Pe 41
Σ DEE*
E po(y) &
log
(15)
Proof. See App. D.
We then construct the Monte Carlo estimator of the gradient using the above theorem, which naturally results in the following unbiased, Rao-Blacke Blackwellized Monte Carlo estimator of the gradient:
M
(m))
log
[Y
(m) Y (m))
Vo log Po (
(m)) <n + Vo log Po (Y
<n
(m))
].
M
m=1
n=1
<n
(16)

---

The above estimator is unbiased, as it is the MC estimator of Eq. (15).
Theorem 5, Assuming [ar[dus] < 0o, Var [d
<0,, have
E |xx - G1²] s E |l0mc GG12]
(17)
Proof. See App. D.I.
Off-policy Gradient. So far, we have discussed how estimate the gradient of KL || q) using samples drawn from the current policy Pe. Crucially, we derive the gradient manually rather than relying on automatic differentiation because the samples depend on 0 through pe- However, in practice and for efficiency reasons, we often collect large batches of samples in parallel with the optimization loop. result, these samples are generated from a slightly outdated version of the policy, denoted Po.s- To compute the KL. divergence using samples from Po.., we first write the KI. as the expectation under Pê.. as
pe (Y)
log Po( PY(Y)
<p 11
E Yopeew
(18)
Pea (Y)
((Y)
Therefore, the MC estimator using samples ( (M) lid. Peu is
old M Po (m()) Psc
Po ()((m))) q(Y (m)
log
(19)
M
m=1 Poas Y (m))
Given the unbiasedness proof of the Rao- Bho-Blaakwellization in Thm. 2, we can similarly write
(m()))
po
YY)
Po Y <n)
KL q)
E Yoreau
E
log
(20)
pe
Y
<n)
Therefore, the Rao Blo-Blackwellizzd MC estimator using samples Y (M) itd.
a0
is
PRB
old
M
pe
( (m)) (Y (m))
(m
E Y
log
PY
(m)
(21)
M
(m))
m=
Pe
Since a and will use samples from the old policy that does not depend on O, we can apply automatic differentiation to compute the estimate of the KL gradient by computing the gradient of pold and pos
Experiments
We use the sentiment control task as the testbed to empirically evaluate our theoretical findings on the KL estimators. Concretely, the reference model, denoted as q, the GPT-IMDB model, i.e., GPT-2 [28] model fine-tuned on IMDB corpus [23]. The goal of the task is to fine-tune this language model such that the samples from it are movie reviews with positive sentiment. The fine-tuned language model is denoted with pe. In the following experiments, we estimate the KL divergence between pe and q. We provide code snippet for implementing the RB estimator in App. F.I.
5.1 Analyzing the KL Estimators
In this experiment, we empirically evaluate the bias, variance, and consistency of various KL estimators. To obtain pe. fine-tune q with direct preference optimization [DPO; 29] sample of 5,000 data points from the IMDB training set. To create the preference data required for DPO training, following Rafailov al. [29], sample responses for each prompt and create pairs per prompt.
"Speciffiallyy https: httpss///hgggnggaee.eeooo htwp://///aa.../o///////////////

---

To determine the preferred response in each pair, we employ binary sentiment classifier. 15 selecting the response with the higher probability of positive senti- ment. Upon successful fine-tuning, Po should assign higher probability mass to movie reviews with posi- tive sentiment while maintaining low KL divergence with q. We then evaluate this KL divergence using our estimators to assess their reliability in measuring distributional shifts induced by fine-tuning
Table Estimated value + empirical stan- dard deviation of different estimators. When aggregating over prompts, Par and Pey fail to significantly reduce the variance of Puc- RB estimator, however, achieves the lowest standard deviation.
M
M 10
6.76 - 0.16
6.76 0.07
6.76 + 0.05
We evaluate on 512 examples from the IMDB dataset. For each review, we randomly select prefix length between and tokens and use it as the prompt. we then generate 4000 samples from Po for each prompt.
Par
6.76 = 0.16
6.76 -.767
6.76 - 0.05
Rev1
6.28 = 2.54
6.28 = 1.13
6.28 + 0.79
6.76 0.16
6.76 = 0.07
6.76 = 0.05
6.76 0.11
6.76 + 0.05
6.76 + 0.03
Using these samples, we compute the MC, control
variate (CV), and weoellaawwellized (RB) estimators and estimate their standard We also implement the Horvitz-Thompson (HT) estimator; see App. A for details. The CV estimator, Pev, is computed twice: once using the optimal estimated from 1000 samples, and once with o to match the setup in Schulman [33].
In Tab. 1, we report the expected KL estimate along with the empirical standard deviation of different estimators evaluated at sample sizes
RB
M = 1, 5, 10. To obtain these estimates, we compute each estimator using M samples, re- peating the process 4000/M M times to estimate both the expected value and the standard devia- tion of the estimates. Our findings confirm that 1), are un-
MC
CV
Cy
HT 11
all estimators except one (Mcv, biased and report an expected KL divergence of 6.76. We also observe that the CV estimators fail to significantly reduce the variance of the standard MC estimator. Importantly, the RB es- timator achieves the lowest standard deviation and offers more robust estimate compared to the MC estimator. Interestingly, we observe that the Mey estimator exhibits a noticeable bias and high standard deviation when a i.e., when is not its optimal value. The bias arises from .(Y)
100
Standard Deviation (log)
Standard deviation of KL estimators across various prompts in the IMDB datasest.
Figure
1:
numerical instability during the computation g(Y) = p(Y) The high variance due to large
values of Var [ar Y)). Specifically, for certain prompts, Var (arlg(()) Y)) can be unbounded. We visualize
the estimates for example prompts in App. E.
Since the robustness of the estimators depends on the choice of the prompt, we further analyze their estimated standard deviations across all prompts. Fig. I presents box plot of standard deviations (in log scale) for each estimator. The Moyy estimator with = shows significant instability for certain prompts, with numerous outliers indicating high variance. In contrast, the pae and the standard pev estimators exhibit comparable standard deviations. In particular, the Boo-Blackwellized estimator con- sistently achieves the lowest standard deviation, suggesting that it provides the most stable estimates.
5.2 KL Estimation and RLHF Training Dynamics
A key application of KL estimation is in the RLHF training loop. From the previous experiment, we observed that the RB estimator significantly reduces the standard deviation of the MC estimator. Therefore, it is natural to ask how this affects RLHF performance when this estimator is used in the training loop. The RLHF objective consists of two terms: (i) the expected rewards for samples generated by the language model Po. which in this case is the samples' score under a sentiment 16 and (ii) the negative KI divergence between the fine-tuned model po and the reference model which represents the language model before fine-tuning.
15 "Specifically, we use https: / httgggggggaaccceeeeee co/ ww//wwerra/distilleeee "" "Specifically, we look the logits of the positive class.
8

---

12
MC
10
Average Reward
Klowlll
100
1.0
50
100
150
Step
Step
Figure 2: Comparison of the Monte Carlo (MC) and Rao-Blackwellized (RB) estimators in the RLHF fine-tuning loop. We perform RLHF with each estimator times and plot the mean and standard deviation (in shades) of the average reward values and the KL at each fine-tuning step. We observe that the MC estimator is not stable the RB estimator and its performance varies significantly across different runs. However, RB estimator reliably offers a good balance between achieving low KL and high reward values in all runs.
We compare the MC and RB estimators for computing the gradient of the KL divergence term in the RLHF training loop. We use the RL algorithm" proposed by Ahmadian et [2], which an improved verfaiion of the REINFORCE algorithm [41]."
First, we empirically test Thm. 5 by measuring the variance of the gradient norm. We sample 40 prompts and compute the gradient of the KL divergence with respect to the model parame- ters using both the MC and RB estimators. To estimate variance empirically, we repeat this pro- cess times. Both estimators are evaluated on the same prompts and model initializations to en- sure a fair comparison. We find that the variance of the gradient norm estimated with the MC esti- mator is 59.90, whereas with the RB estimator is 45. 44. This corresponds to 24.6% reduction in variance, providing direct empirical evidence consistent with our theoretical motivation.
MC
2,5
RB
2-
Average Rewards
1.5
0.5
KL. range
Stativtc
P-value
All
0.78
xx 10-5
0-
(0,6)
0.56
x 10
16
0.5
-0.5
0
10
15
We then proceed with using both estimators in RL fine-tuning. We track two metrics throughout fine-tuning: (i) the average reward associated with samples drawn from pa. and (ii) the KI. divergence between Po and y. The results are visualized in Fig. 2. The purple trace represents the training run where the Mxu is used in the optimization loop to estimate the gradient of the KI. divergence, while the blue trace represents the run using Puc- The x-axis denotes the fine-tuning step. with the left plot showing the evolution of the average reward and the right plot displaying the KL divergence between Po and q over the course of fine-tuning. We repeat each experiment times and report the mean and standard deviation of each metric. Notably, the KL values in the right plot are estimated using the RB estimator. However, we observe the same overall trend when using the MC estimator for evaluation.
Figure 3: Compared to models trained with MC esimator, models trained with RB appear on the Pareto front 78% of the time.
As illustrated in Fig. 2, the performance of the models trained using the MC estimator varies significantly across the 5 experiments, resulting in large standard deviation in both average rewards and the KI. divergence. However, estimator consistantly achieves high rewards and reasonable KI. values across all runs. This observation suggests that the RB estimator makes RLHF runs more stable.
App. D.2 discusses common mistake when Rao- hho-Blackwellizingg the KL estimator in trust-region algorithms. 18 #Specifically, we use the available implementation this algorithm the trl library [40]. App. F reports the hyperparameters used for the algorithu.

---

Finally, we vary the KL coefficient, 3, in [0.01, 0.1] range and fine-tune 18 models with each estimator. For each estimator, we plot the Pareto frontier of average rewards versus KL divergence in Fig. 3, displaying models that do not appear the Pareto front with reduced opacity. Overall, we find that fine-tuning with the estimator is more effective at achieving high rewards while maintaining low KI divergence from the reference model. To quantify this effect, we compute the fraction of RB fine-tuned models that appear on the overall Pareto front-i.e., the frontier obtained when considering all models fine-tuned with either estimator. We then conduct permutation test and report the results in Fig. 3. We find that 78% of the points on the overall Pareto front come from RB fine-tuned models. Restricting models with KL values below this fraction rises to 95%, with both results being statistically significant.
Conclusion
In this paper, we study the problem of estimating the KL divergence between language models. We provide a comprehensive formal analysis of various KL estimators, with a focus on their bias and variance. We introduce the RB estimator, which is provably unbiased and has variance most equal to that of the standard NC estimator. This estimator applies the well-known technique to reduce the variance of the standard MC method. Our empirical results show that the RB estimator significantly reduces the variance compared to the MC estimator, while other estimators fail to achieve meaningful variance reduction or, in some cases, suffer from unbounded variance. Additionally, we find that using our proposed RB estimator makes RLHF more stable and produces models that more frequently lie on the Pareto frontier of reward versus KL, compared models fine-tuned with the MC estimator.
Blao-laawwlllizttionn
Impact Statement
In this paper, investigate the fundamental problem of estimating KL divergence between language models. One key application of KL estimation is in RLHF, which aims to enhance fluency while aligning language models with user preferences. However, RLHF can also be misused by bad actors to optimize models for generating misleading, biased, or harmful content. While our work provides a deeper understanding of KL estimation techniques, it is purely foundational research and does not introduce new risks or directly contribute to harmful applications.
Limitations
In RLHF experiments, evaluating the variance of estimator and comparing it to existing methods requires training a large number of models. For instance, the significance test in $5.2 involves training models. Due to limited computational resources, we used the antrolled-generaaton task as a testbed. Given the strength of both our theoretical and empirical results, we hope future work will adopt the wlakkllawwellized estimator and apply to larger language models and wider variety of RL-baased approaches to LM alignment.
Acknowledgements
We thank Ahmad Beirami and Cristina Pinneri for the insightful discussions throughout the course of this project. We also thank Alexander K. Lew for the valuable feedback on a draft of this paper. Afra Amini supported by the ETH AI Center doctoral fellowship.
References
[1] Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. 2024. On-policy distillation of language models: Learning from self-generated mistakes. In The International Conference Learning Representations.
[2] Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. 2024. Back to basics: Revisiting REINFORCE style optimization for learning from human feedback in 1lms.
10

---

[3] David Blackwell. 1947. Conditional expectation and unbiased sequential estimation. The Annals of Mathematical Statistics, 18(1):105-110.
[4] Nadav Borenstein, Anej Svete, Robin Chan, Josef Valvoda, Franz Nowak, Isabelle Augenstein, Eleanor Chodroff, and Ryan Cotterell. 2024. What languages are easy to language-modee perspective from learning probabilistiic regular languages. In Proceedings of the Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).
[5] George Casella and Christian P. Robert. 1996. Rao-Blackwellisation of sampling schemes. Biometrika.
[6] Paul F. Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. 2017. Deep reinforcement learning from human preferences. In Advances in Neural Information Processing Systems.
[7] Imre Csiszár. 1967. On information-type measure of difference of probability distributions and indirect observations. Studia Scientiarum Mathematicarum Hungarica.
[8] Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Bingxiang He, Wei Zhu, Yuan Ni, Guotong Xie, Ruobing Xie, Yankai Lin, Zhiyuan Liu, and Maosong Sun. 2024. UtraFeedback: Boosting language models with scaled AI feedback. In Proceedings of the International Conference on Machine Learning-
[9] Kevin Du, Vésteinn Snæbjarnarson, Niklas Stochr, Jennifer White, Aaron Schein, and Ryan Cotterell. 2024. Context versus prior knowledge language models. In Proceedings of the Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers).
[10] Tomas Geffner and Justin Domke. 2018. Using large ensembles of control variates for variational inference. In Advances in Neural Information Processing Systems.
[1]] Alan E. Gelfand and Adrian F. M. Smith. 1990. Sampling- approaches to calculating marginal densities. Journal of the American Statisticul Association, 85(410):398-499
[12] Alison L. Gibbs and Francis Edward Su. 2002. On choosing and bounding probability metrics. International Statistical Review Revue Internationale de Statistique.
[13] Alexander Havrilla, Maksym Zhuravinskyi. Duy Phung. Aman Tiwari, Jonathan Tow, Stella Biderman, Quentin Anthony, and Louis Castricato. 2023. trlX: A framework for large scale reinforcement learning from human feedback. In Proceedings of the Conference on Empirical Methods in Natural Language Processing.
[14] John R. Hershey and Peder A. Olsen. 2007. Approximating the Kullback Leibler divergence between Gaussian mixture models. In Proceedings of the IEEE International Conference on Acoustics, Speech and Signal Processing.
[15] Hendrik Jan Hoogeboom. 2015. Undecidable problems for context-free grammars.
[16] D. G. Horvitz and D. J. Thompson. 1952. A generalization of sampling without replacement from finite universe. Journal of the American Statistical Association.
[17] Jian Hu, Xibin Wu, Zilin Zhu, Xianyu, Weixun Wang, Dehao Zhang, and Yu Cao. 2024. OpenRLHF: An easy-to-use, scalable and high-performance RLHF framework. Computing Research Repository.
[18] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang. Timothée Lacroix, and William El Sayed. 2023. Mistral 7b. Computing Research Repository, arXiv:2310.0692)
[19] S. Kullback and R. A. Leibler. 1951. On information and sufficiency. The Annals of Mathemati- cal Statistics.
[20] Daniel J. Lehmann. 1977. Algebraic structures for transitive closure, Theoretical Computer Science.
11

---

[21] E. L. Lehmann and George Casella. 1998. Theory of Point Estimation.
[22] Zhifei Li and Jason Eisner. 2009. First- and second-order expectation semirings with applications to minimum-risk training on translation forests. In Proceedings of the Conference on Empirical Methods Natural Language Processing.
[23] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. 2011. Learning word vectors for sentiment analysis. In Proceedings of the Annual Meeting of the Association for Computational Linguistics: Human Language Technologies.
[24] Luca Malagutti, Andrius Buinovskij. Anej Svete, Clara Meister, Afra Amini, and Ryan Cotterell. 2024. The role of n-gram smoothing in the age of neural networks. In Proceedings of the Conference of the North American Chapter of the Association for Compulational Linguistics: Humun Language Technologies (Volume Long Pupers).
[25] Meta. 2023. Llama 2: Open foundation and fine-tuned chat models. Technical report, Meta.
[26]
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang. Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul F. Christiano, Jan Leike, and Ryan Lowe. 2022. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems.
[27] Pouya Pezeshkpour. 2023. Measuring and modifying factual knowledge in large language models. Computing Research Repository, rcXiv:230600624.
[28] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners.
[29] Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D. Manning, and Chelsea Finn. 2023. Direct preference optimization: Your language model is secretly a reward model. In Advances in Neural Information Processing Systems.
[30] C. R. Rao. 1945. Information and the accuracy attainable in the estimation of statistical parameters. Bulletin of the Calcutta Mathematical Society, 37:81-91.
[31] Christian P. Robert and George Casella. 2004. Monte Carlo Statistical Methods.
[32] Sheldon M. Ross. 2002. Simulation.
[33] John Schulman. 2020. Approximating KL divergence. http:///acch..nt////lgkkkaaaaaa
html.
[34] Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. 2024. HybridFlow: flexible and efficient RLHF framework. Computing Research Repository, arXiv:2409. 19256..
[35] Nisan Stiennon, Long Ouyang, Jeffrey Wu. Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F. Christiano. 2020. Learning to summarize with human feedback. In Advances in Neural Information Processing Systems.
[36] Niklas Stoehr, Kevin Du, Vésteinn Snæejarnarson, Robert West, Ryan Cotterell, and Aaron Schein. 2024. Activation sealing for steering and interpreting language models. In Findings of the Association for Computational Linguistics: EMNLP 2024.
[37] Anej Svete, Nadav Borenstein, Mike Zhou, Isabelle Augenstein, and Ryan Cotterell. 2024. Can transformers learn n-gram language models? In Proceedings of the Conference on Empirical Methods in Natural Language Processing.
[38] Lewis Tunstall, Edward Emanuel Beeching, Nathan Lambert, Nazneen Rajani, Kashif Rasul, Younes Belkada, Shengyi Huang, Leandro Von Werra, Clémentine Fourrier, Nathan Habib, Nathan Sarrazin, Omar Sanseviero, Alexander Rush, and Thomas Wolf. 2024. Zephyr: Direct distillation of LM alignment. In Proceedings of the Conference Language Modeling.
12

---

[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems.
[40] Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, Shengyi Huang, Kashif Rasul, and Quentin Gallouêdec. 2020. Trl: Transformer reinforcement learning. https: wtw.ttuu.cc... com/huugiigfface/trll
[41] R. J. Williams. 1992. Simple statistical gradient-following algorithms for connectionist rein- forcement learning. Machine Learning.
13

---

A Horvitz-Thompson Estimation
When estimating the KL divergence between language models p and q, have access not only to samples from p but also to the probability assigned by p to any string y- This enables the use of more informed estimator, which leverages these probabilities during construction. Notably, this estimator is a specific instance of the Horvitz-Thompson (HT) estimator [16], defined as
ply)
p(y)
ply)
PELT
log
f(y),
(22)
(y)
(3))
DES
WES
where S is the random variable representing the set of all sampled strings. Any sampling design
can be specified to generate the elements of S. The inclusion probability, denoted by TT the probability that particular string is included in S, equivalently, E (y E 5)}-
Proposition Par is an unbiased estimator of the KL divergence, i.e., E [PTT] (L q).
Proof. The bias of the estimator is as follows:
Pur - KL(p 4)
p(y)
log p(y) g
kL(P ll g)
(definition of par)
(23a)
Ts (
p(y)
log ply)
1yy E s}}
-KL 9)
(23b)
It's (y)
(v)
DEM-
p(y)
log p(U)
Σ
E
[ s}}
KL(p 1| g)
(linearity of expectation)
(23c)
(yy)
((y)
3
Σ
p(y)
log
p(y)
xs(rt KL 91
(definition of (s()))
(23d)
(yy)
yET-
0.
(23e)
Similar the MC estimator, the HT estimator does not necessarily return non-negative estimate of the KL. In principle, however, we should prefer Eq. (22) to Eq. (4) because it exploits more information namelly the knowledge p. Whether the HT estimator yields lower variance than the
MC estimator depends on the sampling design used to construct S. In our experiments in App. E, we used the sampling-withe with-replacement design, where T3()) MC estimator, we observed no significant reduction in variance in our experiments.
(1 p(y)) Compared to the
14

---

B Control Variate Monte Carlo Estimation
Proposition 1. Consider the control variate MC estimator Pev defined in Eq. (5), and assume that E[g(Y)) < cc. Then Pey is an unbiased estimator, and its variance is given by
Var [f] + a² Var[gi + 2a Cov [f.
Var[uev]
(6)
M
Proof. Recall that the control variate Monte Carlo estimator is defined as
1 M M mmm
Mcv
Σ f()(m))))))))))))))))))))) + (9( G).
(5)
Note that (y M m=
ltd p. First, look at the expected value of the estimator:
M
Elpevl
ΣE E [((((())))) m==1
+ a (E (E))())))) G) Effl
(24a)
M
0
Therefore, it is unbiased. Next, we manipulate the variance
Varluev] 1 Varls + a (g G)] M
2a
M Var[s]
M Varlgl +
Cov[Sgg],
(24b)
M
where the second equality above stems from well-known variance identities." 20
Proposition 7. Provided that the target KL divergence finite, the estimator pev with is nonnegative with probability one. Furthermore, is the only value for which we can guarantee nonnegativity of pev.
Proof. To prove that Pey is nonnegative with probability one when = 1, we will prove that each
term log (m))
p(Y(m) + (((((m)))
-1) in the summation that defines μev is nonnegative.
(((Y(m))
Let r 2 Y (m))
pm) We identify a for which log + a (r- 1) 2 0 holds for all r >> 0. Equivalently,
we want to find a values for which:
inf logr + a (r 1) 2 0. rS0
(25a)
Next, prove that 1 is the only value satisfies the inequality Eq. (25a). To why, consider the first-order optimality conditions for the minimization over r. These conditions are necessary and sufficient log + a (r 1) is convex for all over the region O. Solving for such that the derivative zero gives us
i
[ --ogrr +a (r- 1)]
(25b)
Dr
1/r+a +
(25c)
1/a.
(25d)
Plugging that value allows us to simplify away the minimization:
0 W inf logr + a (r (r-1) 1)
(25e)
r>0
= =- log(1 (1/g)1/)00 - a (1/a 1)
(25f)
= log(a) + -1-a. a.
(25g)
For each term in Mey to be nonnegative, need log(a) log(a) (x) + 1 a is strictly concave and nonpositive for all a > 0, attaining zero only at a = 1. Thus, when a 1 every term in Mlyv nonnegative. When a + some terms may be nonnegative, leading to the possibility that Pev may be negative.
a However, the function
Let X and be real-valued random variables, let be real scalar. Then. the following hold: Var[X Y] = Var [Xr[X]] + Var[Y]] [V]++ + Cov [X, Y). Var[aX] = Var (X), and Var X + aaj Var[X],
Note that 0 whenever this term is finite.
15

---

Proposition 8. Given the random variable Y ~ ps let f(Y) prove Cov [f g] is zero if and only if q.
log p
and g(
p(Y). We
2(22
y
Proof.
log p(Y)
log p Y
((
Cov[f.g] E
E
(26a)
g(M)
q(
p( YY
log p(x))
g(
E
log
log
1 (( ]
E
log
(26b)
p(Y)
(
((Y)
p
q(
log
IE log 9
E
log
(26c)
p(Y)
(
p(Y)
Y
7( (
P
P log
E
log
(26d)
P
g
(L p) (L q).
(26e)
Since each KI divergence is non-negative and equals zero if and only if the two distributions coincide almost everywhere, we have,
Cov [J =0 - -kL(y p) KL 4) =0 = p=
(26f)
Proposition 9.
Var
(0)) (ov
S Var [Mac]
= a E [min(0, 2a*). max(0, 2a)]
(27)
Proof. We first establish conditions of a for which the variance of Var [ucv] does not increase that of Var[ In other words, seek conditions such that the following inequality holds:
Var
uev
I Var pMc
(28)
Var [f] a Var [g] + 2a Cov [f < Var
(29)
M
M
Var - a² Var + 2a Cov [f, S Var[f]
(30)
a² Var[g] + 2a Cov [f < 0.
(31)
Observe that the function the left-hand side is quadratic in a, and, moreover, it is convex because Covif.g) We can use the quadratic formula
Var[g] 2 The minimum of this quadratic is a*
Var [ to identify the two values of a where it equals 0, i.e., E {0, 2x J. Now, because the quadratic is convex, we have that it is < 0 for values of a E [0nin0, 2a"),, max(0, ). Note that a* can be positive or negative; hence the min and max.
Remark 10. When does Schulman's 's (2020) suboptimal choice of a = not hurt estimator variance? To answer this question, we substitute a 1 into the variance of μcv given in Eq. (6), we have
Var[S] + Var[g] + 2 Cov
Var
(32)
M
Therefore, for Var [ev]]= I Var [M we must have
Var [ʃ] Var [g] + Cov [F1]]
Var[/] M
<
(33a)
Var [f] + Var [g] 2 Covlf.gl Varif]
(33b)
16

---

Var(gl + 2 Cov [f, gl
- 0
(33c)
[f.g] 1
(33d)
Var [g]
2
1
a
2'
(33e)
Therefore, choosing 1 does not hurt the variance if a 2 21 but does otherwise.
17

---

C weoellzac Estimator
Lemma 11. With regard to the estimator Pas the the following two properties hold for all N >
1. Blugg (
(unbiasedness) (yariance reduction)
Var (MRD [N)] S Var
(N)
PMC
Proof. In the proof below, consider the case where M = 1. The proof easily generalizes to M > using the i.i.d. assumption. Additionally, we write Pia (Y), rather than suppressing the argument to the estimator as we do in the main text. We begin by proving the unbiasedness property of the estimator.
)
1 E PRB
E
MC
<n =Y
(definition of (N),
(34a)
N
MMc (8)
Yn <n = =Y'n]
(linearity expectation)
(34b)
N
E
[ (Y)
(law of total expectation)
(34c)
E
Mc
(linearity of expectation)
(34d)
Elume (N)
(definition of M
(34e)
Next, we prove the variance-reduccion property:
Var
(N) PRB
Y 1) <m
27 <m1))
E Y
Σ
MC
Y <n
B (N
Biyh )2
(35a)
(defintion of (N) and unbiasedness)
N
<n 1)
27
E
Σ
E
MM
Y <n
Elugh (2))
(35b)
(each 8 is distributed independently and identically to Y)
N
Y Y211 1)
1))
E
M (Y")
.
Y <N
+
Bl
(N))
(35c)
Y'
(lincority expectation)
E
E
E A.V
Σ PMC (PmH)
V <N
Y
Y Y
(35d)
(Jensen's 's incquality)
N
27
...
E
Σ
M
(
E[paec (N)
(35e)
(law of total expectation)
N
27
E
Σ
(
E ju (2))2
(35f)
18

---

((Y J are lidd)
E
(1)))()))
( (Haic (
(8))
Bi Bisuac 12
(35g)
(detinition of Maic!
Var (N)
(35h)
(definition of variance)
Lemma 12.
Σ MMC = Psc n=1
(36)
Proof.
M
pl (m())
PMc
M
log
g( (Y (m))
(37a)
P(EOS
Y (m))
(mm))
p( (m)
(mm)
Y
[m)
log
log
<n
(37b)
M
(EOS
(V (m)
((m)
M
(m))
Σ
log
(37c)
M
m=1
n=l
(m)
Σ
log
lim
II
(37d)
M
N+00
M
N
lim N
log
(37e)
M
m==1
M
pl
<N
lim
log
(37f)
No
M
m
<N
lim
PMc
(37g)
N
Note that going from Eq. (37b) to Eq. (37c), we use the padding construction given in $3.
Theorem 2. Suppose the MC estimator Pac has finite variance, i.e., Var [ar[) < co. Then the following properties regarding pus hold:
(1) E [RB] = (L || q) (unbiasedness)
(ii) Var I Var [arce (variance reduction)
Proof. In this proof, consider the special case of M 1. The proof easily generalizes M > 1 using the i.i.d. assumption. Additionally, write Mau (Y), ), rather than suppressing the argument to the estimator as we do in the main text. We begin with proving the unbiasedness of the estimator, using Lemma 12 and Lemma 11.
E PAR Y
E
lim
Y
Y <n
(38a)
N
<D
vt
lim
E
M
Y
<m
Y 2n
(38b)
N00
Y
(Tonelli's Theorem)
19

---

lim
[ (x)
(38c)
N +00
(Lemma
N
lim
MM
(38d)
N
1000
n=
(Fubinn's Theoren, >
MMin (T) < o0)
Umcc Y)]
(38e)
(L 9).
(38f)
(unbiasedneess pse))
Finally, prove the variance-reduction property.
Var
E r
lim N-+00 n=
E [ (Y) Ynn v''nl - -KL()ll))
(39a)
(Definition of Men (Y))
N
27
= lim E
Σ
M
Yen=
<n - (k((pll))
(39b)
N-00
n==1
(dominated convergence theorem, Var [122] 00)
27
lim E - Y
Σ MMc (Y) kL(p.g)
(39c)
(Lermna 11. variance reduction)
lim
Σ Msac (Y) - KL(P.9) n=11
(39d)
(dominated convergence theorom, Var [aNas] < 00)
=E (AMc (Y) - KL (l))
(39e)
(Lemmna 12)
-Var(uac)
(39f)
20

---

D wlaelizck Estimator of the Gradient
Theorem 4. Let po and q be two language models over ا and po the prefix probability function of po- Then, the following equality holds
Pe (Y (Y| y)
VOKI.(pO (1 = -E Foly) E log
q y) (Vo (Y y) + Vo log log
(15)
NEΣ*
Proof.
(okl.pp11g)
=E KL(Po(- 8) gt
Vol))) [
v))Vo(o(y))
y) q(- y))
(40a)
yE
Po (T y)
Σ ye
( 3) |l @
(o(
V VoPoy y) log
(40b)
g y)
DEΣ
Pe (T y)
Σ НEΣ.
kL 1) ll at V + (v)
Σ
log
V Vopo v) + Vopo y)
q0 y)
VE
(40c)
(Y y)
E KL a (dll1)))Voool)) + (
log
+1
V 0 log To (Y y))
q(F y)
VEΣ*
(40d)
Σ
(L 3) c (ll11)) (y)
Fe ( E log
(Y
y)
Vo log (Y 5)
(40e)
y)
YEΣ*
log
Po y)
V Vo)o(y) (0 E
log
po Y
y)
Vo log Fo ( 2)
(40f)
QY y)
Y
y)
YEΣ*
= Σ
(y E log
Po Y
y)
Vo loggFo + (o) E log
Po Y
y)
Vo log Po (Y y)
9 Y
g)
y)
yEE-
(40g)
=
Σ BEΣ-
(o(y) E
log
Po (Y Y
y)
(Vo log (Y ()) + V Vologio(y))
(40h)
9 Y
y)
D.1 Variance-Reducee Proof
The proof structure is as follows: We first prove that the inequality holds when constrain Y to have length less than or equal to N. We then generalize to the infinite-lnngth sequences by analyzing as N oc. We begin with defining the truncated MC and RB estimators. Let 6(N2 be the truncated MC estimator of the gradient:
N
(m)
(m)
log
<n
Vo
log Po
(m)) <N)-
(41)
MC
M
(m)
m))
q
Let be the truncated RB estimator:
(N)
Y (m)
8
(NN)
E
log
(Y
<n
Vo (o
( (mm) + Vo +oloogoo
( (m))
RB
M
m))
<n
(42)
21

---

Lemma 13. The truncated MC estimator of the gradient 8 dac. converges to dac as N goes to oc, ie.
1im 6/c)
Ôucc-
Proof.
M Σ
Pe ( (m))
8 MMC
log
Vo log Fo (m))
(43)
M m=1
q ((m)))
M
Po
Y )
(r(m)]
Po PY, Y ((m)
Σ
log
II
Ve log po (m))
(44)
M m=1
(EOS
Y (m)
TY,
Y (m))
M
II Pel (Y
V log Po ( (m))
log
(45)
M m=1
(m))
TY,,,
<n
M Σ
N II
Po (Y,,,
(m))
Vo log po (Y (m)
log lim
SNN)
(46)
M m=1
N 500 n=1 q(
(m))
<a
M N0000
1 M
Σ
lim
log
n=1 (F,
II
Po V
Y
v(m)
<n
Vo log Po Y
.NN
(m)
(47)
<n
No
lim
M
m=1
M
log
Pa
9
Y
SN
m)
<N
(m)
Vo log Po (Y
(m)
S/N)
(48)
lim
8
(N))
MC
(49)
N-
Lemma 14. The following identity holds:
Vo log Po Y sn)
E Y
V Volog sN) Y Sn Y Sn
(50)
for any N n.
Proof. Note that (Y Sa) = ErN- Po "nt') for any N > n. Therefore, we have
Vopo ( sn)
V Vo
(Y sm)=
(51a)
To sn)
VOÑO Y <nD')
(51b)
Po Y sn)
VEEN-N
Pol Y EnD
Vo log po (Y snly')
(51c)
Do Y sn)
V'EX
Po T
Y Y Po Y sn)
Vo log Po (Y snt')
(51d)
Po (Y cn)
Σ (T
Y Sn) Vo log Y zn
(51e)
VEL
E
Ve
Po sN) Y
Ysn
(51
Lemma 15. Let (,N) be the truncated MC estimator of the gradient defined in Eg. (41) and
6 the
truncated RB estimator of the gradient defined in Eg. (42). Define N.N E 5 (N) have
MC = R. B.o ' We
17]
J - G.V
E
R
< E
(52)
22

---

Proof. Without loss of generality, we assume M = 1. The proof generalizes to M >> with the ii.d. assumption.
E
(N)
|
G.N.
(53a)
KB
2-
Σ
E Y
MMc (Y sa) (Vo log cn) co Vo Vologio Volggg Y <n))
G
Y <n = <n
Y
Vo Volog gn)
(53b)
(delinition )
N
I Y
Σ
E
Pai (Y"
Sn Vo Po (Y
5n)
G.
Y
(53c)
cn
(each Y is distributed independently and identically tn Y)
MC
Sn
Vo
log Po Y (a) - G.V
F <N = F <N ... Y
<1
=
Y
<1
(53d)
(lincarity of expectation)
E
E
Σ
n (Y"
Sn)
Vo go (Y sn) G.a
Y <N Y <N
Y
Y
Y
A.V
(53e)
(Jenson's 's incquality)
N
Σ
Moc (Y
Sn Vo log Po Ynn) G.V
(53f)
expectation)
N
E
Σ p
(
Vologgoe sm) G.¹.
(53g)
Y
(V)
are Lid.)
E
Σ
(Y
Sn E
Velog WY Y sn Ys
G.V
(53h)
X
n=1
(Lomma 14)
N
Σ
<N) G. log Po Vo
V
Sn = Y
(53i)
Sn
(linearity expectation)
E
E
V (n) Vo
log
(Y xx) G.
pe
PIN
SN Y YN
Y
Y 11
I
(53)
(lincarity of expectation)
23

---

EE
E
E Y
Σ Mee (Y
Sn Vo logio
(o (Y SN) G.V
Y
SN Y SN
Y 1
Y S1
Y
(53k)
(Jensen's 's incquality)
27
E Y
Σ n=1
Me
Sn Vo Po (x')) G. Po
(531)
Y W
expectation)
N Σ Mc (
Vg (xx)
SN) - G.¹
(53m)
n=1
are Lid.)
E Y
0 (NV)
(
(Y) - G
(53n)
(delinition )
Theorem 5. Assuming Var[ōus] oc, Var[ōmc] 2o, we have
E |800 - GIP] < ||0mc G
(17)
Proof. Without loss of generality, we assume M = 1. The proof generalizes to M with the i.i.d.
assumption.
lim S(N) G.V. |"]
E |kkb
G||² =E
(54a)
-00
(detinition 6sn)
s( G.N Gi1(?)
1p]
27
lim E
(54b)
N N0
(dominated convergence theorem, Var Bea 20)
G.N. (m1
<
E N
lim
(54c)
(Lemma 15)
G.N 27
19]
E
lim 5 5(N) N-000
(54d)
(dominated convergence theorem. Var [wee < oc)
= E |0nec 1000]
(54e)
eemma 13)
D.2 A Note Rao-Blackwellizingg KL in Trust-Region Algorithms
The conventional Monte Carlo estimator of KL(pe || q) used in the PPO algorithm in open-sourced RLHF libraries, e.g., [13, 40], is as follows:
PPO MC
1
Σ
M
D mo(Y(m) (m))
log
Peau
(Y (m))
(55)
M
Pe (y-())
q( (Y (me))
mL
24

---

Pens Notably, the expected value of the above estimator is not equal to where Y
(Y)) Lid
.. kI, g) and
PPO El
po
log
PM.u (Y)
=
log
aL (Y) (
(56)
Y au
Po
(Y)
Yrupa
A natural question at this point is: what is the relationship between M and Jocc, and why is minimizing PM valid proxy for minimizing (mme? Crucially, the KI. divergence between po and can be decomposed into the sum of MC and the KI. divergence between and poau as shown in the following equation:
log
aa
log
Po Y)
pe Y log
KL(Po 11
(57)
Y Y~upo
(
Y Yrupo
Po (Y)
Y Y-up
g (Y)
Elagel
trutt region. wiggio...(pp.... ap
Therefore, minimizing KT.(pe || q) is equivalent to minimizing both Mly and KI || pe Notably, since the KL divergence between the current policy and the old policy, KL(Pe || pe..). is already constrained by PPO's clipping mechanism, the algorithm effectively focuses penalizing only the first term, using PPo.
A naïve approach to Rao Blackwellizing MMC defined in Eq. (55), is as follows:
M Po (Y(())
Po (Yn
(mj) <a
PRRB
PPO
M
Pa
(Y(m))
lim N-+00
E Yar
log
(Y
Y
<n
(58)
Importantly, PPO does not give an unbiased estimate Ex
pa (X) log
i.c.,
q(Y)
E
[
E rpp
lim
log
Po
n2
E Youpo
log
ao (Y)
(59)
Y
<n)
(Y)
Pe
Therefore, caution the reader against using this estimator replacement for PPO in practice.
25

---

Ihia gceat novie!
I've seen several
proopt: 1hia toreible zovie
CVe-1
10
MC
Extimated KL.
10
4-
Sample Size
Sample Size
Sample Size
Figure Comparing the bias, variance, and consistency of the estimators the sample size increases. The Hev estimator with 1 exhibits a higher standard deviation, particularly for neutral and negative prompts, where the variance becomes extremely large. contrast, the RB estimator, Jes. achieves the lowest standard deviation.
E
Additional Experiments
In Fig. 4, we visualize the KL estimates for three different prompts: (left) positive prompt, (middle) a neutral prompt, and (right) negative adversarial prompt. The traces represent the average estimates from all the repetitions, while the shaded regions indicate the standard deviation. Except Hev, 1, all other estimators are unbiased and consistent.
As the sample size increases, the chance of sampling from the tail of pa also increases. These tail samples often correspond negative movie reviews that had high probability under the language model prior to fine-tuning, i.e., g. leading to extremely large values of g(Y) and, consequently, high standard deviation. This effect indeed depends on the prompt and is particularly pronounced for neutral and adversarial prompts.
We conducted an additional experiment to evaluate the estimators random subset of prompts from the UltraFeedback dataset [8]. We compute the KL divergence between Zephyr-7B-Beta [38] and its reference model, Mistral-7B-ve [18]. Zephyr is fine-tuned from Mistral using DPO on UltraFeed- back, and as part of this fine-tuning. is desirable not to diverge significantly from the base model.
We randomly sampled 512 prompts and generated 100 responses per prompt. For estimation, we used subsets of 1, 5, and 10 samples, reserving the maining samples to estimate each method's 's standard deviation. Tab. reports the KL estimate + standard deviation for each estimator. Unlike our GPT-2 ex- periments, we had to use half- malf-precision to perform inference and forward passes on single GPU, which introduces bias in the HT and CV estimators. Con- sistent with our findings on the IMDB dataset, our proposed RB estimator consistently achieves the low- est standard deviation across all settings, reaffirming its stability and reliability.
Table 2:
Estimated value + empirical stan- dard deviation of different estimators. RB es- timator consistently achieves the lowest stan- dard deviation.
M
M
M
PMc
18.05 3.199
18.05 +: 1.63
18.05 0.75
18.05 = 12.64
18.56
+ 5.34
19.24 = 3.62
Pev1
17.17 1 3.17
17.17
1.62
17.17 = 0.75
Pev
17.50
3.118
17.80
1.63
17.80 + 0.75
Pas
3.16 18.05
18.05 1.61
18.05
26

---

Experimental Details
F.1 Code Snippet
def compute_ [rmpute_k1(logrrbbs. ref _oogprbs logits, ref_ sef_logits):
3
Compute KL divergence using two estimators.
Args:
logprobs: Log probabilities of sampled actions from policy rog_logpoob: Log probabilities same actions from reference
logits:
Full distribution logits from policy (all actions)
ref_logits: Full distribution logits from reference (all actions)
Example:
Policy samples action from 1000 possible actions logprobs [-2.3] logits = 10.1, -0.5, ...] # All 1000 action logits (raw)
# Only action 3's log prob
MC: uses only sampled action
RB: uses all 1000 actions for exact expectation
# Monte Carlo: unbiased but higher variance
kl_mc mean(logprobs ref_logprobs)
21
22
# Rao-Blackwell: lower variance, uses full distribution
log_p wal.oootmaxxlooiiii
Normalize to log probs
log_q sos_sootmmxx(ree _lggits)
Normalize log probs
kl_rb mean(sum(exp(Log_D) (log_p log_g)))
26
return kl kl_mc, kl_rb
F.2 RLHF Experiments
In App. F.2, we include the hyperparameters used with the RLOO algorithm for the sentiment control experiment. Each experiment takes approximately minutes on a single rtx_4090 GPU.
Hypterparameter
Value
Optimizer
Adam (e = le-5, 1r= 3e- 6)
Scheduler
Linear
Batch Size
32
B
0.07
k
Number of RLOO Updates Iteration Per Epoch
Clip range
0.2
Sampling Temperature
27

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main contributions are that our proposed estimator is unbiased and has variance less than equal to the variance of the MC estimator, and it improves the stability of RLHF training. All these claims are clearly stated in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the
results can be expected to generalize to other settings. is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We explain the limitations in $6.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification,
asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how
they scale with dataset size.
applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions favor of transparency play important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof? Answer: [Yes]
Justification: Yes, the main theoretical claim is that the RB estimator is unbiased and has variance less than or equal the variance of the MC estimator, this is supported by theorem 2, and the proof is in appendix E.3. The rest of the theoretical claims regarding other estimators are proved in appendix A-F.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
28

---

All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: We provide simple code snippet to implement RB estimator in App. F.1, which also highlights the difference between the RB and MC estimator.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
example, if the contribution
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of sosed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes]
Justification: We have released the code but, to preserve anonymity, we have not included the link to the public repository in the paper. Instead, the code is attached as supplementary material to the submission.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https:/ httplps..cppbblcc//
gui / ssi ssionnooicy myy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
29

---

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nip) http://iippcccppplliiipp cc/publcc/ gui /ubb/ ssionPoli imPoliiyll cy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce experimental results for the new proposed method and basclines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: The experimental details explained in App. F.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is
necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or supplemental material. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: When evaluating estimators in $5.1, we report the standard deviation across runs. Furthermore, for RLHF Pareto plot, we perform a permutation test in $5.2 to test significance of the results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to
library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors). It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1 1-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
Experiments compute resources
Question: For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experi- ments?
Answer: [Yes]
Justification: We provide this in App.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
30

---

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e-g., preliminary or failed experiments that didn't 't make into the paper).
Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code Ethics https://neur:: ips. ccc/publc/Ethhicsuideelllll
Answer: [Yes]
Justification: We have reviewed the NeurIPS code of ethics and we confirm that the paper conforms to the NeurIPS code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consideration due laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss this in $6. Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance). fairness considerations (e.g.. deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to par- ticular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, is legitimate to point out that improvement in the quality of generative models could be used generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results. and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
We do not release any new models or datasets, therefore, the paper poses no such
risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
31

---

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the dataset $5.1. While the original paper doesn't specify a particular license, it emphasizes that the data is intended for academic and non-commercial use.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's
creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We do not release new assets.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/codd/model as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human
subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should included in the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
32

---

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equívalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble). such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non- standard component of the methods in this research? Note that if the L.M is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification:
The core method developed in this paper does involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, or non-standard components. Please refer to our LLM policy pectps://neurlie http:r/s.cccoooeeeeee/////////// for what should or should be described.
33


Reviews:
[{"review_id": "XjJvUDJsws", "reviewer": "Reviewer_Q7xD", "summary": "This paper proposes a new estimator of KL divergence between language models through Rao-Blackwellization. The proposed estimator is unbiased and has less variance than the baseline Monte Carlo estimator that is currently well used in machine learning applications. The paper introduces the Rao–Blackwellized Monte Carlo method and analyzes the gradient updates. The paper also presents experiments on preference finetuning with GPT-2 on the IMDB dataset. The experiments present bias and variances of different baseline estimators along with averaged rewards using different KL estimators. Results show the proposed RB estimator has low variance and achieves pareto fronts in evaluation.", "questions": "- I wonder if there is anyway to quantify the improvements in numerical stability. It could be a plus to show this in the experiment sections.\n- Does the method apply to other divergences in the f-divergence family like Jensen–Shannon divergence?\n- In the authors' opinion, what downstream tasks benefit most from lower variances of KL divergence estimation? It would be helpful to elaborate the potential impact of the proposed method.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths\n- The paper gives detailed preliminaries and derivations of the proposed method such as Control Variate Monte Carlo and Rao–Blackwellization.\n- The paper also analyzes gradient computation and also off-policy gradient for real-world settings.\n- The paper implements and compares with a number of baselines thus making results convincing in experiments.\n\nWeaknesses\n- The experiments only include one preference finetuning task on one small-scale dataset IMDB-sentiment. The experiments are done with GPT-2 model. It would be more helpful to show the advantage of the proposed estimator in more real-world applications. These could be done with more diverse variety of experiments like experiments with other applications such as knowledge distillation which also use KL divergence a lot, or experiments with different tasks in preference finetuning domain, or experiments with different model families.", "was_revised": true, "final_justification": "Additional results strengthen the paper, but a higher score would need more extensive experiments from different real world applications.", "rebuttal": "We thank the reviewer for their thoughtful feedback.\n\n## Experiments on another model family and real-world benchmarks\n\nWe thank all the reviewers for acknowledging our limited computational resources to conduct large-scale experiments. That said, a recurring concern was the lack of evaluation on a different model family and a real-world benchmark (for example, reviewer `#5fbx` specifically suggested experiments on the UltraFeedback dataset)\n\nTo address this, we conducted an additional experiment to evaluate all the discussed KL divergence estimators on a random subset of prompts from the UltraFeedback dataset. We compute the KL divergence between [Zephyr-7B-Beta](https://huggingface.co/HuggingFaceH4/zephyr-7b-beta) and its reference model, [Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1). Zephyr is fine-tuned from Mistral using DPO on UltraFeedback, and as part of this fine-tuning, it is desirable not to diverge significantly from the base model.\n\nWe randomly sampled 512 prompts and generated 100 responses per prompt. For estimation, we used subsets of 1, 5, and 10 samples, reserving the remaining samples to estimate each method’s standard deviation. The table below reports the KL estimate ± standard deviation for each estimator. Unlike our GPT-2 experiments, we had to use half-precision to perform inference and forward passes on a single GPU, which introduces bias in the HT and CV estimators. Consistent with our findings on the IMDB dataset, our proposed RB estimator consistently achieves the lowest standard deviation across all settings, reaffirming its stability and reliability.\n\n\n| Method | M = 1 | M = 5 | M = 10 |\n|--------|--------|--------|---------|\n| HT     | 18.05 ± 12.64 | 18.56 ± 5.34 | 19.24 ± 3.62 |\n| CV1    | 17.17 ± 3.17 | 17.17 ± 1.62 | 17.17 ± 0.75 |\n| CV     | 17.80 ± 3.18 | 17.80 ± 1.63 | 17.80 ± 0.75 |\n| MC     | 18.05 ± 3.19 | 18.05 ± 1.63 | **18.05 ± 0.75** |\n| **RB** | **18.05 ± 3.16** | **18.05 ± 1.61** | **18.05 ± 0.75** |\n\n## Questions\n> I wonder if there is anyway to quantify the improvements in numerical stability. It could be a plus to show this in the experiment sections.\n\nThank you for this suggestion. The numerical instability arises when computing g(Y), and it's partially captured in Table 1 and Figure 1. However, it's more clearly illustrated in Appendix G (Figure 4). This issue typically occurs with samples from the tail of p (often negative movie reviews) that are still likely under q but suppressed during DPO, making p very small. As a result, the ratio q/p can overflow and bias the estimates. Figure 4 shows that this becomes more likely as the sample size increases (since tail samples are more likely to appear), especially for neutral or negative prompts. We will include a more detailed discussion in the final version of the paper.\n\n> Does the method apply to other divergences in the f-divergence family like Jensen–Shannon divergence?\n\nYes. While our focus is on KL divergence due to its common use, the proposed estimator can also be applied to both forward and backward KL, which together can be used to estimate the Jensen–Shannon divergence.\n\n> In the authors' opinion, what downstream tasks benefit most from lower variances of KL divergence estimation? It would be helpful to elaborate the potential impact of the proposed method.\n\nThis estimator is particularly useful when the number of samples available for estimating the divergence is limited, i.e., cases where standard Monte Carlo estimates of KL become highly unreliable. One important example is in RLHF, where GPU memory constraints the batch size, and therefore, the number of samples used for KL estimation is limited. In such settings, a lower-variance estimator can lead to more stable training and better-aligned models."}, {"review_id": "6JYDU3AzSa", "reviewer": "Reviewer_5fbx", "summary": "This paper introduces a Rao-Blackwellized estimator for KL divergence between language models, offering unbiased estimates with lower variance than standard Monte Carlo methods. It improves training stability in RLHF and yields models closer to the reward–KL Pareto frontier. To facilitate fine-tuning practices, this paper also derives a more stable and effective version of the estimator. The arguments are well-supported by both theoretical analysis and experimental results.", "questions": "Q1, It is out of the scope for this paper, but it would be much helpful if the dependency w.r.t. the reward function is clarified. \n\nIn many real-world datasets, the binary reward models, e.g., the Bradley-Terry(BT) model, have been known to be exposed to 'intransitivity' risk because it relies on scalar variables, where all preferences are transitive by assumption. \n- The literature below studied representative preference datasets in the real world, where the 'transitive' relationship between preference annotations may not always hold. \n- https://arxiv.org/abs/2409.19325 (Duan et al, 2017) \n\nWhat is the implication or relationship between this 'intransitivity' and the proposed RB estimator?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. This paper has proposed an RB estimator for KL divergence estimation between LLMs. Lower variance and unbiasedness are theoretically justified, despite the higher estimation cost.\n2. The idea is straightforward, originating from the context of divergence estimation; a code snippet for implementing the RB estimator is provided.\n3. Evaluation is mostly limited to sentiment control.\n\nWeakness:\n1. Broader alignment tasks and certain well-benchmarked datasets such as UltraFeedback remain to be validated.", "was_revised": true, "final_justification": "The authors provided their candid response. The Pareto frontier result is impactful, and I will maintain my score as 'Weak Accept'.", "rebuttal": "We thank the reviewer for their thoughtful feedback. For the comment regarding other model families and datasets, we refer the reviewer to the response [to the reviewer #Q7xD](https://openreview.net/forum?id=um9kHMof0c&noteId=WF9BPFMFcR).\n\nQuestions:\n>  It is out of the scope for this paper, but it would be much helpful if the dependency w.r.t. the reward function is clarified. What is the implication or relationship between this 'intransitivity' and the proposed RB estimator?\n\nThat is an interesting question. As discussed in the cited work, intransitivity in preference data can introduce significant challenges when learning a reliable reward model. In particular, it can lead to a reward function that is noisy or poorly calibrated. When optimizing with a KL-constrained RL objective in this setting, we suspect that the model may compensate for the weak reward signal by relying more heavily on the KL term, since it provides a more consistent optimization direction. However, if the KL estimator used in this process has high variance, this over-reliance can destabilize training. Our proposed RB estimator could help in this case by providing a more stable estimate of the KL divergence."}, {"review_id": "fPyJlo9VMg", "reviewer": "Reviewer_A7gG", "summary": "This paper addresses the problem of estimating the Kullback-Leibler (KL) divergence between language models, which is crucial for applications like reinforcement learning from human feedback (RLHF). The authors identify that while Monte Carlo (MC) estimation is commonly used for this task, it suffers from high variance and can even produce negative estimates of KL divergence, which is theoretically non-negative. To address these issues, the paper introduces a Rao-Blackwellized estimator that applies the classical variance reduction technique of Rao-Blackwellization to KL divergence estimation. The key insight is to decompose the language model estimation into step-wise components and compute exact conditional KL divergences at each step, which are then combined to form the final estimate. The authors prove that their estimator is unbiased and has variance less than or equal to the standard MC estimator, while maintaining the same computational complexity and guaranteeing non-negative estimates.", "questions": "Is there any assumption about the variance?", "limitations": "See the weakness part.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths：\n\n1. The paper provides rigorous theoretical analysis with formal proofs showing that the Rao-Blackwellized (RB) estimator is unbiased and has variance less than or equal to the Monte Carlo estimator. \n\n2. A key strength is that the proposed method achieves variance reduction without additional computational overhead compared to the standard MC estimator. \n\n3. The experiments systematically compare multiple estimators across different metrics (bias, variance, consistency) and demonstrate practical benefits in RLHF training, including improved stability and better Pareto frontier performance.\n\nWeakness：\n\n1. The core contribution applies an existing statistical technique (Rao-Blackwellization) to a specific domain. While the application is novel and valuable, the methodological innovation is somewhat incremental.\n\n2. The empirical evaluation is restricted to relatively small models (GPT-2) and a single task (sentiment control). While this is acknowledged as a limitation due to computational constraints, it raises questions about scalability to modern large language models and diverse applications.\n\n3. The theoretical results assume finite variance for both estimators, but the paper doesn't deeply analyze when these assumptions might be violated in practice or provide guidance for detecting such cases.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful feedback. \n\n- Regarding the comment that our work applies an existing technique and is methodologically incremental, it is true that our method is based on Rao–Blackwellization, a well-known variance reduction technique. Our contribution lies in demonstrating that this simple, principled technique, requiring only minimal changes and reusing logits already computed during the forward pass, can substantially reduce variance at almost no additional cost. In fact, this simplicity could facilitate the integration of this method in RLHF libraries.\n\n- Regarding other model families and datasets, we refer the reviewer to the response [to the reviewer #Q7xD](https://openreview.net/forum?id=um9kHMof0c&noteId=WF9BPFMFcR).\n\n\n- Regarding the assumption of bounded variance, we thank the reviewer for raising this important point. We assume that the true KL divergence is finite. This assumption is generally realistic when p is a model fine-tuned to approximate q, often initialized from a checkpoint of q. However, in cases where q is close to zero while p remains positive, the KL divergence and consequently the variance of our estimators can become unbounded. We agree with the reviewer that this aspect deserves further clarification, and we will address it in the final version of the paper."}, {"review_id": "ZmaaemMmTn", "reviewer": "Reviewer_oszg", "summary": "This paper propose an estimator to the KL divergence and its gradient during RLHF, based on the decomposition of conditional expectation. The proposed estimation is shown to have lower variance, compared with the classicial Montecarlo exstimation, both in KL divergence itself and its gradient. The authors also use the proposed estimation to conduct RLHF and verify it induce a more stable training process.", "questions": "See Weankess, I will raise my score if the author can address my concerns in above.", "limitations": "See Weakness.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nThis paper is well written and discuss an important question. \n\nThis paper has a strong theoretical guarantees of the proposed techniques. \n\nWeakness:\n\nMy major concerns are two folds, respectively about the improvement over other varaince reduction techniques and the empirical evidance on the varaince reduction of gradient on KL divergence. \n\nFor the first, though the author compare their estimation with the vanilla MC method, but the improvement over the other variance reduction method is neglected. For example to comparison with  \"Schulman. 2020. Approximating KL divergence.\"\n\nFor the second, the empirical results only verify the estimated KL divergence has lower variance, but the verification on the gradient of KL divergence is unexplored, which actually decides the dynamics of RLHF.", "was_revised": true, "final_justification": "See above", "rebuttal": "We thank the reviewer for their thoughtful feedback.\n\n> For the first, though the author compare their estimation with the vanilla MC method, but the improvement over the other variance reduction method is neglected. For example to comparison with \"Schulman. 2020. Approximating KL divergence.\"\n\nWe want to clarify that the control variate estimator we discuss (CV with alpha = 1) is the same as the estimator proposed by Schulman (2020), known as the K3 estimator. We will make sure to emphasize this point further in the final version of the paper. In our experiments, we observe that this estimator (CV with alpha=1) can suffer from unbounded variance (figure 1), which limits its reliability compared to our proposed method.\n\n> For the second, the empirical results only verify the estimated KL divergence has lower variance, but the verification on the gradient of KL divergence is unexplored, which actually decides the dynamics of RLHF.\n\nIn our RLHF experiments, we compare the variance of the MC and RB estimators at each training step. Figure 2 (right) shows the standard deviation of both estimators’ KL estimates throughout fine-tuning. We run RLHF five times with each estimator and plot the mean KL values along with their standard deviations (shown as shaded regions). The results clearly show that the MC estimator exhibits higher variance across training steps, while the RB estimator provides more stable and reliable KL estimates. Importantly, this improved stability translates to better downstream performance. As demonstrated in Figure 3, models trained using the RB estimator more consistently appear on the Pareto front balancing reward and KL divergence, highlighting the practical benefits of reduced variance in the training dynamics."}]

Author rebuttal:

Meta-review: This paper proposes a Rao-Blackwellization of the typical sampling-based estimator of KL-divergence and its gradient, particularly in the context of its use in RLHF. They show that the new estimator has indeed smaller MSE theoretically and they also empirically demonstrate that this leads to better estimation in practice. The reviewers appreciate the rigor of the approach and its potential impact.

During the discussions, the chief concern was whether this impact was properly demonstrated to occur in practice. The authors provided additional experiments, which should be part of the paper. The results on assessing the gradient variance reduction and a discussion on the conditions necessary for the new estimator to be well behaved, should also be included. More RLHF experiments would better motivate the community to adopt this new estimator, considering that this context is the main driver of the research, and practitioners would heavily weigh whether the benefits of making this change are worthwhile. Additionally, the authors could explore other contexts where KL estimation is critical and where adopting the new estimator can make a significant difference.
Decision: Accept (poster)