PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding and Enforcing Weight Disentanglement in Task Arithmetic

Shangge Liu, Yuehan Yin, Lei Wang, Qi Fan, Yinghuan Shi, Wenbin Li, Yang Gao, Dacheng Tao

arXiv:2604.17078v1cs.AI

TL;DR

The paper explains task arithmetic through Task-Feature Specialization, which supports weight disentanglement and produces weight vector orthogonality. It proposes OrthoReg to enforce orthogonality on weight updates, with theory and experiments showing consistent performance improvements.

  • Problem

    Task arithmetic lacks a fundamental explanation for why task compositions avoid interference, while conventional joint fine-tuning is costly, data-intensive, and risks catastrophic forgetting.

  • Method

    The paper identifies Task-Feature Specialization as sufficient for weight disentanglement and proposes OrthoReg, which enforces column-wise orthogonality on weight updates as its geometric consequence.

  • Results

    OrthoReg consistently improves performance across baselines and model scales; on ViT-L-14, it raises Non-lin. FT accuracy from 84.07% to 88.23% and TTA accuracy from 86.19% to 87.52%.

  • Takeaways & Limitations

    Weight vector orthogonality provides an observable geometric route for promoting task disentanglement and improving model merging.

  • Takeaways & Limitations

    The theoretical disentanglement guarantee assumes disjoint specialized feature sets, but distinct tasks can share features, producing interference.

Abstract

from arXiv · show

Task arithmetic provides an efficient, training-free way to edit pre-trained models, yet lacks a fundamental theoretical explanation for its success. The existing concept of ``weight disentanglement" describes the ideal outcome of non-interfering task composition but does not reveal its underlying cause. Crucially, what intrinsic properties of the pre-trained model ($θ_0$) or the task vectors ($τ_t$) enable this disentanglement remains underexplored. In this paper, we introduce Task-Feature Specialization (TFS), a model's ability to allocate distinct internal features to different tasks, as the fundamental principle. We first prove that TFS is a sufficient condition for weight disentanglement. More importantly, we find that TFS also gives rise to an observable geometric consequence: weight vector orthogonality. This positions TFS as the common cause for both the desired functional outcome (disentanglement) and a measurable geometric property (orthogonality). This relationship provides the key insight for our method: since the abstract TFS property is intractable to enforce directly, we can instead promote weight disentanglement by shaping its concrete geometric consequence, orthogonality. Therefore, we propose OrthoReg, a simple and effective regularization method that actively enforces an internal orthogonal structure on weight updates ($ΔW$) that constitute $τ_t$ during fine-tuning. And we theoretically prove that OrthoReg promotes disentanglement. Extensive experiments demonstrate that OrthoReg consistently and significantly enhances the performance of various task arithmetic methods. Code is available at \href{https://github.com/RL-MIND/OrthoReg}{https://github.com/RL-MIND/OrthoReg}.

1. Introduction

Task arithmetic offers efficient, training-free model editing, but its success lacks a fundamental explanation. The paper identifies Task-Feature Specialization as the common cause of weight disentanglement and weight vector orthogonality, motivating OrthoReg.

  • Task arithmetic combines task vectors from fine-tuned models without costly joint retraining.Each task vector is a parameter shift from pre-trained weights, and vector arithmetic composes, removes, or relates skills.
  • Weight disentanglement explains ideal task composition by isolating each task vector’s effects to its respective data domain.The paper presents this concept as a partial answer to why task arithmetic works.
  • The paper asks which intrinsic properties of θ0 and construction choices for τt enable task vectors to compose without interference.Standard fine-tuning offers no guarantee that task vectors will compose well.
  • Task-Feature Specialization allocates distinct internal features to tasks, suffices for weight disentanglement, and gives rise to observable weight vector orthogonality.This links the desired functional property and measurable geometric consequence through a common underlying principle.
  • OrthoReg enforces orthogonality on weight updates constituting task vectors because directly enforcing Task-Feature Specialization is intractable.The method is theoretically shown to promote disentanglement.
  • The paper experimentally reports that OrthoReg consistently and significantly improves various task arithmetic methods.

2. Related Work

