Blame for code/notes/src/0201-verifier-binary-reward-mdp.tex

\documentclass[11pt]{article}
\usepackage{common}
\begin{document}
\notehead{0201}{Verifier-in-the-loop as a binary-reward MDP}{Batch M2, learning from a verifier. Companion notebook: \texttt{notebooks/0201-verifier-binary-reward-mdp.ipynb} (fp64, executed).}

\begin{problem}
\question{What exactly is optimised when a proof is only accepted or rejected by the Lean kernel, and is that reward an unbiased signal of what we care about?}
A theorem prover is a language model $\pi_\theta$ that, given a formal statement $x$ (a Lean~4 theorem), writes a proof attempt $y$. We want to \emph{improve} $\pi_\theta$ using only the fact that a proof is either accepted or rejected by the Lean kernel. Three questions must be answered before any algorithm (rejection sampling, GRPO, \dots) can be written down:
\begin{enumerate}[nosep]
\item What is the Markov decision process, and what exactly is optimised?
\item Is the reward the compiler returns an \emph{unbiased} signal of what we care about, so that no learned reward model (and hence no reward-model error to over-optimise) sits in the loop?
\item Where can that claim fail, and what does the optimiser do when it fails?
\end{enumerate}
\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
$x\sim D$ & formal statement (prompt), drawn from a problem distribution $D$ \\
$y=(y_1,\dots,y_T)$ & proof attempt, tokens $y_t\in\mathcal V$; $T$ may vary, EOS terminates \\
$s_t=(x,y_{<t})$ & state before token $t$ (the full prefix) \\
$\pi_\theta(y_t\mid s_t)$ & token policy; $\pi_\theta(y\mid x)=\prod_t\pi_\theta(y_t\mid s_t)$ \\
$G(x,y)\in\{0,1\}$ & ground truth: $y$ elaborates to a term of exactly the type $x$, no \texttt{sorry}, no extra axioms \\
$V(x,y)\in\{0,1\}$ & what the verifier (Lean compile under resource limits) returns; $r:=V$ \\
$p_\theta(x)=\Pr_{y\sim\pi_\theta}[V=1]$ & per-problem solve probability, pass@1 \\
$J(\theta)=\E_{x}\,p_\theta(x)$ & expected return \\
$A_+=\{V=1,G=0\}$ & false accepts (verifier holes) \\
$A_-=\{V=0,G=1\}$ & false rejects (timeouts, memory limits, non-determinism) \\
$s_\theta(y|x)=\score(y|x)$ & score function \\
$b$ & baseline, a function that does not depend on $y$ \\ \bottomrule
\end{tabular}
\end{center}

\section{Derivation}
\subsection{The MDP}
Generation is a deterministic-transition MDP. State $s_t=(x,y_{<t})$; action $a_t=y_t\in\mathcal V$; transition $s_{t+1}=(x,y_{\le t})$ (append the token, no environment randomness); discount $\gamma=1$. The reward is zero at every step except the terminal one, where it is the verifier output:
\[
r_t=0\ (t<T),\qquad r_T=V(x,y).
\]
The return of a trajectory is therefore $\sum_t r_t=V(x,y)\in\{0,1\}$: a \emph{binary, terminal} reward. At the sequence level the same object is a contextual bandit: context $x$, action $y$, reward $V(x,y)$. Both views give the same objective,
\begin{equation}\label{eq:J}
J(\theta)=\E_{x\sim D}\,\E_{y\sim\pi_\theta(\cdot|x)}\,V(x,y)=\E_{x\sim D}\,p_\theta(x).
\end{equation}

\subsection{Score-function gradient}
Fix $x$. Writing the expectation as a sum over the (finite, though astronomically large) set of proofs,
\begin{align}
\nabla_\theta p_\theta(x)&=\nabla_\theta\sum_y\pi_\theta(y|x)V(x,y)=\sum_y V(x,y)\,\nabla_\theta\pi_\theta(y|x)\label{eq:g1}\\
&=\sum_y V(x,y)\,\pi_\theta(y|x)\,\nabla_\theta\log\pi_\theta(y|x)\label{eq:g2}\\
&=\E_{y\sim\pi_\theta}\bigl[V(x,y)\,\nabla_\theta\log\pi_\theta(y|x)\bigr].\label{eq:g3}
\end{align}
Step (\ref{eq:g1}) uses that $V$ does not depend on $\theta$: the verifier is a fixed function, not a differentiable network, so the gradient acts only on the sampling distribution. Step (\ref{eq:g2}) is $\nabla\pi=\pi\nabla\log\pi$ (valid where $\pi>0$; terms with $\pi=0$ contribute $0$ to both sides).

