{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module 27: Model Merging"
   ],
   "id": "cdc96fef-2627-48dc-9796-cfd9b89883ce"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "Every module so far has *made* a model — trained one\n",
    "([m07](../m07_training/lesson.qmd)), aligned one\n",
    "([m12](../m12_alignment/lesson.qmd)), or adapted one cheaply\n",
    "([m25](../m25_peft/lesson.qmd)). This module does something stranger and\n",
    "almost free: it takes **several finished models and combines them into\n",
    "one by arithmetic in weight space** — no gradient step, no data, no\n",
    "change to inference cost.\n",
    "\n",
    "**Model merging** treats a network’s weights as a point in a very\n",
    "high-dimensional space and its fine-tuning as a *direction*. Once you\n",
    "see weights that way, you can average two models, *add* a skill,\n",
    "*subtract* a behavior, or fold a dozen specialists into a single\n",
    "generalist — all by editing numbers, never by training.\n",
    "\n",
    "Why it matters for LLMs:\n",
    "\n",
    "- **It is how the open-weight world actually ships.** Most models near\n",
    "  the top of community leaderboards are *merges* — a base plus a handful\n",
    "  of fine-tunes, combined with a tool like `mergekit`. Merging is the\n",
    "  cheapest known way to compose capabilities.\n",
    "- **A fine-tune is reusable, not disposable.** The difference\n",
    "  `theta_ft − theta_base` is a portable **task vector** you can move\n",
    "  onto another checkpoint, scale up or down, or negate to *un*learn.\n",
    "- **Zero marginal cost.** The merged model has the *same architecture\n",
    "  and size* as its parents. Unlike an ensemble (run N models, average\n",
    "  outputs), a merge runs once.\n",
    "\n",
    "### What You’ll Learn\n",
    "\n",
    "After this module, you can:\n",
    "\n",
    "- Explain why the **weights of same-init fine-tunes can be averaged** —\n",
    "  the **model soup** — and why that is not obvious.\n",
    "- Compute a **task vector** `tau = theta_ft − theta_base` and use task\n",
    "  arithmetic to **add** skills and **negate** behaviors.\n",
    "- Diagnose the **interference** (sign disagreement + redundancy) that\n",
    "  makes a naive sum of many task vectors degrade.\n",
    "- Build **TIES-Merging** from scratch — **trim**, **elect sign**,\n",
    "  **disjoint merge** — and **DARE** (drop-and-rescale), and explain why\n",
    "  dropping 90% of a delta can be lossless.\n",
    "\n",
    "### Prerequisites\n",
    "\n",
    "This module requires familiarity with:\n",
    "\n",
    "- [Module 01: Tensors](../m01_tensors/lesson.qmd) — weights are just\n",
    "  tensors; here we do arithmetic on whole state dicts.\n",
    "- [Module 07: Training](../m07_training/lesson.qmd) — fine-tuning is the\n",
    "  move in weight space we are about to name and reuse.\n",
    "- [Module 25: PEFT](../m25_peft/lesson.qmd) — LoRA already merged an\n",
    "  adapter back into a weight; this module merges *whole models* the same\n",
    "  way.\n",
    "\n",
    "## Intuition: Weights Are a Place\n",
    "\n",
    "Picture the millions of numbers in a model as the coordinates of a\n",
    "single point. Training moves that point. When you fine-tune a shared\n",
    "base model $\\theta_{base}$ on task A, you end at $\\theta_A$; the arrow\n",
    "between them,\n",
    "\n",
    "$$\\tau_A = \\theta_A - \\theta_{base},$$\n",
    "\n",
    "is the **task vector** — everything the model *changed* to learn A. It\n",
    "points from “generic” toward “good at A.”\n",
    "\n",
    "Two facts make this picture powerful. First, **arrows add**: if $\\tau_A$\n",
    "points toward skill A and $\\tau_B$ toward skill B, then\n",
    "$\\theta_{base} + \\tau_A + \\tau_B$ often reaches a model good at *both*.\n",
    "Second, **arrows have a sign**: $-\\tau_A$ walks *away* from A, which\n",
    "turns out to be a clean way to make a model forget a behavior. The\n",
    "stepper below walks these operations on a 2-D toy — the same algebra\n",
    "runs unchanged in a billion dimensions."
   ],
   "id": "26c3acdc-15e9-44f7-bccd-acd5e7e565f0"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "learnllm-ojs-note": true
   },
   "source": [
    "> **Interactive figure**\n",
    ">\n",
    "> This visualization runs in the browser. Open the [HTML lesson](lesson.html) to explore it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Key Insight**\n",
    ">\n",
    "> A trained model is a *point*; a fine-tune is a *vector*. Merging is\n",
    "> linear algebra on those points and vectors — which is why it costs no\n",
    "> training and no extra inference: you are only ever adding and scaling\n",
    "> weights you already have.\n",
    "\n",
    "## Model Soups: Average the Weights Themselves\n",
    "\n",
    "The simplest merge ignores task vectors entirely and averages the\n",
    "*weights* directly. Take several models fine-tuned from the **same**\n",
    "pretrained checkpoint — say, with different learning rates or seeds —\n",
    "and set every parameter to its mean across them. Wortsman et al. (2022)\n",
    "called this a **model soup**, and the surprise is that it *works*: the\n",
    "uniform soup often beats every individual member on both accuracy and\n",
    "robustness.\n",
    "\n",
    "Why should averaging weights be legal at all? Two randomly initialized\n",
    "networks average into garbage. The trick is **shared initialization**:\n",
    "models fine-tuned from one base stay in the *same loss basin*, close\n",
    "enough that the straight line between them stays low-loss. The mean sits\n",
    "near the bottom of that basin — a flat, well-generalizing spot."
   ],
   "id": "a6393d7e-31de-43cb-8215-7f0c8835ebe7"
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "uniform soup: [1.0, 1.0, 1.0, 1.0]\n",
      "weighted 3:1: [1.5, 1.5, 0.5, 0.5]"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from merging import model_soup, task_vector, apply_task_vector\n",
    "\n",
    "# A tiny \"state dict\": parameter name -> tensor. The algorithms are model-agnostic,\n",
    "# so what works here works on a real GPTModel.state_dict().\n",
    "base = {\"w\": torch.zeros(4)}\n",
    "model_a = {\"w\": torch.tensor([2.0, 2.0, 0.0, 0.0])}   # fine-tuned toward skill A\n",
    "model_b = {\"w\": torch.tensor([0.0, 0.0, 2.0, 2.0])}   # fine-tuned toward skill B\n",
    "\n",
    "soup = model_soup([model_a, model_b])\n",
    "print(\"uniform soup:\", soup[\"w\"].tolist())            # the elementwise mean\n",
    "\n",
    "# A soup of one model is just that model; weights need not sum to 1 (they're normalized).\n",
    "print(\"weighted 3:1:\", model_soup([model_a, model_b], weights=[3.0, 1.0])[\"w\"].tolist())"
   ],
   "id": "868cc18b"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Soups need a shared ancestor**\n",
    ">\n",
    "> Averaging weights only makes sense for models that started from the\n",
    "> **same** initialization (a common base checkpoint). Merge two models\n",
    "> trained from scratch with different inits and you land between two\n",
    "> unrelated basins — usually far worse than either. Every technique in\n",
    "> this module assumes a shared base.\n",
    "\n",
    "## The Math: Task Vectors\n",
    "\n",
    "Model soups throw away a useful distinction: *what each fine-tune\n",
    "changed*. Task arithmetic keeps it. Define the **task vector** as the\n",
    "delta a fine-tune applied:\n",
    "\n",
    "$$\\tau_i = \\theta_i - \\theta_{base}.$$\n",
    "\n",
    "Then editing a model is vector algebra. The **merged** model is the base\n",
    "plus a scaled sum of task vectors:\n",
    "\n",
    "$$\\theta_{merged} = \\theta_{base} + \\lambda \\sum_i \\tau_i .$$\n",
    "\n",
    "- **Adding** ($+\\tau_A + \\tau_B$) composes skills into one multi-task\n",
    "  model.\n",
    "- **Negating** ($-\\tau_A$) *removes* a behavior — the unlearning / detox\n",
    "  operator.\n",
    "- **Scaling** by $\\lambda$ dials the strength; with several vectors you\n",
    "  usually need $\\lambda < 1$ so their sum does not overshoot.\n",
    "\n",
    "The bookkeeping identity that makes this exact: applying the task vector\n",
    "back to the base with $\\lambda = 1$ reconstructs the fine-tune. There is\n",
    "nothing to approximate — merging is defined by these equations.\n",
    "\n",
    "## Code: Task Arithmetic From Scratch\n",
    "\n",
    "The whole toolkit lives in `merging.py`. Every function takes and\n",
    "returns a state dict, so it composes freely. Here are the three moves\n",
    "and the identities the tests pin:"
   ],
   "id": "f033e4e6-6d2f-4460-9cb7-57e70b23e192"
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "round-trip matches θ_A: True\n",
      "forget A: [-2.0, -2.0, 0.0, 0.0]\n",
      "compose A+B: [2.0, 2.0, 2.0, 2.0]"
     ]
    }
   ],
   "source": [
    "from merging import negate, merge\n",
    "\n",
    "# --- Round-trip identity: base + (θ_ft − θ_base) == θ_ft ---\n",
    "tau_a = task_vector(base, model_a)\n",
    "recon = apply_task_vector(base, tau_a, 1.0)\n",
    "print(\"round-trip matches θ_A:\", torch.allclose(recon[\"w\"], model_a[\"w\"]))\n",
    "\n",
    "# --- Negate to forget: base + (−τ_A) walks away from skill A ---\n",
    "forget_a = apply_task_vector(base, negate(tau_a), 1.0)\n",
    "print(\"forget A:\", forget_a[\"w\"].tolist())   # negative on A's coordinates\n",
    "\n",
    "# --- Add to compose: base + τ_A + τ_B is good at both ---\n",
    "tau_b = task_vector(base, model_b)\n",
    "both = merge(base, [tau_a, tau_b], coeff=1.0)\n",
    "print(\"compose A+B:\", both[\"w\"].tolist())     # all four coordinates active"
   ],
   "id": "d1dd690f"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice `merge` reached `[2, 2, 2, 2]` — full strength on *both* skills —\n",
    "because A and B touched **disjoint** coordinates. That is the easy case.\n",
    "The interesting question is what happens when task vectors *fight* over\n",
    "the same parameter.\n",
    "\n",
    "## The Interference Problem\n",
    "\n",
    "Real task vectors are dense and they overlap. When you sum many of them,\n",
    "two things go wrong (Yadav et al., 2023):\n",
    "\n",
    "1.  **Sign conflict.** For a given parameter, task A may want it $+0.9$\n",
    "    while task B wants it $-0.1$. Their sum, $+0.8$, is dragged *down*\n",
    "    by a disagreement that was never important to B — the strong,\n",
    "    correct update gets diluted by weak noise pointing the other way.\n",
    "2.  **Redundant interference.** Most of a task vector’s entries are\n",
    "    small and irrelevant. Summed across many tasks, this low-magnitude\n",
    "    bulk piles up into a haze that drowns the few coordinates that\n",
    "    actually carry a skill.\n",
    "\n",
    "A plain sum treats all of this as signal. The fix is to be selective:\n",
    "keep only the entries that matter, and when they disagree, let the\n",
    "*strong* side win.\n",
    "\n",
    "## TIES-Merging\n",
    "\n",
    "**TIES** (**T**rim, **I**nterference resolution via **E**lect **S**ign,\n",
    "disjoint merge) resolves both problems in three steps, applied to the\n",
    "task vectors before they are summed:\n",
    "\n",
    "1.  **Trim** — per task vector, keep only the top-`density` fraction of\n",
    "    entries by magnitude (the paper uses 20%); zero the rest. This\n",
    "    deletes the redundant haze.\n",
    "2.  **Elect sign** — for each parameter, sum the (trimmed) values across\n",
    "    tasks and take the **sign of that sum**. This is a\n",
    "    magnitude-weighted majority vote: the direction the tasks most\n",
    "    strongly agree on.\n",
    "3.  **Disjoint merge** — average only the values whose sign *matches*\n",
    "    the elected sign, dividing by the count of agreeing tasks. The\n",
    "    dissenters are dropped, so a strong correct update is no longer\n",
    "    diluted."
   ],
   "id": "30ce26fb-e786-4222-b92f-7d5889d26576"
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "elected sign: [1.0, 1.0, 1.0]\n",
      "TIES  coord 2: 3.0\n",
      "naive coord 2: 2.0"
     ]
    }
   ],
   "source": [
    "from merging import trim, elect_sign, disjoint_merge, ties_merge\n",
    "\n",
    "# Two fine-tunes with a planted conflict on coordinate 2: A wants +3 (strong),\n",
    "# B wants −1 (weak). Coordinates 0 and 1 are each task's own disjoint skill.\n",
    "base3 = {\"w\": torch.zeros(3)}\n",
    "a = {\"w\": torch.tensor([1.0, 0.0, 3.0])}\n",
    "b = {\"w\": torch.tensor([0.0, 1.0, -1.0])}\n",
    "\n",
    "tvs = [task_vector(base3, a), task_vector(base3, b)]\n",
    "elected = elect_sign(tvs)\n",
    "print(\"elected sign:\", elected[\"w\"].tolist())         # coord 2: 3+(−1)=+2 -> +1\n",
    "\n",
    "ties = ties_merge(base3, [a, b], density=1.0, coeff=1.0)\n",
    "naive = merge(base3, tvs, coeff=1.0)\n",
    "print(\"TIES  coord 2:\", float(ties[\"w\"][2]))          # 3.0 — keeps the strong vote\n",
    "print(\"naive coord 2:\", float(naive[\"w\"][2]))         # 2.0 — the −1 cancels part of it"
   ],
   "id": "57c0df1a"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TIES kept the full $+3$ on the contested coordinate; the naive sum let\n",
    "the weak $-1$ eat a third of it. On a real merge of many models, that\n",
    "difference between “strong side wins” and “everything averages toward\n",
    "mush” is the difference between a merge that works and one that\n",
    "degrades.\n",
    "\n",
    "### Interactive: Watch Sign Election Resolve a Conflict\n",
    "\n",
    "Drag the **trim density** to see the low-magnitude entries vanish first,\n",
    "then watch each parameter’s column resolve to its elected sign and\n",
    "disjoint-merge only the agreeing cells. Green = positive, red =\n",
    "negative, faded = trimmed or dropped."
   ],
   "id": "0a67f0f9-3105-4893-be90-1b1c07fbe3d4"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "learnllm-ojs-note": true
   },
   "source": [
    "> **Interactive figure**\n",
    ">\n",
    "> This visualization runs in the browser. Open the [HTML lesson](lesson.html) to explore it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Try This**\n",
    ">\n",
    "> 1.  **Turn density down to 0.25.** Each task vector keeps only its\n",
    ">     single largest entry — the redundant middle columns fade out, and\n",
    ">     the merged row collapses to the few coordinates that carry real\n",
    ">     signal.\n",
    "> 2.  **Watch param 2.** Two task vectors want it negative and one wants\n",
    ">     it positive; the elected sign is negative, so the lone positive\n",
    ">     cell is *dropped* from the average rather than cancelling the two\n",
    ">     that agree.\n",
    "\n",
    "## DARE: Drop Almost Everything, Lose Nothing\n",
    "\n",
    "TIES trims by magnitude. **DARE** (Yu et al., 2023) shows you can trim\n",
    "*at random* and still be fine — because a task vector is astonishingly\n",
    "redundant. **D**rop **A**nd **RE**scale zeros each delta entry with\n",
    "probability $p$, then multiplies the survivors by $\\frac{1}{1-p}$:\n",
    "\n",
    "$$\\text{DARE}(\\tau)_j =\n",
    "\\begin{cases}\n",
    "\\dfrac{\\tau_j}{1-p} & \\text{with probability } 1-p,\\\\[4pt]\n",
    "0 & \\text{with probability } p.\n",
    "\\end{cases}$$\n",
    "\n",
    "The rescale is the whole trick: it makes DARE\n",
    "**expectation-preserving**, $\\mathbb{E}[\\text{DARE}(\\tau)] = \\tau$\n",
    "entrywise. So *on average* the dropped vector equals the original — and\n",
    "empirically you can drop 90–99% of a fine-tune’s deltas and merge with\n",
    "no measurable loss. Applied before TIES (**DARE-TIES**), it lets you\n",
    "stack many models with far less interference."
   ],
   "id": "b342cafe-d6bd-4728-8219-403f1058f3ce"
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "one draw (p=0.6): [0.0, 0.0, 7.5, 10.0, 12.5]\n",
      "mean of 4000 draws at p=0.9: [1.02, 2.09, 2.87, 3.88, 5.6]\n",
      "original: [1.0, 2.0, 3.0, 4.0, 5.0]"
     ]
    }
   ],
   "source": [
    "from merging import dare\n",
    "\n",
    "tau = {\"w\": torch.arange(1.0, 6.0)}          # the delta to sparsify: [1,2,3,4,5]\n",
    "\n",
    "# p = 0 is the identity; survivors at p = 0.6 are rescaled by 1/(1−0.6) = 2.5\n",
    "g = torch.Generator().manual_seed(0)\n",
    "print(\"one draw (p=0.6):\", dare(tau, 0.6, g)[\"w\"].tolist())\n",
    "\n",
    "# Average many independent draws -> back to the original (expectation preserved)\n",
    "acc = torch.zeros(5)\n",
    "N = 4000\n",
    "for i in range(N):\n",
    "    acc += dare(tau, 0.9, torch.Generator().manual_seed(i))[\"w\"]\n",
    "print(\"mean of 4000 draws at p=0.9:\", [round(x, 2) for x in (acc / N).tolist()])\n",
    "print(\"original:\", tau[\"w\"].tolist())"
   ],
   "id": "caf3a965"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even at $p = 0.9$ — throwing away nine of every ten delta entries per\n",
    "draw — the average lands back on $[1, 2, 3, 4, 5]$. The individual draws\n",
    "are sparse; their expectation is exact.\n",
    "\n",
    "## Interactive Exploration\n",
    "\n",
    "The widget below is driven by real numbers from `demonstrate_merging()`\n",
    "in `merging.py`: two models fine-tuned toward disjoint skills, with a\n",
    "planted sign conflict on one coordinate. Compare how each merge fills\n",
    "the six weight coordinates — and watch the **conflict coordinate**\n",
    "(outlined) separate the naive sum from TIES."
   ],
   "id": "e443da9e-bd42-4669-a647-79cce7aebbee"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "learnllm-ojs-note": true
   },
   "source": [
    "> **Interactive figure**\n",
    ">\n",
    "> This visualization runs in the browser. Open the [HTML lesson](lesson.html) to explore it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Try This**\n",
    ">\n",
    "> 1.  **Flip between “task arithmetic” and “TIES.”** On the disjoint\n",
    ">     skill coordinates (w0–w3) both reach ~1. On the outlined **w4**,\n",
    ">     task arithmetic sits near +2 (the −1 ate part of the +3); TIES\n",
    ">     keeps the full +3.\n",
    "> 2.  **Look at “soup.”** Every coordinate is halved — the average never\n",
    ">     composes skills to full strength the way task arithmetic does.\n",
    ">     Soup is robust; arithmetic is expressive.\n",
    "\n",
    "## Common Pitfalls\n",
    "\n",
    "| Pitfall | Why it bites | Fix |\n",
    "|-----------------------|------------------------------------|--------------|\n",
    "| Merging models with **different initializations** | Weights from unrelated basins average to nonsense | Only merge fine-tunes of a **shared** base checkpoint |\n",
    "| **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$) |\n",
    "| Summing dense task vectors **naively** | Sign conflicts dilute the strong updates | Resolve interference with **TIES** (trim + elect sign) |\n",
    "| Forgetting the **rescale** in DARE | Dropping entries without $\\tfrac{1}{1-p}$ shrinks the delta | Always rescale survivors; that is what preserves the expectation |\n",
    "| Mismatched **keys/shapes** across models | Merging assumes identical parameter tensors | Merge same-architecture models; align tokenizer/embedding sizes first |\n",
    "\n",
    "## Exercises\n",
    "\n",
    "### Exercise 1: The round-trip identity"
   ],
   "id": "606c7e22-efb2-411e-8b34-0332f134d309"
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from merging import task_vector, apply_task_vector\n",
    "\n",
    "base = {\"w\": torch.randn(8)}\n",
    "ft = {\"w\": torch.randn(8)}\n",
    "\n",
    "# Reconstruct `ft` from `base` and its task vector. It should match to float tol.\n",
    "# tv = ...\n",
    "# recon = ...\n",
    "# print(torch.allclose(recon[\"w\"], ft[\"w\"]))"
   ],
   "id": "e612d48c"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2: Negate to forget"
   ],
   "id": "03213a30-8440-49f6-b5a6-7d2c286b3ad6"
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from merging import negate, model_soup\n",
    "\n",
    "# Build a model good at both A and B, then use a NEGATED task vector to remove\n",
    "# skill A while leaving B intact. Which coordinates should change sign?\n",
    "base = {\"w\": torch.zeros(4)}\n",
    "a = {\"w\": torch.tensor([2.0, 2.0, 0.0, 0.0])}\n",
    "b = {\"w\": torch.tensor([0.0, 0.0, 2.0, 2.0])}\n",
    "# both = ...\n",
    "# forget_a = apply_task_vector(both, negate(task_vector(base, a)), 1.0)\n",
    "# print(forget_a[\"w\"].tolist())"
   ],
   "id": "831b5387"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3: How many entries survive a trim?"
   ],
   "id": "4ef2724d-184b-4f58-be75-0f91488254e7"
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from merging import trim\n",
    "\n",
    "tau = {\"w\": torch.randn(100)}\n",
    "# For density in {0.1, 0.2, 0.5}, predict the nonzero count (ceil(100*density)),\n",
    "# then check with trim(...). Confirm the survivors are the largest by |value|.\n",
    "# for d in (0.1, 0.2, 0.5):\n",
    "#     out = trim(tau, d)\n",
    "#     print(d, int((out[\"w\"] != 0).sum()))"
   ],
   "id": "31659872"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4: DARE preserves the mean"
   ],
   "id": "d86280c7-2d70-4c21-8e6f-96d484675783"
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from merging import dare\n",
    "\n",
    "tau = {\"w\": torch.tensor([1.0, 2.0, 3.0, 4.0])}\n",
    "# Average many DARE draws at p=0.8 and confirm the mean approaches tau.\n",
    "# Then verify each nonzero survivor equals value / (1 - p) exactly."
   ],
   "id": "188e8711"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Key takeaways:\n",
    "\n",
    "1.  **Weights are a place; a fine-tune is a vector.** Model merging is\n",
    "    arithmetic on points and directions in weight space — no training,\n",
    "    no extra inference cost.\n",
    "2.  **Model soups average same-init weights** and often beat every\n",
    "    member, because shared-initialization fine-tunes share a loss basin.\n",
    "3.  **Task vectors** $\\tau = \\theta_{ft} - \\theta_{base}$ let you\n",
    "    **add** skills, **negate** behaviors, and **scale** strength via\n",
    "    $\\theta_{base} + \\lambda \\sum_i \\tau_i$.\n",
    "4.  **Interference** — sign conflict and redundancy — is why naive sums\n",
    "    degrade; **TIES** fixes it by **trimming**, **electing a sign**, and\n",
    "    **disjoint-merging** only the agreeing entries.\n",
    "5.  **DARE** shows task vectors are so redundant you can randomly drop\n",
    "    90%+ of their entries and, with a $\\tfrac{1}{1-p}$ rescale, merge\n",
    "    with no loss in expectation.\n",
    "\n",
    "## What’s Next\n",
    "\n",
    "Model merging is the cheapest way to *compose* the models this book\n",
    "taught you to build. It pairs naturally with\n",
    "[PEFT](../m25_peft/lesson.qmd) — LoRA adapters are themselves task\n",
    "vectors, so everything here applies to merging adapters — and with\n",
    "[quantization](../m14_quantization/lesson.qmd) and\n",
    "[evaluation](../m17_evaluation/lesson.qmd), which score a merge exactly\n",
    "as they would any other model. The open frontier is *learned* merging:\n",
    "**Fisher-weighted** averaging (weight each parameter by how much each\n",
    "task relies on it), **RegMean** (a closed-form least-squares merge), and\n",
    "**evolutionary** recipe search over merge coefficients — replacing the\n",
    "hand-tuned $\\lambda$ and density with something the data chooses.\n",
    "\n",
    "### Going Deeper\n",
    "\n",
    "**Core Papers:**\n",
    "\n",
    "- [Model Soups](https://arxiv.org/abs/2203.05482) — Wortsman et\n",
    "  al., 2022. Averaging the weights of multiple fine-tunes of one base\n",
    "  improves accuracy and robustness at no inference cost.\n",
    "- [Editing Models with Task\n",
    "  Arithmetic](https://arxiv.org/abs/2212.04089) — Ilharco et al., 2022.\n",
    "  Introduced **task vectors**: add to learn, negate to forget, combine\n",
    "  by analogy.\n",
    "- [TIES-Merging: Resolving Interference When Merging\n",
    "  Models](https://arxiv.org/abs/2306.01708) — Yadav et al., 2023. Trim →\n",
    "  elect sign → disjoint merge, the interference fix this lesson builds.\n",
    "- [Language Models are Super Mario\n",
    "  (DARE)](https://arxiv.org/abs/2311.03099) — Yu et al., 2023.\n",
    "  Drop-and-rescale sparsification: delta redundancy makes 90%+ dropping\n",
    "  lossless.\n",
    "\n",
    "**Practical Resources:**\n",
    "\n",
    "- [mergekit](https://github.com/arcee-ai/mergekit) — the toolkit the\n",
    "  open-weight community uses to merge models (soups, task arithmetic,\n",
    "  TIES, DARE, SLERP) in practice."
   ],
   "id": "ab8fa468-3b75-4752-a67e-e32ccba9daa8"
  }
 ],
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3"
  },
  "language_info": {
   "name": "python",
   "codemirror_mode": {
    "name": "ipython",
    "version": "3"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.14"
  }
 }
}