Related work studies task arithmetic, weight disentanglement, and orthogonality as ways to understand or improve model merging. This paper frames orthogonal weight updates as a direct mechanism for mitigating task interference.

  • Task arithmetic merges specialized models by algebraically manipulating task vectors, but composing multiple tasks can cause destructive interference.Prior solutions operate either during merging or before merging.
  • Tangent Task Arithmetic shows that fine-tuning in linearized tangent space promotes weight disentanglement.Other work studies attention-only fine-tuning, nonlinear Transformer generalization, and bounds requiring nearly orthogonal task vectors.
  • The paper repurposes orthogonality as a constraint on weight updates and demonstrates its effectiveness for mitigating interference in task arithmetic.

3. Preliminaries and Problem Formulation

The paper formulates task arithmetic through parameter shifts and weighted task-vector addition, defines weight disentanglement, and analyzes it under NTK linearization. Its theory focuses on linear-layer parameters and in-domain interference.

  • 3.1. Basic Setup and Notation: A task vector τt is the parameter shift from pre-trained weights θ0 to fine-tuned weights θt*.
  • 3.1. Basic Setup and Notation: Task addition constructs a multi-task model as θMT = θ0 + Σ_t αtτt, where αt scales each task vector.
  • 3.1. Basic Setup and Notation: The theoretical analysis concentrates on linear-layer parameters, including fully connected layers and attention projections.The paper justifies this focus by their role in modern architectures and model merging.
  • 3.2. The Weight Disentanglement Property: Weight disentanglement requires merged behavior on task-domain data to depend only on that task’s vector and to revert to pre-trained behavior out of domain.
  • 3.2. The Weight Disentanglement Property: The paper investigates both inherent properties of θ0 and task-vector construction methods that can promote disentanglement.
  • 3.3. The NTK Linearization Hypothesis: Under the NTK linearization hypothesis, a small parameter change τ approximates the output change by τ^T∇θf(x;θ0).The Jacobian J(x) is defined as the model-output gradient with respect to parameters.

4. The Proposed Framework

The framework identifies Task-Feature Specialization (TFS) as the sufficient functional condition for weight disentanglement and the source of weight vector orthogonality. Because TFS is difficult to enforce directly and may fail under feature overlap, OrthoReg enforces orthogonality on weight updates to promote disentanglement.

  • Task-Feature Specialization: Task-Feature Specialization allocates distinct internal features to different tasks, such as separate feature sets for car classification and MNIST digit classification.For task t, the specialized feature set contains features whose activations affect the model output on that task’s data domain.
  • Task-Feature Specialization: Under the NTK linearization hypothesis, TFS is sufficient for weight disentanglement between tasks.The result formalizes why functionally dedicating distinct features to distinct tasks can prevent cross-task interference.
  • From TFS to Weight Vector Orthogonality: TFS also implies block orthogonality in model weight matrices, with empirical evidence in CLIP ViT-B/16 approaching column-wise orthogonality.Angles between all column-vector pairs are sharply concentrated near 90°, suggesting decorrelation both across and within task features.
  • Orthogonality as a Clue for Disentanglement: Weight vector orthogonality is a measurable signature of the abstract TFS property, but is not itself presented as the direct cause of disentanglement.The paper uses WVO diagnostically: observing it after effective training increases belief that the model has developed a TFS-like structure and will exhibit disentanglement.
  • Limitations of Ideal TFS: The ideal TFS assumption requires disjoint specialized feature sets, but feature overlap can invalidate the TFS-based disentanglement theorem and leave static pre-trained weights insufficient.When tasks share features, the paper shifts attention to actively constructing task vectors that promote disentanglement.
  • OrthoReg: OrthoReg replaces intractable direct enforcement of TFS with column-wise orthogonality regularization on weight updates, which theory shows promotes disentanglement despite feature overlap.The regularizer adds Lortho to the task objective; its strength is controlled by λ, while Δθ represents the parameter updates constituting the task vector.

5. Experiments

