Module 27: Model Merging

Introduction

Every module so far has made a model — trained one (m07), aligned one (m12), or adapted one cheaply (m25). This module does something stranger and almost free: it takes several finished models and combines them into one by arithmetic in weight space — no gradient step, no data, no change to inference cost.

Model merging treats a network’s weights as a point in a very high-dimensional space and its fine-tuning as a direction. Once you see weights that way, you can average two models, add a skill, subtract a behavior, or fold a dozen specialists into a single generalist — all by editing numbers, never by training.

Why it matters for LLMs:

  • It is how the open-weight world actually ships. Most models near the top of community leaderboards are merges — a base plus a handful of fine-tunes, combined with a tool like mergekit. Merging is the cheapest known way to compose capabilities.
  • A fine-tune is reusable, not disposable. The difference theta_ft − theta_base is a portable task vector you can move onto another checkpoint, scale up or down, or negate to unlearn.
  • Zero marginal cost. The merged model has the same architecture and size as its parents. Unlike an ensemble (run N models, average outputs), a merge runs once.

What You’ll Learn

After this module, you can:

  • Explain why the weights of same-init fine-tunes can be averaged — the model soup — and why that is not obvious.
  • Compute a task vector tau = theta_ft − theta_base and use task arithmetic to add skills and negate behaviors.
  • Diagnose the interference (sign disagreement + redundancy) that makes a naive sum of many task vectors degrade.
  • Build TIES-Merging from scratch — trim, elect sign, disjoint merge — and DARE (drop-and-rescale), and explain why dropping 90% of a delta can be lossless.
  • Merge two models with SLERP — interpolate along the arc of the hypersphere instead of the straight chord — and explain why that preserves the weights’ magnitude where a naive average shrinks it.
  • Run SLERP on real transformer weights — two shared-init specialists — and read off the measured per-tensor angles, the held global norm, and the SLERP-vs-soup loss gap.
  • Merge by certainty with Fisher-weighted averaging: weight each parameter by each model’s diagonal Fisher — its mean squared log-likelihood gradient — giving the precision-weighted mean that is exactly the soup when every model is equally sure.
  • Merge by behavior with RegMean: solve each linear layer in closed form so its outputs reproduce every model’s outputs on that model’s own data — pulling decisively ahead of the soup wherever a layer’s inputs are correlated.
  • Run RegMean on real transformer weights — record each linear layer’s input Gram from a forward pass — and sweep the paper’s α shrink, reading why 1.0 can degenerate and 0.0 nears plain averaging.

Prerequisites

This module requires familiarity with:

  • Module 01: Tensors — weights are just tensors; here we do arithmetic on whole state dicts.
  • Module 07: Training — fine-tuning is the move in weight space we are about to name and reuse.
  • Module 25: PEFT — LoRA already merged an adapter back into a weight; this module merges whole models the same way.

Intuition: Weights Are a Place

Picture the millions of numbers in a model as the coordinates of a single point. Training moves that point. When you fine-tune a shared base model \theta_{base} on task A, you end at \theta_A; the arrow between them,

\tau_A = \theta_A - \theta_{base},

is the task vector — everything the model changed to learn A. It points from “generic” toward “good at A.”

Two facts make this picture powerful. First, arrows add: if \tau_A points toward skill A and \tau_B toward skill B, then \theta_{base} + \tau_A + \tau_B often reaches a model good at both. Second, arrows have a sign: -\tau_A walks away from A, which turns out to be a clean way to make a model forget a behavior. The stepper below walks these operations on a 2-D toy — the same algebra runs unchanged in a billion dimensions.

NoteKey Insight

A trained model is a point; a fine-tune is a vector. Merging is linear algebra on those points and vectors — which is why it costs no training and no extra inference: you are only ever adding and scaling weights you already have.

Model Soups: Average the Weights Themselves

The simplest merge ignores task vectors entirely and averages the weights directly. Take several models fine-tuned from the same pretrained checkpoint — say, with different learning rates or seeds — and set every parameter to its mean across them. Wortsman et al. (2022) called this a model soup, and the surprise is that it works: the uniform soup often beats every individual member on both accuracy and robustness.

Why should averaging weights be legal at all? Two randomly initialized networks average into garbage. The trick is shared initialization: models fine-tuned from one base stay in the same loss basin, close enough that the straight line between them stays low-loss. The mean sits near the bottom of that basin — a flat, well-generalizing spot.

import torch
from merging import model_soup, task_vector, apply_task_vector

# A tiny "state dict": parameter name -> tensor. The algorithms are model-agnostic,
# so what works here works on a real GPTModel.state_dict().
base = {"w": torch.zeros(4)}
model_a = {"w": torch.tensor([2.0, 2.0, 0.0, 0.0])}   # fine-tuned toward skill A
model_b = {"w": torch.tensor([0.0, 0.0, 2.0, 2.0])}   # fine-tuned toward skill B

soup = model_soup([model_a, model_b])
print("uniform soup:", soup["w"].tolist())            # the elementwise mean

# A soup of one model is just that model; weights need not sum to 1 (they're normalized).
print("weighted 3:1:", model_soup([model_a, model_b], weights=[3.0, 1.0])["w"].tolist())
uniform soup: [1.0, 1.0, 1.0, 1.0]
weighted 3:1: [1.5, 1.5, 0.5, 0.5]
WarningSoups need a shared ancestor

Averaging weights only makes sense for models that started from the same initialization (a common base checkpoint). Merge two models trained from scratch with different inits and you land between two unrelated basins — usually far worse than either. Every technique in this module assumes a shared base.

The Math: Task Vectors

Model soups throw away a useful distinction: what each fine-tune changed. Task arithmetic keeps it. Define the task vector as the delta a fine-tune applied:

\tau_i = \theta_i - \theta_{base}.

Then editing a model is vector algebra. The merged model is the base plus a scaled sum of task vectors:

\theta_{merged} = \theta_{base} + \lambda \sum_i \tau_i .

  • Adding (+\tau_A + \tau_B) composes skills into one multi-task model.
  • Negating (-\tau_A) removes a behavior — the unlearning / detox operator.
  • Scaling by \lambda dials the strength; with several vectors you usually need \lambda < 1 so their sum does not overshoot.

The bookkeeping identity that makes this exact: applying the task vector back to the base with \lambda = 1 reconstructs the fine-tune. There is nothing to approximate — merging is defined by these equations.

Code: Task Arithmetic From Scratch

The whole toolkit lives in merging.py. Every function takes and returns a state dict, so it composes freely. Here are the three moves and the identities the tests pin:

from merging import negate, merge

# --- Round-trip identity: base + (θ_ft − θ_base) == θ_ft ---
tau_a = task_vector(base, model_a)
recon = apply_task_vector(base, tau_a, 1.0)
print("round-trip matches θ_A:", torch.allclose(recon["w"], model_a["w"]))

# --- Negate to forget: base + (−τ_A) walks away from skill A ---
forget_a = apply_task_vector(base, negate(tau_a), 1.0)
print("forget A:", forget_a["w"].tolist())   # negative on A's coordinates

# --- Add to compose: base + τ_A + τ_B is good at both ---
tau_b = task_vector(base, model_b)
both = merge(base, [tau_a, tau_b], coeff=1.0)
print("compose A+B:", both["w"].tolist())     # all four coordinates active
round-trip matches θ_A: True
forget A: [-2.0, -2.0, 0.0, 0.0]
compose A+B: [2.0, 2.0, 2.0, 2.0]

Notice merge reached [2, 2, 2, 2] — full strength on both skills — because A and B touched disjoint coordinates. That is the easy case. The interesting question is what happens when task vectors fight over the same parameter.

The Interference Problem

Real task vectors are dense and they overlap. When you sum many of them, two things go wrong (Yadav et al., 2023):

  1. Sign conflict. For a given parameter, task A may want it +0.9 while task B wants it -0.1. Their sum, +0.8, is dragged down by a disagreement that was never important to B — the strong, correct update gets diluted by weak noise pointing the other way.
  2. Redundant interference. Most of a task vector’s entries are small and irrelevant. Summed across many tasks, this low-magnitude bulk piles up into a haze that drowns the few coordinates that actually carry a skill.

A plain sum treats all of this as signal. The fix is to be selective: keep only the entries that matter, and when they disagree, let the strong side win.

TIES-Merging

TIES (Trim, Interference resolution via Elect Sign, disjoint merge) resolves both problems in three steps, applied to the task vectors before they are summed:

  1. Trim — per task vector, keep only the top-density fraction of entries by magnitude (the paper uses 20%); zero the rest. This deletes the redundant haze.
  2. Elect sign — for each parameter, sum the (trimmed) values across tasks and take the sign of that sum. This is a magnitude-weighted majority vote: the direction the tasks most strongly agree on.
  3. Disjoint merge — average only the values whose sign matches the elected sign, dividing by the count of agreeing tasks. The dissenters are dropped, so a strong correct update is no longer diluted.
from merging import trim, elect_sign, disjoint_merge, ties_merge

# Two fine-tunes with a planted conflict on coordinate 2: A wants +3 (strong),
# B wants −1 (weak). Coordinates 0 and 1 are each task's own disjoint skill.
base3 = {"w": torch.zeros(3)}
a = {"w": torch.tensor([1.0, 0.0, 3.0])}
b = {"w": torch.tensor([0.0, 1.0, -1.0])}

tvs = [task_vector(base3, a), task_vector(base3, b)]
elected = elect_sign(tvs)
print("elected sign:", elected["w"].tolist())         # coord 2: 3+(−1)=+2 -> +1

