Some forewords

I’m currently experimenting with different ways to produce content in the hope to be able to produce more than one blog post a year in 2024. Expect some variations in my content style and feel free to share any advices or feedbacks you could have.

Editing Models With Task Arithmetic

Recently I came across mergekit a “tools for merging pretrained large language models”, this picked my curiosity about how exactly does one create a new model by “merging” or “fusing” multiple models into one. Naively I though MoE was the only way to achieve that, but god I was wrong… It appears there are a bunch of different technics available to merge or fuse models into a one.

Today I will cover Editing Models With Task Arithmetic, the first paper I read one the subject will investigating the art of merging/fusing models. No prerequire is needed in order to understand this blog post except having basic knowledges on Deep Learning.

Task vectors

This paper explores the idea of being able to transmit knowledge from multiple specialized models sharing the same architecture (i.e finetuned) into one final model that is capable of leveraging the collective expertise of those diverse domains.

In order to transmit knowledge from a specialized model (i.e fine-tuned) into another target model, the authors create task vector, which serves as vectorized representation encapsulating task-specific knowledge from one model.

The knowledge encapsulated in this task vector can later be integrated into the base model, enhancing and incorporating the acquired information.

To create task vectors $T_t$ the authors simply substracted the base model’s weight parameters $\theta_{b}$ from the fine-tuned model’s weight parameters $\theta_{ft}$ like follow: $$ T_t = \theta_{ft} - \theta_{b} $$

Learning via addition

To transmit the knowledge using a generated task vector to a base mode, one can simply revert the formula as follow: $$ \theta_{b} + T_t = \theta_{ft} $$

But what’s intresting is that one can also do the same with multiple task vectors at once!

Consider two fine-tuned models with weight parameters denoted as $\theta_{ft}^{t1}$ and $\theta_{ft}^{t2}$, fine-tuned on tasks $T_{1}$ and $T_{2}$, respectively. Enhancing the performance of a base model on both $T_{1}$ and $T_{2}$ can be achieved by generating task vectors from $\theta_{ft}^{t1}$ and $\theta_{ft}^{t2}$ and adding them into the weight parameters of the base model, designated as $\theta_{b}$: $$T_{1} = \theta_{ft}^{t1} - \theta_{b}$$ $$T_{2} = \theta_{ft}^{t2} - \theta_{b}$$ $$\theta_{b}^{t_1+t_2} = \theta_{b} + T_{1} + T_{2}$$

Figure 1.c: Learning via addition


Let’s quickly go over one of the examples showcased in the paper where they fuse eight different models into one and see if the resulting models can perform on all those eight tasks.

Each of the eight models is a CLIP model fine-tuned on a task-specific data, encompassing satellite imagery, traffic signs, and even MNIST. Task vectors were generated for each of these models, and subsequently, they were fused to a unified CLIP model.

And you know what’s interesting? The performance of the resulting model improves as more task vectors are incorporated into the base model for each of those eight tasks!

Figure 3: Adding task vectors builds multi-task models

In the depicted plot, you can observe the averaged accuracy of the resulting model as task vectors are progressively added to it. Notably, the more task vectors are added, the better the resulting model performs.

And now, what if I tell you that we can also do the other way around?

Forgetting via negation

Suppose you intend for your model to selectively forget certain abilities, whether for ethical considerations or regulatory compliance. Just as you can teach new tasks to your model using task arithmetic, you can employ the same principle in the opposite direction.

In this case, instead of adding a task vector, you would remove it to achieve the desired forgetting effect:

$$T_{t} = \theta_{ft} - \theta_{b}$$ $$\theta_{b}^{-T_t} = \theta_{b} - T_{t}$$

The remarkable aspect is that this process results in only minimal changes to the control tasks!

Figure 1.b: Forgetting via negation


Let’s revisit the previous example where eight CLIP models were fused into one. However, in this iteration, the goal is not to create a multi-task model. Instead, the objective is to diminish accuracy across all target tasks while preserving the performance on the control task, namely ImageNet. To achieve this, rather than adding each task vector, they opt to subtract it.

To gauge the effectiveness of this technique, they compare it against two alternative approaches:

  • Subtracting random task vectors from the target model.
  • Employing gradient ascent on the target tasks.

Once again, the results speak for themselves:

Table 1: Forgetting image classification tasks via negation

Gradient ascent significantly degrades the performance on the control task, while subtracting random task vectors shows no significant impact on the accuracy of the target tasks. In contrast, the negative task vectors successfully diminishe the accuracy on the target tasks while having a minimal impact on the control task. Once again, the outcomes are truly impressive.

Going even further: Task analogies

The final application demonstrated in this paper introduces task analogies. Consider three tasks, $A$, $B$, $C$, and $D$, where “$A$ is to $B$ as $C$ is to $D$”. In the absence of specific data for task $D$, one can enhance the accuracy of the target model by utilizing tasks $A$, $B$, and $C$ with the following equation:

$$T_{D} = T_{C} + (T_{B} - T_{A})$$

To illustrate, suppose you aim to train your model to detect indoor lions. Acquiring labeled data specifically for indoor lions might be challenging, but obtaining labeled data for another indoor animal, such as dogs, could be more feasible. Therefore, to teach a model to detect indoor lions, one could create a task vector and add it to the base model as follows:

$$T_{lionIndoors} = T_{lionOutdoors} + (T_{dogIndoors} - T_{dogOutdoors})$$ $$\theta_{b}^{lionIndoors} = \theta_{b} + T_{lionIndoors}$$

This concept bears resemblance to the behavior observed when working with embeddings, such as in word2vec. The parallels are quite remarkable, I find it truly amazing!

Figure 1.d: Task analogies


Conclusion

This wrap up the mains ideas behind Editing Models With Task Arithmetic. If you found this blog post intriguing, you can stay updated on my latest content by following me on Twitter. Additionally, for a glimpse into my most recent projects and developments, visit my Github. Your continued interest and support are greatly appreciated!