Experiments evaluate OrthoReg on task addition, task negation, orthogonality, and hyperparameter sensitivity across CLIP Vision Transformers and eight image-classification datasets. OrthoReg consistently improves merging performance, produces more orthogonal task vectors, and strengthens forgetting while preserving control-task accuracy.

  • Experimental Setup: Experiments use eight image-classification datasets and CLIP-pretrained ViT-B-32, ViT-B-16, and ViT-L-14 models.The text encoder is frozen during fine-tuning.
  • Experimental Setup: A single scaling coefficient α is grid-searched from 0.0 to 1.0 for every method to enable direct comparison of task-vector quality.The evaluation does not use task-adaptive coefficients α_t.
  • Task Addition: OrthoReg consistently improves performance across all baselines and model scales.On ViT-L-14, it raises Non-lin. FT accuracy from 84.07% to 88.23% and TTA accuracy from 86.19% to 87.52%.
  • Task Addition: OrthoReg provides a balanced and widespread performance lift across most individual tasks, rather than only improving the average.Figure 4 reports ViT-L-14 per-task accuracy for zero-shot, baseline merging, and OrthoReg-enhanced merging.
  • Task Addition: Norm.Acc. reaches 100.08% for Non-lin. FT and 100.05% for ATT-FT on ViT-L-14, indicating near-total absence of task interference.The paper interprets normalized accuracy at or above 100% as functional realization of ideal weight disentanglement.
  • Task Negation: OrthoReg strengthens task negation across baselines; for ATT-FT on ViT-L-14, it lowers target-task accuracy by an additional 10.18 percentage points without reducing ImageNet control performance.The experiment requires preserving at least 95% of zero-shot accuracy on the control task.
  • Orthogonality and Sensitivity: OrthoReg makes task vectors more orthogonal and maintains higher accuracy across a wide range of merging coefficients α.Figure 5 shows darker cosine-similarity heatmaps after OrthoReg, while Figure 6 compares regularization and merging-coefficient sensitivity.

6. Conclusion

The paper explains task arithmetic through Task-Feature Specialization, linking it to weight disentanglement and weight-vector orthogonality. This connection motivates OrthoReg, which improves performance by enforcing orthogonality on weight updates.

  • Conclusion: Task-Feature Specialization ensures weight disentanglement and creates weight-vector orthogonality as its geometric consequence.The paper identifies TFS as the common principle underlying both outcomes.
  • Conclusion: OrthoReg promotes disentanglement by enforcing orthogonality on weight updates and significantly improves performance by producing more orthogonal task vectors.The paper proposes exploring additional orthogonality constraints in future work.

Supplementary Material

The theoretical analysis concentrates on linear-layer parameters because they dominate model parameters and align with model-merging practice. This focus yields a tractable framework intended to capture core task-arithmetic mechanisms.

  • Scope of Analysis: The analysis focuses on fully connected layers and attention projection matrices while omitting biases and normalization-layer parameters.The paper states that linear layers constitute the vast majority of parameters in modern Transformer models.
  • Scope of Analysis: Restricting the geometric analysis to linear layers is presented as a tractable yet powerful framework for capturing core task-arithmetic mechanisms.The choice also aligns with established model-merging strategies that apply complex methods mainly to linear layers.

B. Justification for Two-Task Simplification

The appendix justifies reducing the proof to pairwise, two-task in-domain disentanglement: linearity makes vanishing pairwise interference sufficient for multi-task composition, while the out-of-domain case follows by the same logic.

  • Two-Task Simplification: The main analysis simplifies weight disentanglement to a two-task, in-domain condition.The displayed condition compares f(x; θ_0 + τ_t + τ_j) with f(x; θ_0 + τ_t) for x in task t’s domain.
  • Two-Task Simplification: The simplification is defended as sufficient because the proofs target pairwise interference terms under the NTK linearization hypothesis.The appendix explicitly presents this as the core disentanglement mechanism.
  • Multi-Task Extension: If each other task has approximately zero interference with task t, linearity makes their summed interference approximately zero in the merged model.The argument extends pairwise cancellation to multiple task vectors and includes their scaling coefficients.
  • Generalization: Therefore, omitting scaling coefficients and proving the two-task interaction does not compromise the generality of the conclusions.The appendix attributes this to linearity of the interference term with respect to task vectors.
  • Out-of-Domain Extension: The out-of-domain preservation condition follows from the same principle: samples outside the task domains should not rely on specialized task features, making each interference term approximately zero.The appendix presents the pairwise in-domain proof as sufficient reasoning for the full weight-disentanglement property.

C. Proof of Lemma 2