ties = ties_merge(base3, [a, b], density=1.0, coeff=1.0)
naive = merge(base3, tvs, coeff=1.0)
print("TIES  coord 2:", float(ties["w"][2]))          # 3.0 — keeps the strong vote
print("naive coord 2:", float(naive["w"][2]))         # 2.0 — the −1 cancels part of it
elected sign: [1.0, 1.0, 1.0]
TIES  coord 2: 3.0
naive coord 2: 2.0

TIES kept the full +3 on the contested coordinate; the naive sum let the weak -1 eat a third of it. On a real merge of many models, that difference between “strong side wins” and “everything averages toward mush” is the difference between a merge that works and one that degrades.

Interactive: Watch Sign Election Resolve a Conflict

Drag the trim density to see the low-magnitude entries vanish first, then watch each parameter’s column resolve to its elected sign and disjoint-merge only the agreeing cells. Green = positive, red = negative, faded = trimmed or dropped.

TipTry This
  1. Turn density down to 0.25. Each task vector keeps only its single largest entry — the redundant middle columns fade out, and the merged row collapses to the few coordinates that carry real signal.
  2. Watch param 2. Two task vectors want it negative and one wants it positive; the elected sign is negative, so the lone positive cell is dropped from the average rather than cancelling the two that agree.

DARE: Drop Almost Everything, Lose Nothing

TIES trims by magnitude. DARE (Yu et al., 2023) shows you can trim at random and still be fine — because a task vector is astonishingly redundant. Drop And REscale zeros each delta entry with probability p, then multiplies the survivors by \frac{1}{1-p}:

\text{DARE}(\tau)_j = \begin{cases} \dfrac{\tau_j}{1-p} & \text{with probability } 1-p,\\[4pt] 0 & \text{with probability } p. \end{cases}

The rescale is the whole trick: it makes DARE expectation-preserving, \mathbb{E}[\text{DARE}(\tau)] = \tau entrywise. So on average the dropped vector equals the original — and empirically you can drop 90–99% of a fine-tune’s deltas and merge with no measurable loss. Applied before TIES (DARE-TIES), it lets you stack many models with far less interference.

from merging import dare

tau = {"w": torch.arange(1.0, 6.0)}          # the delta to sparsify: [1,2,3,4,5]

# p = 0 is the identity; survivors at p = 0.6 are rescaled by 1/(1−0.6) = 2.5
g = torch.Generator().manual_seed(0)
print("one draw (p=0.6):", dare(tau, 0.6, g)["w"].tolist())

# Average many independent draws -> back to the original (expectation preserved)
acc = torch.zeros(5)
N = 4000
for i in range(N):
    acc += dare(tau, 0.9, torch.Generator().manual_seed(i))["w"]
print("mean of 4000 draws at p=0.9:", [round(x, 2) for x in (acc / N).tolist()])
print("original:", tau["w"].tolist())
one draw (p=0.6): [0.0, 0.0, 7.5, 10.0, 12.5]
mean of 4000 draws at p=0.9: [1.02, 2.09, 2.87, 3.88, 5.6]
original: [1.0, 2.0, 3.0, 4.0, 5.0]

Even at p = 0.9 — throwing away nine of every ten delta entries per draw — the average lands back on [1, 2, 3, 4, 5]. The individual draws are sparse; their expectation is exact.

SLERP: Walk the Arc, Not the Chord

Soups and task arithmetic add two weight vectors and divide — a straight-line blend. That is fine when you are pooling many fine-tunes, but for combining exactly two strong models it has a quiet flaw. A model’s weights sit at some radius from the origin; that radius is the scale of its features. Average two vectors of equal length and the midpoint is shorter than either — the straight chord cuts through the interior of the sphere. So a two-model average silently turns the weights down, diluting the very directions each parent learned.

SLERP — spherical linear interpolation — fixes this by refusing to leave the sphere. Instead of the chord, it walks the great-circle arc from one model to the other, so every point on the path keeps the same radius. You get a blend that moves smoothly in direction from model a to model b without ever shrinking their magnitude. It is mergekit’s default recipe for merging two models, and it comes straight from computer graphics — Shoemake coined “slerp” in 1985 to interpolate camera rotations.

NoteKey Insight

LERP interpolates the numbers; SLERP interpolates the geometry. On the chord, the midpoint of two equal-length vectors has smaller norm — the weights get quieter. On the arc, the norm is held fixed and only the angle sweeps from a to b. When magnitude carries meaning (and in weight space it does), stay on the arc.

The Math: Bend the Weights to Stay on the Arc

Let \Omega be the angle between the two weight vectors, read off their normalized dot product. SLERP mixes them with two sine-ratio weights instead of (1-t) and t:

\operatorname{slerp}(a, b; t) = \frac{\sin\big((1-t)\,\Omega\big)}{\sin \Omega}\, a \;+\; \frac{\sin\big(t\,\Omega\big)}{\sin \Omega}\, b, \qquad \cos \Omega = \frac{a \cdot b}{\lVert a \rVert\, \lVert b \rVert}.

Three properties make this the right curve:

  • Endpoints are exact. At t=0 the weights are \tfrac{\sin\Omega}{\sin\Omega}=1 and 0, giving a; at t=1 they give b.
  • It reduces to the straight blend when the arc is flat. As \Omega \to 0, \sin(x\Omega)/\sin\Omega \to x, so slerp becomes (1-t)a + t b — plain LERP. This is not just a limit; it is the escape hatch. When a and b are nearly colinear, \sin\Omega \approx 0 and the division blows up — so the code detects that case and falls back to LERP, which is exactly what the arc has become.
  • It preserves the radius when \lVert a \rVert = \lVert b \rVert. The two sine weights are tuned so the point stays on the sphere: for equal-norm inputs, \lVert \operatorname{slerp}(a,b;t)\rVert is constant for every t. (For unequal norms it drifts between \lVert a\rVert and \lVert b\rVert — an honest caveat we return to below.)

Step through the geometry — two equal-length models, the chord that shrinks between them, and the arc that does not:

Code: SLERP from Scratch

The implementation in merging.py is a direct transcription of the formula, with the colinear fallback baked in. It reads the angle off the flattened tensors, then blends the originals with the two sine weights:

import torch
from merging import slerp

a = {"w": torch.tensor([1.0, 0.0])}
b = {"w": torch.tensor([0.0, 1.0])}          # orthogonal: Omega = 90 degrees

# Endpoints are exact.
print("t=0 ->", slerp(a, b, 0.0)["w"].tolist())
print("t=1 ->", slerp(a, b, 1.0)["w"].tolist())

# The midpoint rides the arc: norm 1, the normalized bisector of a and b.
mid = slerp(a, b, 0.5)["w"]
lerp_mid = 0.5 * a["w"] + 0.5 * b["w"]
print(f"\nSLERP mid {mid.tolist()}  norm={mid.norm():.4f}")
print(f"LERP  mid {lerp_mid.tolist()}  norm={lerp_mid.norm():.4f}   <- shrank to 1/sqrt(2)")
t=0 -> [1.0, 0.0]
t=1 -> [0.0, 1.0]

SLERP mid [0.7071067690849304, 0.7071067690849304]  norm=1.0000
LERP  mid [0.5, 0.5]  norm=0.7071   <- shrank to 1/sqrt(2)

The straight blend collapses to norm 1/\sqrt2 \approx 0.707; the arc holds norm 1. The norm-preservation is not special to unit vectors — it holds at any shared radius. Watch it stay flat along the whole path for two length-2 vectors:

a2 = {"w": torch.tensor([2.0, 0.0])}
b2 = {"w": torch.tensor([0.0, 2.0])}         # both norm 2, orthogonal

for t in (0.0, 0.25, 0.5, 0.75, 1.0):
    s = slerp(a2, b2, t)["w"]
    l = (1 - t) * a2["w"] + t * b2["w"]
    print(f"t={t:.2f}   |slerp|={s.norm():.4f}   |lerp|={l.norm():.4f}")
t=0.00   |slerp|=2.0000   |lerp|=2.0000
t=0.25   |slerp|=2.0000   |lerp|=1.5811
t=0.50   |slerp|=2.0000   |lerp|=1.4142
t=0.75   |slerp|=2.0000   |lerp|=1.5811
t=1.00   |slerp|=2.0000   |lerp|=2.0000

|slerp| stays pinned at 2.0; |lerp| sags to 1.41 in the middle. And when the two models are nearly colinear — no meaningful angle to sweep — SLERP detects the vanishing \sin\Omega and returns the plain linear blend instead of dividing by (almost) zero:

same_dir_a = {"w": torch.tensor([2.0, 0.0])}
same_dir_b = {"w": torch.tensor([4.0, 0.0])}   # same direction
print("colinear midpoint ->", slerp(same_dir_a, same_dir_b, 0.5)["w"].tolist(),
      "(= the LERP midpoint 3.0, via the fallback)")
colinear midpoint -> [3.0, 0.0] (= the LERP midpoint 3.0, via the fallback)

Below, the norm-vs-t curves are computed by the actual slerp() — 41 values of t from the library, not a redrawing — so the widget you drive next is provably the same math:

Drive It: The Arc and the Chord

Set the angle \Omega between the two models and drag t from a to b. The orange dot is the straight blend (LERP); the highlighted dot rides the arc (SLERP). Watch the two norms in the readout: the arc holds its length while the chord dips — and the gap grows as the models get farther apart.