\subsection{Factorising over tokens}
Since $\log\pi_\theta(y|x)=\sum_{t=1}^T\log\pi_\theta(y_t|s_t)$,
\begin{equation}\label{eq:tok}
\nabla_\theta\log\pi_\theta(y|x)=\sum_{t=1}^T\nabla_\theta\log\pi_\theta(y_t|s_t).
\end{equation}
Because the only non-zero reward arrives at the end, the reward-to-go at every step is $G_t=\sum_{u\ge t}r_u=V(x,y)$. Hence
\[
\nabla_\theta p_\theta(x)=\E\Bigl[\sum_{t=1}^T V(x,y)\,\nabla_\theta\log\pi_\theta(y_t|s_t)\Bigr]:
\]
\emph{every token of an accepted proof receives the same credit}, whether it was the key lemma or a filler \texttt{simp}. This is the credit-assignment problem in its purest form; notes 0204 and 0206 are about reducing its variance.

\paragraph{Softmax score identity.} For a tabular softmax $\pi(a|s)=e^{z_a}/\sum_b e^{z_b}$ with logits $z=z(s)$,
\[
\frac{\partial\log\pi(a|s)}{\partial z_b}=\ind[a=b]-\frac{e^{z_b}}{\sum_c e^{z_c}}=\ind[a=b]-\pi(b|s),
\]
i.e.\ $\nabla_z\log\pi(a|s)=e_a-\pi(\cdot|s)$. The notebook uses this hand formula for the ``naive'' gradient and compares it with autograd.

\subsection{Baselines keep the gradient unbiased}
Differentiate the normalisation $\sum_y\pi_\theta(y|x)=1$ (true for every $\theta$): $0=\sum_y\nabla\pi_\theta(y|x)=\sum_y\pi_\theta(y|x)\nabla\log\pi_\theta(y|x)$, that is
\begin{equation}\label{eq:zero}
\E_{y\sim\pi_\theta}\bigl[\score(y|x)\bigr]=0.
\end{equation}
Hence for any $b=b(x)$ not depending on $y$, $\E[(V-b)\,\score]=\E[V\score]-b\cdot 0=\nabla p_\theta(x)$. (If $b$ depended on $y$ the factorisation $\E[b(y)\score]=b\,\E[\score]$ would fail; this is the whole unbiasedness condition, and notes 0204 and 0206 re-use it.)

\subsection{Variance of the one-sample estimator and the optimal constant baseline}
Let $g=(V-b)\,s$ with $s=\score(y|x)$ (a vector). Then $\E g=\nabla p$ for every $b$ and
\[
\Var[g]:=\E\|g\|^2-\|\E g\|^2=\E\bigl[(V-b)^2\|s\|^2\bigr]-\|\nabla p\|^2 .
\]
Set the derivative in $b$ to zero: $\partial_b\Var=-2\,\E[(V-b)\|s\|^2]=0$, so
\begin{equation}\label{eq:bstar}
b^\star=\frac{\E[V\,\|s\|^2]}{\E[\|s\|^2]},\qquad
\Var_{\min}=\E[V\|s\|^2]-\frac{\E[V\|s\|^2]^2}{\E\|s\|^2}-\|\nabla p\|^2,
\end{equation}
using $V^2=V$ for binary $V$. It is a \emph{score-norm-weighted} success rate. It lies near $p$ but is not equal to it. In the notebook: $p=0.3045$, $b^\star=0.3382$, and the exact variances are $0.599$ ($b=0$), $0.394$ ($b=p$), $0.392$ ($b=b^\star$): almost all of the gain comes from using \emph{any} sensible baseline.

