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