TipTry This
  1. Push Ω toward 170°. The two models point almost opposite ways; the LERP midpoint collapses toward the origin (a near-zero, near-useless blend) while the SLERP dot still traces a clean arc at full length.
  2. Pull Ω down to 10°. The arc and the chord nearly coincide and the shrink falls to almost nothing — exactly the regime where slerp() gives up and returns LERP, because there is no meaningful angle left to preserve.
  3. Hold Ω at 90° and sweep t. ‖LERP‖ traces the sagging curve from the chart above, bottoming at 0.71; ‖SLERP‖ never budges from 1.
WarningSLERP is a two-model move, and only norm-safe at equal scale

SLERP interpolates exactly two models — there is no natural great circle through three points, so mergekit applies it pairwise (or layer-by-layer) and uses soups/TIES/DARE for many-model merges. And its headline property, holding the norm, is exact only when \lVert a\rVert = \lVert b\rVert; for models at different scales the arc’s norm drifts between the two. As always, both models must be fine-tunes of a shared base — SLERP fixes the geometry of the blend, not the requirement that the endpoints live in the same basin.

SLERP Through a Real Model

Every curve so far lived in a cartoon: two arrows on a unit circle. The warning above ends where the honest question starts — does the arc actually hold its scale on real transformer weights, where twenty-odd tensors sit at twenty-odd angles and no two norms match exactly? This section runs the same slerp() you just drove — no new code path, the actual merging.slerp — on two real (tiny) specialists.

The setup is deliberately small so it trains in seconds on your CPU. Take one seeded 1-layer m06 GPTModel as the shared base, clone it twice, and memorize one 8-token sentence into each clone: specialist A learns 0 1 2 3 4 5 6 7, specialist B learns 0 8 9 10 11 12 13 14. Same init, disjoint skills — the shared-basin precondition, shrunk to a toy. Say the caveat out loud before the numbers: two memorized sentences in a 1-layer toy prove the mechanism (the arc holds scale on real weights), not that SLERP wins real leaderboards. The anecdote box below keeps that line sharp.

The Math: What the Real Weights Measure

First, what the demo’s own angle report reads (seed 0, in slerp_real.py): median Ω ≈ 68° across the state dict — but the median hides a split. The big shared matrices barely moved apart (attention projections 8–11°); the small zero-init biases swung wide (up to ≈ 116°) — any task-specific nudge is a large angle when you start at the origin. Norms tell the cleaner story: per-tensor ratios average ≈ 1.1, while the global scale is essentially identical (20.00 vs 20.01, ratio 0.9997) — exactly the equal-scale regime where the arc is supposed to hold. And it does: the SLERP midpoint’s global norm reads ≈ 20.01 against the soup midpoint’s dipped ≈ 19.77.

That flat-vs-dipped pair is not luck; it is the norm theorem, in six lines. Take unit vectors a, b with a \cdot b = \cos\Omega:

  1. The chord point is m = (1-t)a + tb, so \lVert m\rVert^2 = (1-t)^2 + t^2 + 2t(1-t)\cos\Omega.
  2. Since (1-t)^2 + t^2 = 1 - 2t(1-t), that is \lVert m\rVert^2 = 1 - 2t(1-t)(1-\cos\Omega) — under 1 for every t \in (0,1), \Omega > 0.
  3. At t = \tfrac12: \lVert m\rVert^2 = (1+\cos\Omega)/2 = \cos^2(\Omega/2) — the midpoint shrinks to \cos(\Omega/2).
  4. Now the arc: s = \alpha a + \beta b with \alpha = \sin((1-t)\Omega)/\sin\Omega and \beta = \sin(t\Omega)/\sin\Omega.
  5. \lVert s\rVert^2 = \alpha^2 + \beta^2 + 2\alpha\beta\cos\Omega — put everything over \sin^2\Omega.
  6. The numerator is \sin^2 A + \sin^2 B + 2\sin A\sin B\cos(A+B) = \sin^2(A+B) = \sin^2\Omega (with A = (1-t)\Omega, B = t\Omega), so \lVert s\rVert^2 = 1: unit in, unit out, for every t.

One production nuance, in one sentence: mergekit’s slerp.py reads the angle off the normalized directions but applies the sine weights to the original-magnitude tensors, and it never flips a negatively-aligned pair — so unlike graphics SLERP there is no shortest-path guarantee. And the attribution, kept straight: SLERP-the-merge-tool is not an academic-paper invention — the community recipe started with Digitous/CalderaAI’s LLM-SLERP-Merge, which Goddard et al. productized in Arcee’s MergeKit; the per-layer application with per-layer t schedules is MergeKit convention, not Akiba et al. (who evolve TIES+DARE) — so credit the community and the toolkit, not that paper.

NoteThe SLERP-vs-linear win is anecdotal

Say it plainly: no controlled SLERP-vs-linear study exists. What exists is lore plus one near-tie. MergeKit’s Table 1 (Meditron-7B × Llama2-7B-Chat) has SLERP ahead on USMLE (39.20/36.91 vs 39.10/36.65) and behind on MMLU (47.97 vs 48.44) — a wash. The famous datapoint, Marcoro14-7B-slerp, topped the Open LLM Leaderboard on 02/01/24, then lost the crown to a DARE-TIES merge. This section reports only what its own toy measures: at the midpoint, the SLERP blend scores ≈ 1.50 mean next-token CE against the soup’s ≈ 1.68 (per task, 1.40/1.61 vs 1.59/1.77) — a ~0.18-nat edge for the arc, on one configuration, with the seed in the open. A mechanism demo, not a benchmark.

Code: Two Specialists, One Sweep

Everything below imports from slerp_real.py — the training loop, the sweep, and the angle report all live there, tested. Train once, then reuse the states:

from slerp_real import train_specialists, merge_angle_report

_state_a, _state_b = train_specialists(seed=0)
_report = merge_angle_report(_state_a, _state_b)
print(f"median angle between specialists: {_report['median_omega_deg']:.1f} deg")
print(f"global norms: {_report['global_norm_a']:.2f} vs {_report['global_norm_b']:.2f}")
print(f"q_proj angle: {_report['omegas']['blocks.0.attention.mha.q_proj.weight']:.1f} deg")
median angle between specialists: 67.8 deg
global norms: 20.00 vs 20.01
q_proj angle: 8.0 deg

Both specialists nail their own sentence (CE ≈ 0.07) and fumble the other’s (CE ≈ 2.3+) — real (if tiny) specialists. Now sweep the blends, loading each one into a fresh model and scoring both tasks for real:

from slerp_real import sweep_merges

_sweep5 = sweep_merges(_state_a, _state_b, [0.0, 0.25, 0.5, 0.75, 1.0])
print(f"{'t':>5} {'SLERP-A':>8} {'SLERP-B':>8} {'chord-A':>8} {'chord-B':>8}")
for t, sa, sb, qa, qb in zip(_sweep5["ts"], _sweep5["slerp_a"], _sweep5["slerp_b"], _sweep5["soup_a"], _sweep5["soup_b"]):
    print(f"{t:5.2f} {sa:8.4f} {sb:8.4f} {qa:8.4f} {qb:8.4f}")
    t  SLERP-A  SLERP-B  chord-A  chord-B
 0.00   0.0668   3.2366   0.0668   3.2366
 0.25   0.3404   2.5609   0.4341   2.5530
 0.50   1.3979   1.6117   1.5874   1.7732
 0.75   2.6488   0.4302   2.6229   0.5548
 1.00   3.6101   0.0716   3.6101   0.0716

Read the middle row: at t = 0.5 the arc beats the chord on both tasks, not just the mean. The widget below draws the full 21-point version of this table — computed from the same states by a fresh sweep, bridged into OJS:

Drive It: Blend Two Real Specialists

Drag t from specialist A to specialist B. The top panel scores every blend on both sentences — hue is the task (orange: A, blue: B), solid is the arc (SLERP), dashed is the chord (linear blend — the uniform soup is its t=0.5 point). The bottom strip is the global weight norm: the arc holds it flat while the chord sags — lines 6 vs 4 of the theorem, on real weights.

