Dataless Weight Disentanglement in Task Arithmetic via Kronecker-Factored Approximate Curvature

AI summary

Overview

• Research area: Task arithmetic (TA) and model merging for foundation models, combined with second-order optimization tools (generalized Gauss-Newton curvature, Kronecker-Factored Approximate Curvature, KFAC).
• Technical level: Advanced.
• Scope: The paper derives a data-free regularizer, called TAK, that improves weight disentanglement in linearized fine-tuning so that task vectors can be added or subtracted without cross-task interference.

What This Paper Is About

Combining multiple fine-tuned task vectors into one model (or subtracting them to remove a capability) often degrades performance because task-specific updates interfere with each other, a failure mode called representation drift. Prior fixes regularize against this drift but need access to the other tasks' training data, which breaks modularity and conflicts with privacy constraints. This paper reframes representation drift regularization as a curvature-approximation problem solvable with pre-computed curvature factors instead of data.

Key Contributions

1. A data-free regularizer for task arithmetic (TAK). The authors show that under model linearization, representation drift becomes a quadratic form of the network Jacobian's Gramian, which is an instance of the generalized Gauss-Newton (GGN) matrix. They approximate this with KFAC to obtain a practical regularizer that needs no external task data.
2. Constant-complexity regularization merging. Because a naive KFAC regularizer must store Kronecker factors per task, costing O(T) memory and run time in the number of tasks T, they propose an aggregation scheme (Eq. 8) that merges per-task factors into a single surrogate, giving O(1) complexity and storage.
3. State-of-the-art results in task addition and negation. TAK improves weight disentanglement on the 8 Vision benchmark with CLIP backbones and on six language tasks with T5-base, in both linearized and (when paired with attention-only fine-tuning) non-linear regimes.
4. Robustness to task-vector rescaling. The method's performance is insensitive to the scaling coefficient α, eliminating the need for held-out tuning or a cross-task validation set.

Main Findings

• Task addition on 8 Vision improves over the unregularized baseline. With coefficients fixed at α = 1.0 (no tuning), TAK reaches absolute/normalized accuracy of 85.8/97.6 on ViT-B/32, 88.3/97.9 on ViT-B/16, and 91.6/99.3 on ViT-L/14. The best results with tuned α are 86.0/97.8, 88.3/98.1, and 91.6/99.3 respectively.
• Comparable to the data-hungry competitor without using data. τJp (Yoshida et al., 2025), which is not dataless, achieves 85.0/97.4 (ViT-B/32), 88.2/98.3 (ViT-B/16), and 90.9/98.3 (ViT-L/14) at α = 1.0, with best values of 85.6/98.2, 88.6/98.7, and 91.1/98.5. TAK matches or exceeds this on most entries while remaining dataless.
• Stronger curvature approximations help. A diagonal GGN approximation (Porrello et al., 2025) reaches 80.1/92.3, 82.9/93.2, and 87.9/96.3 at α = 1.0, better than plain linear fine-tuning but below KFAC.
• Plain linear fine-tuning is a weak baseline. Linear FT reaches 76.7/87.2 (ViT-B/32), 80.2/88.9 (ViT-B/16), and 88.0/94.8 (ViT-L/14) at α = 1.0, with best values 78.8/89.9, 82.0/90.9, and 88.0/94.8.
• Benefits extend to the non-linear regime when paired with attention-only fine-tuning. Attention-only FT alone reaches 22.5/23.3, 22.8/23.4, and 66.2/69.7 at α = 1.0 (best: 78.2/86.3, 80.4/87.1, 88.2/93.8). Adding TAK yields best accuracy of 83.1/91.3, 84.3/91.0, and 89.9/95.9.
• Task negation (unlearning) works better with TAK. Using ImageNet as a control task and reporting minimum target accuracy while preserving at least 95% of pretrained accuracy on control tasks, TAK achieves target/control scores of 3.4/62.4 (ViT-B/32), 3.4/66.4 (ViT-B/16), and 3.5/72.6 (ViT-L/14). τJp gets 6.7/60.8, 4.7/66.0, and 3.7/73.0; TaLoS gets 11.0/60.7, 10.6/66.1, and 10.7/73.6.
• Language results with T5-base favor TAK. Across six natural language tasks (SNLI, MultiNLI, SICK, SciTail, RTE, QNLI), TAK reaches 78.7 absolute / 98.9 normalized, versus 81.3/100 for τJp, 76.3/93.4 for TaLoS, 76.9/92.8 for linear FT, 72.9/85.2 for attention-only FT, and 85.9 for individually fine-tuned models. The authors note that using other tasks' data (τJp) still yields additional gains, suggesting textual domains may benefit from even more accurate curvature estimation.
• The merging heuristic costs nothing in accuracy. Comparing naive multi-task FT (O(T)) against accumulated regularization (TAK, O(1)): naive reaches 86.5/98.4 (ViT-B/32), 88.0/97.5 (ViT-B/16), 78.5/97.0 (T5-base) at α = 1.0, with best values 86.6/98.5, 88.1/97.6, and 78.5/97.0. TAK reaches 85.8/97.6, 88.3/97.9, and 78.6/98.7, with best values 86.0/97.8, 88.3/98.1, and 78.7/98.9.
• Robustness across the α sweep. In an α-sweep over [0, 2] on ViT-B/32, KFAC-regularized models with simple task-vector summation show the best peak performance and stable accuracy across a wide interval, unlike post-hoc merging strategies TIES (Yadav et al., 2023), TSV (Gargiulo et al., 2025), and ISO (Marczak et al., 2025). The authors note TAK also avoids the SVD computations required by ISO and TSV.
• Task localization. Figure 5 reports the distribution of ‖J_θ f(x, θ₀) τ_t‖₂² for inputs from a task's own training distribution (inliers) versus other tasks (outliers), under regularized and non-regularized fine-tuning.

