Blame for code/notes/src/0205-model-averaging-weight-interpolation.tex

\documentclass[11pt]{article}
\usepackage{common}
\begin{document}
\notehead{0205}{Model averaging and weight interpolation}{Batch M2. Companion notebook: \texttt{notebooks/0205-model-averaging-weight-interpolation.ipynb} (fp64, executed).}

\begin{problem}
\question{Why should interpolating the weights back toward the base keep the pass@1 gain while restoring pass@$k$, and when is there an interior optimum?}
Later stages of SFT and RL raise pass@1 and lower pass@$N$ for large $N$: the model gets sharper and loses the rare proofs that carry coverage (note 0203). Goedel-Prover-V2 repairs this \emph{without retraining} by interpolating the trained weights back toward the base: with base parameters $\theta_0$ and fine-tuned $\theta$ it uses $(1-\alpha)\theta_0+\alpha\theta$, $\alpha\in(0,1)$, after SFT (the averaged model starts RL) and again after RL \cite{gps}. We want to know to first order \emph{why} this keeps pass@1 gains while restoring pass@$k$, when an interior optimum $\alpha^\star$ exists, and what the approximation error is.

\textbf{Convention warning.} \cite{gps}, Sec.~2.3 defines the average as $(1-\alpha)\theta_0+\alpha\theta$ ($\alpha$ weights the fine-tuned model), but Sec.~3.6 says ``$\alpha$ is the weight of the base model'' and then reports that ``a higher proportion of the base model (i.e., lower $\alpha$) leads to lower pass@1, but pass@$N$ first rises and then falls as $\alpha$ increases''. The clause ``a higher proportion of the base model (i.e., lower $\alpha$)'' is consistent only with $\alpha$ weighting the \emph{fine-tuned} model, contradicting the preceding ``$\alpha$ is the weight of the base model'', so Sec.~3.6 is internally inconsistent as well as inconsistent with Sec.~2.3. Throughout this note $\alpha$ is the weight of the \emph{fine-tuned} model, $\alpha=0$ is the base.
\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
$\theta_0,\ \theta_1$ ($=\theta$) & base and fine-tuned parameters; $d=\theta_1-\theta_0$ \\
$\theta_\alpha=\theta_0+\alpha d$ & interpolated parameters \\
$z(\theta)=z_\theta(s,\cdot)$ & next-token logits at state $s$; $\pi_\theta=\mathrm{softmax}(z_\theta)$ \\
$g=z_1-z_0$ & logit change ("tilt"), per context \\
$\Cov_\alpha,\ \E_\alpha$ & moments under $\pi_{\theta_\alpha}$ \\
$p_x(\alpha)$, $\ \Delta_x$ & solve probability of prompt $x$ along the path; log-odds slope $\Delta_x=\mathrm{logit}\,p_x(1)-\mathrm{logit}\,p_x(0)$ \\
$\kappa=d^\top Hd$ & curvature of the loss along the path \\ \bottomrule
\end{tabular}
\end{center}

