Model Merging & Composition

The whole is greater than the sum of its parts, especially when the parts are neural networks that learned different things.

Distill, Synergy-Seeking AI Agent

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 (for an overview of these model families, see Chapter 07), or even two Llama models fine-tuned from different pretrained versions. The common ancestry is what ensures the weight spaces are compatible. Figure 16.2.1 visualizes these task vectors in parameter space.

Why does model merging work at all? This is one of the most surprising results in recent LLM research. The intuition relies on the concept of linear mode connectivity: models fine-tuned from the same base checkpoint tend to occupy the same low-loss basin in the parameter space. Because they share the same initialization, the paths from the base model to each fine-tuned variant do not cross high-loss barriers. Averaging (or otherwise combining) their weights produces a point that also lies within this shared basin, inheriting capabilities from both. This explains why merging fails for models trained from different random initializations: they occupy entirely separate loss basins with no linear path between them. The practical consequence is that model merging is only reliable when all constituent models descend from the same base checkpoint.

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:

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.

A tiny numeric example clarifies the effect. Consider a single parameter from two fine-tuned models:

# Linear merge of a single weight with different alpha values
w_a, w_b = 1.4, 0.6 # same parameter in models A and B

for alpha in [0.3, 0.5, 0.7]:
 w_merged = alpha * w_a + (1 - alpha) * w_b
 print(f"alpha={alpha:.1f} -> W_merged = {alpha}*{w_a} + {1-alpha}*{w_b} = {w_merged:.2f}")
# alpha=0.3 -> W_merged = 0.84 (closer to model B)
# alpha=0.5 -> W_merged = 1.00 (midpoint)
# alpha=0.7 -> W_merged = 1.16 (closer to model A)

alpha=0.3 -> W_merged = 0.3*1.4 + 0.7*0.6 = 0.84
alpha=0.5 -> W_merged = 0.5*1.4 + 0.5*0.6 = 1.00
alpha=0.7 -> W_merged = 0.7*1.4 + 0.3*0.6 = 1.16

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.

Figure 16.2.3 walks through these three TIES steps visually.

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

Figure 16.2.5 provides an at-a-glance comparison of the merging methods covered above. The table below adds quantitative detail on each method's capabilities and trade-offs.

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

Code Fragment 16.2.2 demonstrates this approach.

# Implement model merging by interpolating weight tensors
# Linear interpolation (LERP/SLERP) blends capabilities from multiple models
import torch
from transformers import AutoModelForCausalLM

def compute_task_vector(base_model_id, finetuned_model_id):
 """Compute task vector: delta = finetuned - base."""
 # Load pretrained causal LM (decoder-only architecture)
 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
)

Merge complete. Test with: transformers AutoModelForCausalLM.from_pretrained('./merged-output')

Code Fragment 16.2.3 demonstrates model merging with a TIES configuration file.

# 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

Run the merge from the command line or use the MergeKit Python API for programmatic control. Code Fragment 16.2.3 shows both approaches.

# pip install mergekit
# Option 1: CLI (most common)
# mergekit-yaml merge_ties.yml ./merged-output --cuda --copy-tokenizer

# Option 2: Python API for programmatic merging
from mergekit.config import MergeConfiguration
from mergekit.merge import MergeOptions, run_merge

config = MergeConfiguration.from_yaml("merge_ties.yml")
run_merge(
 config,
 out_path="./merged-output",
 options=MergeOptions(
 cuda=True,
 copy_tokenizer=True,
 lazy_unpickle=True, # lower RAM usage
 ),
)
print("Merge complete. Test with: transformers AutoModelForCausalLM.from_pretrained('./merged-output')")

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. Figure 16.2.4 depicts this optimization loop.