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.

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.

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.

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

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)]))

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.
  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.

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.