| \documentclass[11pt]{article} |
| \usepackage{common} |
| \begin{document} |
| \notehead{0103}{Importance sampling for off-policy updates; the product-of-token-ratios explosion}{Policy-optimisation lineage, batch M1. Companion notebook: \texttt{notebooks/0103-importance-sampling-offpolicy.ipynb}.} |
|
|
| \begin{problem} |
| \question{Can we take many optimiser steps on one batch of samples without biasing the gradient, and what does that correction cost?} |
| Generating responses is the expensive part of RL for LLMs; a gradient step is cheap. Practical systems therefore take many optimiser steps on one batch: DAPO samples $512\times16$ responses per rollout step and takes 16 mini-batch updates (arXiv:2503.14476 Sec.~4.1), and MiniMax-M1 runs ``16 rounds of off-policy updates per generation batch'' (arXiv:2506.13585 Sec.~3.1). After the first step the batch was drawn from $\pi_{\theta_{\rm old}}$ while we differentiate $J(\theta)$ for a different $\theta$; the on-policy estimator of 0102 is then biased. Importance sampling (IS) repairs the bias exactly. The price, which this note quantifies, is a variance that grows \emph{exponentially} in the number of tokens. |
| \end{problem} |
|
|
| \section{Notation} |
| \begin{center}\small\begin{tabular}{@{}ll@{}}\toprule |
| symbol & meaning\\\midrule |
| $q=\pi_{\theta_{\rm old}},\ p=\pi_\theta$ & behaviour (sampling) and target distribution over responses\\ |
| $w(o)=p(o)/q(o)$ & sequence importance weight\\ |
| $r_t(\theta)=\pi_\theta(o_t\mid s_t)/\pi_{\theta_{\rm old}}(o_t\mid s_t)$ & token ratio; $w_T=\prod_{t\le T}r_t$ is the weight of the first $T$ tokens\\ |
| $\chi^2(p\|q)=\sum_o(p-q)^2/q$ & chi-square divergence; $\chi^2_t(s)$ is its conditional (per-state) version\\ |
| $\ESS=(\sum_iw_i)^2/\sum_iw_i^2$ & effective sample size of $n$ weighted samples\\ |
| $R,b,s(o)$ & reward, baseline, score $\nabla_\theta\log\pi_\theta(o)$ (as in 0102)\\\bottomrule\end{tabular}\end{center} |
|
|
| \section{Derivation} |
| \subsection{The IS identity} |
| Let $f$ be any function of the response. If $q(o)>0$ wherever $p(o)f(o)\ne0$ (absolute continuity), |
| \begin{equation}\E_p[f]=\sum_op(o)f(o)=\sum_oq(o)\frac{p(o)}{q(o)}f(o)=\E_q\big[w\,f\big].\label{eq:is}\end{equation} |
| This is an algebraic identity, hence unbiased; the only assumption is the support condition. Taking $f\equiv1$: $\E_q[w]=1$. |
| \ver{$\E_q[wf]=\E_p[f]$ and $\E_q[w]=1$ hold to $0$ (fp64) on all 121 sequences.} |
|
|
| \subsection{Variance of the weight} |
| \begin{align} |
| \Var_q[wf]&=\E_q[w^2f^2]-(\E_pf)^2=\E_p[w f^2]-(\E_pf)^2,\label{eq:isvar}\\ |
| \Var_q[w]&=\E_q[w^2]-1=\sum_o\frac{p^2}{q}-1=\sum_o\frac{(p-q)^2}{q}=\chi^2(p\|q)=\E_p[w]-1.\label{eq:chi2} |
| \end{align} |
| Eq.~\eqref{eq:chi2}: $\E_q[w^2]=\sum_o q\,p^2/q^2=\sum_op^2/q=\E_p[w]$, and expanding $(p-q)^2/q$ gives $p^2/q-2p+q$, whose sum is $\sum p^2/q-1$. Jensen links this to KL: $\KL(p\|q)=\E_p\log w\le\log\E_pw=\log(1+\chi^2)$, so large KL forces large $\chi^2$. |
| \ver{Eq.~\eqref{eq:isvar} and Eq.~\eqref{eq:chi2} hold to $\le10^{-13}$ on the toy ($\chi^2=0.0$ cancels exactly in fp64).} |
|
|
| \paragraph{Effective sample size.} With $n$ i.i.d.\ draws from $q$, the law of large numbers gives $\sum_iw_i\approx n\E_qw=n$ and $\sum_iw_i^2\approx n\E_qw^2$, so |
| \begin{equation}\ESS\approx\frac{n}{1+\chi^2(p\|q)}.\label{eq:ess}\end{equation} |
| Self-normalised IS, $\hat\mu=\sum_iw_if_i/\sum_iw_i$, drops the need to know $q,p$ normalisers but is biased at order $1/n$ (ratio of two random sums). |
| \ver{MC: mean ESS at $n=256$ is $252.20$ vs.\ $n/(1+\chi^2)=252.17$ (rel.\ error $7\times10^{-5}$); SNIS bias $-0.0107\,(n{=}4)\to-0.0028\,(16)\to-0.0010\,(64)\to-0.0002\,(256)$ with standard errors $0.0016,\,0.0008,\,0.0008,\,0.0004$: the $O(1/n)$ shrinkage is statistically resolved from $n=4$ to $16$ only; the two largest-$n$ values are within $1.3\sigma$ of zero.} |
|
|
| \subsection{Off-policy policy gradient} |
| With $b$ a baseline independent of the sample, 0102 gives $\nabla J(\theta)=\sum_op(o)(R-b)\nabla\log p(o)$. Insert $p=wq$: |
| \begin{equation}\nabla_\theta J(\theta)=\E_{o\sim\pi_{\theta_{\rm old}}}\Big[w(o)\,(R(o)-b)\,\nabla_\theta\log\pi_\theta(o)\Big],\qquad w(o)=\prod_{t=1}^{|o|}r_t(\theta).\label{eq:offpg}\end{equation} |
| The product form is exactly the chain rule of 0101, Eq.~(1). This is the \emph{sequence-level} estimator written in arXiv:2510.26788 Eq.~5. For an outcome reward $R$ the weight cannot be shortened: the reward depends on every token, so every token's ratio enters. |
| \ver{Eq.~\eqref{eq:offpg} equals the autograd gradient of $J(\theta)$ to $4\times10^{-17}$ at $\delta=0$ and $2.8\times10^{-17}$ at $\|\theta-\theta_{\rm old}\|=0.32$.} |
|
|
| \subsection{Why the ratio explodes with length} |
| \paragraph{Exact recursion.} By the tower rule on the last token, |
| \begin{equation}\E_q[w_T^2]=\E_q\Big[w_{T-1}^2\,\E_{o_T\mid s_T}[r_T^2]\Big]=\E_q\Big[w_{T-1}^2\big(1+\chi_T^2(s_T)\big)\Big],\label{eq:rec}\end{equation} |
| using Eq.~\eqref{eq:chi2} at the single state $s_T$. If every state has $\chi^2_t(s)\ge c$ then $\E_q[w_T^2]\ge(1+c)^T$, i.e.\ $\ESS/n\lesssim(1+c)^{-T}$ (an \emph{upper} bound on the effective sample fraction). For i.i.d.\ token ratios it is an equality: |
| \begin{equation}\Var_q[w_T]=(1+\chi_1^2)^T-1.\label{eq:geom}\end{equation} |
| \ver{brute-force enumeration of all $3^T$ token strings equals $(1+\chi_1^2)^T$ to $\le4.3\times10^{-16}$ (relative) for $T=1,3,6,9$.} |
|
|
| \paragraph{Mean one, typical value zero.} The mean $\E_q[w_T]=1$ is exactly preserved for every $T$ (the weights are a martingale). But for the typical sample, |
| \begin{equation}\E_q[\log w_T]=-\sum_{t\le T}\E_{s_t\sim q}\KL\big(\pi_{\rm old}(\cdot|s_t)\,\|\,\pi_\theta(\cdot|s_t)\big)\le0,\label{eq:typ}\end{equation} |
| which follows by taking $\E_q$ of $\sum_t\log r_t$ and the tower rule. If $\log r_t$ are approximately i.i.d.\ $\mathcal N(-\sigma^2/2,\sigma^2)$ (the Gaussian that makes $\E r=1$), then $\log w_T\sim\mathcal N(-T\sigma^2/2,T\sigma^2)$: the median weight is $e^{-T\sigma^2/2}\to0$, $\E w_T^2=e^{T\sigma^2}$, and $\ESS/n\approx e^{-T\sigma^2}$. The unit mean is carried by exponentially rare samples, and any finite batch sees only the ones with $w\approx0$ plus a few enormous outliers. |
| \ver{on the 8191-sequence toy: $\E_q[\log w_T]$ equals minus the tree-recursion sum of conditional KLs to $\le1.9\times10^{-16}$ (checked at $T=3,8$); $\chi^2_T$ grows monotonically ($\times13.3$ from $T=1$ to $12$), while $\E_q[w_T]=1\pm2\times10^{-16}$.} |
|
|
| \paragraph{Magnitude.} Take a tiny per-token $\chi^2_1=0.01$ (our toy, i.i.d.\ tokens, so Eq.~\eqref{eq:geom} is an equality): $\ESS/n\approx(1.01)^{-T}$ gives $0.91$ at $T=10$, $0.37$ at $T=100$, $4.8\times10^{-5}$ at $T=1000$, $5\times10^{-18}$ at $T=4000$. With $10^{4}$-token reasoning traces even $\chi_1^2=10^{-3}$ gives $e^{-10}\approx4.5\times10^{-5}$. This is the argument behind abandoning the exact sequence weight in practice. |
|
|
| \subsection{The token-level approximation (what PPO/GRPO do) and its error} |
| Replace $w(o)=\prod_sr_s$ by the \emph{single} token ratio of the token being differentiated: |
| \begin{equation}\hat g_{\rm tok}(\theta)=\E_q\Big[\sum_tr_t(\theta)\,(R-b)\,\nabla\log\pi_\theta(o_t\mid s_t)\Big]=\nabla_\theta\E_q\Big[\sum_tr_t(\theta)(R-b)\Big].\label{eq:tok}\end{equation} |
| The discarded factor for token $t$ is $\prod_{s\ne t}r_s$: the other tokens' ratios, i.e.\ (for $s<t$) the mismatch of the \emph{state distribution} and (for $s>t$) the mismatch of the future. Both factors equal $1$ at $\theta=\theta_{\rm old}$, so Eq.~\eqref{eq:tok} reproduces the true gradient there; it is wrong at $\theta\ne\theta_{\rm old}$ at first order in $\|\theta-\theta_{\rm old}\|$ for the gradient and second order for the surrogate value (0104 turns this into a bound). Measured on the toy: |
| \begin{center}\small\begin{tabular}{@{}rcc@{}}\toprule |
| $\|\theta-\theta_{\rm old}\|$ & $\|\hat g_{\rm tok}-\nabla J\|$ & $|\Delta L-\Delta J|$\\\midrule |
| 0.00 & $6\times10^{-17}$ & $3\times10^{-17}$\\ |
| 0.02 & $2.9\times10^{-4}$ & $3.9\times10^{-7}$\\ |
| 0.08 & $1.2\times10^{-3}$ & $6.2\times10^{-6}$\\ |
| 0.32 & $4.7\times10^{-3}$ & $1.0\times10^{-4}$\\\bottomrule\end{tabular}\end{center} |
| \ver{log--log slopes $1.00$ (gradient) and $2.01$ (value); both errors vanish at $\theta_{\rm old}$ to $\le6\times10^{-17}$.} |
| GSPO (arXiv:2507.18071) argues from the same observation that a single-sample token weight cannot play the distribution-correcting role of Eq.~\eqref{eq:is} (Sec.~3, Eq.~4) and instead uses the length-normalised sequence ratio $s_i=(\pi_\theta/\pi_{\rm old})^{1/|y_i|}$ (Sec.~4.1, Eq.~7), whose log is the \emph{mean} (not sum) of token log-ratios, which keeps it near $1$ at the price of changing the objective. |
|
|
| \subsection{Implementation: log-space, not a running product} |
| Naive: multiply token ratios. Stable: $\log w=\sum_t(\log\pi_\theta-\log\pi_{\rm old})_t$ and exponentiate once (or never, if clipping happens in log space). With $T=4000$, $\sigma=0.5$ the true $\log w=-521.5$; a float32 running product returns exactly $0$, while the log-space value is $3.4\times10^{-227}$ (finite in float64). Both agree to within $10^{-12}$ relative error at $T=50$. |
|
|
| \section{Boxed result} |
| \begin{result} |
| \begin{align*} |
| &\E_p[f]=\E_q[wf],\qquad \Var_q[w]=\chi^2(p\|q),\qquad \ESS\approx\frac{n}{1+\chi^2}\\ |
| &\E_q[w_T^2]=\E_q\big[w_{T-1}^2\,(1+\chi_T^2(s_T))\big]\ \Longrightarrow\ \Var_q[w_T]=(1+\chi_1^2)^T-1\quad\text{(i.i.d.\ tokens)}\\ |
| &\nabla_\theta J(\theta)=\E_{o\sim\pi_{\rm old}}\Big[\prod_{t=1}^{|o|}r_t(\theta)\,(R-b)\,\nabla_\theta\log\pi_\theta(o)\Big] |
| \end{align*} |
| \end{result} |
|
|
| \section{What breaks} |
| \begin{breaks}\small\begin{itemize} |
| \item \textbf{Support violations.} Eq.~\eqref{eq:is} needs $q>0$ wherever $p\,f\ne0$. Any sampler that truncates the vocabulary (top-$k$, top-$p$, min-$p$) has $\mu(a)=0$ for tokens the trainer's $\pi(a)>0$; those tokens can never be corrected. \emph{Violated:} absolute continuity (cf.\ 0108). |
| \item \textbf{Heavy tails / infinite variance.} $\chi^2=\infty$ if $p^2/q$ is not summable; with finite $\chi_1^2$ the variance still grows like $(1+\chi_1^2)^T$. |
| \item \textbf{ESS collapse.} Eq.~\eqref{eq:ess} is a first-order law-of-large-numbers statement; for $n\ll\ESS^{-1}$ the empirical ESS is far below the prediction because rare huge weights are unobserved. |
| \item \textbf{SNIS bias} is $O(1/n)$ and not bounded uniformly in $T$. |
| \item \textbf{Token-level surrogate is not the gradient.} Eq.~\eqref{eq:tok} drops $\prod_{s\ne t}r_s$; the error grows linearly in $\|\theta-\theta_{\rm old}\|$ (measured slope $1.00$). \emph{Violated:} $\pi_\theta$ stays close to $\pi_{\rm old}$ (the trust-region assumption of 0104). |
| \item \textbf{Stale advantages.} $R-b$ was computed at $\theta_{\rm old}$; for a learned $b$ or value function it is also stale. |
| \item \textbf{Numerics.} Running products underflow or overflow; clip in log space. |
| \end{itemize}\end{breaks} |
|
|
| \nextnote{The weights above are harmless only if $\pi_\theta$ stays near $\pi_{\rm old}$. Note 0104 turns ``stay near'' into a theorem (a lower bound on improvement, TRPO), then derives the PPO clipped objective as a cheap pessimistic proxy for that bound and tabulates exactly which tokens it silences.} |
|
|
| \section*{References} |
| {\small\begin{itemize}[leftmargin=1.2em] |
| \item Yu et al., \emph{DAPO}, arXiv:2503.14476, Sec.~4.1 (512 prompts $\times$16 samples, mini-batch 512, i.e.\ 16 updates per rollout step). |
| \item MiniMax, \emph{MiniMax-M1}, arXiv:2506.13585, Sec.~3.1 (Eq.~1: PPO ratio $r_{i,t}$; ``16 rounds of off-policy updates per generation batch''). |
| \item Zheng et al., \emph{Group Sequence Policy Optimization}, arXiv:2507.18071, Sec.~3 (Eq.~4, IS needs many samples; token weight is single-sample noise), Sec.~4.1 Eqs.~5--7. |
| \item Qi et al., arXiv:2510.26788, Sec.~2.1, Eq.~5 (sequence-level IS policy gradient). |
| \item The IS identity, $\chi^2$ variance formula, ESS, and the $(1+\chi^2)^T$ recursion are derived above from first principles. |
| \end{itemize}} |
| \end{document} |
|
|