Under NTK linearization, weight disentanglement between tasks is equivalent to the interference from one task vector vanishing on the other task’s data domain. The proof derives this by linearizing the two-task composition and cancelling the shared terms.

  • Conclusion: Therefore, under NTK linearization, functional non-interference is equivalent to task-vector orthogonality with the model’s gradient Jacobian on the other task’s data.The result links the abstract disentanglement property to a geometric condition involving τj and J(x).
  • Proof setup: Weight disentanglement requires f(x; θ0 + τt + τj) ≈ f(x; θ0 + τt) for every x in task t’s domain.This is the simplified two-task functional condition used as the proof’s starting point.
  • Linearization: First-order NTK expansions express the composed model as f(x; θ0) + τt^T J(x) + τj^T J(x), while the task-t model contains only τt^T J(x).The left-hand side uses the total perturbation τt + τj; the right-hand side uses τt alone.
  • Equivalent condition: Cancelling the shared terms yields the equivalent condition τj^T J(x) = 0 for all x in task t’s data domain.Thus, task j’s interference term must be approximately zero on task t inputs.

D.2. Supporting Proposition for Theorem 3

The supporting proposition shows that TFS leaves task-vector components near zero for features outside a task’s specialized set. Combining this with NTK stability makes every cross-task interference term vanish, proving weight disentanglement.

  • Supporting proposition: If feature k is outside task j’s specialized set, its task-vector component (τj)k is approximately zero under NTK linearization and TFS.The task-vector component is the accumulated fine-tuning update for the corresponding weight column.
  • Fine-tuning updates: The proof represents (τj)k as the sum of gradient updates across fine-tuning steps and shows that each expected gradient is approximately zero.The argument uses the chain rule: the loss and input terms are bounded, while output sensitivity to an unused feature is negligible.
  • NTK stability: NTK stability keeps the fine-tuning trajectory near θ0, preserving negligible sensitivity to features outside task j’s specialized set throughout training.This condition is applied for every update step s = 0,1,…,S−1.

E.2. Detailed proof of Corollary 2

The proof connects TFS to block-orthogonal weights through feature-activation covariance. With whitened layer inputs, the activation covariance equals W^T W, and cross-task feature decorrelation makes this matrix block-diagonal.

  • Weight–covariance link: The proof first relates the weight Gram matrix W^T W to the covariance matrix Σz of feature activations.For a linear layer, z = W^T In(x), so input whitening reduces the activation covariance to the weight Gram matrix.
  • Whitening assumption: Normalization is assumed to whiten layer inputs, giving ΣIn ≈ I_m and μIn ≈ 0.Under these conditions, the input second-moment matrix is approximately I_m.
  • Weight–covariance link: This yields Σz = W^T W, making block-orthogonality equivalent to a block-diagonal activation covariance.The Gram-matrix entry for two columns is their inner product.
  • Block structure: TFS makes activations of features specialized for different tasks approximately decorrelated, so cross-task covariance entries are near zero.The theory does not require features specialized for the same task to be decorrelated.
  • Conclusion: Consequently, weight columns assigned to different tasks span approximately orthogonal subspaces, while within-task blocks may remain non-diagonal.This is the block-orthogonality conclusion of Corollary 2.

F. Bayesian Analysis of the Relationship between TFS, WVO, and WD

The Bayesian analysis treats TFS as the common cause of weight disentanglement and weight-vector orthogonality. Because WVO raises the posterior belief in TFS, it also raises the posterior probability of WD under the stated assumptions.

  • Event definitions: The analysis defines TFS as Event A, weight disentanglement as Event B, and weight-vector orthogonality as Event C.TFS is the underlying feature-allocation property, WD is the non-interfering composition outcome, and WVO is the measurable geometric property.
  • Causal structure: TFS is assumed sufficient for both WD and WVO, giving P(B|A) = 1 and P(C|A) = 1.The logical implication is A ⇒ (B ∧ C).
  • Posterior WD: Under the additional assumption that WD and WVO are approximately conditionally independent when TFS is absent, WVO increases the probability of WD: P(B|C) > P(B).The assumption is represented by P(B|¬A,C) ≈ P(B|¬A), with the residual probability denoted q < 1.
  • Bayesian update: Conditioning on WVO increases belief in TFS because WVO is treated as a non-trivial geometric structure rather than an arbitrary-network default.The analysis assumes P(C) < 1 and derives P(A|C) > P(A).
  • Implication: The result supports using WVO as a diagnostic for whether a model is suitable for task arithmetic.This conclusion follows from WVO providing evidence for the abstract functional property of WD.

G.1. Proof of Theorem 4

