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.

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.

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

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.

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.

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. The open frontier is learned merging: Fisher-weighted averaging (weight each parameter by how much each task relies on it), RegMean (a closed-form least-squares merge), and evolutionary recipe search over merge coefficients — replacing the hand-tuned \(\lambda\) and density with something the data chooses.

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.