TATR: Trust Region Task Merging

• Task Arithmetic in Trust Regions is a method that fuses task deltas with trust region constraints to prevent negative transfer during multi-task model merging.
• It leverages unlabeled exemplars and per-layer projections to restrict harmful updates, ensuring that fine-tuned parameter changes do not degrade performance on other tasks.
• Empirical evaluations demonstrate robust gains—improving accuracy by up to 10 percentage points—and confirm TATR's plug-and-play compatibility with various merging strategies.

Task Arithmetic in Trust Region (TATR) is a training-free model merging approach designed to address knowledge conflicts when integrating multiple fine-tuned neural network models. Building on Task Arithmetic (TA)—the linear combination of parameter deltas from individually fine-tuned task models—TATR introduces the notion of a trust region in parameter space. This region constrains the directions in which task vectors are merged, preserving task performance by mitigating harmful interference between tasks. TATR is compatible with a broad range of TA-based methods and delivers robust empirical improvements across diverse benchmarks (Sun et al., 25 Jan 2025).

1. Task Arithmetic and Model Merging

Task Arithmetic (TA) begins with a pre-trained model parameterized by $\theta_{\text{pre}}$. For each task $i$, fine-tuning yields parameters $\theta_i = \theta_{\text{pre}} + \Delta\theta_i$, where $\Delta\theta_i$ denotes the task delta or task vector. The standard variant of TA merges $K$ tasks by a simple linear sum:

$$\theta_{\text{merge}} = \theta_{\text{pre}} + \lambda \sum_{i=1}^K \Delta\theta_i$$

where $\lambda$ is a scalar weight, often set to $1/K$ or tuned on held-out data. This paradigm enables training-free multi-task model fusion, circumventing the computational and storage overhead of joint or sequential fine-tuning. However, naive linear combination can introduce negative transfer when the directions of task deltas conflict.

2. Knowledge Conflicts in TA-Based Merging

Knowledge conflicts arise when different task deltas $\Delta\theta_i$ encode incompatible parameter updates, resulting in one task's improvement causing loss increases for others. This phenomenon is quantifiable: letting $L_j(\theta)$ be the loss on task $j$, a conflict manifests if the loss increases sharply upon merging $\Delta\theta_i$:

$$\Delta L_j = L_j(\theta_{\text{pre}} + \Delta\theta_i) - L_j(\theta_{\text{pre}}) \gg 0$$

Empirically, such conflicts are evidenced by decreased accuracy or increased loss on one or more tasks after merging. The underlying mechanism is the alignment of certain components of $\Delta\theta_i$ with the gradient of $L_j$ at $\theta_{\text{pre}}$, leading to substantial first-order loss increments (Sun et al., 25 Jan 2025).

3. Trust Region Derivation and Enforcement

TATR introduces a projection step to restrict each task delta to a trust region within the parameter space. The trust region is defined such that only directions inducing small changes to all tasks' losses are retained. Formally, the trust region for a delta $\Delta$ is:

$$\text{TR} = \left\{ \Delta : |\nabla_\theta L_j(\theta_{\text{pre}})^\top \Delta| \leq \epsilon_j \ \forall j \right\}$$

where $\epsilon_j$ bounds the first-order loss change for each task $j$.

Practical Trust Region Approximation

Instead of exact gradient or Hessian computation, TATR leverages unlabeled exemplars $D_j$ from each task to extract data-dependent features $X_j$. For every linear layer $l$ and task $k$, a removal basis $B_k^l$ is computed by maximizing the exclusion of directions likely to harm other tasks:

$$\max_{B} \sum_{i \neq k} \|X_k^l W_i^l B B^\top\|_F^2 - \alpha\|X_i^l W_i^l B B^\top\|_F^2$$

Projecting $\Delta\theta_i^l$ onto the orthogonal complement of $B_k^l$ guarantees that merged updates do not move along harmful directions. For normalization or bias parameters, TATR applies a data-driven mask $m_k^l \in \{0,1\}^d$, optimized by maximizing the same separation objective in a coordinate-wise fashion.

