Model Merging & Composition

1. Why Model Merging Works

Model merging exploits a remarkable property of neural network loss landscapes: models fine-tuned from the same base occupy a connected region of parameter space where linear interpolations between them tend to perform well. When two models are fine-tuned from the same pretrained checkpoint, their weight differences from the base represent "task vectors" in parameter space. These task vectors can be combined arithmetically because the underlying loss landscape in the neighborhood of the pretrained model is approximately convex.

The key constraint is that all models being merged must share the same architecture and originate from the same base pretrained checkpoint. You cannot merge a Llama model with a Mistral model, or even two Llama models fine-tuned from different pretrained versions. The common ancestry is what ensures the weight spaces are compatible.

2. Merging Methods

2.1 Linear (Weighted Average)

The simplest merging method computes a weighted average of model weights. Given models A and B with weights W_A and W_B, the merged model has weights:

W_merged = α · W_A + (1 - α) · W_B

The mixing coefficient α controls the balance between models. Setting α = 0.5 gives equal weight to both. Linear merging is fast and simple, but it can produce mediocre results when the models have very different weight distributions because opposing parameter changes can cancel each other out.

2.2 SLERP (Spherical Linear Interpolation)

SLERP treats weight vectors as points on a high-dimensional sphere and interpolates along the geodesic (shortest path on the sphere surface) rather than in a straight line through the interior. This preserves the magnitude of weight vectors better than linear interpolation, which tends to shrink weights toward zero when the models diverge.

SLERP is applied layer by layer (or parameter by parameter) and produces consistently better results than linear averaging. It is the recommended default for merging two models.

2.3 TIES (TRIM, Elect Sign, merge)

TIES-Merging (2023) addresses the interference problem that plagues simple averaging. When two fine-tuned models modify the same parameter in opposite directions, averaging cancels out both changes. TIES handles this through three steps:

1. TRIM: Remove small-magnitude changes (below a threshold) that are likely noise rather than meaningful task knowledge.
2. Elect Sign: For each parameter, take a majority vote across models on the sign of the change. Parameters where models disagree on direction are resolved by the majority.
3. Merge: Average only the values that agree with the elected sign, zeroing out conflicting contributions.

2.4 DARE (Drop And REscale)

DARE (2024) takes a different approach to reducing interference: it randomly drops a large fraction (typically 90-99%) of the delta parameters before merging, then rescales the remaining parameters to compensate. The intuition is that fine-tuning changes are highly redundant, and a small random subset of changes captures the essential task knowledge. By keeping only a sparse set of changes from each model, the probability of destructive interference is dramatically reduced.

DARE can be combined with other merging methods (DARE+TIES is a popular combination) and tends to work especially well when merging many models or when the models have been fine-tuned for very different tasks.

2.5 Model Stock

Model Stock (2024) draws on ideas from portfolio theory in finance. Instead of merging all models uniformly, it selects an optimal subset and weighting based on the geometric properties of the models in weight space. It computes the distances between models and the base, then assigns weights that maximize diversity while minimizing deviation from the base. This produces more robust merges, particularly when some of the input models are lower quality or overfitted.

3. Merging Method Comparison

4. Task Arithmetic

Task arithmetic provides the theoretical framework for understanding model merging. A "task vector" is defined as the element-wise difference between a fine-tuned model and its pretrained base: τ = W_ft - W_base. Task arithmetic shows that these vectors can be manipulated algebraically:

• Addition: W_base + τ_A + τ_B creates a model with both capabilities
• Negation: W_base - τ_A removes capability A (useful for unlearning)
• Scaling: W_base + λτ_A controls the strength of adaptation

import torch
from transformers import AutoModelForCausalLM

def compute_task_vector(base_model_id, finetuned_model_id):
    """Compute task vector: delta = finetuned - base."""
    base = AutoModelForCausalLM.from_pretrained(base_model_id)
    finetuned = AutoModelForCausalLM.from_pretrained(finetuned_model_id)

    task_vector = {}
    for name in base.state_dict():
        task_vector[name] = (
            finetuned.state_dict()[name].float()
            - base.state_dict()[name].float()
        )
    return task_vector

