Blame for code/notes/src/0203-pass-at-k.tex

\documentclass[11pt]{article}
\usepackage{common}
\begin{document}
\notehead{0203}{pass@$k$: the unbiased estimator, its variance, and why pass@1 and pass@$k$ disagree}{Batch M2. Companion notebook: \texttt{notebooks/0203-pass-at-k.ipynb} (exact rational arithmetic and fp64, executed).}

\begin{problem}
\question{How do we estimate pass@$k$ without bias from $n\ge k$ samples, and why is optimising pass@$k$ not the same as optimising pass@1?}
A prover is judged by pass@$k$: the probability that at least one of $k$ sampled proofs is accepted. We sample $n\ge k$ proofs per problem (a few hundred is typical), count the $c$ accepted ones, and need a number. The obvious guess $1-(1-c/n)^k$ is biased. And once pass@$k$ is a training or model-selection target it is \emph{not} the same target as pass@1: later-stage RL typically raises pass@1 while pass@$N$ for large $N$ falls \cite{gps,yue}. This note (i) derives the unbiased estimator and its variance, (ii) shows precisely where the two objectives differ, and (iii) derives an unbiased \emph{gradient} estimator for pass@$k$.
\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
$p=p_\theta(x)$ & per-problem pass@1 (note 0201) \\
$n,\ k,\ c$ & samples drawn ($n\ge k$), attempts in the metric, number accepted; $c\sim\mathrm{Bin}(n,p)$ \\
$\rho_{n,k}(c)=1-\binom{n-c}{k}\big/\binom nk$ & the estimator of pass@$k$ \\
$J_k(\theta)=\E_x[1-(1-p_\theta(x))^k]$ & pass@$k$ objective ($J_1=J$) \\
$w_k(p)=k(1-p)^{k-1}$ & per-problem weight of $J_k$ relative to $J_1$ \\
$\rho_{-i}$ & unbiased pass@$(k{-}1)$ estimate from the $n-1$ samples other than $i$ \\
$s_i=\nabla_\theta\log\pi_\theta(y_i|x)$ & score of sample $i$ \\ \bottomrule
\end{tabular}
\end{center}

\section{Derivation}
\subsection{Definition}
Draw $k$ i.i.d.\ proofs from $\pi_\theta(\cdot|x)$. All fail with probability $(1-p)^k$, so
\begin{equation}\label{eq:passk}
\pk(x)=1-(1-p)^k,\qquad \pk=\E_x\bigl[1-(1-p_\theta(x))^k\bigr]=J_k .
\end{equation}
$p$ is unknown; we have only $c\sim\mathrm{Bin}(n,p)$.

\subsection{The plug-in estimator is biased}
$\hat p=c/n$ is unbiased for $p$, but $u\mapsto(1-u)^k$ is convex, so by Jensen $\E(1-\hat p)^k\ge(1-p)^k$ and $1-(1-\hat p)^k$ \emph{under}estimates pass@$k$. For $k=2$ the bias is exact: $\E(1-\hat p)^2=(1-p)^2+\Var(\hat p)=(1-p)^2+p(1-p)/n$, so the estimate is low by $p(1-p)/n$. Notebook, $n{=}8$, $p{=}0.3$: bias $-0.026$ ($k{=}2$), $-0.075$ ($k{=}4$), $-0.104$ ($k{=}8$), as exact rationals. Codex notes the same underestimate (Appendix A, Fig.~13) \cite{codex}.

\subsection{The unbiased estimator}
\begin{proposition}
For $n\ge k$, $\ \rho_{n,k}(c)=1-\binom{n-c}{k}/\binom nk$ satisfies $\E_{c\sim\mathrm{Bin}(n,p)}\rho_{n,k}(c)=1-(1-p)^k$ for all $p$.
\end{proposition}
\emph{Proof 1 (probabilistic).} Condition on $c$. Choose a uniformly random $k$-subset of the $n$ samples. It contains no correct sample with probability $\binom{n-c}{k}/\binom nk$ (hypergeometric). Unconditionally, a uniformly random $k$-subset of $n$ i.i.d.\ samples is itself $k$ i.i.d.\ samples, which are all wrong with probability $(1-p)^k$. By the tower property $\E_c[\binom{n-c}{k}/\binom nk]=(1-p)^k$. $\square$

