Blame for code/notes/src/0202-rejection-sampling-expert-iteration-em.tex

\documentclass[11pt]{article}
\usepackage{common}
\begin{document}
\notehead{0202}{Rejection sampling and expert iteration as policy improvement}{Batch M2. Companion notebook: \texttt{notebooks/0202-rejection-sampling-expert-iteration-em.ipynb} (fp64, executed).}

\begin{problem}
\question{What objective does ``sample, keep the accepted, fine-tune, repeat'' optimise, how does it relate to the policy gradient, and when does it stop working?}
Note 0201 gave us a binary reward $V(x,y)$ and the policy gradient $\nabla J=\E[(V-b)\,\score]$. Running that gradient online needs a rollout engine inside the training loop, a baseline, clipping, and careful stability engineering. Rejection-sampling fine-tuning (RFT), STaR and expert iteration instead do something much simpler: \emph{sample from a frozen policy, keep what the verifier accepts, run ordinary supervised fine-tuning on the kept samples, repeat.} The question: what objective does that procedure optimise, how does it relate to the policy gradient, does it improve the policy monotonically, and when does it stop working?
\end{problem}
\stepcounter{section}% the problem box counts as section 1: keeps every \S and notebook cross-reference stable

\section{Notation}
\begin{center}
\begin{tabular}{@{}p{0.36\linewidth}p{0.60\linewidth}@{}}
\toprule
Symbol & Meaning \\ \midrule
$\pi_t=\pi_{\theta_t}$ & frozen sampling policy at iteration $t$ \\
$V(x,y)\in\{0,1\}$ & verifier, note 0201 \\
$p_t(x)=\sum_y\pi_t(y|x)V(x,y)$ & per-problem solve probability under $\pi_t$ \\
$q_t(y|x)=\pi_t(y|x)V(x,y)/p_t(x)$ & reward-weighted ("accepted") distribution, defined where $p_t(x)>0$ \\
$O$ & binary optimality variable with $\Pr[O{=}1\,|\,x,y]=V(x,y)$ \\
$\mathcal L(\theta,q)$ & evidence lower bound (ELBO) on $\log p_\theta(x)$ \\
$F(\theta)=\frac1{|S|}\sum_{x\in S}\log p_\theta(x)$ & EM objective over the solvable set $S=\{x:\exists y,\,V(x,y)=1\}$ \\
$L_t(\theta)=\E_x\E_{y\sim\pi_t}[V\log\pi_\theta]$ & pooled RFT objective \\
$n,\ c$ & samples per problem, number accepted \\ \bottomrule
\end{tabular}
\end{center}

\section{Derivation}
\subsection{The reward-weighted distribution and the RFT objective}
Pooled RFT draws $y\sim\pi_t(\cdot|x)$, keeps pairs with $V=1$ and minimises the token-level cross-entropy on them. In the infinite-sample limit its objective is
\[
L_t(\theta)=\E_x\sum_y\pi_t(y|x)V(x,y)\log\pi_\theta(y|x)=\E_x\,p_t(x)\sum_y q_t(y|x)\log\pi_\theta(y|x),
\]
using $\pi_tV=p_tq_t$. Since $\sum_yq_t\log\pi_\theta=-H(q_t)-\KL(q_t\|\pi_\theta)$ and $H(q_t)$ does not depend on $\theta$,
\begin{equation}\label{eq:rft}
\argmax_\theta L_t(\theta)=\argmin_\theta\ \E_x\,p_t(x)\,\KL\bigl(q_t(\cdot|x)\,\|\,\pi_\theta(\cdot|x)\bigr).
\end{equation}
RFT is \emph{forward-KL projection of the policy onto the accepted distribution, weighted by $p_t(x)$}.

\subsection{First-order equivalence with the policy gradient}
Differentiate $L_t$ at $\theta=\theta_t$:
\[
\nabla_\theta L_t\big|_{\theta_t}=\E_x\sum_y\pi_t(y|x)V(x,y)\,\nabla_\theta\log\pi_\theta(y|x)\big|_{\theta_t}=\E_x\E_{y\sim\pi_t}\bigl[V\,\score\bigr]=\nabla_\theta J(\theta_t),
\]
by Eq.~(4) of note 0201. So an \emph{infinitesimal} RFT step is exactly a REINFORCE step with zero baseline; a finite RFT step (training to convergence on the fixed data) is a large, off-policy step along the same direction at its start. STaR writes the same policy gradient with an indicator reward as its Eqs.~(1)--(2) and remarks that ``the indicator function discards the gradient for all sampled rationales that do not lead to the correct answer'' \cite{star}; it then says STaR only \emph{approximates} $J$ (greedy decoding, which biases exploration, and several gradient steps on one batch), which is exactly the finite-step, off-policy caveat above. Notebook check: the autograd gradient of $L_t$ at $\theta_t$ matches $\nabla J$ to $3\times10^{-18}$.