\subsection{Why the verifier reward is unbiased for what we want}
Define the error indicators. By definition $V-G=\ind_{A_+}-\ind_{A_-}$, hence, exactly,
\begin{equation}\label{eq:bias}
J_V(\theta)=J_G(\theta)+\E_x\,\pi_\theta(A_+|x)-\E_x\,\pi_\theta(A_-|x).
\end{equation}
Two cases give the two halves of the claim.
\begin{itemize}[nosep]
\item \emph{Soundness} ($A_+=\varnothing$): the verifier never accepts a non-proof. Then $V\le G$ pointwise, $J_V\le J_G$, the reward can only err on the \emph{pessimistic} side, and there is no region of output space where a high reward is unearned.
\item \emph{Independent timeouts}: if (with a sound verifier) rejection of a true proof happens with probability $\delta$ independently of $y$ \emph{and of $x$}, then $J_V=(1-\delta)J_G$ and $\nabla J_V=(1-\delta)\nabla J_G$: a positive rescaling of the true gradient, which a learning-rate schedule absorbs. The gradient \emph{direction} is unbiased. If $\delta=\delta(x)$ depends on the prompt (long statements time out more), the result is $\nabla J_V=\E_x[(1-\delta(x))\nabla p^G_\theta(x)]$: still $y$-independent, but a re-weighting of prompts, not a global rescaling.
\end{itemize}
Compare a learned reward model $r_\phi=G+\varepsilon(x,y)$, whose error $\varepsilon$ is a function of $y$ that the policy is free to exploit: $J_\phi-J_G=\E_\pi\varepsilon$. Choosing $y$ to maximise $G+\varepsilon$ selects \emph{for} large $\varepsilon$ (the winner's curse). For $n$ candidates, each correct w.p.\ $p$, and i.i.d.\ Gaussian error of scale $\sigma$,
\[
\Pr[\text{RM pick correct}]=\int n\,p\,f_1(u)\bigl[p\,F_1(u)+(1-p)F_0(u)\bigr]^{n-1}du,
\qquad \text{verifier: }1-(1-p)^n,
\]
($f_j,F_j$: pdf/cdf of $\mathcal N(j,\sigma^2)$), and the score the RM reports for its pick exceeds the truth by an amount bounded by $\E\max_i\varepsilon_i\le\sigma\sqrt{2\ln n}$. Numbers from the notebook (quadrature matched by Monte Carlo): $n{=}16$, $p{=}0.1$: verifier $0.815$, RM with $\sigma=1$: $0.337$. The verifier loses \emph{only} when no sample is correct (coverage), never to noise. Gao et al.\ measure the same effect in RLHF: gold reward as a function of $d=\sqrt{\mathrm{KL}}$ follows $d(\alpha-\beta d)$ for best-of-$n$ and $d(\alpha-\beta\log d)$ for RL, rising then falling \cite{gao2022}.

\subsection{What the optimiser does with a hole}
Take a tabular softmax over three output classes $k\in\{\text{wrong},\text{proof},\text{hole}\}$ with $V=(0,1,1)$, $G=(0,1,0)$. From $J=\sum_k\pi_kr_k$ and the softmax derivative,
\[
\frac{\partial J}{\partial z_k}=\pi_k\Bigl(r_k-\sum_j\pi_jr_j\Bigr)=\pi_k(r_k-J).
\]
Gradient flow $\dot z_k=\pi_k(r_k-J)$ gives, for the two accepted classes (both $r=1$),
\[
\frac{d}{dt}\bigl(z_{\text{proof}}-z_{\text{hole}}\bigr)=(\pi_{\text{proof}}-\pi_{\text{hole}})(1-J).
\]
Whichever accepted class is \emph{more reachable} gains on the other, truth plays no role. Starting from $\pi=(0.96,0.01,0.03)$ the notebook integrates this to $E[V]=0.99995$ while $E[G]=0.0008$ and $\pi(\text{hole})=0.9992$: the true-proof share of the accepted mass falls from $0.25$ to $0.001$. A sound reward is an exact reward only on the set the verifier actually checks.

\section{Boxed result}
\begin{result}
\begin{align*}
J(\theta)&=\E_{x}\,\E_{y\sim\pi_\theta}V(x,y),\qquad
\nabla J=\E_{x,y}\Bigl[\bigl(V(x,y)-b(x)\bigr)\sum_{t=1}^{T}\nabla_\theta\log\pi_\theta(y_t|s_t)\Bigr]\ \ \text{for every } b(x),\\[3pt]
J_V&=J_G+\E\,\pi_\theta(A_+)-\E\,\pi_\theta(A_-),\qquad
b^\star=\frac{\E[V\|s\|^2]}{\E[\|s\|^2]} .
\end{align*}
With a sound verifier ($A_+=\varnothing$) and timeouts that are independent of both $y$ and $x$ ($A_-$ a fixed fraction $\delta$ of true proofs), $\nabla J_V=(1-\delta)\nabla J_G$.
\end{result}

\section{What breaks}
\begin{breaks}
\begin{tabular}{@{}p{0.30\linewidth}p{0.30\linewidth}p{0.34\linewidth}@{}}
\toprule
Failure & Assumption violated & Consequence \\ \midrule
Verifier hole ($A_+\neq\varnothing$): \texttt{sorry}-like escapes, extra axioms, unsafe/\texttt{native} evaluation, kernel or elaborator bugs & Soundness, $V\le G$ & $J_V=J_G+\pi(A_+)$. The optimiser climbs toward the most reachable accepted output; $E[G]$ can collapse while $E[V]\to1$ (\S3.7) \\
Statement mismatch: formal $x$ does not say what the informal problem says & $G$ measures the intended property & Perfect verification of the wrong theorem; outside the verifier's remit \\
Timeouts correlated with $y$ (long or heavy proofs time out more) & $\delta$ independent of $y$ & $\nabla J_V=\nabla J_G-\nabla\pi(A_-)$: pressure toward cheap-to-check proofs \\
Non-deterministic verifier (tactic randomness, load-dependent limits) & $V$ is a function of $(x,y)$ & Reward becomes random; still unbiased for $\E V$ but noisier and not reproducible \\
Reward too sparse ($p_\theta\approx0$ on most $x$) & non-degenerate $p_\theta$ & $\nabla p=0$ at $p\in\{0,1\}$, see note 0206 \\
Credit assignment over long $y$ & --- (not a bias, a variance) & all $T$ tokens share one scalar; $\E\|s\|^2$ is a sum of $T$ score terms and grows with $T$; cf.\ VinePPO \cite{vineppo} \\
Reward-model substitution (for e.g.\ informal reasoning) & $\varepsilon\perp y$ & Winner's curse above; gold reward peaks then declines \cite{gao2022} \\ \bottomrule
\end{tabular}
\end{breaks}

\nextnote{Note 0202 spends the binary reward in the cheapest possible way. Sampling then discarding unaccepted proofs and fine-tuning on the rest is the M-step of an EM algorithm whose E-step is ``sample and verify''. We show that, to first order, one such step \emph{is} the policy-gradient step derived here, and say when it can and cannot improve the policy.}

\begin{thebibliography}{9}
\bibitem{gao2022} L.~Gao, J.~Schulman, J.~Hilton. Scaling Laws for Reward Model Overoptimization. arXiv:2210.10760, Sec.~1 (functional forms $R_{\mathrm{bon}}(d)=d(\alpha-\beta d)$, $R_{\mathrm{RL}}(d)=d(\alpha-\beta\log d)$).
\bibitem{rloo} A.~Ahmadian et al. Back to Basics: Revisiting REINFORCE Style Optimization. arXiv:2402.14740, Sec.~2.3 (leave-one-out baseline) and Eq.~8 (REINFORCE with baseline).
\bibitem{vineppo} A.~Kazemnejad et al. VinePPO: Refining Credit Assignment in RL Training of LLMs. arXiv:2410.01679 (value-network credit assignment; motivation).
\bibitem{gps} Y.~Lin et al. Goedel-Prover-V2. arXiv:2508.03613, Sec.~2.3--2.4 (Lean compiler as verifier; RL on verifier reward).
\bibitem{dsmath} Z.~Shao et al. DeepSeekMath. arXiv:2402.03300, Sec.~4.1 (GRPO with outcome reward).
\end{thebibliography}
\end{document}