Blame for code/notes/src/0109-kl-estimators.tex

\documentclass[11pt]{article}
\usepackage{common}
\begin{document}
\notehead{0109}{KL regularisation estimators $k_1,k_2,k_3$; why recent verifier-RL sets $\beta=0$}{Policy-optimisation lineage, batch M1. Companion notebook: \texttt{notebooks/0109-kl-estimators.ipynb}.}

\begin{problem}
\question{What does each KL estimator estimate as a value, what does it optimise as a loss, and should the penalty be there at all for an exact verifier?}
RLHF maximises $\E[R]-\beta\,\KL(\pi_\theta\|\pi_{\rm ref})$ so that the policy stays where a learned reward model is accurate. The KL over a space of sequences cannot be summed, so it is estimated from samples, and one of three Monte-Carlo estimators (Schulman's $k_1,k_2,k_3$) is used, either as a number or as a loss to differentiate. Three separate questions must not be conflated: (i) as a \emph{value}, how biased and how noisy is each estimator? (ii) as a \emph{loss}, what does differentiating it actually optimise? (iii) for RL with an exact verifier (theorem proving), should the penalty be there at all? Recent recipes answer (iii) with $\beta=0$; we say why, and what that choice costs.
\end{problem}

\section{Notation}
\begin{center}\small\begin{tabular}{@{}ll@{}}\toprule
symbol & meaning\\\midrule
$q=\pi_\theta$ (sampling), $p=\pi_{\rm ref}$ & policy and frozen reference; samples $x\sim q$\\
$r=p(x)/q(x)$, $\ell=\log r$ & ratio, log-ratio; $\E_q[r]=1$ if $\mathrm{supp}\,p\subseteq\mathrm{supp}\,q$\\
$\KL(q\|p)=\E_q[-\ell]$ & reverse KL (the regulariser); $\KL(p\|q)$ the forward KL\\
$k_1,\,k_2,\,k_3$ & $-\log r,\ \tfrac12(\log r)^2,\ (r-1)-\log r$ (Schulman's note)\\
$\sigma^2$ & for Gaussians $q=\mathcal N(0,1)$, $p=\mathcal N(\sigma,1)$: $\KL=\sigma^2/2$\\
$\lambda=-\ell=\log(q/p)$ & log-ratio in the \emph{policy} direction; $k_1=\lambda$, $k_2=\tfrac12\lambda^2$\\
$s_t,\ \lambda_t$ & token score and token log-ratio $\lambda_t=\log(\pi_\theta/\pi_{\rm ref})(o_t|s_t)$\\\bottomrule\end{tabular}\end{center}
\textbf{Sign convention.} Secs.~2.1--2.2 (values) use $\ell=\log(p/q)=-\lambda$; Secs.~2.3--2.4 (gradients, sequences) use $\lambda=\log(q/p)$, so that $\nabla\lambda=+\nabla\log q$; in the boxed result the first line (a value) has $\ell$, the last line (gradients) has $\lambda$.

\section{Derivation}
\subsection{Values: unbiasedness}
\begin{itemize}
\item $k_1$: $\E_q[-\log r]=\sum q\log\frac qp=\KL(q\|p)$. Unbiased by definition.
\item $k_3$: $\E_q[r-1]=\sum p-1=0$ (a control variate with coefficient $c=1$), so $\E_q[k_3]=\KL$. It is $\ge0$ pointwise since $\log x\le x-1$ (the vertical distance between $\log$ and its tangent at $1$, a Bregman divergence).
\item $k_2$: $\E_q[\tfrac12\ell^2]$ is the $f$-divergence with $f=\tfrac12\log^2$, and $f''(1)=1$ equals KL's, so $k_2$ and KL agree to second order near $p\approx q$ but not beyond. The exact bias: $\E_q[e^\ell]=\E_q[r]=1$ gives $\sum_{n\ge1}\E[\ell^n]/n!=0$, i.e.
\begin{equation}\KL(q\|p)=-\E[\ell]=\sum_{n\ge2}\frac{\E_q[\ell^n]}{n!},\qquad\E_q[k_2]-\KL=-\sum_{n\ge3}\frac{\E_q[\ell^n]}{n!}.\label{eq:k2bias}\end{equation}
\end{itemize}
\ver{on three random categorical pairs ($K=6$): $\E[k_1]=\E[k_3]=\KL$ to $\le2.9\times10^{-16}$; the series in Eq.~\eqref{eq:k2bias} (40 terms) reproduces both KL and the bias of $k_2$ to $\le2.8\times10^{-16}$; the sign of $k_2$'s bias is not fixed ($-0.019,\ +0.014,\ +0.155$ for KL $=0.106,\ 0.286,\ 0.612$); $k_3,k_2\ge0$ on a $10^5$-point log-grid of $r\in[e^{-12},e^{12}]$ while $k_1<0$ on half of it.}

\subsection{Values: variance, in closed form for Gaussians}
Let $q=\mathcal N(0,1)$, $p=\mathcal N(\sigma,1)$. Then $\ell=m+\sigma Z$ with $m=-\sigma^2/2$, $Z\sim\mathcal N(0,1)$ under $q$, and $\KL=\sigma^2/2$.
\begin{align*}
\Var(k_1)&=\Var(\ell)=\sigma^2=2\KL.\\
\E[k_2]&=\tfrac12\E\ell^2=\tfrac12(\sigma^2+\tfrac{\sigma^4}4),\ \text{so}\ \text{bias}(k_2)=\tfrac{\sigma^4}8=\tfrac12\KL^2.\\
\Var(k_2)&=\tfrac14\Var(\ell^2)=\tfrac14(4m^2\sigma^2+2\sigma^4)=\tfrac{\sigma^6}4+\tfrac{\sigma^4}2.\\
\Var(k_3)&=\Var(r)+\Var(\ell)-2\Cov(r,\ell)=(e^{\sigma^2}-1)+\sigma^2-2\sigma^2=e^{\sigma^2}-1-\sigma^2,
\end{align*}
using $\E_q[r^2]=e^{\sigma^2}$ and $\Cov_q(r,\ell)=\E_p[\ell]-\E_q[\ell]=\sigma^2$. Hence, relative to the true value,
\[\frac{\mathrm{std}(k_1)}{\KL}=\sqrt{\tfrac{2}{\KL}},\qquad\frac{\mathrm{std}(k_2)}{\KL}=\sqrt{2+2\KL},\qquad\frac{\mathrm{std}(k_3)}{\KL}=\frac{\sqrt{e^{2\KL}-1-2\KL}}{\KL},\qquad\frac{\text{bias}(k_2)}{\KL}=\frac{\KL}2.\]
For small KL, $k_2$ and $k_3\approx\ell^2/2$ have standard deviation $\approx\sqrt2\,\KL$ while $k_1$ has $\sqrt{2\KL}\gg\KL$: that is why $k_1$ is unusable as a value when the policy is near the reference.
\ver{Gauss--Hermite quadrature (120 nodes) matches every closed form: variances to $\le2.2\times10^{-12}$ (relative), biases to $\le1.2\times10^{-16}$ at $\KL=0.005,\,0.05,\,0.5$. At $\KL=0.005$: std/KL $=20.0,\,1.418,\,1.417$ for $k_1,k_2,k_3$ and bias$(k_2)/\KL=0.0025$; at $\KL=0.5$: $2.0,\,1.732,\,1.695$ and $0.25$ (Schulman's table gives $20,\,1.42,\,1.42$ and $2,\,1.73,\,1.7$, $0.25$).}

\paragraph{Heavy tails of $k_3$.} $k_3\approx r$ for large $r$, so $\Var(k_3)\sim\E_q[r^2]=1+\chi^2(p\|q)$: infinite when the policy gives vanishing mass to something the reference likes. Example: $q=(0.001,0.999)$, $p=(0.3,0.7)$: $\KL=0.350$; std of $k_1=0.19$, $k_2=0.51$, $k_3=9.27$ ($r=300$ on the rare token, $\chi^2=89.5$). So $k_3$'s ``low variance'' holds near $q\approx p$ and fails exactly when the policy has drifted away from reference support.
\ver{$\E_q[r^2]-1=\chi^2(p\|q)=89.5$ exactly; stds as above.}

\subsection{Gradients: what each estimator optimises when differentiated}
With $p$ frozen, $\nabla\lambda=\nabla\log q$ (note the sign: $\ell=\log p-\log q$ would give $\nabla\ell=-\nabla\log q$). Samples $x\sim q_\theta$ are drawn \emph{on-policy and detached}; the gradient flows through $q_\theta$ inside $r$ (``$k_n$ as loss''), or the estimator is used as a detached coefficient on the score (``$k_n$ in reward'' as in an RL reward penalty).
\begin{align*}
\nabla k_1&=\nabla\log q\ \Rightarrow\ \E_q[\nabla k_1]=0\quad\text{(Lemma 1 of 0102: no regularising gradient)}\\
\nabla k_2&=\lambda\,\nabla\log q\ \Rightarrow\ \E_q[\nabla k_2]=\E_q[\lambda\nabla\log q]=\nabla\KL(q\|p)\\
\nabla k_3&=(1-r)\nabla\log q\ \Rightarrow\ \E_q[\nabla k_3]=0-\textstyle\sum_xp\,\nabla\log q=\nabla\KL(p\|q)
\end{align*}
Derivations: $\nabla\KL(q\|p)=\sum\nabla q\,\lambda+\sum q\nabla\lambda=\E_q[\lambda\nabla\log q]+\E_q[\nabla\log q]=\E_q[\lambda\nabla\log q]$. For $k_3$: $\nabla r=-r\nabla\log q$, so $\nabla k_3=\nabla r\,(1-1/r)=(1-r)\nabla\log q$, and $\E_q[r\nabla\log q]=\sum p\nabla\log q$, while $\nabla\KL(p\|q)=-\sum p\nabla\log q$.
\textbf{Conclusions} (Liu et al., arXiv:2510.01555, framework in Sec.~4, gradient analysis in Secs.~5.1--5.3, off-policy correction in App.~G): (a) $k_1$ is an unbiased \emph{value} (noisy near $\pi_{\rm ref}$, Sec.~2.2) and a useless \emph{loss}; (b) $k_2$ as a loss and $k_1$ as a detached coefficient (reward penalty $\beta\,\mathrm{sg}(\log q/p)$) both give exactly the gradient of the reverse KL; (c) $k_3$ as a loss (GRPO's choice, DeepSeekMath Eqs.~3--4) has the gradient of the \emph{forward} KL $\KL(p\|q)$, not of the quantity it is an unbiased estimate of. The coefficient $1-r$ is the first-order Taylor expansion of the principled $-\log r$ at $r=1$, so the two gradients agree to first order and differ at second.
\ver{exact on-policy enumeration: $\E[\nabla k_1]=5.6\times10^{-17}$; $\E[\nabla k_2]$ and the detached-coefficient form equal $\nabla\KL(q\|p)$ to $5.6\times10^{-17}$; $\E[\nabla k_3]$ equals $\nabla\KL(p\|q)$ to $8.3\times10^{-17}$; the two gradients differ by $21.5\%$ in relative norm at one distance, and the gap halves as $q\to p$: $15.3\%,\,9.2\%,\,5.1\%,\,2.6\%,\,1.35\%$ for distance $0.4,\,0.2,\,0.1,\,0.05,\,0.025$ (first-order agreement).}
Off-policy (several updates per batch) all ``as loss'' forms additionally need the importance weight (0103); Liu et al.\ App.~G give the corrected form. The same gradient lens applies to the \emph{sequence} view below.

\subsection{Sequences: per-token estimators ignore how earlier tokens move later contexts}
The sequence KL gradient is $\nabla\KL_{\rm seq}=\E[(\sum_t\lambda_t)\,s(o)]=\E\big[\sum_t(\sum_{t'\ge t}\lambda_{t'})\,s_t\big]$ (0102: reward-to-go; ``$k_1$ in reward''). A per-token loss $\sum_tk_2(\lambda_t)$ instead gives $\E[\sum_t\lambda_ts_t]$, which drops the dependence of the token-$t'$ log-ratio on tokens $t<t'$. The exact difference is
\begin{equation}\nabla\KL_{\rm seq}-\E\Big[\sum_t\lambda_ts_t\Big]=\sum_t\E\Big[\KL_t(s_t)\,\nabla\log\pi_\theta(o_{<t})\Big],\qquad\KL_t(s)=\KL\big(\pi_\theta(\cdot|s)\|\pi_{\rm ref}(\cdot|s)\big),\label{eq:missing}\end{equation}
the covariance between the prefix's score and the KL incurred afterwards (the tower rule on $\E[\lambda_t\,\nabla\log\pi(o_{<t})]=\E[\KL_t\nabla\log\pi(o_{<t})]$). Likewise per-token $k_3$-as-loss has the gradient of $\sum_sP_\pi(s)\,\KL(\pi_{\rm ref}(\cdot|s)\|\pi_\theta(\cdot|s))$ with the context weights $P_\pi(s)$ \emph{frozen}. As \emph{values}, the per-token estimators are still fine: $\E[\sum_tk_3]=\KL_{\rm seq}$.
\ver{$\E[(\sum_t\lambda_t)s]$ and the reward-to-go form equal $\nabla\KL_{\rm seq}$ to $8.3\times10^{-17}$; per-token $k_2$-as-loss deviates by $3.8\%$ in relative norm and Eq.~\eqref{eq:missing} holds to $9.7\times10^{-17}$; per-token $k_3$-as-loss equals the frozen-context forward-KL gradient to $5.9\times10^{-17}$; $\KL_{\rm seq}=0.01885=\E\sum_tk_1=\E\sum_tk_3$ (to $6\times10^{-17}$) while $\E\sum_tk_2=0.02103$.}

\subsection{Numerics}
$k_3=(r-1)-\log r$ cancels catastrophically when $|\log r|\ll1$: with $x=\log r$ the result is $\approx x^2/2$ but is built from terms of size $|x|$, so each subtraction loses $\approx\log_{10}(2/|x|)$ digits. Compute $x=\log p-\log q$ from log-probabilities, never from a ratio of probabilities. In fp32 at $x=10^{-4}$ the naive form $(e^x-1)-x$ gives $1.66\times10^{-8}$ against the true $5.00\times10^{-9}$ (relative error $2.3$, i.e.\ $230\%$); at $x=10^{-5}$ it is wrong by a factor $\approx270$.
Using $\mathrm{expm1}(x)-x$ removes the cancellation \emph{inside} $e^x-1$ but not the one in the final subtraction, so it is only a partial fix: rounding $\mathrm{expm1}(x)$ to $u=2^{-24}$ gives a worst-case relative error of $\approx2u/|x|$ ($1.2\times10^{-4}$ for $|x|\sim10^{-3}$, $10^{-3}$ at $10^{-4}$, $9\times10^{-3}$ at $10^{-5}$, $0.12$ at $10^{-6}$). A complete fix switches to the Taylor series $\tfrac{x^2}2(1+\tfrac x3+\tfrac{x^2}{12})$ for $|x|<10^{-2}$ (truncation error $\approx x^3/60$ relative), which is accurate to $\approx2\times10^{-7}$ down to $10^{-6}$.
\ver{expm1 form: $3\times10^{-4}$ at the four tested points $|x|=10^{-1}\dots10^{-4}$, but the worst case over $4\times10^5$ random points per decade (both signs) is $1.0\times10^{-6},\,9.3\times10^{-6},\,1.2\times10^{-4},\,9.7\times10^{-4},\,9.0\times10^{-3},\,0.115$ for decades starting at $10^{-1},\dots,10^{-6}$, always $\le2u/|x|_{\min}$ (max ratio $0.97$); the Taylor/expm1 hybrid has worst-case relative error $9.3\times10^{-6}$ over $|x|\in[10^{-6},1)$.}

\subsection{Why $\beta=0$ in recent verifier-reward RL, and what it costs}
\paragraph{What the papers say.} Dr.\ GRPO (arXiv:2503.20783 Sec.~3): the KL regulariser helps ``prevent $\pi_\theta$ from deviating too far from the distribution where the reward model is accurate'', but ``RL-tuning reasoning models typically employs rule-based verifiers as $r$, eliminating the concerns of distributional shift. This allows us to remove the KL term, which not only saves the memory and computation required by $\pi_{\rm ref}$ during training, but also potentially leads to better performance'' (they assume $\beta=0$). DAPO (arXiv:2503.14476 Sec.~2.3): the model distribution ``can diverge significantly from the initial model, thus this restriction is not necessary''. MiniMax-M1: ``There is no KL penalty term in CISPO similar to other recent works'' (arXiv:2506.13585 Sec.~3.1). The Leanstral~1.5 announcement gives no KL setting (it names CISPO but no hyper-parameters); the Leanstral technical report (Sec.~2.3) writes the CISPO objective with no KL term in it, and does not report the truncation value, so for that system $\beta=0$ is read from the displayed objective rather than from a stated hyper-parameter.
\paragraph{The argument, in three parts.}
(1) \emph{Reason for the penalty disappears.} KL exists to keep the policy inside the region where a \emph{learned} reward is trustworthy. A proof checker is not learned; its verdict on a given string does not degrade with distance from the SFT model.
(2) \emph{Cost.} The reference needs a second set of weights and one extra forward pass per sample. If a backward pass costs twice a forward (the assumption stated in arXiv:2510.26788 Sec.~1), a forward--backward step costs $1+2=3$ forward-units and the reference forward adds $1$, i.e.\ $+33\%$ of that step (generation not included), plus a second model in memory. Qi et al.\ quote $\approx25\%$ for their extra $\pi_{\rm old}$ forward under the same assumption, so their baseline step must be counted as $\approx4$ units (it already contains another forward); the overhead is therefore $\approx25$--$33\%$ depending on what is counted, not a single number.
(3) \emph{It can hurt.} A KL term pulls towards $\pi_{\rm ref}$, which for an SFT model is exactly what RL is trying to improve upon; with $\beta>0$ the optimum of the regularised objective is $\pi^\star\propto\pi_{\rm ref}\,e^{R/\beta}$, which for $R$ the pass indicator keeps probability on every reference proof, including wrong ones.
\paragraph{What it costs, honestly.} Nothing in the objective then keeps the policy near its initialisation; drift, entropy collapse and repetition are controlled only by the other mechanisms in 0106--0108. ProRL (Liu et al., cited in arXiv:2510.01555 Sec.~2) is reported to keep a $k_2$-as-loss KL penalty and periodically reset the reference model, and to find KL important for long-run stability: so $\beta=0$ is a design choice with evidence on both sides, not a theorem. And a KL penalty is \emph{not} a defence against a flaw in the verifier: a verifier hole reachable by small steps from $\pi_{\rm ref}$ is still reachable under $\beta>0$.

\section{Boxed result}
\begin{result}
\begin{align*}
&\E_q[k_1]=\E_q[k_3]=\KL(q\|p),\qquad\E_q[k_2]-\KL=-\textstyle\sum_{n\ge3}\E_q[\ell^n]/n!\quad(=\tfrac12\KL^2\text{ for Gaussians})\\
&\text{std}/\KL\ \text{(Gaussian)}:\ \ k_1:\sqrt{2/\KL},\quad k_2:\sqrt{2+2\KL},\quad k_3:\sqrt{e^{2\KL}-1-2\KL}/\KL\\
&\E_q[\nabla k_1]=0,\quad\E_q[\nabla k_2]=\nabla\KL(q\|p),\quad\E_q[\nabla k_3]=\nabla\KL(p\|q)\\
&\E_q[\mathrm{sg}(\lambda)\nabla\log q]=\nabla\KL(q\|p),\qquad\nabla\KL_{\rm seq}-\E\textstyle\sum_t\lambda_ts_t=\sum_t\E[\KL_t(s_t)\nabla\log\pi(o_{<t})]
\end{align*}
\end{result}

\section{What breaks}
\begin{breaks}\small\begin{itemize}
\item \textbf{Support.} $\E_q[k_3]=\KL$ needs $\mathrm{supp}\,p\subseteq\mathrm{supp}\,q$ (so that $\E_qr=1$). If the policy has assigned $0$ to something the reference allows, $k_3$ is biased and its variance infinite.
\item \textbf{Value $\ne$ gradient.} A good value estimator ($k_3$) as a loss optimises the wrong KL direction; $k_1$ as a loss optimises nothing. \emph{Violated:} ``a good estimator of a number is a good loss''.
\item \textbf{Per-token vs sequence.} Per-token losses hold the context distribution fixed (Eq.~\eqref{eq:missing}); correct treatment needs reward-to-go on the log-ratios.
\item \textbf{Off-policy.} All ``as loss'' forms are on-policy gradients; with 16 updates per batch (0103) they need importance weights, and clipping interacts with them.
\item \textbf{Variance when drifted.} The estimator that is best near $q\approx p$ ($k_3$) is worst far from it.
\item \textbf{Mismatch.} $\log\pi_{\rm ref}$ and $\log\pi_\theta$ come from the trainer but the samples from the sampler (0108): a spurious non-zero KL even at $\beta$-fixed point.
\item \textbf{$\beta=0$ is not free:} no anchor; it relies on the verifier being exact and on the other stabilisers (clip, caps, dynamic sampling).
\end{itemize}\end{breaks}

\nextnote{With $\beta=0$ the only signal left is the verifier's verdict. Note 0110 analyses that signal: binary rewards, dead groups and their exactly-zero gradient, overlong penalties as reward (not loss), potential-based shaping, and what verifier errors do to the objective.}

\section*{References}
{\small\begin{itemize}[leftmargin=1.2em]
\item Schulman, \emph{Approximating KL Divergence} (blog, \texttt{joschu.net/blog/kl-approx.html}; vault note \texttt{approximating-kl-divergence}): definitions of $k_1,k_2,k_3$, the control-variate construction, the table at $\KL=0.005$ and $0.5$. (Not an arXiv paper.)
\item Shao et al., arXiv:2402.03300, Sec.~4.1.1, Eqs.~3--4 ($k_3$ with $r=\pi_{\rm ref}/\pi_\theta$ used as the GRPO penalty).
\item Liu et al., \emph{Rethinking KL Regularization in RLHF: From Value Estimation to Gradient Optimization}, arXiv:2510.01555v2 (vault title ``Technical Report''): abstract; Sec.~2 (ProRL remark); Sec.~4 (framework of ``$k_n$ as loss'' vs ``$k_n$ in reward''); Secs.~5.1--5.3 (gradient equivalences, Eqs.~17--18); App.~G (off-policy correction). Their conclusions (a)--(c) are re-derived and verified above.
\item Liu et al., arXiv:2503.20783, Sec.~3 (removing KL); Yu et al., arXiv:2503.14476, Sec.~2.3; MiniMax, arXiv:2506.13585, Sec.~3.1.
\item Mistral AI, Leanstral technical report (\texttt{LeanstralReport.pdf}, Sec.~2.3; not on arXiv): CISPO objective without a KL term; vault note \texttt{in-this-work-we-present-leanstral-a-series-of-generalist-code-agent-models-for}.
\item Qi et al., arXiv:2510.26788, Sec.~1 (backward $\approx2\times$ forward cost).
\item Eqs.~\eqref{eq:k2bias}, the Gaussian variances and Eq.~\eqref{eq:missing} are derived here and verified by quadrature / exhaustive enumeration.
\end{itemize}}
\end{document}