| \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} |
|
|
| \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} |
|
|