\subsection{The EM view}
Let $\Pr[O{=}1|x]=\sum_y\pi_\theta(y|x)V(x,y)=p_\theta(x)$. For any distribution $q(\cdot|x)$ over $y$ that puts mass wherever $V\pi_\theta>0$ (so that the importance weight $V\pi_\theta/q$ is defined and $\E_q[V\pi_\theta/q]=p_\theta(x)$), Jensen's inequality ($\log$ concave) gives
\begin{align}
\log p_\theta(x)&=\log\E_{y\sim q}\Bigl[\frac{V\,\pi_\theta}{q}\Bigr]\ \ge\ \E_{y\sim q}\Bigl[\log\frac{V\,\pi_\theta}{q}\Bigr]=\E_q[\log V]-\KL(q\|\pi_\theta)=:\mathcal L(\theta,q).\label{eq:elbo}
\end{align}
(For $q$ supported on $\{V=1\}$, $\log V=0$. If $q$ misses part of $\{V\pi_\theta>0\}$ the first equality becomes ``$\ge$'' and the bound $\log p_\theta\ge\mathcal L$ still holds. If $q$ puts mass on $\{V=0\}$ then $\mathcal L=-\infty$.) Write $q^\star_\theta=\pi_\theta V/p_\theta$. Then the gap is exactly a KL:
\begin{equation}\label{eq:gap}
\log p_\theta(x)-\mathcal L(\theta,q)=\KL(q\|q^\star_\theta).
\end{equation}
\emph{Proof.} $\KL(q\|q^\star_\theta)=\E_q[\log q-\log\pi_\theta-\log V+\log p_\theta]=-\mathcal L(\theta,q)+\log p_\theta$ (when $q$ is supported on $\{V=1\}$). $\square$

\textbf{E-step.} Maximise $\mathcal L(\theta_t,q)$ over $q$: by (\ref{eq:gap}) the maximiser is $q^\star_{\theta_t}=q_t$ and the bound is tight, $\mathcal L(\theta_t,q_t)=\log p_{\theta_t}(x)$. Sampling from $\pi_t$ and discarding rejects \emph{is} sampling from $q_t$ (rejection sampling).

\textbf{M-step.} Maximise $\mathcal L(\theta,q_t)$ over $\theta$, equivalently minimise $\KL(q_t\|\pi_\theta)$, i.e.\ maximum likelihood on the accepted samples. Summing over problems with equal weights gives the exact EM M-step for the per-problem ELBOs, which is the form in which \cite{rest} derives it (their ELBO chain). Note that the objective ReST$^{\mathrm{EM}}$ actually states for non-negative rewards, their Eq.~(3), $\E_x\E_{y\sim p_t}[r\log p_\theta]$, and the loss in their Algorithm~1 pool all accepted samples, so they are the \emph{pooled} objective $L_t$ of the previous subsection (weight $p_t(x)$ per problem), not the equal-weight one.

\subsection{Monotone improvement (what is and is not guaranteed)}
\begin{proposition}
If the M-step is solved exactly, $F(\theta_{t+1})\ge F(\theta_t)$.
\end{proposition}
\emph{Proof.} For each $x\in S$ (with $q_t(\cdot|x)\ll\pi_{\theta_{t+1}}(\cdot|x)$, which holds for the full-support softmax policies used here):
\[
\log p_{\theta_{t+1}}(x)\ \overset{(\ref{eq:elbo})}{\ge}\ \mathcal L(\theta_{t+1},q_t)\qquad\text{and}\qquad \mathcal L(\theta_t,q_t)\ \overset{(\ref{eq:gap})}{=}\ \log p_{\theta_t}(x)
\]
hold pointwise. The middle step, $\mathcal L(\theta_{t+1},q_t)\ge\mathcal L(\theta_t,q_t)$, holds only \emph{summed} over $x\in S$, because the exact M-step maximises $\sum_{x\in S}\mathcal L(\theta,q_t(\cdot|x))$ over the shared $\theta$ and $\theta_t$ is a feasible point. Averaging the chain over $S$ gives $F(\theta_{t+1})\ge F(\theta_t)$. $\square$

