Download code/notes/src/0207-test-time-scaling-cost-model.tex from viai957/align: direct link, hf CLI and curl.
- Browser
- Download file 14.9 kB
-
https://huggingface.co/viai957/align/resolve/main/code/notes/src/0207-test-time-scaling-cost-model.tex
- Command line
-
hf download hf://viai957/align/code/notes/src/0207-test-time-scaling-cost-model.tex
-
curl -L -o 0207-test-time-scaling-cost-model.tex https://huggingface.co/viai957/align/resolve/main/code/notes/src/0207-test-time-scaling-cost-model.tex
14.9 kB
| \documentclass[11pt]{article} | |
| \usepackage{common} | |
| \begin{document} | |
| \notehead{0207}{Test-time scaling: a cost model for samples versus self-correction}{Batch M2. Companion notebook: \texttt{notebooks/0207-test-time-scaling-cost-model.ipynb} (exact rationals, quadrature and Monte Carlo; executed).} | |
| \begin{problem} | |
| \question{With a perfect verifier, how should inference compute be split between fresh samples and revisions, and when should sampling stop?} | |
| At inference time a prover with a verifier can spend compute in two ways: more independent samples (pass@$n$), or longer sequential chains in which failed attempts are revised using compiler feedback (note 0204). With a perfect verifier the accounting changes qualitatively compared with chat-style best-of-$n$: nothing is lost to selection, and sampling can stop at the first success. This note builds the cost model: coverage versus $n$ over a population of problems, the expected cost with early stopping, the optimal stopping rule, the break-even between a fresh sample and a revision, and what an imperfect verifier would do instead. | |
| \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$ & solve probability of one sample on a given problem; across problems $p\sim\mathrm{Beta}(a,b)$ \\ | |
| $n$ & number of attempts (budget cap); $N\sim\mathrm{Geom}(p)$ the index of the first success \\ | |
| $\mathrm{cov}(n)=\E_p[1-(1-p)^n]$ & coverage (best-of-$n$ with a perfect verifier) \\ | |
| $c$ & cost of one attempt (generation plus verification); $U$ value of a solved problem \\ | |
| $h_j=\Pr[\text{success at attempt }j{+}1\mid j\text{ failures}]$ & posterior hazard \\ | |
| $p_1,\ q$ & success of a fresh attempt; success of a revision given earlier failures \\ | |
| $\rho$ & cost of a revision relative to a fresh attempt \\ | |
| $\phi$ & false-accept rate of an imperfect verifier \\ \bottomrule | |
| \end{tabular} | |
| \end{center} | |
| \section{Derivation} | |
| \subsection{Coverage with a perfect verifier} | |
| A problem with solve probability $p$ is solved by $n$ attempts with probability $1-(1-p)^n$ and, because the verifier is exact, \emph{the solving attempt can be identified}: the selection loss of note 0201 is zero. Over a population with density $\propto p^{a-1}(1-p)^{b-1}$, | |
| \begin{equation}\label{eq:cov} | |
| \E_p(1-p)^n=\int_0^1\frac{p^{a-1}(1-p)^{b+n-1}}{B(a,b)}dp=\frac{B(a,b+n)}{B(a,b)},\qquad \mathrm{cov}(n)=1-\frac{B(a,b+n)}{B(a,b)} . | |
| \end{equation} | |
| Using $B(a,b+n)=\Gamma(a)\Gamma(b+n)/\Gamma(a+b+n)$ and $\Gamma(b+n)/\Gamma(a+b+n)\sim n^{-a}$, | |
| \begin{equation}\label{eq:tail} | |
| 1-\mathrm{cov}(n)\ \sim\ \frac{\Gamma(a+b)}{\Gamma(b)}\,n^{-a}\qquad(n\to\infty). | |
| \end{equation} | |
| The uncovered fraction is a power law in $n$ whose exponent $a$ is the order of the density of problems near $p=0$: the heavier the tail of hard problems (small $a$), the slower the gain from more samples. Notebook: (\ref{eq:cov}) matches Monte Carlo ($2\times10^6$ problems) within $1.3$ standard errors for $n\in\{1,\dots,256\}$; the ratio of exact uncovered fraction to (\ref{eq:tail}) is $0.939$ at $n=10$, $0.9933$ at $100$ and $1.0000$ at $10^5$. Large Language Monkeys observes coverage rising log-linearly over four orders of magnitude and models it as $c\approx\exp(ak^b)$ (Fig.~5), and notes some curves do not follow it closely \cite{lm}. Fitting that law to Beta-mixture curves on $k\in[1,10^4]$ gives mean absolute errors $0.006$, $0.001$, $0.006$ (maxima $0.026,0.012,0.053$) for three Beta populations, comparable to the errors in that paper's fits; we do not claim the Beta model is the true generative process, only that a mixture of per-problem geometric laws reproduces the shape. | |
| \subsection{Early stopping: the expected cost with a perfect verifier} | |
| Sample until the first verified success or $n$ attempts. Then the attempts used are $\min(N,n)$ and | |
| \begin{equation}\label{eq:early} | |
| \E\min(N,n)=\sum_{j=0}^{n-1}\Pr[N>j]=\sum_{j=0}^{n-1}(1-p)^j=\frac{1-(1-p)^n}{p}, | |
| \end{equation} | |
| exactly (rational arithmetic; also by direct enumeration of the truncated geometric distribution). Averaging over the population, $\sum_{j<n}B(a,b+j)/B(a,b)$, which grows like $n^{1-a}$ for $a<1$: \emph{sublinear} in the cap $n$. For $\mathrm{Beta}(0.4,2)$ and $n=64$ the coverage is $0.767$ while the expected cost is $23.05$ attempts, not $64$ (Monte Carlo agreement $0.5$ s.e.). At equal expected cost $64$ the adaptive policy can run with cap $n_{\max}=322$ (expected cost $63.997$; $n_{\max}=323$ costs $64.12$) and reaches coverage $0.877$ against $0.767$ for a fixed $64$. Early stopping needs a verifier that tells you \emph{when} to stop; majority voting or a learned scorer does not offer this. | |
| \subsection{When to stop: posterior hazard and the threshold rule} | |
| Given $p\sim\mathrm{Beta}(a,b)$ and $j$ failures, Bayes gives the posterior $\propto p^{a-1}(1-p)^{b+j-1}$, i.e.\ $\mathrm{Beta}(a,b+j)$, so | |
| \begin{equation}\label{eq:hazard} | |
| h_j=\E[p\mid j\text{ failures}]=\frac{B(a+1,b+j)}{B(a,b+j)}=\frac a{a+b+j}. | |
| \end{equation} | |
| The longer a problem resists, the less likely the next attempt is to succeed. Suppose each attempt costs $c$ and a solve is worth $U$. Write $W_j$ for the optimal value after $j$ failures: $W_j=\max\{0,\ h_j(U-c)-(1-h_j)c+(1-h_j)W_{j+1}\}$ (stop, or pay $c$ and continue). Because $h_j$ is decreasing, if $h_jU<c$ at some $j$ then it stays so for every later $j$, and continuing can never pay: the one-step-lookahead rule is optimal, and the rule is | |
| \begin{equation}\label{eq:stop} | |
| \text{continue while}\quad h_jU\ge c\ \iff\ j<\frac{aU}{c}-(a+b). | |
| \end{equation} | |
| The notebook solves the dynamic programme backward for $\mathrm{Beta}(0.4,2)$: stopping indices $78$, $398$, $398$ for $(U,c)=(10,0.05),(10,0.01),(50,0.05)$ against the thresholds $77.6$, $397.6$, $397.6$. The effort scales with $U/c$ and the exponent $a$: for hard populations (small $a$) the optimal effort per problem is small unless the solve is valuable relative to cost. This is a Bayesian-optimal version of the idea behind compute-optimal allocation across problems of different difficulty \cite{snell} (Sec.~3.1--3.2 there choose a per-difficulty-bin strategy empirically): equalise the marginal value $h_jU$ across problems and stop where it falls to $c$. Snell et al.\ do not derive this rule. | |
| \subsection{Fresh sample or revision?} | |
| A fresh attempt costs one unit and succeeds with $p_1$. A revision costs $\rho$ units and succeeds with $q$ given that all previous attempts in the chain failed. The failure probability of a chain with $R$ revisions is $(1-p_1)(1-q)^R$ at cost $1+\rho R$, and $n=\lfloor B/(1+\rho R)\rfloor$ independent chains fit in a budget $B$, so | |
| \[ | |
| -\ln\Pr[\text{fail}]=n\bigl(a_1+Ra_2\bigr),\qquad a_1=-\ln(1-p_1),\ a_2=-\ln(1-q),\qquad n\approx\frac B{1+\rho R}. | |
| \] | |
| The log-failure removed per unit cost is $f(R)=B\,\dfrac{a_1+Ra_2}{1+\rho R}$, and | |
| \[ | |
| \frac{df}{dR}\propto a_2(1+\rho R)-\rho(a_1+Ra_2)=a_2-\rho a_1 . | |
| \] | |
| It does not depend on $R$: $f$ is monotone, so the optimum is at an extreme (no revisions, or the maximum allowed). Revise iff | |
| \begin{equation}\label{eq:rev} | |
| \frac{\ln(1-q)}{\ln(1-p_1)}\ \ge\ \rho\qquad\Bigl(\approx q/p_1\ge\rho\ \text{for small probabilities}\Bigr). | |
| \end{equation} | |
| Brute-force enumeration over integer $(n,R)$ with $R\le12$ and $n(1+\rho R)\le600$ on a grid of $p_1\in\{0.02,0.05,0.1\}$, $q\in\{0.01,0.03,0.06,0.12\}$, $\rho\in\{1,1.5,2.5\}$ agrees with the sign rule at $36/36$ points (for example $p_1{=}0.05$, $q{=}0.06$: with $\rho=1$ the optimum is $50$ chains of $11$ revisions, with $\rho=1.5$ it is $600$ fresh samples). The cap $R\le12$ matters: the objective is monotone in $R$ in the continuous relaxation, so without a cap the integer optimum would be a single chain of $599$ revisions (log-failure $37.1$ against $36.6$ for $50\times11$); $11$ rather than $12$ comes from the integrality of $n=\lfloor600/(1+R)\rfloor$. In practice the cap is set by the context window and by $q_j$ falling with $j$ (see the table below). A revision must be both informative and not much more expensive than a resample. If each revision re-reads all earlier attempts and feedback, then $\rho$ grows with the round and the test must be applied per round with $\rho_j$ and $q_j$. | |
| \textbf{A published data point, token accounting only.} Goedel-Prover-V2 reports on miniF2F that self-correction at pass@32 with the context extended to 128k tokens (YaRN) and up to five revisions reaches $92.7\%$, above the $92.2\%$ of the non-correcting 32B model at pass@8192 (Sec.~3.5; Table~4) \cite{gps}. Its main-table self-correction mode (initial attempt plus two revisions) is capped at 40,000 tokens in total and the first round of either mode at 30,000 (Sec.~3.2); that mode reaches only $90.4\%$ at pass@32, below $92.2\%$ (Table~4). By the caps, the token ratio for the $92.7\%$ configuration is $8192\cdot30\text{k}/(32\cdot128\text{k})=60\times$ (it would be $192\times$ with the 40k cap, but that configuration does not beat $92.2\%$). A ratio of caps is neither an upper nor a lower bound on the ratio of tokens actually used. In the cost model's terms, a 40k cap for a chain of one attempt and two revisions against 30k for one fresh attempt means a chain costs at most $\tfrac43$ of a fresh attempt at the caps, $1+2\rho\le\tfrac43$, i.e.\ $\rho\le\tfrac16$ if the caps were costs, so (\ref{eq:rev}) would hold whenever $q/p_1\gtrsim\tfrac16$ under that accounting. The cost model explains why such a trade can pay when (\ref{eq:rev}) holds; it does not predict that it does, because $q$ must be measured. | |
| \subsection{Imperfect verifiers cap the benefit of $n$} | |
| Suppose the verifier accepts every correct proof but also accepts an incorrect one with probability $\phi$. The first accepted sample is correct with probability | |
| \begin{equation}\label{eq:prec} | |
| \frac{p}{p+(1-p)\phi}, | |
| \end{equation} | |
| \emph{independent of the budget $n$}: a precision cap. (Monte Carlo: $0.8414$ vs $0.8403$, $0.8335$ vs $0.8333$, $0.8315$ vs $0.8347$.) Selecting by a learned scorer is worse still (note 0201: pick-correct probability $0.337$ vs $0.815$ for the exact verifier at $n=16$, $p=0.1$, $\sigma=1$). Large Language Monkeys reports exactly this gap on MATH with a Llama-3-8B-Instruct: coverage rises from $82.9\%$ to $98.44\%$ between 100 and 10{,}000 samples while majority-vote and reward-model selection reach only about $40.5\%\to41.4\%$ \cite{lm}. Holes amplify with $n$: if a sample hits a verifier hole with probability $h$, the chance that \emph{some} sample among $n$ does is $1-(1-h)^n$; for $h=10^{-4}$ and $n=8192$ that is $0.56$. Large-$n$ best-of-$n$ therefore \emph{requires} the audit of note 0201 more, not less. | |
| \subsection{Wall-clock and cost per attempt} | |
| One attempt costs $c=c_{\mathrm{gen}}L+c_{\mathrm{ver}}$ with $L$ generated tokens and a verification (compile) cost; with early stopping the expected cost per problem is $\E\min(N,n)\cdot c$ and with $W$ parallel workers the latency is roughly $\lceil n/W\rceil c$ (without early stopping) or the time to the first success among $W$ parallel streams (with it). Heavier proofs raise both terms, so $c$ is not constant across attempts; treat $c$ as a random variable correlated with $p$. | |
| \section{Boxed result} | |
| \begin{result} | |
| \begin{align*} | |
| \mathrm{cov}(n)&=1-\frac{B(a,b+n)}{B(a,b)},\qquad 1-\mathrm{cov}(n)\sim\frac{\Gamma(a+b)}{\Gamma(b)}n^{-a},\qquad \E\min(N,n)=\frac{1-(1-p)^n}{p},\\ | |
| h_j&=\frac a{a+b+j},\qquad \text{continue iff } j<\tfrac{aU}{c}-(a+b),\\ | |
| &\text{revise iff }\frac{\ln(1-q)}{\ln(1-p_1)}\ge\rho,\qquad | |
| \Pr[\text{correct}\mid\text{accepted}]=\frac p{p+(1-p)\phi}. | |
| \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 | |
| Attempts are correlated (shared prefix cache, seed reuse, low temperature) & i.i.d.\ attempts given $p$ & Coverage grows more slowly than $1-(1-p)^n$; hazards $h_j$ fall faster than $a/(a+b+j)$ \\ | |
| Beta prior misspecified, or $p$ drifts with context length & exchangeable Beta population & Stopping threshold (\ref{eq:stop}) and tail exponent $a$ are wrong; fit $a$ on held-out problems \\ | |
| Revisions are not independent of the failed attempt (anchoring on a wrong approach) & $q$ independent of history & $q_j$ decreases with $j$; (\ref{eq:rev}) must be applied per round \\ | |
| Revision cost grows with accumulated context & constant $\rho$ & $\rho_j$ increases linearly with the round; break-even moves toward fresh samples \\ | |
| Verifier holes & $V\le G$ (note 0201) & $1-(1-h)^n$ false successes; the larger $n$, the more audit is needed \\ | |
| Learned scorer instead of exact verifier & perfect verifier & Precision cap $p/(p+(1-p)\phi)$; no early stopping; selection error that does not vanish with $n$ \\ | |
| Heterogeneous verification cost (timeouts on long proofs) & constant $c$ & Expected cost under early stopping is larger for hard problems, where long proofs concentrate \\ | |
| Reporting a curve selected on the test set & held-out budget choice & Optimistic stopping thresholds and break-evens; see note 0208 \\ \bottomrule | |
| \end{tabular} | |
| \end{breaks} | |
| \nextnote{Every number in this note is an estimate from a finite benchmark. Note 0208 gives the statistics needed to report them honestly: exact McNemar for paired comparisons, Wilson intervals, cluster bootstrap over problems, and the power analysis that says what a $+5/244$ improvement can and cannot support.} | |
| \begin{thebibliography}{9} | |
| \bibitem{lm} B.~Brown et al. Large Language Monkeys: Scaling Inference Compute with Repeated Sampling. arXiv:2407.21787, Fig.~5 (coverage vs.\ samples, $c=\exp(ak^b)$) and the comparison with majority voting and reward-model selection (MATH, Llama-3-8B-Instruct). | |
| \bibitem{snell} C.~Snell et al. Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters. arXiv:2408.03314, Sec.~3.1--3.2 (compute-optimal strategy, difficulty bins), Sec.~6 (refining the proposal distribution; sequential revisions vs.\ parallel best-of-$N$, Fig.~5). | |
| \bibitem{gps} Y.~Lin et al. Goedel-Prover-V2. arXiv:2508.03613, Sec.~3.2 (30K-token first round; two additional self-correction rounds, 40K total), Sec.~3.5 (128k context via YaRN, up to 5 revisions: pass@32 $92.7\%$ vs.\ $92.2\%$ at pass@8192 without correction), Table~4 (miniF2F scaling). | |
| \bibitem{codex} M.~Chen et al. Evaluating Large Language Models Trained on Code. arXiv:2107.03374 (repeated sampling: $70.2\%$ at 100 samples, abstract; pass@$k$, Sec.~2.1). | |
| \bibitem{cobbe} K.~Cobbe et al. Training Verifiers to Solve Math Word Problems. arXiv:2110.14168 (learned verifiers for best-of-$n$; motivation for the imperfect-verifier case). | |
| \end{thebibliography} | |
| \end{document} | |