\section{Derivation}
\subsection{Weight averaging is logit interpolation, up to a named second-order error}
Expand the logits of each context in $\alpha$ (second-order Taylor in the parameters, $J=\partial z/\partial\theta$, $H_z=\partial^2z/\partial\theta^2$):
\[
z(\theta_\alpha)=z_0+\alpha Jd+\tfrac12\alpha^2d^\top H_zd+O(\|d\|^3),\qquad z_1=z_0+Jd+\tfrac12d^\top H_zd+O(\|d\|^3).
\]
Therefore the interpolated logits $(1-\alpha)z_0+\alpha z_1=z_0+\alpha Jd+\tfrac12\alpha\,d^\top H_zd$ differ from the true ones by
\begin{equation}\label{eq:err}
z(\theta_\alpha)-\bigl[(1-\alpha)z_0+\alpha z_1\bigr]=\tfrac12\alpha(\alpha-1)\,d^\top H_zd+O(\|d\|^3).
\end{equation}
The error vanishes at the endpoints, is largest at $\alpha=\tfrac12$ (where it is $-\tfrac18d^\top H_zd$), and is zero for all $\alpha$ if the logits are linear in $\theta$ (the regime of a linearised/lazy network and of the notebook's log-linear policies). Under that approximation
\begin{equation}\label{eq:geo}
\pi_\alpha(y_t|s_t)=\frac{\pi_0(y_t|s_t)^{1-\alpha}\,\pi_1(y_t|s_t)^{\alpha}}{\sum_v\pi_0(v|s_t)^{1-\alpha}\pi_1(v|s_t)^{\alpha}}:
\end{equation}
a \emph{geometric} (product-of-experts) mixture at each token, renormalised per context. It is not the ensemble $(1-\alpha)\pi_0+\alpha\pi_1$: the maximal difference between them is $0.256$ in the notebook. The sequence law is a product of per-token normalised factors, not a normalised product of sequence probabilities. Notebook: evaluating the model at $\theta_\alpha$ (naive) and interpolating the two logit tables (efficient) agree to $4\times10^{-16}$, and both equal (\ref{eq:geo}).

\subsection{Exponential tilt: $\alpha$ is an inverse temperature on the update}
Equation (\ref{eq:geo}) says that if $\pi_1=\pi_0e^{g}/Z$ then $\pi_\alpha=\pi_0e^{\alpha g}/Z_\alpha$. \textbf{Scope.} In this subsection and the next, ``$\pi$'' is a single softmax over complete outputs $y$ (one categorical per prompt, which is what the notebook's log-linear families are), so that $g=z_1-z_0$ is one number per output and $Z_\alpha$ one number per prompt. For an autoregressive policy the tilt acts token by token and $Z_\alpha(s_t)$ depends on the context; the corrected statement is given at the end of this subsection. For a KL-regularised RL solution $\pi^\star\propto\pi_0e^{r/\beta}$ ($g=r/\beta$) of such a single-categorical policy, interpolation moves along the family of regularised optima with effective strength $\beta/\alpha$. Differentiating the normalised expectation, with $Z_\alpha=\E_0e^{\alpha g}$:
\begin{equation}\label{eq:cov}
\frac{d}{d\alpha}\E_\alpha[r]=\frac{d}{d\alpha}\frac{\E_0[re^{\alpha g}]}{\E_0[e^{\alpha g}]}=\E_\alpha[rg]-\E_\alpha[r]\E_\alpha[g]=\Cov_\alpha(g,r).
\end{equation}
For $g=r/\beta$ this is $\Var_\alpha(r)/\beta\ge0$, so pass@1 is non-decreasing in $\alpha$; for binary $r$ and $g=cr$ the odds simply multiply by $e^{\alpha c}$ and
\begin{equation}\label{eq:logistic}
p(\alpha)=\sigma\bigl(\mathrm{logit}\,p_0+\alpha c\bigr),
\end{equation}
exactly (notebook: error $0$ at five values of $\alpha$; and (\ref{eq:cov}) matches autograd to $5\times10^{-18}$).

\textbf{Autoregressive correction.} For a sequence $y=(y_1,\dots,y_T)$ with $\pi_\alpha(y)=\prod_t\pi_0(y_t|s_t)\,e^{\alpha g_t}/Z_\alpha(s_t)$, $g_t=g(s_t,y_t)$, one has $\frac{d}{d\alpha}\log\pi_\alpha(y)=\sum_t\bigl[g_t-\E_{\pi_\alpha(\cdot|s_t)}g\bigr]=:\tilde g_\alpha(y)$, the score of the family, and therefore
\[
\frac d{d\alpha}\E_\alpha[r]=\Cov_\alpha\bigl(\tilde g_\alpha,r\bigr)\neq\Cov_\alpha\Bigl(\sum_tg_t,\,r\Bigr)\quad\text{in general}.
\]
(The identity is the score-function identity of note 0201 applied to the path $\alpha$; $\E_\alpha\tilde g_\alpha=0$ is its zero-mean property.) For a single categorical $\tilde g_\alpha=g-\E_\alpha g$ and the two covariances coincide, which is why (\ref{eq:cov}) is exact there. In a two-token toy ($|\mathcal V|=3$, random logit tables) the finite-difference derivative is $-0.1919$, $\Cov_\alpha(\tilde g_\alpha,r)$ gives $-0.1919$ (error $5\times10^{-11}$) and $\Cov_\alpha(\sum_tg_t,r)$ gives $-0.1769$ (an $8\%$ error). Everywhere below, $g$ inside $\Cov_\alpha(g,\cdot)$ means $\tilde g_\alpha$ for autoregressive policies, and the second, log-odds-linear form of (\ref{eq:dpk}) is a per-prompt modelling assumption about $p_x(\alpha)$ that does not depend on this distinction.

\subsection{pass@1 versus pass@$k$ along the path}
Combine (\ref{eq:cov}) with $\pk=\E_x[1-(1-p_x)^k]$:
\begin{equation}\label{eq:dpk}
\frac{d}{d\alpha}\pk=\E_x\Bigl[k(1-p_x)^{k-1}\,\Cov_\alpha(g_x,r_x)\Bigr]\ \overset{\text{log-odds linear}}{=}\ \E_x\bigl[k\,p_x(1-p_x)^k\,\Delta_x\bigr].
\end{equation}
The first form is exact for any logit-linear path of a single-categorical policy (use $\tilde g_\alpha$ for autoregressive ones) and needs no assumption on the shape of $p_x(\alpha)$; the second assumes $\mathrm{logit}\,p_x(\alpha)=l_x+\alpha\Delta_x$ (true for binary tilts, a first-order approximation otherwise). Notebook: the second form matches autograd to $3\times10^{-17}$ for $k\in\{1,4,32\}$. For $k=1$ the weight is $p(1-p)$, peaked at $p=\tfrac12$; for large $k$ it is $kp(1-p)^k$, peaked at $p=\tfrac1{k+1}$ ($0.03$ for $k=32$). Whether pass@$k$ improves or degrades along the path is therefore decided by the sign of $\Delta_x$ on the \emph{hard tail}, not on the bulk:
\begin{itemize}[nosep]
\item RL tilts toward modes the model already finds, so $\Delta_x>0$ on prompts with moderate $p_x$ (large effect on pass@1) and can be $\le0$ on tail prompts where the useful modes are rare (small effect on pass@1, large on pass@$k$, $k\gg1$).
\item $\alpha^\star_k$ solves $\E_x[kp_x(1-p_x)^k\Delta_x]=0$. An \emph{interior} global optimum exists \emph{if} the derivative is positive at $\alpha=0$ (RL helps the base) and negative at $\alpha=1$ (RL has overshot): then neither end point can be the maximum. This is sufficient but not necessary. Counterexample: half of the prompts with $l=-3$, $\Delta=-1$ (declining) and half with $l=-8$, $\Delta=+10$ (late bloomers whose $p$ rises from $3\times10^{-4}$ to $0.88$), $k=8$. Then $\frac d{d\alpha}\mathrm{pass@}8=-0.115$ at $\alpha=0$ and $-0.062$ at $\alpha=1$, yet pass@8 is $0.162$ at $\alpha=0$, $0.568$ at $\alpha=1$ and $0.579$ at $\alpha^\star=0.80$. The path of pass@$k$ need not be unimodal, so the sign test at the end points only certifies the ``if'' direction (equivalence holds when pass@$k(\alpha)$ is unimodal).
\end{itemize}
Synthetic two-population model (65\% ``learned'' prompts with $l\sim\mathcal N(-0.5,1.5^2)$, $\Delta\sim\mathcal N(1.5,0.3^2)$; 35\% ``tail'' prompts with $l\sim\mathcal N(-3,0.5^2)$, $\Delta\sim\mathcal N(-1.2,0.3^2)$): $\alpha^\star_1=1$ ($0.278\to0.436$); $\alpha^\star_8=0.525$ (pass@8 $0.6717$ at $\alpha=0$, $0.6762$ at $\alpha^\star$, $0.6738$ at $\alpha=1$: a shallow optimum); $\alpha^\star_{32}=0$ ($0.900$ vs $0.787$ at $\alpha=1$).

On 30 random log-linear ``SFT $\to$ RL'' families (96 prompts, 24 candidates, 6 features; SFT = 80 steps and RL = 400 further steps of pass@1 ascent, exact pass@$k$): pass@1 is non-decreasing in $\alpha$ in all 30; the end-point derivative (\ref{eq:dpk}) at $\alpha=1$ is negative for pass@8 and pass@64 in all 30 (interpolating back toward the base \emph{always} helps locally); the derivative at $\alpha=0$ is positive in 14/30 (pass@8) and 3/30 (pass@64), which are exactly the families with an interior optimum. Means: pass@1 $0.091\to0.135$ (base $\to$ RL), pass@64 $0.779\to0.466$. So in these toys averaging with the base \emph{recovers} diversity, but a strictly interior optimum requires that the RL stage still improves pass@$k$ at the base, which is what Goedel-Prover-V2 reports for its checkpoints (pass@$N$ first rises and then falls as $\alpha$ grows) \cite{gps}. We take that as the paper's empirical finding; nothing in this derivation proves it holds for their models.

\subsection{Why a straight line in weight space is a sensible path}
Along $\theta_\alpha=\theta_0+\alpha d$ a loss $\ell$ satisfies, to second order,
\[
\ell(\theta_\alpha)-\bigl[(1-\alpha)\ell(\theta_0)+\alpha\ell(\theta_1)\bigr]=\tfrac12\alpha(\alpha-1)\,\kappa,\qquad\kappa=d^\top Hd,
\]
exactly for a quadratic (the same algebra as (\ref{eq:err})). If $\kappa\ge0$ the interpolated model beats the chord: no barrier. The \emph{barrier} $\max_\alpha[\ell(\theta_\alpha)-\text{chord}]$ is the object studied as linear mode connectivity: Frankle et al.\ find that networks become stable to SGD noise, i.e.\ end in a linearly connected region, early in training \cite{lmc}; the soups paper observes that models fine-tuned independently from one pretrained initialisation are connected, with an error barrier mainly when fine-tuning with high learning rates (Sec.~3.3, takeaway 2; App.~J.1, Fig.~J.1) \cite{soups}. Notebook (tiny tanh MLP, 2-D spirals, fp64): two fine-tunes from one pretrained checkpoint have barrier $0.0000$; two networks trained from independent initialisations have barrier $+3.19$ (endpoint losses $\le10^{-3}$ in both cases). For the convex log-linear NLL, the barrier is non-positive by convexity (checked: $0$).

For cross-entropy along linear logits the curvature is explicit: $L(\alpha)=-\E_x\log\mathrm{softmax}(z_\alpha)_{y^\star}$ gives
\begin{equation}\label{eq:curv}
\frac{d^2L}{d\alpha^2}=\E_x\Var_{y\sim\mathrm{softmax}(z_\alpha)}\bigl[\Delta z_y\bigr]\ \ge0
\end{equation}
(second derivative of a log-partition function; notebook error $4\times10^{-16}$). Wortsman et al.\ give the general soup-versus-logit-ensemble approximation $L^{\mathrm{soup}}_\alpha-L^{\mathrm{ens}}_\alpha\approx\frac{\alpha(1-\alpha)}2\bigl[-\frac{d^2}{d\alpha^2}L^{\mathrm{soup}}_\alpha+\beta^2\E_x\Var_{Y\sim p_{\mathrm{softmax}}(\beta f(x;\theta_\alpha))}[\Delta f_Y(x)]\bigr]$ (their Eq.~1, Sec.~4): for logit-linear models $f^{\mathrm{ens}}=f^{\mathrm{soup}}$, both sides are $0$ and (\ref{eq:curv}) shows the two bracket terms cancel. The notebook confirms the soup loss equals the logit-ensemble loss to machine precision in that regime. Lin et al.\ argue, for RLHF, that averaging improves the alignment-forgetting Pareto front by increasing feature diversity on layers whose feature spaces overlap \cite{aligntax}; WiSE-FT interpolates zero-shot and fine-tuned weights to trade in-distribution accuracy against robustness \cite{wiseft}.

\section{Boxed result}
\begin{result}
\begin{align*}
&z(\theta_\alpha)=(1-\alpha)z_0+\alpha z_1+\tfrac12\alpha(\alpha-1)\,d^\top H_zd+O(\|d\|^3),\qquad
\pi_\alpha\propto\pi_0^{1-\alpha}\pi_1^{\alpha}\ \ (\text{per token, logit-linear}),\\
&\frac{d}{d\alpha}\E_\alpha[r]=\Cov_\alpha(g,r),\qquad
\frac{d}{d\alpha}\pk=\E_x\bigl[k(1-p_x)^{k-1}\Cov_\alpha(g_x,r_x)\bigr]\overset{\text{log-odds lin.}}{=}\E_x\bigl[kp_x(1-p_x)^k\Delta_x\bigr],\\
&\ell(\theta_\alpha)-\text{chord}=\tfrac12\alpha(\alpha-1)\,d^\top Hd,\\
&\alpha^\star_k\ \text{interior}\ \Leftarrow\ \tfrac{d}{d\alpha}\pk\big|_{0}>0\ \text{and}\ \tfrac{d}{d\alpha}\pk\big|_{1}<0\quad(\text{sufficient; iff if }\pk(\alpha)\text{ is unimodal}).
\end{align*}
\small (Covariances are those of the sequence score $\tilde g_\alpha$ for autoregressive policies; $\tilde g_\alpha=g-\E_\alpha g$ for a single categorical.)
\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
Large update $\|d\|$ (long RL, high LR) & small second-order term in (\ref{eq:err}) & Logit interpolation no longer predicts $z(\theta_\alpha)$; barrier $\propto\kappa$ \\
Models from different pretraining or independent inits & same basin (linear mode connectivity) & Barrier is large ($+3.19$ in the toy): the average is worse than both parents \\
High learning-rate fine-tunes & connectivity of fine-tunes & Error barrier between fine-tuned models reported at high LR \cite{soups} \\
Averaging LoRA factors separately & averaging is linear in the \emph{update} $\Delta W=BA$ & $\bar B\bar A\ne\overline{BA}$; relative error $0.62$ for random rank-4 factors; merge then scale $\Delta W$ instead \\
Interpolating in bf16 & update $\gg$ bf16 resolution (unit roundoff $2^{-8}$, spacing $2^{-7}$ to $2^{-8}$ of the weight) & Fine-tuning delta $10^{-3}\times$ weight scale: $69\%$ of weights are \emph{bit-identical} to the base in bf16; compute in fp32 and cast once \\
No interior optimum & $\tfrac{d}{d\alpha}\pk|_0>0$ & If also $\pk(\alpha)$ is unimodal the base is the pass@$k$ optimum (14/30 and 3/30 families have an interior one for $k=8,64$); averaging then only trades pass@1 for pass@$k$ monotonically \\
Choosing $\alpha$ on the reporting set & $\alpha$ selected on held-out data & Selection bias on the benchmark; select on a dev split (note 0208) \\
Convention slip ($\alpha$ vs $1-\alpha$) & single definition & The two sections of the source report opposite conventions; state $\alpha$'s meaning in every table \\ \bottomrule
\end{tabular}
\end{breaks}

\nextnote{Averaging repairs diversity after the fact. Note 0206 looks at the other lever, the training data: which problems carry gradient at all. With binary rewards the group-relative advantage is nonzero only when a problem is neither always nor never solved; we derive the signal as a function of the solve rate $p$, show it peaks at $p=\tfrac12$ and vanishes at $0$ and $1$, and connect it to the pass-rate band used for dynamic sampling.}

\begin{thebibliography}{9}
\bibitem{gps} Y.~Lin et al. Goedel-Prover-V2. arXiv:2508.03613, Sec.~2.3 (``Model averaging for enhanced diversity''; $(1-\alpha)\theta_0+\alpha\theta$), Sec.~2.4 (averaging after SFT and after RL), Sec.~3.6 and Fig.~8 ($\alpha\in\{0.6,\dots,0.9\}$, pass@1 vs pass@$N$).
\bibitem{soups} M.~Wortsman et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. arXiv:2203.05482, Sec.~4 (Eq.~1: soup vs.\ ensemble), Sec.~3.3 takeaway~2 and App.~J.1/Fig.~J.1 (error barriers between fine-tunes at high learning rates).
\bibitem{wiseft} M.~Wortsman et al. Robust fine-tuning of zero-shot models. arXiv:2109.01903, Eq.~(1) (WiSE-FT: weight-space ensemble $f(x,(1-\alpha)\theta_0+\alpha\theta_1)$ of zero-shot and fine-tuned weights).
\bibitem{lmc} J.~Frankle et al. Linear Mode Connectivity and the Lottery Ticket Hypothesis. arXiv:1912.05671, Sec.~3 (instability analysis of unpruned networks; networks become stable to SGD noise early in training).
\bibitem{aligntax} Y.~Lin et al. Mitigating the Alignment Tax of RLHF. arXiv:2309.06256, Sec.~5 (Proposition~5.1, proof in App.~F.2; feature diversity on layers whose feature spaces overlap; model averaging).
\bibitem{leash} F.~Wu et al. The Invisible Leash. arXiv:2507.14843 (support shrinkage under RLVR).
\end{thebibliography}
\end{document}