def apply_task_vectors(base_model_id, task_vectors, scaling_factors):
    """Apply scaled task vectors to base model."""
    model = AutoModelForCausalLM.from_pretrained(base_model_id)
    state_dict = model.state_dict()

    for tv, scale in zip(task_vectors, scaling_factors):
        for name in state_dict:
            state_dict[name] = state_dict[name].float() + scale * tv[name]

    model.load_state_dict(state_dict)
    return model

# Example: combine code and medical capabilities
code_tv = compute_task_vector("base-model", "code-finetuned")
medical_tv = compute_task_vector("base-model", "medical-finetuned")

merged = apply_task_vectors(
    "base-model",
    [code_tv, medical_tv],
    [0.7, 0.5],   # Scaling factors per task
)

5. Model Soups

Model soups (Wortsman et al., 2022) are a specialized form of model merging where you average multiple checkpoints from the same training run or from runs with different hyperparameters. The key insight is that checkpoints along a training trajectory, or from runs that differ only in learning rate or data ordering, lie in a connected low-loss basin. Averaging these checkpoints produces a model that is more robust and generalizes better than any individual checkpoint.

import torch
from pathlib import Path

def create_model_soup(checkpoint_dirs: list[str], model_class):
    """Average weights from multiple training checkpoints."""

    # Load all checkpoint state dicts
    state_dicts = []
    for ckpt_dir in checkpoint_dirs:
        model = model_class.from_pretrained(ckpt_dir)
        state_dicts.append(model.state_dict())
        del model  # Free memory

    n = len(state_dicts)
    print(f"Creating soup from {n} checkpoints")

    # Uniform average
    soup_state = {}
    for key in state_dicts[0]:
        soup_state[key] = sum(sd[key].float() for sd in state_dicts) / n

    # Load into fresh model
    model = model_class.from_pretrained(checkpoint_dirs[0])
    model.load_state_dict(soup_state)
    return model

# Example: soup from checkpoints at different training stages
soup_model = create_model_soup(
    ["./checkpoints/epoch-2", "./checkpoints/epoch-3",
     "./checkpoints/epoch-4", "./checkpoints/epoch-5"],
    AutoModelForCausalLM,
)

6. MergeKit: The Standard Tool

MergeKit is the most widely used tool for model merging, supporting all major algorithms and providing a YAML-based configuration interface. It handles the complex details of loading models, computing task vectors, applying merge operations, and saving results in standard Hugging Face format.

# Install: pip install mergekit

# SLERP merge configuration (merge_slerp.yml)
slices:
  - sources:
      - model: models/code-llama-7b
        layer_range: [0, 32]
      - model: models/medical-llama-7b
        layer_range: [0, 32]

merge_method: slerp
base_model: models/code-llama-7b
parameters:
  t:
    - filter: self_attn
      value: [0, 0.5, 0.3, 0.7, 1]  # Gradient across layers
    - filter: mlp
      value: [1, 0.5, 0.7, 0.3, 0]
    - value: 0.5                     # Default for other params
dtype: bfloat16

# TIES merge configuration (merge_ties.yml)
models:
  - model: models/code-llama-7b
    parameters:
      density: 0.5    # Keep top 50% of changes
      weight: 1.0
  - model: models/medical-llama-7b
    parameters:
      density: 0.5
      weight: 0.8
  - model: models/math-llama-7b
    parameters:
      density: 0.5
      weight: 0.6

merge_method: ties
base_model: meta-llama/Llama-3-8B
parameters:
  normalize: true
dtype: bfloat16

7. Evolutionary Model Merging

Evolutionary model merging (Sakana AI, 2024) automates the search for optimal merge configurations using evolutionary algorithms. Instead of manually selecting merge methods, weights, and layer-specific parameters, an evolutionary optimizer explores the space of possible merges and evaluates each candidate against a benchmark suite. This approach has produced merged models that significantly outperform manually configured merges.

The search space includes merge method selection, per-layer interpolation weights, density parameters (for TIES/DARE), and even layer permutations. The evolutionary algorithm (typically CMA-ES or NSGA-II) optimizes these parameters against multiple objectives: performance on target benchmarks, retention of base model capabilities, and overall coherence.