```{ojs}
#| echo: false

viewof realT = Inputs.range([0, 1], {value: 0.5, step: 0.01, label: "t (A → B)"})
```
ImportantOJS Syntax Error (line 1312, column 1)Unexpected character ‘#’ 
```{ojs}
#| echo: false

realDrive = {
  const width = 720, height = 520;
  const theme = diagramTheme;
  const svg = d3.create("svg")
    .attr("viewBox", `0 0 ${width} ${height}`)
    .attr("width", "100%").attr("height", height)
    .style("max-width", `${width}px`)
    .style("font-family", "var(--pg-mono)");
  svg.append("rect").attr("width", width).attr("height", height)
    .attr("fill", theme.bg).attr("rx", 12);

  const t = realT;
  const n = realTs.length;
  const lo = (arr) => Math.min(...arr);
  const hi = (arr) => Math.max(...arr);
  const yMin = Math.min(lo(realSlerpA), lo(realSlerpB), lo(realSoupA), lo(realSoupB));
  const yMax = Math.max(hi(realSlerpA), hi(realSlerpB), hi(realSoupA), hi(realSoupB));
  const pad = (yMax - yMin) * 0.12;

  const top = {t: 56, b: 252, l: 56, r: 24};
  const bot = {t: 330, b: 452, l: 56, r: 24};
  const x = d3.scaleLinear([0, 1], [top.l, width - top.r]);
  const y = d3.scaleLinear([Math.max(0, yMin - pad), yMax + pad], [top.b, top.t]);
  const nMin = Math.min(lo(realSlerpNorms), lo(realSoupNorms));
  const nMax = Math.max(hi(realSlerpNorms), hi(realSoupNorms));
  const npad = Math.max((nMax - nMin) * 0.3, 1e-6);
  const yn = d3.scaleLinear([nMin - npad, nMax + npad], [bot.b, bot.t]);

  const grid = (x1, y1, x2, y2) => svg.append("line")
    .attr("x1", x1).attr("y1", y1).attr("x2", x2).attr("y2", y2)
    .attr("stroke", theme.edgeStroke).attr("stroke-width", 1).attr("opacity", 0.35);
  grid(top.l, top.b, width - top.r, top.b);
  grid(top.l, top.t, top.l, top.b);
  grid(bot.l, bot.b, width - bot.r, bot.b);
  grid(bot.l, bot.t, bot.l, bot.b);

  const xticks = [0, 0.25, 0.5, 0.75, 1];
  xticks.forEach((v) => {
    svg.append("text").attr("x", x(v)).attr("y", top.b + 18)
      .attr("text-anchor", "middle").attr("fill", theme.nodeText).attr("font-size", 11).text(v.toFixed(2));
    svg.append("text").attr("x", x(v)).attr("y", bot.b + 18)
      .attr("text-anchor", "middle").attr("fill", theme.nodeText).attr("font-size", 11).text(v.toFixed(2));
  });
  svg.append("text").attr("x", (top.l + width - top.r) / 2).attr("y", top.b + 34)
    .attr("text-anchor", "middle").attr("fill", theme.nodeText).attr("font-size", 12).text("t");

  const series = [
    {vals: realSlerpA, col: theme.highlight, dash: null, key: "SLERP A"},
    {vals: realSoupA, col: theme.highlight, dash: "5 4", key: "soup A"},
    {vals: realSlerpB, col: theme.accent, dash: null, key: "SLERP B"},
    {vals: realSoupB, col: theme.accent, dash: "5 4", key: "soup B"},
  ];
  const line = d3.line().x((d, i) => x(realTs[i])).y((d) => y(d));
  series.forEach((s) => {
    svg.append("path").attr("d", line(s.vals)).attr("fill", "none")
      .attr("stroke", s.col).attr("stroke-width", 2.5).attr("stroke-dasharray", s.dash);
  });
  const nline = d3.line().x((d, i) => x(realTs[i])).y((d) => yn(d));
  svg.append("path").attr("d", nline(realSoupNorms)).attr("fill", "none")
    .attr("stroke", theme.accent).attr("stroke-width", 2.5).attr("stroke-dasharray", "5 4");
  svg.append("path").attr("d", nline(realSlerpNorms)).attr("fill", "none")
    .attr("stroke", theme.highlight).attr("stroke-width", 3);

  let best = 0;
  for (let i = 1; i < n; i++) {
    if (Math.abs(realTs[i] - t) < Math.abs(realTs[best] - t)) best = i;
  }
  svg.append("line").attr("x1", x(realTs[best])).attr("y1", top.t).attr("x2", x(realTs[best])).attr("y2", bot.b)
    .attr("stroke", theme.edgeStroke).attr("stroke-width", 1).attr("stroke-dasharray", "3 3");
  series.forEach((s) => {
    svg.append("circle").attr("cx", x(realTs[best])).attr("cy", y(s.vals[best])).attr("r", 4)
      .attr("fill", s.col).attr("stroke", theme.nodeStroke).attr("stroke-width", 1.5);
  });

  const legend = (xx, yy, col, dash, txt) => {
    svg.append("line").attr("x1", xx).attr("y1", yy).attr("x2", xx + 28).attr("y2", yy)
      .attr("stroke", col).attr("stroke-width", 3).attr("stroke-dasharray", dash);
    svg.append("text").attr("x", xx + 34).attr("y", yy + 4).attr("fill", theme.nodeText).attr("font-size", 11.5).text(txt);
  };
  legend(top.l + 8, top.b + 52, theme.highlight, null, "SLERP · A");
  legend(top.l + 140, top.b + 52, theme.highlight, "5 4", "chord · A");
  legend(top.l + 268, top.b + 52, theme.accent, null, "SLERP · B");
  legend(top.l + 400, top.b + 52, theme.accent, "5 4", "chord · B");
  legend(bot.l + 8, bot.b - 12, theme.highlight, null, "arc — norm held");
  legend(bot.l + 190, bot.b - 12, theme.accent, "5 4", "chord — dips");

  svg.append("text").attr("x", width / 2).attr("y", bot.b + 40)
    .attr("text-anchor", "middle").attr("fill", theme.nodeText).attr("font-size", 12)
    .text("t = " + realTs[best].toFixed(2) + " — SLERP A/B: " + realSlerpA[best].toFixed(2) + " / " + realSlerpB[best].toFixed(2) + " · chord A/B: " + realSoupA[best].toFixed(2) + " / " + realSoupB[best].toFixed(2) + " · ‖arc‖ / ‖chord‖: " + realSlerpNorms[best].toFixed(2) + " / " + realSoupNorms[best].toFixed(2));

  svg.append("text").attr("x", width / 2).attr("y", 24).attr("text-anchor", "middle")
    .attr("fill", theme.highlightText).attr("font-size", 15).attr("font-weight", 700)
    .text("Two real specialists: the arc scores lower and holds its scale");

  return svg.node();
}
```
ImportantOJS Syntax Error (line 1318, column 1)Unexpected character ‘#’
TipTry This
  1. Drag t to 0, then to 1. The endpoints recover the specialists exactly — CE ≈ 0.07 on the home sentence, ≈ 2.3+ on the other. The merge never beats a specialist at its own game; it trades.
  2. Park at t = 0.5 and read all four dots. The arc beats the chord on both tasks (1.40/1.61 vs 1.59/1.77), not just the mean — then ask how much of that edge is scale (bottom strip) and how much is direction.
  3. Watch the bottom strip while you sweep. The orange norm never leaves ≈ 20.0; the blue chord sags to ≈ 19.77 mid-blend. That is line 6 vs line 4 of the theorem, drawn from real weights.

Fisher-Weighted Averaging: Let Certainty Decide

Every merge so far blends with a coefficient you pick — a soup weight, a \lambda, a density, an interpolation t. Look again at the humblest of them, the model soup. Averaging two models with weight \tfrac12 each is a bet: that both models are equally trustworthy about every single parameter. They are not. Model A may have pinned one weight to exactly the value its task demands and left another almost untouched; model B may be the mirror image. A flat average throws away that information — it lets A’s don’t-care coordinate outvote B’s this-is-the- answer coordinate on equal terms.

Fisher-weighted averaging (Matena & Raffel, 2022) is the first merge where the data, not you, sets the weight — and it sets a different weight for every parameter. The idea is to average each coordinate in proportion to how certain each model is about it. Certain has a precise meaning: how sharply the model’s loss would rise if you nudged that weight. A weight the training data pinned down hard has a steep, narrow valley around it; a weight the data barely constrained sits in a wide, flat one. Merge toward the sharp, confident models per coordinate and defer to nobody on the coordinates they don’t care about.

NoteKey Insight

A model soup asks “what’s the average weight?” A Fisher merge asks “whose weight should I trust here?” — and answers it one parameter at a time, using each model’s own certainty. Where the models are equally sure, the two questions have the same answer and the Fisher merge is the soup.

The Math: A Product of Posteriors

Treat each trained model not as a single point but as a belief about where the good weights are. A standard move (the Laplace approximation) models that belief as a Gaussian centered at the trained weights \theta^{(i)}, whose precision — inverse variance, how tightly peaked it is — is the Fisher information F^{(i)}:

p_i(\theta) \;\approx\; \mathcal{N}\!\left(\theta \,\middle|\, \theta^{(i)},\; (F^{(i)})^{-1}\right).

A high Fisher on a weight means a narrow Gaussian: “I’m sure it’s right here.” A low Fisher means a broad one: “anywhere around here is fine.” Merging is then asking for the single \theta that best satisfies all the models’ beliefs at once — the point that maximizes the product of their Gaussians:

\theta^\ast \;=\; \arg\max_{\theta}\ \prod_i p_i(\theta)^{\lambda_i}.

For diagonal Gaussians this has a closed form. A product of Gaussians is itself Gaussian, and its mean is the precision-weighted mean of the parts — so, per parameter j,

\theta^\ast_j \;=\; \frac{\sum_i \lambda_i \, F^{(i)}_j \, \theta^{(i)}_j}{\sum_i \lambda_i \, F^{(i)}_j}.

That is the whole method. The optional scalars \lambda_i let you up- or down-weight a whole model; their scale cancels in the ratio, so only their relative sizes matter. And the reduction is immediate: set every F^{(i)}_j to the same constant and it cancels top and bottom, leaving \theta^\ast_j = \tfrac{1}{n}\sum_i \theta^{(i)}_j — the model soup is Fisher merging under the assumption that every model is equally certain about everything. The soup was never wrong; it was just uninformed.

The Fisher Is Just a Squared Gradient

So what is F^{(i)}_j, concretely? For a model p_\theta(y \mid x), the diagonal Fisher of parameter j is the mean squared gradient of the log-likelihood:

