Distilling Linearized Behavior for Effective Task Arithmetic

2026-05-18

AI Summary

Existing linear fine-tuning methods struggle to balance representational capacity with inference efficiency and fail to support efficient task vector arithmetic. This work proposes a novel approach that distills the hidden representations of a curvature-regularized linear teacher model into a nonlinear student model by imposing constraints in activation space, thereby emulating linear behavior in parameter space at the activation level for the first time. This enables the nonlinear student to preserve compositional task vectors under standard fine-tuning, supporting zero-overhead task addition and subtraction. Experiments demonstrate that the method achieves strong performance on both vision and language benchmarks while maintaining task arithmetic capabilities and deployment efficiency.

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.

Problem

• task arithmetic
• linear fine-tuning
• model merging
• non-linear fine-tuning
• inference overhead

Innovation

• task arithmetic
• linear fine-tuning
• knowledge distillation
• activation space constraint
• model merging