Experiments: Editing Models with Task Arithmetic

Experiments: Editing Models with Task Arithmetic#

Editing Models with Task Arithmetic (2023) defined “task vectors” as vectors in parameter space that can be added or subtracted from the parameters of a pre-trained model to make it perform better (or worse) on a given target task without affecting its performance on other tasks. This approach is different from fine-tuning on a task, which improves performance on that task but adversely worsens the performance on other tasks. Task vectors for different tasks can be added together to form a new vector that improves the performance of a model for multiple tasks.

While the idea is simple and effective, it rests on a key assumption: task vectors are added or subtracted to a base model that was used to determine the task vector in the first place.

The implication is a task vector calculated for one model A, will not necessarily work for another model B with a different set of parameters. This is expected since two neural networks that are exactly the same function can have different sets of parameters. Models with many parameter combinations for the same function are called “singular models”. Therefore, task vectors for the same task may point towards different directions for model A and model B

Our hypothesis therefore is: “Given two networks with the same architecture and trained on the same dataset but with different parameters will not necessarily share the same task vector for a given task”

If this is the case, then this is a problem because it means we cannot reuse the same task vector on different models with different parameters, even if they share the same architecture.

Experiment Setup#

We train a ResNet18 base model on CIFAR10. We then finetune on MNIST and FashionMNIST to calculate the task vectors. We add the two task vectors and add the resulting vector to the parameters of the base model. We then quantify the accuracy after adding the task vectors. We expect the accuracy on MNIST and FashionMNIST to improve.

We then trained another ResNet18 base model on CIFAR10. We then add the same task vector to see if the accuracy improves. Because neural networks are singular models, we expect a different parameter set for this second base model and therefore we hypothesize that the same task vector won’t work for the second base model.

This limitation has indeed been noted in the reference paper mentioned above.

TLDR Results#

The task vectors worked to improve the performance of the 1st base model on MNIST and FashionMNIST. However, as expected, it did not work for the 2nd base model. These are the accuracies on the train set. (Note we are not concerned with test performance at this point. We only care about the applicability of task vectors across networks)

Model: CIFAR10
CIFAR10: 0.984375
MNIST: 0.078125
Fasion MNIST: 0.078125

Model: CIFAR10 (diff init)
CIFAR10: 0.96875
MNIST: 0.0625
Fasion MNIST: 0.078125

Model: CIFAR10 + MNIST task vector + FashionMNIST task vector
CIFAR10: 0.859375
MNIST: 0.296875
Fasion MNIST: 0.421875

Model: CIFAR10 (diff init) + MNIST task vector + FashionMNIST task vector
CIFAR10: 0.71875
MNIST: 0.046875
Fasion MNIST: 0.078125

We compared the parameters of the two base networks and their cosine similarity is only 0.09. Therefore they are very different and we don’t expect the same task vector to work for both models

Computing and applying task vector on the same model#

Applying task vector on the same base model used to compute the task vector First, for the baseline experiment, we follow the setup in the paper, but we use a simpler network and simpler dataset.

We download ResNet and train it on CIFAR 100 Dataset with data augmentation. We call this \(\theta_0\)
We finetune the network on upright MNIST dataset to get \(\theta_A\)
We finetune the network on upright FashionMNIST dataset to get \(\theta_B\)
We compute the task vectors
- \(T_{A} = \theta_A - \theta_0\)
- \(T_{B} = \theta_B - \theta_0\)
- \(T_{C} = T_A + T_B\)
We add the task vectors to \(\theta_0\) where \(\lambda\) is determined via validation set optimization
- \(\theta_{0A} = \theta_0 + \lambda_A * T_A\)
- \(\theta_{0B} = \theta_0 + \lambda_B * T_B\)
- \(\theta_{0C} = \theta_0 + \lambda_C * T_C\)
We compare the performances of the different models on the following tasks:
- Target tasks (MNIST and FashionMNIST)
- Control task (Imagenet)

from torchvision import transforms 
from torch.utils.data import DataLoader, Subset
from torchvision.datasets import CIFAR10, MNIST, FashionMNIST
import matplotlib.pyplot as plt
import numpy as np

EPOCHS = 20
LR=0.001

# CIFAR 10
transform = transforms.Compose([
    # transforms.RandomCrop(32, padding=4),
    # transforms.RandomHorizontalFlip(),
    # transforms.RandomRotation(45),
    transforms.ToTensor(),
    transforms.Normalize((0.3907, 0.3807, 0.3520), (0.2938, 0.2880, 0.2904))
])

train_data_cifar = CIFAR10(root='../data/cifar10', train=True, download=True, transform=transform)
test_data_cifar = CIFAR10(root='../data/cifar10', train=False, download=True)

# # Temp. only get 256 images for testing
# indices = list(range(64))
# train_data_cifar = Subset(train_data_cifar, indices)
train_loader_cifar = DataLoader(train_data_cifar, batch_size=64, shuffle=True, num_workers=4)