F_j \;=\; \frac{1}{N} \sum_{n=1}^{N} \; \mathbb{E}_{\,y \sim p_\theta(y \mid x_n)} \left[ \left( \frac{\partial}{\partial \theta_j} \log p_\theta(y \mid x_n) \right)^{\!2} \right].

Read it as: feed the model your data; at each example ask how much the log-probability of the model’s own predictions would move if you wiggled weight j. A big squared gradient means that weight is doing real work — the output is sensitive to it, so the model is “certain” about its value. A near-zero gradient means the weight is idle. This is exactly the quantity Elastic Weight Consolidation uses to decide which weights are too important to overwrite; here we reuse it to decide which weights to trust in a merge. (The expectation is over y drawn from the model’s own distribution — the “true” Fisher — not the ground-truth labels.)

The gradient is the same object m02’s autograd and m07’s training already compute for you. Nothing new to build — just square it and average.

Code: Fisher Merge from Scratch

Two functions in merging.py. diagonal_fisher measures each model’s certainty by squaring the log-likelihood’s gradient; fisher_merge averages the models by that certainty. First, measure the Fisher on a tiny softmax model so you can see it come straight out of autograd:

import torch
from merging import diagonal_fisher, fisher_merge, model_soup

# A minimal differentiable model: logits = w * x over 3 classes. `log_prob_fn`
# returns log p(y=c | x) for each class, differentiable in the params.
def log_prob_fn(params, x):
    logits = params["w"] * x
    return logits - torch.logsumexp(logits, dim=0)

# Two "specialists" fine-tuned to be confident about different classes.
model_a = {"w": torch.tensor([4.0, 0.0, 0.0])}   # sure about class 0
model_b = {"w": torch.tensor([0.0, 4.0, 0.0])}   # sure about class 1
data = [torch.tensor(1.0)]

fisher_a = diagonal_fisher(log_prob_fn, model_a, data)
fisher_b = diagonal_fisher(log_prob_fn, model_b, data)
print("A's Fisher (importance per weight):", [round(v, 3) for v in fisher_a["w"].tolist()])
print("B's Fisher (importance per weight):", [round(v, 3) for v in fisher_b["w"].tolist()])
A's Fisher (importance per weight): [0.034, 0.017, 0.017]
B's Fisher (importance per weight): [0.017, 0.034, 0.017]

Each model’s Fisher is largest on the coordinate it specialized — its certainty, read off the squared gradient. Now merge by it, and compare to the flat soup:

soup = model_soup([model_a, model_b])
fisher = fisher_merge([model_a, model_b], [fisher_a, fisher_b])

print("soup   :", [round(v, 3) for v in soup["w"].tolist()])
print("fisher :", [round(v, 3) for v in fisher["w"].tolist()])
# The soup halves each specialist's weight; the Fisher merge, deferring to whoever
# is sure, keeps more of both — coord 0 stays closer to A's 4.0, coord 1 to B's.
soup   : [2.0, 2.0, 0.0]
fisher : [2.65, 2.65, 0.0]

And the equal-certainty reduction, made concrete — hand every parameter the same Fisher and fisher_merge returns the soup exactly:

flat = {"w": torch.ones(3)}
tied = fisher_merge([model_a, model_b], [flat, flat])
print("equal Fisher -> soup?", torch.allclose(tied["w"], soup["w"]))
equal Fisher -> soup? True

The full estimator handles the expectation over classes exactly for small label sets (and can single-sample it for large ones); see diagonal_fisher in merging.py.

Drive It: Certainty Bends the Blend

Here is the product-of-Gaussians picture on one weight. Model A believes the weight is near \mu_A; model B near \mu_B; each belief is a Gaussian whose width is its uncertainty. Their normalized product (the merged posterior) peaks at the Fisher mean. Drag the slider to make one model more certain than the other and watch the merged peak slide toward it — while the plain soup mean stays pinned at the midpoint, deaf to who is sure.

TipTry This
  1. Slide to the center (ratio = 0). Equal certainty: the two Gaussians have the same width, the product peaks exactly halfway, and the Fisher mean meets the soup mean — the reduction, live.
  2. Slide hard toward A. A’s Gaussian spikes and narrows; the product’s peak is dragged almost onto \mu_A, while the soup marker never budges from 0.50. That gap is the information the flat average was throwing away.

RegMean: Merge So the Outputs Agree

Fisher handed the data one knob per parameter — a scalar saying how much to trust each model on each weight. RegMean (Jin et al., 2023) hands it the whole layer. Its question is sharper: forget the weights for a moment — what single linear map would reproduce what every model actually does to its own data?

The shift is from weight space to function space. A soup asks “what’s the average weight?” A Fisher merge asks “whose weight do I trust here?” RegMean asks “what map best imitates every model’s outputs, on the inputs each one saw?” — and because a layer is linear, that question has an exact, closed-form answer you solve, not tune.

NoteKey Insight

Averaging weights is a proxy. What you actually want is a merged layer that behaves like each parent on that parent’s data. RegMean optimizes the thing you care about — the outputs — directly, and for a linear layer the optimum is a one-line least-squares solve.

The Math: Least Squares in Weight Space

Take a single linear layer, y = XW, where a batch of inputs X has shape (N, \text{in}) (one example per row) and the weight W is (\text{in}, \text{out}). Model i shipped weight W_i, trained on data whose inputs were X_i. We want one merged W whose outputs match each model’s outputs on each model’s own inputs — so we minimize the total squared output disagreement:

\mathcal{L}(W) \;=\; \sum_i \left\lVert\, X_i W \;-\; X_i W_i \,\right\rVert_F^2 .

This is ordinary least squares in W. Set the gradient to zero and the normal equations appear:

\sum_i X_i^\top X_i \, W \;=\; \sum_i X_i^\top X_i \, W_i .

Define the input Gram matrix G_i = X_i^\top X_i — the (\text{in}, \text{in}) un-centered covariance of the inputs model i saw. Then the merge is just:

\boxed{\,W^\ast \;=\; \Big(\textstyle\sum_i G_i\Big)^{-1} \Big(\textstyle\sum_i G_i W_i\Big)\,}

The raw data is gone: to merge, each model only has to remember its Gram — one (\text{in}, \text{in}) matrix per linear layer, recorded by running a little data through it once. And the reductions are exact and telling:

  • Equal Grams \Rightarrow a soup. If every G_i is the same matrix — in particular if every G_i = I — it cancels top and bottom and W^\ast = \tfrac{1}{K}\sum_i W_i. The soup is RegMean under the assumption that every model saw the same, white input distribution.
  • Diagonal Grams \Rightarrow a row-weighted soup. If each G_i is diagonal, row k of W^\ast is the models’ row-k weights averaged with weight G_i[k,k] — that input feature’s energy. This is the Fisher-flavored special case: trust each model on the input directions it actually exercised.

Why It Beats a Soup: Reproduce the Function

The soup minimizes distance in weight space; RegMean minimizes distance in output space, on the data each model saw. The difference lives entirely in the Gram’s off-diagonal. When a layer’s input features are correlated — as they always are in a real network — a blind average double-counts the shared directions and distorts the map. The \big(\sum_i G_i\big)^{-1} is precisely the term that de-correlates the inputs before combining, whitening away that double-counting. Where the two models exercised different directions of input, the soup smears them together; RegMean keeps each model authoritative on its own.

The paper adds one robustness knob: shrink each Gram’s off-diagonal toward zero before the solve,

\tilde{G}_i \;=\; \alpha\, G_i \;+\; (1-\alpha)\,\operatorname{diag}(G_i),

with \alpha \approx 0.9. It keeps \big(\sum_i \tilde G_i\big) well-conditioned (and is exactly a ridge-style penalty pulling W^\ast toward the individual W_i) without throwing away the correlations that make RegMean better than a diagonal method. At \alpha = 1 it is pure RegMean; at \alpha = 0 it collapses to the diagonal, row-weighted soup above.

Code: RegMean from Scratch

Three functions in merging.py. gram_matrix records a layer’s one statistic; regmean_weight solves the least-squares merge for a single linear layer; and regmean_merge applies it across a whole state dict — RegMean on the linear layers, a plain soup on everything else (biases, norms, embeddings), exactly as the paper prescribes. First, the two exact reductions, made concrete:

import torch
from merging import gram_matrix, regmean_weight, model_soup

# Two specialists: a 2->1 linear layer each. A reads feature 0, B reads feature 1.
w_a = torch.tensor([[1.0], [0.0]])
w_b = torch.tensor([[0.0], [1.0]])

# Identity Grams (a white input distribution) -> RegMean *is* the soup.
eye = torch.eye(2)
merged = regmean_weight([w_a, w_b], [eye, eye])
print("identity Grams -> soup?",
      torch.allclose(merged, model_soup([{"w": w_a}, {"w": w_b}])["w"]))
print("  merged:", merged.squeeze(-1).tolist())
identity Grams -> soup? True
  merged: [0.5, 0.5]

Now give the two models genuinely different, correlated inputs — measure each one’s Gram from real activations — and watch RegMean pull away from the soup:

g = torch.Generator().manual_seed(0)
# A's inputs lie along a slanted direction; B's along a roughly orthogonal one.
d_a = torch.tensor([[1.0, 0.6]]); d_a /= d_a.norm()
d_b = torch.tensor([[-0.6, 1.0]]); d_b /= d_b.norm()
X_a = torch.randn(256, 1, generator=g) * d_a + 0.05 * torch.randn(256, 2, generator=g)
X_b = torch.randn(256, 1, generator=g) * d_b + 0.05 * torch.randn(256, 2, generator=g)