This is the monotonicity statement ReST$^{\mathrm{EM}}$ quotes ($\mathcal L(\theta_{t+1},q_{t+1})\ge\mathcal L(\theta_t,q_{t+1})\ge\mathcal L(\theta_t,q_t)$ \cite{rest}); as item~2 below shows, it applies to the equal-weight M-step and not literally to the pooled objective their Eq.~(3) and Algorithm~1 optimise. Three things it does \emph{not} say:
\begin{enumerate}[nosep]
\item It is a statement about $F=\E_x\log p_\theta$, not about $J=\E_xp_\theta$ (pass@1). The gradients differ per problem by the factor $1/p$: $\nabla\log p=\nabla p/p$. EM weights hard problems more than the policy gradient does.
\item The pooled objective actually run in practice has weight $p_t(x)$ on each problem (\ref{eq:rft}); the guarantee holds for the equal-weight (per-problem normalised) version.
\item Pointwise improvement $p_{t+1}(x)\ge p_t(x)$ is not guaranteed when $\theta$ is shared across problems.
\end{enumerate}
In the notebook (log-linear policy, $d{=}5$ features, $K{=}12$ outputs, 8 prompts, M-step solved by L-BFGS to gradient norm $1.2\times10^{-9}$): $F$ rises monotonically from $-1.733$ to $-1.361$ in 12 iterations with zero decreases. With the \emph{pooled} (unnormalised) weights, the mean solve rate on solvable prompts rises from $0.26$ (after the first iteration) to $0.42$ (after the sixth), yet $F$ falls from $-1.45$ to $-2.35$ over the same span and in 20 of the 42 (prompt, iteration) pairs the prompt's own $p$ decreased: the update trades hard prompts for easy ones.

\subsection{Full-capacity limit and the coverage theorem}
If the model class contains $q_t$ (tabular softmax), the exact M-step sets $\pi_{t+1}(\cdot|x)=q_t(\cdot|x)$ for every $x$ with $p_t(x)>0$. Then
\[
p_{t+1}(x)=\sum_yq_t(y|x)V(x,y)=\frac{\sum_y\pi_tV^2}{p_t}=1 .
\]
One iteration solves every problem the current policy can solve with positive probability. For $p_t(x)=0$, $q_t$ is $0/0$: there is no accepted sample to imitate and the iteration leaves the problem untouched (notebook, six prompts: $p_t=(0.261,0.366,0.180,0.168,2\!\times\!10^{-4},0)\ \to\ p_{t+1}=(1,1,1,1,1,0)$). With finite $n$ the problem contributes at all only with probability $1-(1-p_t)^n$: the \emph{coverage frontier} is $p_t(x)\gtrsim1/n$. Any progress on $p_t=0$ problems comes from generalisation across problems, not from the algorithm.

\subsection{The finite-sample estimator}
Fix $x$, draw $y_1,\dots,y_n\sim\pi_t$ i.i.d., $c=\sum_iV_i\sim\mathrm{Bin}(n,p)$. Condition on $c=k\ge1$ and on \emph{which} indices are correct: the correct samples are i.i.d.\ from $q_t$ (conditioning an i.i.d.\ $\pi_t$ sample on $V=1$). Therefore for any $f$,
\begin{equation}\label{eq:ratio}
\E\Bigl[\frac1c\sum_{i}V_if(y_i)\,\ind[c>0]\Bigr]=\Pr[c\ge1]\ \E_{q_t}f=\bigl(1-(1-p)^n\bigr)\,\E_{q_t}f,
\end{equation}
and the \emph{pooled} average satisfies $\E[\tfrac1n\sum_iV_if(y_i)]=p\,\E_{q_t}f$ by linearity. The filtered mean is exactly unbiased conditional on $c\ge1$; its only error is the missing event $c=0$. Verified by enumerating all $4^n$ tuples for $n\in\{1,2,3,5\}$: errors below $2.3\times10^{-15}$.

\subsection{How the variants differ}
\begin{center}
\begin{tabular}{@{}p{0.40\linewidth}p{0.30\linewidth}p{0.22\linewidth}@{}}
\toprule
Procedure & Gradient on problem $x$ at $\theta_t$ & Per-problem weight on $\nabla p_x$ \\ \midrule
Pooled RFT / STaR / ReST$^{\mathrm{EM}}$ Eq.~(3) (\ref{eq:rft}) & $\nabla p_t(x)$ & $1$ \\
Per-problem-normalised EM (exact M-step of the per-problem ELBOs; not what Eq.~(3) of \cite{rest} runs) & $\nabla\log p_t(x)=\nabla p_t/p_t$ (notebook check) & $1/p_t(x)$ \\
Cap or dedupe per problem (RFT distinct-path filter \cite{rft}) & not derived here; changes the effective per-problem count & depends on cap \\ \bottomrule
\end{tabular}
\end{center}
RFT and STaR also retrain from the original base model each round \cite{rft,star}, which for a convex, exactly solved M-step is irrelevant but for finite training acts as a regulariser toward the base. Note 0203 adds the third objective, pass@$k$, whose weight is $k(1-p)^{k-1}$.