# MNIST
transform = transforms.Compose([
    # transforms.RandomCrop(32, padding=4),
    # transforms.RandomHorizontalFlip(),
    # transforms.RandomRotation(45),
    transforms.Grayscale(num_output_channels=3), # repeat channel
    transforms.ToTensor(),    
    transforms.Normalize((0.1307, 0.1307, 0.1307), (0.3081, 0.3081, 0.3081)),
])

train_data_mnist = MNIST(root='../data/mnist', train=True, download=True, transform=transform)
test_data_mnist = MNIST(root='../data/mnist', train=False, download=True)

# # Temp. only get 256 images for testing
# indices = list(range(64))
# train_data_mnist = Subset(train_data_mnist, indices)
train_loader_mnist = DataLoader(train_data_mnist, batch_size=64, shuffle=True, num_workers=4)

# FashionMNIST
transform = transforms.Compose([
    # transforms.RandomCrop(32, padding=4),
    # transforms.RandomHorizontalFlip(),
    # transforms.RandomRotation(45),
    transforms.Grayscale(num_output_channels=3), # repeat channel
    transforms.ToTensor(),    
    # transforms.Normalize((0.1307, 0.1307, 0.1307), (0.3081, 0.3081, 0.3081)),
])

train_data_fmnist = FashionMNIST(root='../data/fashionmnist', train=True, download=True, transform=transform)
test_data_fmnist = FashionMNIST(root='../data/fashionmnist', train=False, download=True)

# # Temp. only get 256 images for testing
# indices = list(range(64))
# train_data_fmnist = Subset(train_data_fmnist, indices)
train_loader_fmnist = DataLoader(train_data_fmnist, batch_size=64, shuffle=True, num_workers=4)

import torch
import torchvision.models as models

device = 'mps'
model_cifar = models.resnet18(weights=None)
model_cifar.fc = torch.nn.Linear(model_cifar.fc.in_features, 10)
model_cifar = model_cifar.to(device)

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model_cifar.parameters(), lr=LR)

model_cifar.train()
for epoch in range(EPOCHS):
    for images, labels in train_loader_cifar:
        optimizer.zero_grad()        

        images = images.to(device)
        labels = labels.to(device)
        outputs = model_cifar(images)
        
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

images = images.to(device)
labels = labels.to(device)
model_cifar.eval()
sum(model_cifar(images).argmax(axis=1) == labels)/len(labels)

tensor(1., device='mps:0')

import copy

# copy base pretrained model
model_mnist = copy.deepcopy(model_cifar)
model_fmnist = copy.deepcopy(model_cifar)

# finetune on mnist
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model_mnist.parameters(), lr=LR)

model_mnist.train()
for epoch in range(EPOCHS):
    for images, labels in train_loader_mnist:
        optimizer.zero_grad()        

        images = images.to(device)
        labels = labels.to(device)
        outputs = model_mnist(images)
        
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

# finetune on fasionmnist
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model_fmnist.parameters(), lr=LR)

model_fmnist.train()
for epoch in range(EPOCHS):
    for images, labels in train_loader_fmnist:
        optimizer.zero_grad()        

        images = images.to(device)
        labels = labels.to(device)
        outputs = model_fmnist(images)
        
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

def print_eval(model):
    model.eval()

    # original model on cifar
    images, labels = next(iter(train_loader_cifar))
    images = images.to(device)
    labels = labels.to(device)
    accuracy = sum(model(images).argmax(axis=1) == labels)/len(labels)
    accuracy = accuracy.cpu().detach().numpy()
    print("CIFAR10:", accuracy)

    # performance of original model on mnist
    images, labels = next(iter(train_loader_mnist))
    images = images.to(device)
    labels = labels.to(device)
    accuracy = sum(model(images).argmax(axis=1) == labels)/len(labels)
    accuracy = accuracy.cpu().detach().numpy()
    print("MNIST:", accuracy)

    # performance of original model on fashionmnist
    images, labels = next(iter(train_loader_fmnist))
    images = images.to(device)
    labels = labels.to(device)
    accuracy = sum(model(images).argmax(axis=1) == labels)/len(labels)
    accuracy = accuracy.cpu().detach().numpy()
    print("Fasion MNIST:", accuracy)

# get the task vector by subtracting parameters
mnist_params = [i for i in model_mnist.parameters()]
fmnist_params = [i for i in model_fmnist.parameters()]
cifar_params = [i for i in model_cifar.parameters()]
task_vector_mnist = [] # task vector
task_vector_fmnist = [] # task vector
for i in range(len(mnist_params)):
    task_vector_mnist += [mnist_params[i]-cifar_params[i]]
    task_vector_fmnist += [fmnist_params[i]-cifar_params[i]]