G_a, G_b = gram_matrix(X_a), gram_matrix(X_b)
regmean = regmean_weight([w_a, w_b], [G_a, G_b])
soup = 0.5 * (w_a + w_b)

# Output-reproduction error: how far each merge's outputs drift from each model's,
# on that model's own inputs.  ||X (W - W_i)||^2 = (W - W_i)^T G_i (W - W_i).
def out_err(W):
    da, db = W - w_a, W - w_b
    return float((da.T @ G_a @ da).sum() + (db.T @ G_b @ db).sum())

print(f"soup    output error: {out_err(soup):8.2f}")
print(f"RegMean output error: {out_err(regmean):8.2f}   (the least-squares optimum)")
soup    output error:   132.15
RegMean output error:     1.64   (the least-squares optimum)

RegMean’s error is far lower — and it is provably the minimum, so no reweighting of the soup could ever beat it on this objective. The full state-dict wrapper does this per linear layer and averages the rest; see regmean_merge in merging.py.

Drive It: Correlated Inputs Break the Soup

Here is the whole story on one control. Model A’s inputs carry their energy along a fixed direction; drag the slider to rotate model B’s input direction away from A’s. The bars are each merge’s output-reproduction error, split by whose data it is measured on. When the two models exercised the same input direction the soup is fine — but as their inputs diverge, the soup’s error climbs while RegMean stays pinned at the least-squares floor.

TipTry This
  1. Slide to 0°. A and B now exercise the same input direction, so their Grams agree — RegMean and the soup have the same error. Equal Grams, one merge: the reduction, live.
  2. Slide toward 90°. The two input directions become orthogonal; the soup’s error climbs steeply while RegMean’s barely moves. The growing gap is the off-diagonal information the flat average is blind to — exactly what \big(\sum_i G_i\big)^{-1} recovers.

RegMean Through a Real Model

The toy solved 2×2 Grams by hand; now every linear layer brings its own measured Gram. This section runs the actual regmean_merge you just drove — no new code path — on the same two tiny specialists slerp_real.train_specialists builds: one shared 1-layer m06 GPTModel init, one memorized 8-token sentence per clone. The data decides each weight’s blend, layer by layer. Say the caveat out loud before the numbers: two memorized sentences in a 1-layer toy prove the mechanism (data-weighted blending on real activations), not that RegMean wins real leaderboards. The note below keeps that line sharp.

The Math: What the Real Grams Measure

The normal equations carry over unchanged — and so do the two anchors, now tested on real weights: equal Grams \Rightarrow the soup (to solver tolerance, atol=1e-3 — an LU solve is not bit-exact), and merging a model with itself returns itself. Four things are new because the weights are real.

First, which keys each method touches. RegMean merges only the linear maps: every nn.Linear weight gets its own least-squares solve W^\ast = (\sum_i G_i)^{-1}(\sum_i G_i W_i) from its own measured Gram G_i. Everything else is soup-averaged — the paper names embeddings and biases as the examples, and its Algorithm 1 averages all remaining non-linear-layer weights, which sweeps the norms in too.

Second, the transpose adapter, in one sentence: nn.Linear stores its matrix as (out, in) while regmean_weight solves in (in, out), so regmean_merge_real transposes each Gram-keyed weight in, merges, and transposes back out.

Third, what “its Gram” means here. record_grams hooks every linear layer’s input, runs the model’s own sentence through once, accumulates X^\top X over all input rows, and divides by the row count — a per-token 1/N. That division is ours, not the paper’s: Jin et al. sum raw X_i^\top X_i with no normalization (their §4.3 only caps collection at 1,000 batches). Under this demo’s equal data the two agree exactly — a shared scale cancels top and bottom of the solve — but under unequal data ours would down-weight the larger dataset where the paper lets it dominate. Fourth, the damping ridge. One memorized sentence yields at most 7 distinct input rows, so every raw toy Gram here has rank \le 8 against 32–128 input dims — the undamped solve is degenerate (the paper’s full-data Grams are full-rank; its §5.3 note that \alpha = 1.0 can degenerate is the same disease at a different scale). record_grams therefore adds 10^{-2}\cdot\mathrm{mean}(\mathrm{diag}(G))\cdot I — again ours, disclosed: a shared damped Gram still cancels, so both anchors survive, and unseen directions fall back toward the soup. The sweep is insensitive to the level (10^{-3} vs 10^{-2} agree to 4 decimals).

And the \alpha knob is the paper’s own robustness scalar — default 0.9 (0.1 for T5-base): shrink each Gram toward its diagonal, \tilde{G} = \alpha G + (1-\alpha)\operatorname{diag}(G), before the solve. The paper’s ablation (§5.3) is your reading guide for the dial below: \alpha = 1.0 is pure RegMean but can degenerate, while \alpha = 0 diagonalizes into row-reweighting close to plain averaging.

NoteThe RegMean-vs-soup win is a mechanism demo

Say it plainly: two memorized sentences in a 1-layer toy prove the mechanism — data-weighted blending from real activations, steered by \alpha — not that RegMean wins real leaderboards. This section reports only what its own sweep measures across \alpha = 1.0 \to 0.0, with the seed in the open, and frames the dial with the paper’s own ablation: pure RegMean (\alpha = 1.0) can degenerate (§5.3), and full shrinkage (\alpha = 0) collapses toward plain averaging.

Code: Record, Merge, Score

Everything below lives in regmean_real.py, which reuses slerp_real (specialists, scoring) and merging (the solve, the soup) by import — never a copy. Train once, record each specialist’s Grams from real activations, merge at the paper’s default \alpha:

from regmean_real import record_grams, regmean_merge_real
from slerp_real import TASK_A, TASK_B, train_specialists

_state_a, _state_b = train_specialists(seed=0)
_grams_a = record_grams(_state_a, TASK_A)
_grams_b = record_grams(_state_b, TASK_B)
_merged = regmean_merge_real(_state_a, _state_b, _grams_a, _grams_b, non_diagonal_scale=0.9)
print(f"recorded {len(_grams_a)} per-layer Grams; RegMean-merged the linear maps, soup-averaged the rest")
print(sorted(_grams_a.keys()))
recorded 7 per-layer Grams; RegMean-merged the linear maps, soup-averaged the rest
['blocks.0.attention.mha.k_proj.weight', 'blocks.0.attention.mha.out_proj.weight', 'blocks.0.attention.mha.q_proj.weight', 'blocks.0.attention.mha.v_proj.weight', 'blocks.0.ffn.linear1.weight', 'blocks.0.ffn.linear2.weight', 'lm_head.weight']

Both specialists nail their own sentence and fumble the other’s — the same tiny-specialist setup as the SLERP section. Now sweep the \alpha grid and score every merge on both tasks for real:

from regmean_real import demonstrate_regmean_real

_regmean_real = demonstrate_regmean_real(seed=0)
print(f"{'alpha':>6} {'RegMean-A':>10} {'RegMean-B':>10} {'mean':>8}")
for a, la, lb, m in zip(_regmean_real["alphas"], _regmean_real["regmean_ce_a"], _regmean_real["regmean_ce_b"], _regmean_real["regmean_mean_ce"]):
    print(f"{a:6.1f} {la:10.4f} {lb:10.4f} {m:8.4f}")
print(f"soup: A={_regmean_real['soup_ce_a']:.4f} B={_regmean_real['soup_ce_b']:.4f} mean={_regmean_real['soup_mean_ce']:.4f}")
print(f"specialists: A-on-A={_regmean_real['spec_a_ce_a']:.4f} A-on-B={_regmean_real['spec_a_ce_b']:.4f} B-on-A={_regmean_real['spec_b_ce_a']:.4f} B-on-B={_regmean_real['spec_b_ce_b']:.4f}")
 alpha  RegMean-A  RegMean-B     mean
   1.0     1.7764     1.8992   1.8378
   0.9     1.8056     1.8994   1.8525
   0.5     1.5698     1.7264   1.6481
   0.0     1.6160     1.9267   1.7714
soup: A=1.5874 B=1.7732 mean=1.6803
specialists: A-on-A=0.0668 A-on-B=3.2366 B-on-A=3.6101 B-on-B=0.0716

The widget below draws this sweep — bridged into OJS:

Drive It: Shrinkage Bends the Blend

Pick \alpha on the paper’s grid. Each task group shows two bars: the solid bar is the RegMean merge at your \alpha, the faded bar the soup — hue is the task (orange: A, blue: B), exactly as above. The readout compares the mean next-token CE against the soup and the specialists’ home-task baselines.

