Model Merging › Implementation
Language AI HandbookAll chapters Model Merging › ImplementationOne section · Sections in this chapter6 / 9 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let's implement these four merging methods from scratch using PyTorch. We will work with small toy models to keep the computation tractable, but the same code logic scales directly to full-size language models. The key operations, computing task vectors, applying masks, electing signs, and combining vectors, are all element-wise or reduction operations over flat parameter vectors that work identically regardless of model size.
Setup
We start by importing the necessary libraries and defining a simple linear model architecture that we can fine-tune on toy tasks. Using a small model lets us inspect the behavior of each method clearly without waiting for large-scale training.
import numpy as np
import torch
torch.manual_seed(42)
np.random.seed(42)Toy Model and Tasks
We define a small feedforward network and two simple regression tasks that require different parameter configurations. Each task is defined by a random projection of the input, so the two tasks require different linear functions of the same input space. A model that perfectly solves task A will not generally solve task B, making this a concrete test of each merging method's ability to combine capabilities.
import torch.nn as nn
class ToyModel(nn.Module):
def __init__(self, input_dim=16, hidden_dim=32, output_dim=1):
super().__init__()
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.act = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
return self.fc2(self.act(self.fc1(x)))
def make_task_data(n=200, input_dim=16, seed=0):
rng = np.random.default_rng(seed)
X = rng.standard_normal((n, input_dim)).astype(np.float32)
# Task-specific random projection defines the target
proj = rng.standard_normal(input_dim).astype(np.float32)
y = (X @ proj).reshape(-1, 1)
return torch.from_numpy(X), torch.from_numpy(y)
# Pretrained model: a random initialization representing a base model
pretrained_model = ToyModel()
# Task A and Task B data: different projections require different parameter configurations
X_a, y_a = make_task_data(seed=1)
X_b, y_b = make_task_data(seed=2)Fine-Tuning Helper
A simple training loop fine-tunes the pretrained model on each task and returns the resulting fine-tuned model. The key point is that both fine-tunes start from the same pretrained weights, which is the essential condition for model merging to work well.
import copy
import torch.optim as optim
def finetune(base_model, X, y, epochs=300, lr=1e-2):
model = copy.deepcopy(base_model)
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss()
for _ in range(epochs):
optimizer.zero_grad()
loss = criterion(model(X), y)
loss.backward()
optimizer.step()
return model
model_a = finetune(pretrained_model, X_a, y_a)
model_b = finetune(pretrained_model, X_b, y_b)Extracting Weights as Flat Vectors
All merging operations work on flat vectors of parameters. Converting model parameters to a flat vector is the first step in any merging pipeline. We define helpers to convert between flat vectors and model instances.
def get_flat_weights(model):
return torch.cat([p.detach().flatten() for p in model.parameters()])
def set_flat_weights(model, flat_weights):
model_copy = copy.deepcopy(model)
offset = 0
with torch.no_grad():
for p in model_copy.parameters():
numel = p.numel()
p.copy_(flat_weights[offset : offset + numel].reshape(p.shape))
offset += numel
return model_copy
theta_pretrained = get_flat_weights(pretrained_model)
theta_a = get_flat_weights(model_a)
theta_b = get_flat_weights(model_b)The flat vector representation makes the merging code clean and general: all the operations work on 1D tensors regardless of the specific architecture. When scaling to real models, you would typically work with state dictionaries and apply operations layer by layer to avoid loading all parameters into memory simultaneously, but the logic is identical.
Weight Averaging
The simplest merge: compute the element-wise mean of all fine-tuned parameter vectors.
def weight_average(pretrained_weights, finetuned_weights_list):
stacked = torch.stack(finetuned_weights_list)
return stacked.mean(dim=0)
theta_averaged = weight_average(theta_pretrained, [theta_a, theta_b])
model_averaged = set_flat_weights(pretrained_model, theta_averaged)Task Arithmetic
Task arithmetic computes task vectors, scales them, and adds them back to the pretrained model. The scaling coefficient is the key hyperparameter controlling how much of the fine-tuned direction to incorporate.
def task_arithmetic(pretrained_weights, finetuned_weights_list, lam=0.4):
task_vectors = [fw - pretrained_weights for fw in finetuned_weights_list]
summed_tv = torch.stack(task_vectors).sum(dim=0)
return pretrained_weights + lam * summed_tv
theta_ta = task_arithmetic(theta_pretrained, [theta_a, theta_b], lam=0.4)
model_ta = set_flat_weights(pretrained_model, theta_ta)TIES Merging
The TIES implementation follows the three steps: trim by magnitude percentile, elect sign by weighted majority vote, and disjoint merge using only agreeing entries.
def ties_merge(
pretrained_weights, finetuned_weights_list, lam=0.4, trim_pct=20.0
):
task_vectors = [fw - pretrained_weights for fw in finetuned_weights_list]
# Step 1: Trim - keep top trim_pct% by absolute magnitude
trimmed = []
for tv in task_vectors:
threshold = torch.quantile(tv.abs(), 1.0 - trim_pct / 100.0)
mask = tv.abs() >= threshold
trimmed.append(tv * mask.float())
# Step 2: Elect sign via magnitude-weighted majority vote
stacked = torch.stack(trimmed) # shape: (k, d)
pos_magnitude = (stacked > 0).float() * stacked.abs()
neg_magnitude = (stacked < 0).float() * stacked.abs()
elected_sign = torch.sign(pos_magnitude.sum(0) - neg_magnitude.sum(0))
elected_sign[elected_sign == 0] = 1.0 # break ties toward positive
# Step 3: Disjoint merge - average only values agreeing with elected sign
agree_mask = (stacked * elected_sign.unsqueeze(0)) > 0 # (k, d)
masked_sum = (stacked * agree_mask.float()).sum(0)
agree_count = agree_mask.float().sum(0).clamp(min=1.0)
merged_tv = masked_sum / agree_count
return pretrained_weights + lam * merged_tv
theta_ties = ties_merge(
theta_pretrained, [theta_a, theta_b], lam=0.4, trim_pct=20.0
)
model_ties = set_flat_weights(pretrained_model, theta_ties)Notice that agree_count.clamp(min=1.0) prevents division by zero in cases where no model has a nonzero value at some position after trimming. This can happen if all task vectors had small values at a particular position that were all trimmed away.
DARE Merging
DARE applies random masking with rescaling to each task vector independently, then combines with simple averaging.
def dare_merge(
pretrained_weights, finetuned_weights_list, lam=0.4, drop_rate=0.9, seed=42
):
rng = torch.Generator()
rng.manual_seed(seed)
dare_task_vectors = []
for i, fw in enumerate(finetuned_weights_list):
tv = fw - pretrained_weights
mask = (torch.rand(tv.shape, generator=rng) > drop_rate).float()
scaled_tv = tv * mask / (1.0 - drop_rate)
dare_task_vectors.append(scaled_tv)
# After DARE preprocessing, merge with simple averaging
summed = torch.stack(dare_task_vectors).mean(0)
return pretrained_weights + lam * summed
theta_dare = dare_merge(
theta_pretrained, [theta_a, theta_b], lam=0.4, drop_rate=0.9
)
model_dare = set_flat_weights(pretrained_model, theta_dare)The rescaling by 1.0 - drop_rate ensures that the surviving entries carry the same expected magnitude as the original task vector, compensating for the entries that were dropped.
Evaluating All Methods
We evaluate each merged model on both tasks and compare their losses. This gives a concrete picture of how each method balances performance across multiple tasks.
def eval_mse(model, X, y):
with torch.no_grad():
preds = model(X)
return nn.MSELoss()(preds, y).item()
results = {
"Pretrained": (
eval_mse(pretrained_model, X_a, y_a),
eval_mse(pretrained_model, X_b, y_b),
),
"Fine-tuned A": (eval_mse(model_a, X_a, y_a), eval_mse(model_a, X_b, y_b)),
"Fine-tuned B": (eval_mse(model_b, X_a, y_a), eval_mse(model_b, X_b, y_b)),
"Avg": (
eval_mse(model_averaged, X_a, y_a),
eval_mse(model_averaged, X_b, y_b),
),
"Task Arith": (eval_mse(model_ta, X_a, y_a), eval_mse(model_ta, X_b, y_b)),
"TIES": (eval_mse(model_ties, X_a, y_a), eval_mse(model_ties, X_b, y_b)),
"DARE": (eval_mse(model_dare, X_a, y_a), eval_mse(model_dare, X_b, y_b)),
}Method Task A MSE Task B MSE Avg MSE ------------------------------------------------------- Pretrained 12.9741 15.2229 14.0985 Fine-tuned A 0.0008 45.2677 22.6343 Fine-tuned B 46.2103 0.0015 23.1059 Avg 11.2517 11.6108 11.4312 Task Arith 11.3515 12.2900 11.8207 TIES 11.3731 14.8002 13.0866 DARE 14.8266 16.7858 15.8062
The results illustrate the core trade-off in model merging. Individual fine-tuned models excel on their own task but perform close to the pretrained baseline on the other task: they have specialized at the expense of generality. Each merging strategy balances performance across both tasks differently, with more sophisticated methods generally achieving better combined performance by reducing interference between task-specific updates. The pretrained model is a lower bound on merge quality: any merge that performs worse than the pretrained baseline on a task has failed its purpose.
Visualizing Task Vectors
The magnitude distribution of task vector entries reveals the core property that makes DARE and TIES effective: most fine-tuning changes are small and concentrated near zero, while only a minority of parameters receive large updates.
The distribution you see here is typical of fine-tuned models at any scale. The spike near zero represents parameters that fine-tuning barely touched. The long tail represents the relatively small number of parameters that were significantly adjusted to accommodate the task. Trimming in TIES targets this structure explicitly: it keeps the tail and discards the spike. DARE samples randomly from both regions, but because the spike contains far more parameters, most of what it keeps are parameters from the tail when the drop rate is high enough.
Merging Method Comparison
Visualizing the Effect of Lambda
The scaling coefficient controls how aggressively task vectors are applied. Too small, and the merged model barely departs from the pretrained baseline. Too large, and the combined updates push the model into territory where neither task can perform well. Scanning across lambda values reveals the operating range and the optimal point for each method.
lambdas = np.linspace(0.0, 1.0, 21)
lambda_results = {"Task Arith": {"a": [], "b": []}, "TIES": {"a": [], "b": []}}
for lam in lambdas:
t = task_arithmetic(theta_pretrained, [theta_a, theta_b], lam=float(lam))
m = set_flat_weights(pretrained_model, t)
lambda_results["Task Arith"]["a"].append(eval_mse(m, X_a, y_a))
lambda_results["Task Arith"]["b"].append(eval_mse(m, X_b, y_b))
t = ties_merge(theta_pretrained, [theta_a, theta_b], lam=float(lam))
m = set_flat_weights(pretrained_model, t)
lambda_results["TIES"]["a"].append(eval_mse(m, X_a, y_a))
lambda_results["TIES"]["b"].append(eval_mse(m, X_b, y_b))
The curves reveal that both Task Arithmetic and TIES initially improve as lambda grows from zero, with performance gains accelerating as the task vectors contribute more of their signal. However, at some threshold, further increases in lambda cause degradation because the combined updates push the model beyond the region where it performs well on both tasks simultaneously. TIES tends to be more stable at higher lambda values because sign conflicts have been resolved before the merge, reducing the chance that a large lambda amplifies destructive interference.
The practical takeaway is that lambda requires tuning. The optimal value depends on the number of tasks, the degree of overlap between task vectors, and the scale of the task vector magnitudes. A small validation set covering the tasks you care about is usually sufficient to find a good lambda with a simple grid search.
Visualizing DARE Drop Rate Sensitivity
drop_rates = np.linspace(0.0, 0.98, 20)
dare_a_mses = []
dare_b_mses = []
for dr in drop_rates:
t = dare_merge(
theta_pretrained, [theta_a, theta_b], lam=0.4, drop_rate=float(dr)
)
m = set_flat_weights(pretrained_model, t)
dare_a_mses.append(eval_mse(m, X_a, y_a))
dare_b_mses.append(eval_mse(m, X_b, y_b))
The plot confirms that performance is relatively stable across a wide range of drop rates before degrading sharply as the drop rate approaches 1.0. This plateau demonstrates why DARE can aggressively prune task vectors: most entries are redundant, and the rescaling step ensures the surviving entries carry the same expected signal as the full task vector. In practice, drop rates of 0.9 are standard, and rates as high as 0.95 often work well, particularly when combining DARE with TIES.