# add mnist and fmnist params then add to the cifar model
task_vector_sum = []
lambda1 = 0.1
lambda2 = 0.2
for i in range(len(mnist_params)):
    task_vector_sum += [task_vector_mnist[i] * lambda1 + task_vector_fmnist[i] * lambda2]

# add the task vector sum to cifar
model_cifar_add = copy.deepcopy(model_cifar)
with torch.no_grad():
    for i, param in enumerate(model_cifar_add.parameters()):
        # add the task vector
        param.add_(task_vector_sum[i])

print("Model: CIFAR10")
model = copy.deepcopy(model_cifar)
print_eval(model)
print()
print("Model: CIFAR10 --> Finetuned MNIST")
model = copy.deepcopy(model_mnist)
print_eval(model)
print()
print("Model: CIFAR10 --> Finetuned FashionMNIST")
model = copy.deepcopy(model_fmnist)
print_eval(model)

Model: CIFAR10
CIFAR10: 0.953125
MNIST: 0.078125
Fasion MNIST: 0.109375

Model: CIFAR10 --> Finetuned MNIST
CIFAR10: 0.234375
MNIST: 1.0
Fasion MNIST: 0.09375

Model: CIFAR10 --> Finetuned FashionMNIST
CIFAR10: 0.140625
MNIST: 0.0625
Fasion MNIST: 0.984375

print("Model: CIFAR10 + MNIST task vector + FashionMNIST task vector")
model = copy.deepcopy(model_cifar_add)
print_eval(model)

Model: CIFAR10 + MNIST task vector + FashionMNIST task vector
CIFAR10: 0.796875
MNIST: 0.390625
Fasion MNIST: 0.25

Computing the task vector on one base model and applying it on a different model#

# new cifar model
model_cifar_2 = models.resnet18(weights=None)
model_cifar_2.fc = torch.nn.Linear(model_cifar_2.fc.in_features, 10)
model_cifar_2 = model_cifar_2.to(device)

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model_cifar_2.parameters(), lr=LR)

model_cifar_2.train()
for epoch in range(EPOCHS):
    for images, labels in train_loader_cifar:
        optimizer.zero_grad()        

        images = images.to(device)
        labels = labels.to(device)
        outputs = model_cifar_2(images)
        
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

# check performance of cifar model 2
print("Model: CIFAR10 (diff init)")
model = copy.deepcopy(model_cifar_2)
print_eval(model)

Model: CIFAR10 (diff init)
CIFAR10: 0.9375
MNIST: 0.078125
Fasion MNIST: 0.078125

# add the task vector sum to cifar model 2
model_cifar_add_2 = copy.deepcopy(model_cifar_2)
with torch.no_grad():
    for i, param in enumerate(model_cifar_add_2.parameters()):
        # add the task vector
        param.add_(task_vector_sum[i])

We see that the task vectors work for the first model trained on CIFAR10, but doesn’t work for the 2nd model trained on CIFAR10 even if they have comparable performance

print("Model: CIFAR10")
model = copy.deepcopy(model_cifar)
print_eval(model)
print()
print("Model: CIFAR10 (diff init)")
model = copy.deepcopy(model_cifar)
print_eval(model)
print()
print("Model: CIFAR10 + MNIST task vector + FashionMNIST task vector")
model = copy.deepcopy(model_cifar_add)
print_eval(model)
print()
# check performance of cifar model 2
print("Model: CIFAR10 (diff init) + MNIST task vector + FashionMNIST task vector")
model = copy.deepcopy(model_cifar_add_2)
print_eval(model)

Model: CIFAR10
CIFAR10: 0.984375
MNIST: 0.078125
Fasion MNIST: 0.078125

Model: CIFAR10 (diff init)
CIFAR10: 0.96875
MNIST: 0.0625
Fasion MNIST: 0.078125

Model: CIFAR10 + MNIST task vector + FashionMNIST task vector
CIFAR10: 0.859375
MNIST: 0.296875
Fasion MNIST: 0.421875

Model: CIFAR10 (diff init) + MNIST task vector + FashionMNIST task vector
CIFAR10: 0.71875
MNIST: 0.046875
Fasion MNIST: 0.078125

Comparing parameters with the same architecture and training configuration#

There is very low cosine similarity between the parameters. Therefore we don’t expect the same task vector for one model with one set of parameters to apply to another model with a different set of parameters

# compare the similarity of the parameters of cifar models 1 and 2
params_cifar1 = torch.cat([i.view(-1) for i in model_cifar.parameters()]).cpu().detach().numpy()
params_cifar2 = torch.cat([i.view(-1) for i in model_cifar_2.parameters()]).cpu().detach().numpy()

# cosine sim
dot_product = np.dot(params_cifar1, params_cifar2)
# l2 norm
norm_A = np.linalg.norm(params_cifar1)
norm_B = np.linalg.norm(params_cifar2)

# Calculate cosine similarity
cosim = dot_product / (norm_A * norm_B)

print("Cosine Similarity:", cosim)

Cosine Similarity: 0.09187909