\section{Boxed result}
\begin{result}
\begin{align*}
q_t(y|x)&=\frac{\pi_t(y|x)V(x,y)}{p_t(x)},\\
\theta_{t+1}&=\argmax_\theta\ \E_{x,\,y\sim\pi_t}\bigl[V\log\pi_\theta\bigr]=\argmin_\theta\E_xp_t\,\KL(q_t\|\pi_\theta),\\
\nabla_\theta L_t\big|_{\theta_t}&=\nabla_\theta J(\theta_t),\qquad
F(\theta_{t+1})\ge F(\theta_t)\quad\text{(exact M-step, equal weights)},\\
\E\Bigl[\tfrac1c\textstyle\sum_iV_if(y_i);\,c>0\Bigr]&=\bigl(1-(1-p)^n\bigr)\E_{q_t}f,\\
p_t(x)&=0\ \Rightarrow\ \text{no datum, no gradient (tabular: }p_{t+1}(x)=0).
\end{align*}
\end{result}

\section{What breaks}
\begin{breaks}
\begin{tabular}{@{}p{0.29\linewidth}p{0.29\linewidth}p{0.36\linewidth}@{}}
\toprule
Failure & Assumption violated & Consequence \\ \midrule
Coverage: $np_t(x)\ll1$ & support of $\pi_t$ meets $\{V=1\}$ at $x$ & No accepted sample; EM is silent on $x$. Hardest problems never improve except by transfer \\
Finite samples, many epochs & M-step is a population argmin & The empirical $\hat q$ has support on $k$ distinct proofs; fitting it drives $H(\pi)\to\le\ln k$ (diversity collapse; see 0203, 0205) \\
Shared parameters & per-prompt independent M-step & $F$ monotone, pointwise $p$ not; 20/42 prompt-steps decreased in the notebook \\
Pooled weights $\propto p_t(x)$ & equal-weight EM objective & Easy problems dominate the data; hard problems trade away \\
Verifier holes & $V\le G$ (note 0201) & $q_t$ puts mass on false accepts and SFT \emph{distils the exploit} directly into the weights \\
Inexact M-step (few steps, LR schedule) & exact argmax & Monotonicity proof does not apply; behaves as a damped policy-gradient step \\
Only positives used & --- & Rejected samples carry gradient information the baseline form $(V-b)$ uses and RFT discards (variance, note 0206) \\ \bottomrule
\end{tabular}
\end{breaks}

\nextnote{Both the EM objective and the policy gradient turn per-problem solve rates $p(x)$ into a training signal, but they weight problems differently ($1/p$ vs $1$). The metric we report, pass@$k$, is yet another per-problem transform, $1-(1-p)^k$. Note 0203 derives its unbiased estimator, its variance, and why optimising pass@1 and pass@$k$ pull in different directions.}

\begin{thebibliography}{9}
\bibitem{rest} A.~Singh et al. Beyond Human Data: Scaling Self-Training for Problem-Solving with Language Models. arXiv:2312.06585, Sec.~3 (ELBO Eq.~2; E-step/M-step; monotone ELBO chain; ``EM with non-negative rewards'', Eq.~3, which is the pooled objective; Algorithm~1 and the ``Improve (M-step)'' loss $J(\theta)=\E_{(x,y)\sim\mathcal D_i}[r\log p_\theta]$).
\bibitem{rft} Z.~Yuan et al. Scaling Relationship on Learning Mathematical Reasoning with Large Language Models. arXiv:2308.01825, Sec.~3.3 (RFT; distinct-equation-list filter; Algorithm~1 is in the appendix).
\bibitem{star} E.~Zelikman et al. STaR: Bootstrapping Reasoning with Reasoning. arXiv:2203.14465, Sec.~3 (Eqs.~1--2: indicator reward and its gradient; retrain from the original model).
\bibitem{exit} T.~Anthony, Z.~Tian, D.~Barber. Thinking Fast and Slow with Deep Learning and Tree Search. arXiv:1705.08439 (expert iteration: planner = search, apprentice = imitation).
\bibitem{gps} Y.~Lin et al. Goedel-Prover-V2. arXiv:2508.03613, Sec.~2.4 (expert-iteration SFT stages S1--S3 before RL).
\end{thebibliography}
\end{document}