\emph{Proof 2 (algebraic, nothing skipped).} First the identity
\[
\binom nc\binom{n-c}{k}\Big/\binom nk=\frac{n!}{c!\,(n-c)!}\cdot\frac{(n-c)!}{k!\,(n-c-k)!}\cdot\frac{k!\,(n-k)!}{n!}=\frac{(n-k)!}{c!\,(n-c-k)!}=\binom{n-k}{c}.
\]
Then
\[
\E_c\frac{\binom{n-c}{k}}{\binom nk}=\sum_{c=0}^{n-k}\binom nc\frac{\binom{n-c}{k}}{\binom nk}p^c(1-p)^{n-c}=\sum_{c=0}^{n-k}\binom{n-k}{c}p^c(1-p)^{n-c}=(1-p)^k\underbrace{\sum_{c=0}^{n-k}\binom{n-k}{c}p^c(1-p)^{n-k-c}}_{=1},
\]
where terms $c>n-k$ vanish because $\binom{n-c}{k}=0$ there. $\square$

The notebook verifies this as an \emph{equality of rational numbers} for all $n\le40$, all $1\le k\le n$ and four values of $p$ (error exactly $0$), and the binomial identity for all $n<30$.

\subsection{A numerically stable form}
Cancel factorials:
\[
\frac{\binom{n-c}{k}}{\binom nk}=\frac{(n-c)!}{(n-c-k)!}\cdot\frac{(n-k)!}{n!}=\frac{(n-k)!/(n-k-c)!}{n!/(n-c)!}=\prod_{i=n-c+1}^{n}\frac{i-k}{i}=\prod_{i=n-c+1}^{n}\Bigl(1-\frac ki\Bigr),
\]
and the whole quantity is $0$ when $n-c<k$ (then $\rho=1$). Every factor lies in $(0,1)$; no factorial is formed. This is the code of Codex Fig.~3 \cite{codex}. In the notebook the naive float ratio raises \texttt{OverflowError} for $n{=}6000$, $k{=}3000$, while the product form agrees with exact rational arithmetic to $8\times10^{-17}$ on a grid including $(n,k)=(1000,999)$.

