Distilling Linearized Behavior for Effective Task Arithmetic
This work bridges the gap between linear and non-linear fine-tuning for task arithmetic, making the benefits of linearized models practically accessible without computational drawbacks.
The authors propose a method to distill linearized behavior from a teacher model into a non-linear student, enabling effective task arithmetic without inference-time overhead. They achieve strong performance across vision and language benchmarks, matching or exceeding linearized models.
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.