Download code/notes/src/0203-pass-at-k.tex from viai957/align: direct link, hf CLI and curl.
- Browser
- Download file 16 kB
-
https://huggingface.co/viai957/align/resolve/main/code/notes/src/0203-pass-at-k.tex
- Command line
-
hf download hf://viai957/align/code/notes/src/0203-pass-at-k.tex
-
curl -L -o 0203-pass-at-k.tex https://huggingface.co/viai957/align/resolve/main/code/notes/src/0203-pass-at-k.tex
16 kB
| \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} | |