Model Merging Applications › Implementation
Language AI HandbookAll chapters Model Merging Applications › ImplementationOne section · Sections in this chapter6 / 10 ↓
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
This section provides a practical implementation of multi-task merging and evaluation using PyTorch. We use small simulated parameter vectors suitable for a local machine, but the same patterns apply at any scale with real model checkpoints.
Setup and Data Preparation
We start by defining simulated model weights representing a base model and two independently fine-tuned models. The simulation uses sparse task vectors (each fine-tune modifies only a subset of parameters) to mimic the behavior of real fine-tunes, where gradient updates are concentrated on the parameters most relevant to each task.
import numpy as np
import torch
np.random.seed(42)
torch.manual_seed(42)
# Simulate base model weights as a simple flat parameter vector
# In practice, this would be loaded from a checkpoint
base_dim = 1000
theta_base = torch.randn(base_dim) * 0.1
# Simulate two task-specific fine-tunes: sentiment and NLI
# Each fine-tune shifts a different subset of parameters
sentiment_delta_indices = np.random.choice(base_dim, size=300, replace=False)
nli_delta_indices = np.random.choice(base_dim, size=300, replace=False)
# Task 1: Sentiment fine-tune
tau_sentiment = torch.zeros(base_dim)
tau_sentiment[sentiment_delta_indices] = torch.randn(300) * 0.3
# Task 2: NLI fine-tune
tau_nli = torch.zeros(base_dim)
tau_nli[nli_delta_indices] = torch.randn(300) * 0.3
# Fine-tuned models
theta_sentiment = theta_base + tau_sentiment
theta_nli = theta_base + tau_nliSentiment task vector norm: 5.563 NLI task vector norm: 5.141 Overlapping parameters: 89 of 300 Cosine similarity: 0.0396
The cosine similarity between the two task vectors indicates how much interference to expect. A near-zero value suggests the two fine-tunes modified largely orthogonal parameters, which is favorable for merging. The overlap count tells you how many parameters both fine-tunes modified simultaneously, which is the region where sign conflicts can occur.
Task Arithmetic Merging
We implement the basic task arithmetic merge across a range of scaling coefficients to trace the performance tradeoff. The evaluate_performance function acts as a proxy for real benchmark evaluation by measuring how close the merged model is to the target fine-tune in parameter space. In a real workflow, you would replace this with actual benchmark scoring.
def task_arithmetic_merge(theta_base, tau_1, tau_2, lambda_1, lambda_2):
"""Merge two task vectors into the base model."""
return theta_base + lambda_1 * tau_1 + lambda_2 * tau_2
def evaluate_performance(theta_merged, theta_target, noise_std=0.05):
"""
Proxy for task performance: similarity to the target fine-tune.
In practice, this would be replaced by a real benchmark evaluation.
"""
diff = torch.norm(theta_merged - theta_target).item()
max_diff = torch.norm(theta_target - theta_base).item()
similarity = max(0.0, 1.0 - diff / max_diff)
return similarity
# Sweep lambda from 0.1 to 1.0 with equal weights for both tasks
lambda_values = np.linspace(0.1, 1.0, 20)
sentiment_scores = []
nli_scores = []
for lam in lambda_values:
theta_merged = task_arithmetic_merge(
theta_base, tau_sentiment, tau_nli, lam, lam
)
s_score = evaluate_performance(theta_merged, theta_sentiment)
n_score = evaluate_performance(theta_merged, theta_nli)
sentiment_scores.append(s_score)
nli_scores.append(n_score)
sentiment_scores = np.array(sentiment_scores)
nli_scores = np.array(nli_scores)Best combined lambda: 0.48 Sentiment score at best lambda: 0.3298 NLI score at best lambda: 0.2798 Score at lambda=1.0 (sentiment): 0.0758, NLI: 0.0000
The best combined lambda balances the two tasks. At lambda = 1.0, the merge uses the full task vectors, which may cause interference depending on overlap between the fine-tunes. Notice that the optimal lambda is often not 1.0: scaling the task vectors down slightly reduces interference while retaining most of the capability contribution from each fine-tune.
TIES-Style Conflict Resolution
Next, we implement a simplified version of TIES merging, which resolves sign conflicts before summing task vectors. The three-step procedure (trim, elect, disjoint merge) handles the case where multiple task vectors modify the same parameter in opposite directions.
def ties_merge(theta_base, task_vectors, lambda_val, trim_fraction=0.2):
"""
Simplified TIES merging:
1. Trim: zero out small-magnitude parameters in each task vector
2. Elect: resolve sign conflicts by majority vote
3. Disjoint merge: sum only the elected parameters
"""
trimmed = []
for tau in task_vectors:
# Step 1: Trim - zero out bottom trim_fraction by magnitude
threshold = torch.quantile(torch.abs(tau[tau != 0]), trim_fraction)
tau_trimmed = tau.clone()
tau_trimmed[torch.abs(tau_trimmed) < threshold] = 0.0
trimmed.append(tau_trimmed)
# Step 2: Elect - for each parameter, determine dominant sign
stacked = torch.stack(trimmed, dim=0) # shape: (num_tasks, base_dim)
sign_sum = torch.sign(stacked).sum(dim=0)
elected_sign = torch.sign(sign_sum)
elected_sign[sign_sum == 0] = 0.0
# Step 3: Disjoint merge - sum task vectors that agree with elected sign
merged_delta = torch.zeros_like(theta_base)
for tau in trimmed:
agrees = (torch.sign(tau) == elected_sign) | (tau == 0)
contribution = tau * agrees.float()
merged_delta += contribution
return theta_base + lambda_val * merged_delta
# Evaluate TIES across lambda values
ties_sentiment_scores = []
ties_nli_scores = []
for lam in lambda_values:
theta_merged_ties = ties_merge(theta_base, [tau_sentiment, tau_nli], lam)
s_score = evaluate_performance(theta_merged_ties, theta_sentiment)
n_score = evaluate_performance(theta_merged_ties, theta_nli)
ties_sentiment_scores.append(s_score)
ties_nli_scores.append(n_score)
ties_sentiment_scores = np.array(ties_sentiment_scores)
ties_nli_scores = np.array(ties_nli_scores)TIES best lambda: 0.48 TIES best combined score: 0.5919 Task arithmetic best combined: 0.6096 TIES improvement: -0.0177
TIES merging reduces interference by eliminating conflicting signals before aggregation. The improvement is most pronounced when the two task vectors have substantial parameter overlap with opposing signs. When the task vectors are nearly orthogonal (low overlap), TIES and task arithmetic tend to produce similar results because there are few conflicts to resolve.
Normalized Performance Evaluation
We now compute the normalized performance score introduced earlier, which expresses the merged model's performance as a fraction of the gap between the base and fine-tuned models. This normalization makes it easy to compare results across different evaluation setups.
def normalized_score(score_merged, score_base, score_ft):
"""Compute normalized performance retaining the fine-tune's gain."""
if score_ft == score_base:
return 0.0
return (score_merged - score_base) / (score_ft - score_base)
# Base model performance (no task vectors)
base_sentiment_score = evaluate_performance(theta_base, theta_sentiment)
base_nli_score = evaluate_performance(theta_base, theta_nli)
# Full fine-tune performance (ceiling)
ft_sentiment_score = evaluate_performance(theta_sentiment, theta_sentiment)
ft_nli_score = evaluate_performance(theta_nli, theta_nli)
# Normalized scores for best task arithmetic and TIES configurations
ta_norm_sentiment = normalized_score(
sentiment_scores[best_combined_idx],
base_sentiment_score,
ft_sentiment_score,
)
ta_norm_nli = normalized_score(
nli_scores[best_combined_idx], base_nli_score, ft_nli_score
)
ties_norm_sentiment = normalized_score(
ties_sentiment_scores[ties_best_idx],
base_sentiment_score,
ft_sentiment_score,
)
ties_norm_nli = normalized_score(
ties_nli_scores[ties_best_idx], base_nli_score, ft_nli_score
)Normalized Performance Scores (1.0 = fine-tune level, 0.0 = base level) Method Sentiment NLI Average -------------------------------------------------------- Task Arithmetic 0.330 0.280 0.305 TIES Merge 0.321 0.271 0.296
A normalized score of 1.0 means the merged model is as good as training a separate fine-tune for that task. Scores above 0.8 are generally considered acceptable for production use cases. The average normalized score across tasks provides a single figure of merit for comparing merge configurations, and it can be extended straightforwardly to three or more tasks by simply averaging over all task scores.
Key Parameters
The key parameters for multi-task merging are:
- lambda (): The scaling coefficient for task vectors. Higher values apply more of each fine-tune's changes. Typical range: 0.3 to 0.8 for two-task merges.
- trim_fraction (TIES): The fraction of low-magnitude parameters to zero out before merging. Higher values reduce interference but may discard useful small adaptations. Typical range: 0.1 to 0.3.
- density (DARE): The fraction of task vector parameters to retain after random sparsification. Lower values reduce interference but require higher lambda to compensate. Typical range: 0.5 to 0.9.