The proof shows that internally orthogonal task updates promote weight disentanglement by bounding task-vector norms and driving inter-task angles toward orthogonality. Under the stated assumptions, this makes expected task interference approximately zero.

  • Theorem 4: Theorem 4 states that approximately internally orthogonal task-update matrices actively promote weight disentanglement, even without Task-Feature Specialization.The result is established under the NTK linearization hypothesis.
  • Proof strategy: The proof targets the interference term τ_j^⊤J(x) on task t’s data domain and decomposes its magnitude into norm and angle factors.The decomposition is followed by directional-alignment, norm-control, angle-control, and synthesis steps.
  • Directional alignment: Task-vector direction aligns with the average Jacobian over the task domain, and typical sample Jacobians align with that average under a data-consistency assumption.This replaces the angle between τ_j and J(x) with the angle between τ_j and τ_t.
  • Norm control: OrthoReg bounds task-vector magnitude by constraining the deviation of each update matrix from orthogonality.The bound follows from controlling the update matrices’ Frobenius norms and therefore the norm of the composed task vector.
  • Angle control: Orthogonal update structure makes task vectors statistically nearly orthogonal, with E[|cos ∠(τ_t, τ_j)|] ≈ 0.Polar decomposition and Stiefel-manifold arguments explain why layerwise orthogonality propagates to whole task vectors.
  • Completing the proof: Because task-vector norms are bounded and inter-task cosine similarity is driven near zero, expected interference is negligible and weight disentanglement follows.The proof concludes with E[|τ_j^⊤J(x)|] ≈ 0 and invokes Lemma 2.

H.2. Experimental Performance Comparison and Analysis

The experiments compare OrthoReg with task-arithmetic baselines across architectures, tasks, and computational settings. OrthoReg generally improves task composition, produces more orthogonal vectors, strengthens negation, and costs less than TTA.

  • Experimental setup: OrthoReg is evaluated across eight image-classification tasks and ViT-B-32, ViT-B-16, and ViT-L-14 architectures.The comparison includes zero-shot, baseline merged models, and orthogonally regularized variants.
  • Computational efficiency: 177.04 minutes and 44500.27 MB are required by Non-lin. FT + OrthoReg, compared with TTA’s 280.86 minutes and 68031.34 MB, while reaching 88.23% versus 86.19% Abs.Acc. on ViT-L-14.OrthoReg adds moderate cost to standard fine-tuning and remains more efficient than TTA.
  • Computational efficiency: 132.96 minutes and 36976.50 MB yield 90.41% for ATT-FT + OrthoReg, versus 126.28 minutes, 36591.06 MB, and 87.81% for ATT-FT.The parameter-efficient setting shows a modest resource increase alongside a considerable performance increase.
  • Evaluation metric: Norm.Acc. measures merged-model performance relative to individually fine-tuned specialists, with 100% indicating average parity and minimal negative interference.The metric averages performance ratios across all merged tasks.
  • Performance comparison: Across most tasks, methods, and architectures, OrthoReg produces a larger performance footprint than the corresponding baselines.This pattern is reported in per-task analyses for ViT-L-14 and in broader radar-chart comparisons.
  • Task negation: Relaxing the control-task accuracy requirement to 90% or 80% strengthens task negation, while OrthoReg further improves forgetting within the threshold.The reported target-accuracy reductions are summarized in Tables 4 and 5.
  • Orthogonality analysis: OrthoReg consistently produces more orthogonal task vectors, shown by darker off-diagonal cosine-similarity heatmap values across methods and architectures.The heatmaps compare baseline task vectors with their OrthoReg-enhanced counterparts.

J.5.2. results

OrthoReg improves performance across most evaluated LoRA module configurations, with especially strong gains in attention-related and larger tunable subspaces. The MLP-only configuration is the sole reported exception, slightly reducing accuracy on smaller models.

  • LoRA module configurations: Up to +4.17 points on ViT-B-16, OrthoReg delivers its largest gains in attention-related LoRA modules such as qkvo.Performance improves consistently across configurations except MLP-only.
  • LoRA module configurations: Full-layer tuning benefits substantially, indicating that larger tunable subspaces help orthogonality constraints isolate task-relevant directions.
  • LoRA module configurations: Q,V-only tuning shows stable improvements when combined with OrthoReg, extending gains to an efficient tuning configuration.
  • LoRA module configurations: The MLP-only setup slightly reduces accuracy on smaller models, making it the only reported configuration without improvement.The authors suggest orthogonality can occasionally restrict useful shared representations in this setting.
  • Overall result: Overall, OrthoReg enhances components responsible for task-discriminative behavior, producing more accurate task vectors and more reliable task arithmetic.
Loading 2604.17078v1…