\subsection{Variance}
$\rho$ is a U-statistic: $1-\rho=\binom nk^{-1}\sum_{S}h(S)$, $h(S)=\ind[\text{all of }S\text{ fail}]$, sum over $k$-subsets. Draw two independent uniform $k$-subsets $S,S'$ of $[n]$; $\E[(1-\rho)^2]=\E_{S,S'}\E[h(S)h(S')]$. If $|S\cap S'|=j$ both fail iff the $2k-j$ distinct samples all fail: probability $(1-p)^{2k-j}$. The overlap is hypergeometric, $w_j=\Pr[|S\cap S'|=j]=\binom kj\binom{n-k}{k-j}/\binom nk$, $\sum_jw_j=1$. Hence
\begin{equation}\label{eq:var}
\Var(\rho)=\sum_{j=0}^{k}w_j(1-p)^{2k-j}-(1-p)^{2k}=\sum_{j=1}^{k}w_j\Bigl[(1-p)^{2k-j}-(1-p)^{2k}\Bigr].
\end{equation}
For $n\gg k^2$, $w_1\simeq k^2/n$ and the $j=1$ term dominates:
\begin{equation}\label{eq:var1}
\Var(\rho)\simeq\frac{k^2\,p\,(1-p)^{2k-1}}{n}.
\end{equation}
This is also the delta-method variance of the plug-in estimator, $\bigl(k(1-p)^{k-1}\bigr)^2p(1-p)/n$: the unbiased estimator costs nothing asymptotically. Notebook: formula (\ref{eq:var}) equals the exact binomial-sum variance in exact rational arithmetic (five $(n,k)$ pairs $\times$ four $p$, difference $0$); for $n{=}200,k{=}8,p{=}0.3$ exact $4.80\times10^{-4}$ vs leading term $4.56\times10^{-4}$ (5\% off). For a benchmark of $N{=}244$ problems with $p_x\sim\mathrm{Beta}(0.5,1)$, $n{=}16$, $k{=}8$, the within-problem variance of the mean is $\tfrac1{N^2}\sum_x\Var_x=1.323\times10^{-4}$ against $1.336\times10^{-4}$ by $4\times10^4$ Monte Carlo replications (1\%), and the mean is unbiased ($0.9$ standard errors). The between-problem component, which dominates, is the subject of note 0208.

\subsection{Why pass@1 and pass@$k$ pull in different directions}
Differentiate (\ref{eq:passk}) with the chain rule:
\begin{equation}\label{eq:gradk}
\nabla J_k=\E_x\bigl[k(1-p_\theta(x))^{k-1}\,\nabla p_\theta(x)\bigr]=\E_x\bigl[w_k(p)\,\nabla p\bigr].
\end{equation}
Within one problem, $1-(1-p)^k$ is increasing in $p$, so the two objectives agree on what is better. They disagree \emph{across} problems, through the weights $1$ versus $w_k(p)$: $w_k$ is decreasing in $p$, so pass@$k$ ignores problems already solved reliably and emphasises problems with small $p$. In logit coordinates $p=\sigma(\ell)$, $dp=p(1-p)\,d\ell$, so a unit logit shift is worth $p(1-p)$ for pass@1 (peak at $p=\tfrac12$) but $kp(1-p)^k$ for pass@$k$ (peak at $p=\tfrac1{k+1}$, e.g.\ $0.03$ for $k=32$).

\emph{A model that cannot tell problems apart makes this a conflict.} Take 12 problems and a context-free policy $\pi$ over four strategies. Strategy~0 is a generalist solving each of 8 ``easy'' problems with probability $0.6$ and no hard problem; strategies $1,2,3$ each solve their own hard problems with probability $0.3$ (strategy~1 problems 9 and 10, strategy~2 problem 11, strategy~3 problem 12) and each easy problem with probability $0.1$. Since $J_k$ is concave in the mixing weights $\pi$ (a concave function of a linear map), the gradient-ascent optima below are global. Maximising $J_1$ by gradient ascent in fp64 puts $\pi=(1,0,0,0)$ (entropy $0$): $J_1=0.400$, $J_{32}=0.667$. Maximising $J_{32}$ gives $\pi=(0.111,0.341,0.274,0.274)$ (entropy $1.32$): $J_{32}=0.981$ but $J_1=0.134$. Each optimum is strictly worse than the other on the other's metric. This is the structure of diversity collapse: sharpening onto the dominant strategy raises pass@1 and removes the rare strategies that carry pass@$k$ at large $k$. It matches the empirical pattern of RLVR \cite{yue,leash}, and Goedel-Prover-V2's observation that later SFT/RL stages raise pass@1 while pass@32 declines \cite{gps}.

Two further points. (a) The average over problems is not a function of the average solve rate: by concavity of $p\mapsto1-(1-p)^k$, $\pk\le1-(1-\bar p)^k$; heterogeneity matters. (b) Verifier holes (note 0201) are amplified by $k$: a per-sample false-accept rate $h$ adds up to $\approx kh$ to pass@$k$ but only $h$ to pass@1.

\subsection{An unbiased pass@$k$-aware policy gradient}
Optimising $J_k$ directly needs (\ref{eq:gradk}) from samples. Plugging $\hat p=c/n$ into $w_k(\hat p)\hat g$ is biased for the same Jensen reason as above. Instead apply the score identity to the \emph{joint} of $k$ samples. For i.i.d.\ $y_1,\dots,y_k$, $\Pr[\max_iV_i=1]=1-(1-p)^k$, so
\[
\nabla J_k(x)=\E\Bigl[\max_{i\le k}V_i\ \sum_{i\le k}s_i\Bigr].
\]
With $n\ge k$ samples, average over all $k$-subsets (a U-statistic, hence unbiased):
\begin{equation}\label{eq:pk}
\hat g_k=\binom nk^{-1}\sum_{|S|=k}\Bigl(\max_{i\in S}V_i\Bigr)\sum_{i\in S}s_i=\sum_{i=1}^nw_i\,s_i,\qquad w_i=\binom nk^{-1}\sum_{S\ni i}\max_{j\in S}V_j .
\end{equation}
Closed form for binary $V$. Each $i$ lies in $\binom{n-1}{k-1}$ subsets. If $V_i=1$ every such subset has max $1$: $w_i=\binom{n-1}{k-1}/\binom nk=k/n$. If $V_i=0$, the subset has max $1$ iff its other $k-1$ members contain a correct sample, i.e.\ $\binom{n-1}{k-1}-\binom{n-1-c}{k-1}$ subsets (here $c=c_{-i}$ is the number correct among the others). Using $\binom{n-1}{k-1}/\binom nk=k/n$,
\[
w_i=\frac kn\Bigl[V_i+(1-V_i)\,\rho_{-i}\Bigr]=\frac kn\Bigl[\rho_{-i}+V_i\,(1-\rho_{-i})\Bigr],\qquad
\rho_{-i}=1-\binom{n-1-c_{-i}}{k-1}\Big/\binom{n-1}{k-1}.
\]
$\rho_{-i}$ is the unbiased pass@$(k{-}1)$ estimate (Section 3.3) built from the other $n-1$ samples; it is independent of $y_i$. Therefore the term $\tfrac kn\rho_{-i}$ is a baseline: $\E[\rho_{-i}s_i]=\E[\rho_{-i}]\,\E[s_i]=0$ by Eq.~(6) of note 0201. Removing it leaves
\begin{equation}\label{eq:pkloo}
\hat g_k^{\mathrm{LOO}}=\frac kn\sum_{i=1}^nV_i\,(1-\rho_{-i})\,s_i .
\end{equation}
Only accepted samples receive gradient, with weight equal to the (estimated) probability that the other $k-1$ draws would \emph{all} fail, i.e.\ the marginal contribution of $y_i$ to a $k$-set. Check: $\E[1-\rho_{-i}]=(1-p)^{k-1}$ so the expected coefficient of an accepted sample's score is $\tfrac kn(1-p)^{k-1}$; summing the $np$ accepted samples gives $k(1-p)^{k-1}\nabla p$, which is (\ref{eq:gradk}). For $k=1$, $\rho_{-i}=0$ and (\ref{eq:pkloo}) is REINFORCE with weights $V_i/n$. The same transformation, applied to the batch of rewards, is the pass@$k$ policy optimisation of Walder and Karkhanis (Theorems~1--2) \cite{pkpo}; Chen et al.\ use pass@$k$ itself as the reward (Eq.~5) \cite{passk-train}.

Notebook (tabular softmax, $K{=}4$ outcomes, $n{=}5$, $k{=}3$, all $4^5{=}1024$ sample tuples enumerated): the closed-form weights equal the naive subset average on every tuple (max difference $4.4\times10^{-16}$); the expectation of (\ref{eq:pk}), of the closed form and of (\ref{eq:pkloo}) all equal the autograd gradient of $1-(1-p)^k$ to $\le6.7\times10^{-16}$; the total variance of the leave-one-out form is $0.140$ versus $0.441$ without the baseline.

\subsection{No unbiased estimator exists for $k>n$}
An estimator $f(c)$ is unbiased for $g(p)$ iff $g(p)=\sum_{c=0}^nf(c)\binom ncp^c(1-p)^{n-c}$, a polynomial of degree at most $n$ (Kolmogorov, quoted as Theorem~8 of \cite{pkpo}). $1-(1-p)^k$ has degree $k$. The notebook projects it on the degree-$n$ Bernstein basis: residual $\le3\times10^{-15}$ for $k\le n$, and $4\times10^{-3}$ ($n{=}4,k{=}5$), $4\times10^{-2}$ ($n{=}4,k{=}8$) for $k>n$.

\section{Boxed result}
\begin{result}
\begin{align*}
\pk&=\E_x\bigl[1-(1-p)^k\bigr],\qquad
\hatpk=\E_{\text{problems}}\Bigl[1-\tbinom{n-c}{k}\big/\tbinom nk\Bigr]\ \ (\text{unbiased, }n\ge k),\\
&\tbinom{n-c}{k}\big/\tbinom nk=\prod_{i=n-c+1}^{n}\bigl(1-\tfrac ki\bigr),\\
\Var\bigl(\rho_{n,k}\bigr)&=\sum_{j=1}^{k}\frac{\binom kj\binom{n-k}{k-j}}{\binom nk}\bigl[(1-p)^{2k-j}-(1-p)^{2k}\bigr]\ \simeq\ \frac{k^2p(1-p)^{2k-1}}{n},\\
\nabla J_k&=\E_x\bigl[k(1-p)^{k-1}\nabla p\bigr],\\
\hat g_k^{\mathrm{LOO}}&=\frac kn\sum_iV_i\bigl(1-\rho_{-i}\bigr)\nabla\log\pi_\theta(y_i|x)\quad(\text{unbiased for i.i.d.\ samples and }n\ge k).
\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
Samples are not i.i.d.\ (shared prefix cache, fixed seed, beam/top-$k$ without replacement, deduplication) & $c\sim\mathrm{Bin}(n,p)$ & $\rho$ is no longer unbiased; deduplicated samples \emph{over}state diversity, correlated ones understate it \\
$n<k$ & $k\le n$ & No unbiased estimator exists (degree argument) \\
Averaging $p$ before applying $1-(1-\cdot)^k$ & per-problem transform & Overestimates pass@$k$ by concavity (heterogeneous $p_x$) \\
Plug-in $1-(1-\hat p)^k$ & unbiasedness & Biased low by about $\binom k2p(1-p)^{k-1}/n$ (leading order) \\
Verifier holes (note 0201) & $V\le G$ & False accepts are amplified roughly $k$-fold in pass@$k$ \\
Training on pass@1 with a context-free bottleneck & objectives separable per problem & Weight $1$ vs $w_k(p)$: collapse onto the generalist strategy; pass@1 up, pass@$k$ down \\
Pass@$k$ gradient at high $k$ & fixed $k$ & $w_k(p)\to0$ for $p\gtrsim1/k$: easy problems stop training and pass@1 can fall; PKPO's Tables~1--2 show pass@1 lower by $1.55$ points (Gemma-2-9B, MATH: $22.24\to20.69$) and $2.52$ points (Llama-3.1-8B: $51.15\to48.63$) at $k_{\mathrm{opt}}=8$ than at $k_{\mathrm{opt}}=1$, and Sec.~5.2.3 anneals $k$ to recover \cite{pkpo} \\
Sampling-distribution mismatch (temperature, top-$p$) & metric uses the training distribution & pass@$k$ is a property of the \emph{sampling} distribution; changing temperature changes it \\ \bottomrule
\end{tabular}
\end{breaks}

\nextnote{Note 0204 moves from one-shot sampling to multi-turn self-correction: after a failed attempt the verifier's error message is fed back and the model revises. That is a multi-turn MDP whose policy-gradient needs a loss mask (feedback tokens are not sampled by the policy) and per-turn baselines; we derive when each choice keeps the gradient unbiased.}

\begin{thebibliography}{9}
\bibitem{codex} M.~Chen et al. Evaluating Large Language Models Trained on Code. arXiv:2107.03374, Sec.~2.1 (Eq.~1, Fig.~3 stable implementation), Appendix~A (bias of $1-(1-\hat p)^k$, Fig.~13).
\bibitem{pkpo} C.~Walder, D.~Karkhanis. Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems. arXiv:2505.15201, Theorem~1 (unbiased pass@$k$ estimator), Theorem~2 (unbiased gradient), Theorem~8 and Corollary~2 (Kolmogorov; App.~A.1: pass@$k$ is unbiasedly estimable iff $n\ge k$), Sec.~5.2.3 (annealing $k$), Tables~1--2 (pass@1 at $k_{\mathrm{opt}}=8$ vs $1$).
\bibitem{passk-train} Z.~Chen et al. Pass@k Training for Adaptively Balancing Exploration and Exploitation of Large Reasoning Models. arXiv:2508.10751, Sec.~2.2 (Eq.~5, pass@$k$ as reward).
\bibitem{yue} Y.~Yue et al. Does Reinforcement Learning Really Incentivize Reasoning Capacity in LLMs Beyond the Base Model? arXiv:2504.13837 (pass@$k$ at large $k$: RLVR narrows rather than expands).
\bibitem{leash} F.~Wu et al. The Invisible Leash: Why RLVR May or May Not Escape Its Origin. arXiv:2507.14843 (empirical-support shrinkage vs expansion).
\bibitem{gps} Y.~Lin et al. Goedel-Prover-V2. arXiv:2508.03613, Sec.~2.3 (``Model averaging for enhanced diversity'': pass@1 increases, pass@N declines in later stages).
\bibitem{lm} B.~Brown et al. Large Language Monkeys. arXiv:2407.21787, Fig.~5 (coverage vs.\ samples).
\end{thebibliography}
\end{document}