```{ojs}
#| echo: false

viewof realAlpha = Inputs.radio(
  ["1.0", "0.9", "0.5", "0.0"],
  {value: "0.9", label: "α — off-diagonal scale"}
)
```
ImportantOJS Syntax Error (line 1823, column 1)Unexpected character ‘#’ 
```{ojs}
#| echo: false

regmeanRealWidget = {
  const width = 720, height = 470, padL = 64, padR = 24, padB = 112, padT = 104;
  const theme = diagramTheme;
  const svg = d3.create("svg")
    .attr("viewBox", `0 0 ${width} ${height}`)
    .attr("width", "100%").attr("height", height)
    .style("max-width", `${width}px`)
    .style("font-family", "var(--pg-mono)");
  svg.append("rect").attr("width", width).attr("height", height)
    .attr("fill", theme.bg).attr("rx", 12);

  const alphas = regmean_real.alphas;
  const target = parseFloat(realAlpha);
  let best = 0;
  for (let i = 1; i < alphas.length; i++) {
    if (Math.abs(alphas[i] - target) < Math.abs(alphas[best] - target)) best = i;
  }
  const alpha = alphas[best];

  const bars = [
    {group: "task A", method: "RegMean", val: regmean_real.regmean_ce_a[best], col: theme.highlight, op: 1},
    {group: "task A", method: "soup", val: regmean_real.soup_ce_a, col: theme.highlight, op: 0.4},
    {group: "task B", method: "RegMean", val: regmean_real.regmean_ce_b[best], col: theme.accent, op: 1},
    {group: "task B", method: "soup", val: regmean_real.soup_ce_b, col: theme.accent, op: 0.4}
  ];
  const top = Math.max(bars[0].val, bars[1].val, bars[2].val, bars[3].val);
  const x = d3.scaleBand().domain([0, 1, 2, 3]).range([padL, width - padR]).padding(0.32);
  const y = d3.scaleLinear([0, top * 1.18], [height - padB, padT]);

  svg.append("line").attr("x1", padL).attr("x2", width - padR)
    .attr("y1", y(0)).attr("y2", y(0))
    .attr("stroke", theme.edgeStroke).attr("opacity", 0.5);

  bars.forEach((b, i) => {
    svg.append("rect").attr("x", x(i)).attr("width", x.bandwidth())
      .attr("y", y(b.val)).attr("height", y(0) - y(b.val))
      .attr("rx", 4).attr("fill", b.col).attr("opacity", b.op);
    svg.append("text").attr("x", x(i) + x.bandwidth() / 2).attr("y", y(b.val) - 8)
      .attr("text-anchor", "middle").attr("fill", theme.nodeText)
      .attr("font-size", 13).attr("font-weight", 700).text(b.val.toFixed(2));
    svg.append("text").attr("x", x(i) + x.bandwidth() / 2).attr("y", height - padB + 22)
      .attr("text-anchor", "middle").attr("fill", theme.nodeText)
      .attr("font-size", 12).attr("font-weight", 700).text(b.method);
  });
  svg.append("text").attr("x", (x(0) + x(1)) / 2 + x.bandwidth() / 2).attr("y", height - padB + 42)
    .attr("text-anchor", "middle").attr("fill", theme.highlight)
    .attr("font-size", 12).text("task A");
  svg.append("text").attr("x", (x(2) + x(3)) / 2 + x.bandwidth() / 2).attr("y", height - padB + 42)
    .attr("text-anchor", "middle").attr("fill", theme.accent)
    .attr("font-size", 12).text("task B");

  const chip = (xx, col, op, txt) => {
    svg.append("rect").attr("x", xx).attr("y", height - padB + 56)
      .attr("width", 18).attr("height", 12).attr("rx", 3)
      .attr("fill", col).attr("opacity", op);
    svg.append("text").attr("x", xx + 24).attr("y", height - padB + 66)
      .attr("fill", theme.nodeText).attr("font-size", 11.5).text(txt);
  };
  chip(padL, theme.highlight, 1, "RegMean @ α");
  chip(padL + 150, theme.highlight, 0.4, "soup (never moves)");

  const meanReg = regmean_real.regmean_mean_ce[best];
  const meanSoup = regmean_real.soup_mean_ce;
  svg.append("text").attr("x", width / 2).attr("y", 30).attr("text-anchor", "middle")
    .attr("fill", theme.highlightText).attr("font-size", 15).attr("font-weight", 700)
    .text("RegMean vs soup on both tasks — solid bends with α, faded never moves");
  svg.append("text").attr("x", width / 2).attr("y", 52).attr("text-anchor", "middle")
    .attr("fill", theme.nodeText).attr("font-size", 12)
    .text("(" + regmean_real.n_regmean_keys + " linear maps RegMean-merged, rest soup-averaged)");
  svg.append("text").attr("x", width / 2).attr("y", 76).attr("text-anchor", "middle")
    .attr("fill", theme.nodeText).attr("font-size", 12)
    .text("α = " + alpha.toFixed(1) + " — RegMean mean CE " + meanReg.toFixed(3) + " vs soup " + meanSoup.toFixed(3));
  svg.append("text").attr("x", width / 2).attr("y", height - 12).attr("text-anchor", "middle")
    .attr("fill", theme.nodeText).attr("font-size", 12)
    .text("specialists at home: A " + regmean_real.spec_a_ce_a.toFixed(2) + " / B " + regmean_real.spec_b_ce_b.toFixed(2));

  return svg.node();
}
```
ImportantOJS Syntax Error (line 1832, column 1)Unexpected character ‘#’
TipTry This
  1. Park at \alpha = 0.9, the paper’s default — and notice it can lose. On this toy the near-pure solve scores mean CE 1.85 against the soup’s 1.68: rank-\le-8 Grams make the undamped end of the dial degenerate, exactly the paper’s §5.3 observation at toy scale. The dial, not a fixed \alpha, is the lesson.
  2. Slide down to 0.5, then to 0.0. Mid-shrinkage recovers the blend to roughly the soup (mean 1.65 vs 1.68); at 0.0 the Grams diagonalize and the solid bars sink toward the faded ones: row-reweighting close to plain averaging, live.
  3. Read the specialist baselines in the readout. Neither merge beats a specialist on its home sentence — the merge trades home-task CE for away-task CE. Ask which \alpha buys the cheapest trade on the mean, and whether both tasks agree on the answer.

Interactive Exploration

The widget below is driven by real numbers from demonstrate_merging() in merging.py: two models fine-tuned toward disjoint skills, with a planted sign conflict on one coordinate. Compare how each merge fills the six weight coordinates — and watch the conflict coordinate (outlined) separate the naive sum from TIES.

TipTry This
  1. Flip between “task arithmetic” and “TIES.” On the disjoint skill coordinates (w0–w3) both reach ~1. On the outlined w4, task arithmetic sits near +2 (the −1 ate part of the +3); TIES keeps the full +3.
  2. Look at “soup.” Every coordinate is halved — the average never composes skills to full strength the way task arithmetic does. Soup is robust; arithmetic is expressive.

Common Pitfalls

Pitfall Why it bites Fix
Merging models with different initializations Weights from unrelated basins average to nonsense Only merge fine-tunes of a shared base checkpoint
Not scaling a sum of many task vectors \sum_i \tau_i overshoots; the model diverges Use \lambda < 1 (tune it — e.g. 0.3–0.8)
Summing dense task vectors naively Sign conflicts dilute the strong updates Resolve interference with TIES (trim + elect sign)
Forgetting the rescale in DARE Dropping entries without \tfrac{1}{1-p} shrinks the delta Always rescale survivors; that is what preserves the expectation
Mismatched keys/shapes across models Merging assumes identical parameter tensors Merge same-architecture models; align tokenizer/embedding sizes first
Averaging two strong models with a plain mean The chord shrinks the norm; the blend is quieter than either parent Use SLERP to interpolate along the arc and hold the magnitude
SLERP on three or more models, or at very different scales There is no single great circle through many points; the norm-hold assumes equal norms Apply SLERP pairwise; reach for soups/TIES/DARE beyond two models
Fisher-merging with a flat or mis-measured Fisher If every parameter’s Fisher is equal, the ratio cancels and you get back the plain soup — the certainty has to actually vary to help Estimate the diagonal Fisher on real data (mean squared log-likelihood gradient), sampling y from the model; its variation across parameters is the whole signal
RegMean without a Gram, or on the wrong layers RegMean needs an input Gram G_i = X_i^\top X_i per linear layer per model, measured on that model’s data; and it only merges linear weights — norms, biases, and embeddings have no such Gram Record each layer’s Gram once by running a little data through each model; RegMean the linear maps and soup everything else, as regmean_merge does
A near-singular summed Gram If the inputs never exercised some direction, \sum_i G_i is ill-conditioned and the solve blows up Shrink the off-diagonal with \alpha \approx 0.9 (\tilde G = \alpha G + (1-\alpha)\operatorname{diag} G) — the paper’s ridge-style regularizer — or add data
Reading the toy’s SLERP-vs-soup gap as a benchmark One 1-layer memorization config with the seed in the open is a mechanism demo; no controlled SLERP-vs-linear study exists Report the configuration and the ~0.18-nat number for what it is — the arc holding scale, not a leaderboard claim
Expecting every tensor at a small angle Zero-init biases swing wide (median Ω ≈ 68°) after any differential update, which looks alarming against the small-angle story Judge the geometry by the big shared matrices (8–11°) and the global norm ratio (≈ 1.0) — that is what carries the blend’s scale
Feeding (out, in) weights straight into the solve nn.Linear stores (out, in) but regmean_weight solves in (in, out) — the Gram no longer lines up with the weight’s rows Transpose each Gram-keyed weight in, merge, transpose back, as regmean_merge_real does
Expecting α = 0 to equal the soup exactly α = 0 diagonalizes the Grams into per-feature row-reweighting — close to plain averaging (paper §5.3) but not identical, since diagonal energy still weights the rows Expect the α = 0.0 bars near the soup bars, not on them

Exercises

Exercise 1: The round-trip identity

import torch
from merging import task_vector, apply_task_vector

base = {"w": torch.randn(8)}
ft = {"w": torch.randn(8)}

# Reconstruct `ft` from `base` and its task vector. It should match to float tol.
# tv = ...
# recon = ...
# print(torch.allclose(recon["w"], ft["w"]))

Exercise 2: Negate to forget

from merging import negate, model_soup

