Distilling Linearized Behavior into Non-linear Fine-Tuning for Effective Task Arithmetic

Task vector composition has emerged as a promising paradigm for editing pre-trained models, enabling model merging through addition and unlearning through subtraction. Fine-tuning in the tangent space of a pre-trained model (linear fine-tuning) has proven effective, as it produces task vectors that are naturally disentangled and resistant to interference. However, linearized models suffer from limited expressivity during training and incur higher computational costs at inference time, which restrict their practical applicability. In this work, we bridge the gap between linear and standard non-linear fine-tuning. We show that linearity with respect to weight perturbations, a property defined in parameter space, can be enforced through constraints in activation space during training. Concretely, we distill hidden representations from a curvature-regularized linearized teacher into a non-linear student trained via conventional fine-tuning. We find that the resulting model inherits key properties of linearized models for task arithmetic, enabling effective composition of task vectors and achieving strong performance across vision and language benchmarks without incurring any inference-time overhead.

Modern AI models are becoming so large that training a new model, or even adapting an existing one, can be very expensive. A promising alternative is to reuse already trained models by combining or modifying their learned skills, much like adding or subtracting ingredients in a recipe. However, this only works reliably when the updates learned for different tasks do not interfere with each other, which is often not the case with standard fine-tuning. In this work, we study how to make model updates easier to combine while keeping models efficient to use. Previous methods achieve good combinability by training models in a special “linearized” regime, but this makes them slower and can limit their flexibility. We instead train a standard model to imitate the internal behavior of a linearized model, so that it learns updates that are both expressive and easier to compose. We further guide this process using regularization and by distilling behavior along the training path, rather than at a single point. Our results show that strict linearity is not necessary: models only need sufficiently local and approximately linear updates. This makes it possible to combine, remove, and reuse learned abilities more reliably, while retaining the efficiency of standard models at deployment.