4. TATR Algorithm and Merging Rule

Clipping is applied per (task, layer) pair as follows. For linear layers:

$$T_i^l \leftarrow T_i^l - (T_i^l B_k^l)(B_k^l)^\top \quad \forall k \neq i$$

For normalization and bias layers, a corresponding mask is applied:

$T_i^l \leftarrow T_i^l \circ (1 - m_k^l)$

The merged model parameters are constructed by summing the projected, trusted task vectors:

$$\theta_\text{TATR} = \theta_{\text{pre}} + \lambda \sum_{i=1}^K T_i$$

where $T_i$ denotes the post-projection, per-layer task vectors. The merge weight $\lambda$ is typically set to $1/K$ or tuned empirically.

A summary of the TATR merging process is as follows:

Step	Operation	Notes
Feature Extraction	Compute features $X_k^l$ for each task and layer	Driven by $D_k$, $W_\text{pre} + T_k^l$
Basis/Mask Computation	Calculate $B_k^l$ (linear) or $m_k^l$ (normalization/bias) for each task/layer	Maximization objective as above
Trust-Region Projection	Project $T_i^l$ w.r.t. $B_k^l$ or $m_k^l$ for all $i \neq k$	Enforces trust region constraints
Model Merge	Aggregate trusted task vectors via $$\theta_\text{TATR} = \theta_{\text{pre}} + \lambda \sum_i T_i$$

5. Empirical Evaluation

TATR was evaluated on eight diverse classification benchmarks fine-tuned from CLIP, using ViT-B/32 and ViT-L/14 backbones. Datasets include SUN397, Stanford Cars, RESISC45, EuroSAT, SVHN, GTSRB, MNIST, and DTD. Baselines span pre-trained (no adaptation), individually fine-tuned, traditional MTL, test-time merging (AdaMerging/AdaMerging++, Surgery Merging), and training-free approaches (Weight Averaging, Fisher Merging, RegMean, Task Arithmetic, TIES-Merging, Consensus Merging).

On ViT-B/32, TATR achieves an average accuracy of 72.8, substantially outperforming vanilla Task Arithmetic (69.1). Gains persist when TATR is combined with other mergers: TIES-Merging + TATR achieves 73.3, AdaMerging++ + TATR attains 82.5, and Surgery + TATR 82.4. Fully TATR-merged models achieve 76.1. For ViT-L/14, TATR improves over vanilla Task Arithmetic by approximately 1 pp (84.5→85.3) and likewise boosts AdaMerging++ from 91.0 to 91.5. The approach exhibits high robustness to the number of exemplars per task, with average accuracy only slightly decreasing ($\sim$1 pp) from 1 to 32 exemplars per task (Sun et al., 25 Jan 2025).

6. Plug-and-Play Compatibility and Component Analysis

TATR is architected as a modular “clipping” operation within the TA model merging pipeline. Because it solely constrains the projection of task vectors, it can be inserted into any task-arithmetic-style merger, including AdaMerging++, Surgery Merging, and TIES-Merging. Across all such configurations, empirical results demonstrate that the addition of TATR consistently yields performance increases, often in the range of +1 to +10 pp over the underlying method.

Ablation studies isolate the individual contributions of linear-layer basis projection and normalization/bias masking, confirming that both components are crucial for conflict mitigation. Together, they recover most or all of the multi-task performance lost in naive merges.

7. Significance and Distinctions

TATR preserves the core simplicity and training-free paradigm of Task Arithmetic while introducing strong safeguards against negative transfer via an efficient, data-driven trust region. By learning where each task vector may move in parameter space—based on a handful of unlabeled exemplars—TATR nearly eliminates the major pathological failure mode of TA. This suggests that parameter-space projections, rather than brute-force averaging, are fundamental for reliable multi-task model merging in the absence of joint training (Sun et al., 25 Jan 2025).