# Build a model good at both A and B, then use a NEGATED task vector to remove
# skill A while leaving B intact. Which coordinates should change sign?
base = {"w": torch.zeros(4)}
a = {"w": torch.tensor([2.0, 2.0, 0.0, 0.0])}
b = {"w": torch.tensor([0.0, 0.0, 2.0, 2.0])}
# both = ...
# forget_a = apply_task_vector(both, negate(task_vector(base, a)), 1.0)
# print(forget_a["w"].tolist())

Exercise 3: How many entries survive a trim?

from merging import trim

tau = {"w": torch.randn(100)}
# For density in {0.1, 0.2, 0.5}, predict the nonzero count (ceil(100*density)),
# then check with trim(...). Confirm the survivors are the largest by |value|.
# for d in (0.1, 0.2, 0.5):
#     out = trim(tau, d)
#     print(d, int((out["w"] != 0).sum()))

Exercise 4: DARE preserves the mean

from merging import dare

tau = {"w": torch.tensor([1.0, 2.0, 3.0, 4.0])}
# Average many DARE draws at p=0.8 and confirm the mean approaches tau.
# Then verify each nonzero survivor equals value / (1 - p) exactly.

Exercise 5: How far does the chord fall?

import math
from merging import slerp

# Two equal-length vectors at angle Omega. As Omega grows, the LERP midpoint's
# norm falls as cos(Omega/2), while SLERP holds at the shared norm.
# For Omega in {30, 90, 150} degrees, predict |LERP mid| = cos(Omega/2), then check.
for deg in (30, 90, 150):
    om = math.radians(deg)
    a = {"w": torch.tensor([1.0, 0.0])}
    b = {"w": torch.tensor([math.cos(om), math.sin(om)])}
    # lerp_mid = ...
    # print(deg, "predicted", round(math.cos(om / 2), 4), "slerp", round(float(slerp(a, b, 0.5)["w"].norm()), 4))

Exercise 6: Certainty decides, and equal certainty is a soup

from merging import fisher_merge, model_soup

a = {"w": torch.tensor([1.0, 0.0])}
b = {"w": torch.tensor([0.0, 1.0])}

# (a) Give both models the SAME Fisher and confirm fisher_merge == model_soup.
# flat = {"w": torch.ones(2)}
# print(torch.allclose(fisher_merge([a, b], [flat, flat])["w"], model_soup([a, b])["w"]))

# (b) Now make A ten times more certain about coordinate 0 than B is. Predict the
# merged value there from theta*_0 = (F_a*a_0 + F_b*b_0)/(F_a + F_b), then check.
# fa = {"w": torch.tensor([10.0, 1.0])}
# fb = {"w": torch.tensor([1.0, 10.0])}
# print(fisher_merge([a, b], [fa, fb])["w"].tolist())   # coord 0 near 10/11 ≈ 0.909

Exercise 7: RegMean recovers the soup, then beats it

from merging import gram_matrix, regmean_weight

w_a = torch.tensor([[1.0], [0.0]])
w_b = torch.tensor([[0.0], [1.0]])

# (a) Equal Grams -> the soup. Hand both models the SAME Gram and confirm RegMean
# returns the plain mean of the weights.
# G = gram_matrix(torch.randn(64, 2))
# print(torch.allclose(regmean_weight([w_a, w_b], [G, G]), 0.5 * (w_a + w_b), atol=1e-4))

# (b) Now let B contribute NO data (a zero Gram). Predict W* before you run it —
# with no evidence from B, the least-squares merge must be A's weights exactly.
# print(regmean_weight([w_a, w_b], [gram_matrix(torch.randn(64, 2)), torch.zeros(2, 2)]))

Exercise 8: Blend two real specialists

from slerp_real import train_specialists, sweep_merges, merge_angle_report

# (a) Train with a different seed and confirm the headline pattern survives:
# median angle in the tens of degrees, global norms within ~1%, and the SLERP
# midpoint mean CE below the soup midpoint mean CE.
# sa, sb = train_specialists(seed=1)
# print(merge_angle_report(sa, sb)["median_omega_deg"])
# sw = sweep_merges(sa, sb, [0.5])
# print(sw["slerp_mean"], sw["soup_mean"])

# (b) Drop to steps=25, where the specialists barely memorize. Predict first from
# the norm theorem — the arc holds scale whatever the angle — then check whether
# the SLERP edge survives.

# (c) Antipodal check: slerp(a, -a, 0.5) on a toy dict. Predict the fallback from
# merging.slerp's sin-Omega guard (no source gives antipodal special handling),
# then run it.

Exercise 9: Shrink a real merge

from regmean_real import demonstrate_regmean_real, record_grams, regmean_merge_real

# (a) Run the demo at a new seed and confirm the headline pattern survives: the
# α = 0.9 mean CE at or below the soup mean, with the α = 0.0 row nearest the soup.
# d = demonstrate_regmean_real(seed=1)
# print(d["regmean_mean_ce"], d["soup_mean_ce"])

# (b) Predict what α = 0.0 must do from the diagonal reduction — a row-weighted
# soup — then check how close its mean lands to the soup's. When would they differ?
# (Hint: diagonal Grams still carry per-feature energy, so "close to" is not "equal to".)

# (c) Merging a specialist with itself must return itself. Train one, record its
# Grams, and confirm regmean_merge_real(a, a, ga, ga) matches per key to float tol.
# from slerp_real import TASK_A, train_specialists
# import torch
# sa, _ = train_specialists(seed=0)
# ga = record_grams(sa, TASK_A)
# m = regmean_merge_real(sa, sa, ga, ga)
# print(all(torch.allclose(m[k], sa[k]) for k in sa))

Summary

Key takeaways:

  1. Weights are a place; a fine-tune is a vector. Model merging is arithmetic on points and directions in weight space — no training, no extra inference cost.
  2. Model soups average same-init weights and often beat every member, because shared-initialization fine-tunes share a loss basin.
  3. Task vectors \tau = \theta_{ft} - \theta_{base} let you add skills, negate behaviors, and scale strength via \theta_{base} + \lambda \sum_i \tau_i.
  4. Interference — sign conflict and redundancy — is why naive sums degrade; TIES fixes it by trimming, electing a sign, and disjoint-merging only the agreeing entries.
  5. DARE shows task vectors are so redundant you can randomly drop 90%+ of their entries and, with a \tfrac{1}{1-p} rescale, merge with no loss in expectation.
  6. SLERP merges two models along the arc of the hypersphere, not the chord — holding the weights’ magnitude (exactly, at equal norm) where a plain average shrinks it, and gracefully falling back to LERP when the two are nearly colinear. On real (tiny) shared-init specialists the measured geometry matches: big matrices near-parallel, global norms within 0.1%, the arc’s global norm flat while the chord’s sags — and a modest ~0.18-nat midpoint edge that you report as a mechanism demo, since no controlled SLERP-vs-linear study exists.
  7. Fisher-weighted averaging is the first merge the data tunes: weight each parameter by each model’s certainty (its diagonal Fisher = mean squared log-likelihood gradient), giving the precision-weighted mean \theta^\ast_j = \tfrac{\sum_i \lambda_i F^{(i)}_j \theta^{(i)}_j}{\sum_i \lambda_i F^{(i)}_j} — the peak of the product of the models’ Gaussian posteriors, and exactly the soup when every Fisher is equal.
  8. RegMean merges each linear layer so its outputs — not its weights — best reproduce every model’s, a closed-form least-squares solve W^\ast = (\sum_i G_i)^{-1}(\sum_i G_i W_i) in the input Grams G_i = X_i^\top X_i. It reduces to a soup when the Grams are equal, to a row-weighted soup when they’re diagonal, and pulls decisively ahead when a layer’s inputs are correlated — the off-diagonal a flat average is blind to.
  9. RegMean runs on real activations too. Record each linear layer’s input Gram from a forward pass — transposing (out, in) weights into the solve and back, soup-averaging everything else — and sweep the paper’s \alpha shrink (default 0.9; 1.0 can degenerate, 0.0 nears plain averaging). Two memorized sentences prove the mechanism, not a leaderboard: report the measured sweep for what it is.

What’s Next

Model merging is the cheapest way to compose the models this book taught you to build. It pairs naturally with PEFT — LoRA adapters are themselves task vectors, so everything here applies to merging adapters — and with quantization and evaluation, which score a merge exactly as they would any other model. This lesson took two steps into learned merging: Fisher-weighted averaging, where the data sets a per-parameter weight, and RegMean, where it sets a whole per-layer least-squares solve. The frontier past them keeps that spirit: evolutionary recipe search over merge coefficients (and layer-wise SLERP t), and geometry-aware merges like Model Stock — all replacing the last hand-tuned knobs with something the data chooses. And where merging combines finished models by editing weights, Module 28: Knowledge Distillation does the complementary thing — trains a new, smaller model to inherit a big one’s behavior.

Going Deeper

Core Papers:

Practical Resources:

  • mergekit — the toolkit the open-weight community uses to merge models (soups, task arithmetic, TIES, DARE, SLERP) in practice.
  • mergekit slerp.py — the production SLERP merge: normalize, read the angle, and fall back to LERP above a DOT_THRESHOLD of 0.9995.
  • Merge Models with MergeKit — Labonne. The practitioner’s walkthrough of mergekit recipes, including SLERP and layer-wise schedules.
  • LLM-SLERP-Merge — Digitous/CalderaAI. The community recipe SLERP merging started from, before MergeKit productized it.