Methodology in Plain English

The authors start from the observation that linearizing a network around its pre-trained weights makes task arithmetic work better. Under that linearization, the "representation drift" caused by adding one task vector while another task is active reduces to a simple squared norm: the Jacobian of the model times the new task vector. Averaging the squared Jacobian over data gives a Gram matrix, which turns the drift penalty into a quadratic form in the task vector.

That Gram matrix is far too large to build, but the authors recognize it as a special case of the generalized Gauss-Newton (GGN) matrix — the same object studied in second-order optimization — because it corresponds to using the squared loss as the criterion. This lets them borrow KFAC, which approximates the GGN as block-diagonal, one block per layer, with each block written as a Kronecker product of two small matrices: the input covariance A and the output-gradient covariance B. Instead of sharing data, one can pre-compute and share these small factors.

To handle many tasks, the paper notes that storing one set of factors per task costs O(T). It proposes merging the factors from all other tasks into a single Kronecker product — summing the B factors and summing the λ-weighted A factors — which keeps cost constant while empirically matching the unmerged formulation. The task-specific regularization strength λ_t is set proportional to dataset size, and β controls overall regularization strength. The final recipe is a linearized fine-tuning loop that minimizes the task loss plus the merged curvature penalty, optionally using the Monte-Carlo KFAC variant (M backpropagations per datum, typically M = 1) rather than the exact variant (C backpropagations per datum) to reduce cost.

Why This Matters

• Impact on research: It connects task arithmetic to the second-order optimization literature, showing that a well-studied curvature approximation (KFAC) can substitute for data sharing in multi-task model composition. It also provides evidence that better curvature approximations produce larger gains in task arithmetic, and that in-training regularization is complementary to post-hoc merging methods.
• Privacy-preserving model customization: Organizations that cannot share customer or regulated data can exchange only pre-computed curvature factors to build disentangled task vectors.
• Decentralized or federated fine-tuning: Constant complexity in the number of tasks makes the approach practical when many participants each contribute a task without centralizing data.
• Machine unlearning: Task negation results show selective removal of a capability (target accuracy drops to 3.4–3.5 on the CLIP backbones) while retaining at least 95% of control-task accuracy, without needing to store or transfer a large control dataset such as ImageNet.
• Deployment without validation data: Insensitivity to α means practitioners can simply add task vectors with α = 1 without holding out a cross-task validation set.
• Industry relevance: Foundation-model vendors and platforms that ship modular adapters benefit from composing them reliably without retraining or exchanging training data; the constant-memory cost also matters when serving many adapters at once.

Future Directions

• Improving curvature estimation for language models: the authors observe that τJp's access to other tasks' data still yields gains on T5-base, implying textual domains may need more accurate curvature estimates than KFAC provides.
• A more principled treatment of the non-linear regime: the regularizer is derived under linearization, and its use with attention-only fine-tuning is justified empirically rather than theoretically.
• Tighter analysis of the accumulated-regularizer heuristic in Eq. 8, which merges per-task Kronecker factors; the paper states it matches the unmerged formulation empirically but does not derive it exactly.
• Further study of task localization, using the inlier/outlier distributions of ‖J_θ f(x, θ₀) τ_t‖₂² shown in Figure 5, to characterize when distinct task vectors occupy separate regions of function space.

Target Audience

Researchers and engineers working on model merging, task arithmetic, and multi-task adaptation of foundation models; practitioners in privacy-sensitive or decentralized settings who need modular, retraining-free model composition; and readers with a background in second-order optimization or Kronecker-factored curvature approximations who want to see them applied outside of training.

Authors’ abstract

Task Arithmetic yields a modular, scalable way to adapt foundation models. Combining multiple task vectors, however, can lead to cross-task interference, causing representation drift and degraded performance. Representation drift regularization provides a natural remedy to disentangle task vectors; however, existing approaches typically require external task data, conflicting with modularity and data availability constraints (e.g., privacy requirements). We propose a dataless approach by framing regularization against representation drift as a curvature matrix approximation problem. This allows us to leverage well-established techniques; in particular, we adopt Kronecker-Factored Approximate Curvature and obtain a practical regularizer that achieves state-of-the-art results in task addition and negation. Our method has constant complexity in the number of tasks and promotes robustness to task vector rescaling, eliminating the need for held-out tuning.