Distilling Linearized Behavior for Effective Task Arithmetic

TL;DR

This paper introduces a method to distill linearized behavior into non-linear models, enabling effective task vector composition for model editing without additional inference costs.

Contribution

It proposes a novel distillation approach that enforces linearity properties during training, bridging linear and non-linear fine-tuning for better task arithmetic.

Findings

Models inherit linearized task composition properties.

Achieves strong performance on vision and language benchmarks.

No inference-time overhead introduced.

Abstract

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.