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 via addition and task removal via subtraction. Fine-tuning in the tangent space of a pre-trained model (linearized fine-tuning) has proven particularly effective in this setting, as it yields task vectors that are naturally disentangled and less prone to interference. However, linearized models suffer from reduced expressivity during training and increased computational cost at inference time, limiting their practical applicability. In this work, we propose to bridge linearized and standard non-linear fine-tuning through knowledge distillation. Specifically, we distill hidden representations from a linearized, curvature-regularized teacher into a non-linear student trained with conventional fine-tuning. By doing so, the goal is to transfer the disentanglement properties of the linearized regime, biasing optimization toward solutions that are composable by design. We show that the resulting task vectors can be composed using naïve Task Arithmetic, achieving strong results across vision and language benchmarks without incurring any inference-time overhead.