Source-linked AI summary
Localizing Task Information for Improved Model Merging and Compression
Ke Wang, Nikolaos Dimitriadis, Guillermo Ortiz-Jimenez, François Fleuret, Pascal Frossard
TL;DR
Multi-task model merging is attractive for combining fine-tuned checkpoints, but it often performs worse than individual models. This paper uses TALL-masks to localize preserved task information and Consensus Merging to remove weights that hinder fusion. Consensus Merging improves existing approaches, while mask-based compression retains 99.7% of performance as storage falls from 57Gb to 8.2Gb.
Problem
Method
TALL-masks localize task-specific parameters in a multi-task vector, and Consensus Merging uses mask agreements to remove weights deemed important by none or only one task.
Source: §1 · pp. 1–2Results
Consensus Merging consistently improves prior merging methods; compression reduces 20 fine-tuned models from 57Gb to 8.2Gb while retaining 99.7% of performance.
Source: §1 · p. 2Takeaways & Limitations
Task-specific information remains in the merged vector and can be localized to recover individual fine-tuned performance without storing each checkpoint.
Source: §7 · p. 9The interpretation assumes task-specific information is preserved in the multi-task vector and treats parameters included in a task mask as relevant to that task.
Source: §3 · p. 3, §5.2 · p. 5
Abstract
from arXiv · showhide
Model merging and task arithmetic have emerged as promising scalable approaches to merge multiple single-task checkpoints to one multi-task model, but their applicability is reduced by significant performance loss. Previous works have linked these drops to interference in the weight space and erasure of important task-specific features. Instead, in this work we show that the information required to solve each task is still preserved after merging as different tasks mostly use non-overlapping sets of weights. We propose TALL-masks, a method to identify these task supports given a collection of task vectors and show that one can retrieve >99% of the single task accuracy by applying our masks to the multi-task vector, effectively compressing the individual checkpoints. We study the statistics of intersections among constructed masks and reveal the existence of selfish and catastrophic weights, i.e., parameters that are important exclusively to one task and irrelevant to all tasks but detrimental to multi-task fusion. For this reason, we propose Consensus Merging, an algorithm that eliminates such weights and improves the general performance of existing model merging approaches. Our experiments in vision and NLP benchmarks with up to 20 tasks, show that Consensus Merging consistently improves existing approaches. Furthermore, our proposed compression scheme reduces storage from 57Gb to 8.2Gb while retaining 99.7% of original performance.
1. Introduction
Weight interpolation can fuse fine-tuned checkpoints, but multi-task models still suffer performance drops relative to individual models. The paper introduces TALL-masks to localize task-specific information and Consensus Merging to improve fusion, while using the masks to compress checkpoints.
1. Introduction
- Prior weight-space fusion methods still show significant performance drops compared with individual fine-tuned models.The introduction situates this problem amid efforts to improve foundation-model use and address their computational and storage requirements.Source: §1 · p. 1
- Mask-agreement analysis identifies catastrophic weights deemed important by no task and selfish weights deemed important exclusively by one task.Consensus Merging removes these weights and keeps general weights.Source: §1 · p. 2
- Consensus Merging improves existing methods, with a 4.9% gain in absolute average accuracy on a 20-task vision benchmark and a 6.3% improvement over TIES on an 8-task NLP benchmark.It is complementary to approaches such as Task Arithmetic and TIES.Source: §1 · p. 2
- 20 fine-tuned models are compressed from 57Gb to 8.2Gb while retaining 99.7% of performance.The compression stores the zero-shot model, merged task vector, and binary masks rather than the individual fine-tuned checkpoints.Source: §1 · p. 2
- The paper reports compression retaining more than 99% of performance across vision settings and task counts, while merging methods can approach zero-shot performance.The compression evaluation spans small-to-large ViTs and benchmarks with 8 to 20 tasks.Source: §1 · p. 2
- TALL-masks localize task-specific information preserved in the multi-task vector by deactivating irrelevant parameters for each task.The method constructs task-specific binary masks by approximating each original task vector from the merged vector.Source: §1 · pp. 1–2
2. Related Work
Related work spans weight interpolation and merging, alongside methods aimed at reducing foundation-model complexity. The paper positions its approach as addressing performance loss while compressing multiple checkpoints.
Weight Interpolation and Model Merging
- Weight interpolation has been studied as a route to low-loss paths and weight ensembles, with work extending to multi-task tradeoffs and pre-trained-model task arithmetic.Other merging approaches address parameter redundancy and sign disagreement, preserve parameters using Fisher information, or learn merging weights from unlabeled test data.Source: §2 · p. 3
Reducing complexity for foundation models
- Foundation-model complexity methods include parameter-efficient fine-tuning, post-training quantization, and sparsification or quantization of fine-tuning residuals for expert merging.Task Arithmetic combines diverse models but can incur severe performance degradation; BYOM sparsifies residuals before merging, with performance depending heavily on sparsity level.Source: §2 · p. 3
- Mixture-of-Experts models specialize sub-networks across tasks or input regions, but training them from scratch can be complex.The paper explores pre-trained models and task-specific information localization as an alternative route to creating expert models.Source: §2 · p. 3
3. Task interference causes performance degradation
Task arithmetic combines task vectors from fine-tuned checkpoints, but its merged model can lose performance. A controlled experiment indicates that this decline can occur even when task information is preserved and weight interference is absent.
3. Task interference causes performance degradation
- Task arithmetic forms task vectors by subtracting the pretrained weights from fine-tuned checkpoints, then sums them and adds the result to the pretrained model.A validation-set scaling factor is applied to the summed vector.Source: §3 · p. 3
- The controlled experiment assigns each task an equal-sized, non-overlapping subset of weights, preserving task information while eliminating parameter interference.The comparison is against standard fine-tuning on an eight-task vision benchmark.Source: §3 · p. 3
- 2.7% worse normalized accuracy than standard fine-tuning was observed for task arithmetic in the controlled setting.The comparison averaged performance over eight tasks.Source: §3 · p. 4, Table 1 · p. 4
- The result suggests task interference, rather than weight interference, contributes to the decline: changes in other tasks alter discriminating features and perturb the input-to-output mapping.The task-specific parameters themselves remain constant in this account.Source: §3 · p. 4
4. TALL-masks: Localizing task-specific information in multi-task vector
TALL-masks constructs task-specific binary masks that extract relevant weights from a merged multi-task vector. On an eight-task vision benchmark, applying the masks restores full performance by deactivating irrelevant parameter subsets.
4. TALL-masks: Localizing task-specific information in multi-task vector
- TALL-masks constructs binary masks that extract relevant weights from the multi-task vector so the extracted weights approximate each original task vector.The masks deactivate irrelevant weights while retaining task-specific information.Source: §4 · p. 4
- The mask threshold λ_t controls how many parameters are selected: smaller values select more, and each task's value is chosen using its validation accuracy.Source: §4 · p. 4
- Full performance is retained on the eight-task vision benchmark by applying task-specific masks to deactivate irrelevant parameter subsets.Figure 2 plots selected parameter percentages and normalized validation accuracy against the task arithmetic baseline.Source: §4 · p. 4, Figure 2 · p. 4
5. Applications
TALL-masks supports two applications: compressing individual checkpoints and improving model merging. It localizes task-relevant information in the multi-task vector, which can be selectively stored or filtered.
5.1. Compressing Task Vectors
- For 20-task ViT-L/14, the masked-checkpoint scheme uses 13.7% of the storage required for individual checkpoints while maintaining most fine-tuned performance.The compressed representation stores the pretrained model, the multi-task vector, and task-specific binary masks.Source: §5.1 · p. 5
5.2. Improving Model Merging
- Consensus Merging targets task interference by removing catastrophic weights used by no task and selfish weights used by only one task.At k = 2, it keeps weights relevant to at least two tasks; the threshold depends on the task combination and merging method.Source: §5.2 · p. 5, Figure 3 · p. 5
- Compression and Consensus Merging are complementary to existing merging approaches and operate on the multi-task vector, allowing practitioners to choose an application for their usage scenario.Unlike compression, Consensus Merging requires no extra storage.Source: §5.2 · pp. 5–6
6. Experiments
Across NLP and vision benchmarks, Consensus Merging improves on existing merging methods, while TALL Mask compression preserves fine-tuned performance with reduced storage. The benefits persist across task counts and combinations, with threshold choices depending on the merging method.
6.1. Model Merging
- Consensus Merging improves on Task Arithmetic and TIES across NLP settings and is best in 5 of 6 vision scenarios; for 20-task ViT-L/14, gains are 4.9% over Task Arithmetic and 1.1% over TIES.For the 8-QA benchmark, the methods improve absolute accuracy over Task Arithmetic by 4.8% and over TIES by 6.3%, respectively.Source: §6.1 · p. 6, Table 3 · p. 7
6.2. Compression
- TALL Mask compression retains almost all fine-tuned performance across vision settings, with TALL Mask + TIES exceeding 99% normalized accuracy in every test case.For ViT-B/32, TALL Mask + TIES achieves around 100% normalized accuracy in all test scenarios.Source: §6.2 · p. 7, Figure 4 · p. 8
6.3. Individual-task performance
- On individual vision tasks, TALL Mask + Task Arithmetic matches fine-tuned models, while merging baselines degrade as task counts grow; Consensus Task Arithmetic suffers the least and leads among merging methods in almost all 20-task datasets.Figure 4 compares individual-task absolute accuracy for computer-vision benchmarks with ViT-B/32.Source: §6.3 · p. 8, Figure 4 · p. 8
6.4. Performance with varying task combinations
- Across sampled task combinations and task counts, TALL Mask + Task Arithmetic matches individually fine-tuned performance, while Consensus Task Arithmetic generally outperforms other merging methods as task counts rise.The study averages results over 8 representative combinations for each task count and reports near-100% normalized accuracy with heavily reduced storage for the mask-based approach.Source: Figure 5 · p. 8, §6.4 · p. 9
6.5. Effect of weight-pruning threshold
- Consensus Task Arithmetic usually performs best at k = 2, whereas Consensus TIES performs best at k = 1; TIES with only catastrophic weights removed consistently outperforms TIES.The different optimal thresholds are attributed to the methods' mask-agreement profiles; removing selfish weights from TIES can reduce salient weights.Source: §6.5 · p. 9, Figure 6 · p. 9
7. Conclusion
The paper presents TALL-masks for checkpoint compression and model merging, arguing that task-specific information survives in the multi-task vector and can be localized with binary masks. It also reports that removing harmful mask categories improves merging performance.
7. Conclusion
- TALL-masks localizes and extracts task-specific information retained in the multi-task vector, supporting recovery of fine-tuned performance and checkpoint compression.The compressed collection consists of the zero-shot model, the merged vector, and task-specific masks; Consensus Merging eliminates harmful weights and improves benchmark performance.Source: §7 · p. 9
A. Experimental Details
- Experiments used four NVIDIA V100 GPUs, each with 32 GB of memory.Source: A. Experimental Details · p. 14
A.1. Fine-tuning
- Fine-tuning followed Ilharco et al. (2023): 2,000 iterations, batch size 128, learning rate 1e-5, cosine annealing with 200 warm-up steps, and a frozen classification layer.The same pretrained CLIP checkpoint was fine-tuned with AdamW.Source: §A.1 · p. 14
A.2. Hyper-parameter tuning
- Normalized accuracy is measured relative to individually fine-tuned accuracy, and depends on both fine-tuning and merging methods.The mask sparsity factor λ was selected per task from {0.2, 0.3, 0.4, 0.5, 0.6} using validation performance; the task-vector scaling factor α was selected from 0.0 to 1.0 in 0.1 increments using average validation performance.Source: §A.1 · p. 14, §A.2 · p. 14
A.3.1. COMPUTER VISION
The experiments cover computer-vision benchmarks with 8, 14, and 20 tasks and two NLP benchmark collections. The larger vision settings extend the 8-task setup with additional datasets.
A.3.1. COMPUTER VISION
- The vision scenarios contain 8, 14, and 20 tasks; the 14-task setup adds six datasets to the original eight, and the 20-task setup adds six more.Beyond eight tasks, datasets come from torchvision; the 14-task selection aims for task diversity, while six tasks for the 20-task setting are selected by zero-shot CLIP versus linear-probe ResNet-50 performance after removing MNIST variants.Source: §A.3.1 · p. 14
A.3.2. NATURAL LANGUAGE PROCESSING
- The NLP benchmark descriptions cover a 7-task collection and an 8-task QA collection, with the latter containing eight listed datasets.The 7-task collection includes QASC, QuaRTz, PAWS, Story Cloze, WikiQA, Winogrande, and WSC; the QA list includes CosmosQA, QASC, QuAIL, QuaRTz, PAWS, ROPES, SocialIQA, and Wiki QA.Source: §A.3.2 · pp. 14–15
B. Derivation of Equation 4
The appendix details the mask-criterion derivation and storage accounting, then examines catastrophic-weight removal and the effect of pruning thresholds. The best threshold depends on the merging method.
B. Derivation of Equation 4
- Independence across sub-problems is used to aggregate the derivation into an optimal mask criterion.The section identifies this step as the transition from Equation 13 to Equation 14.Source: B. Derivation of Equation 4 · p. 15
C. Storage cost calculation
- Fine-tuned models cost 32(TP′+F), while Task arithmetic, Ties-merging, Consensus variants, and zero-shot each store one model at 32P.Magnitude pruning and magnitude masking have separate accounting, and their stated costs include stored models, task vectors, or masks.Source: C. Storage cost calculation · p. 15
- TALL Mask with Task arithmetic or Ties-merging costs (64 + T)P′ + 32F, including the zero-shot model, multi-task vector, and T binary masks.The storage calculation assumes one float parameter takes 32 bits.Source: C. Storage cost calculation · p. 15
D.1. Performance for removing only catastrophic weights
- This section studies Consensus Merging after removing only catastrophic weights—those beneficial for none of the tasks—with image-classification experiments using ViT-B/32, ViT-L/14, and ViT-B/16.Results are presented in Tables 4 and 5.Source: §D.1 · p. 16, Table 4 · p. 16, Table 5 · p. 16
D.2. Effect of weight-pruning threshold for 8 and 14 tasks
- Consensus Task Arithmetic performs best at threshold 2, while Consensus TIES performs best at threshold 1; further pruning reduces performance toward zero-shot.The comparison covers 8-task and 14-task image-classification benchmarks with ViT-B/32.Source: §D.2 · p. 16
D.3. Performance with Parameter-efficient fine-tuning methods
On (IA)3 NLP models, removing only catastrophic weights consistently improves Consensus Merging over Task Arithmetic and TIES, whereas removing selfish weights degrades performance. The appendix also supplies mask-agreement and additional vision results.
D.3. Performance with Parameter-efficient fine-tuning methods
- For 7 NLP tasks using (IA)3, Consensus TA gains 3.8% over Task Arithmetic and Consensus TIES gains 5.3% over TIES when only catastrophic weights are removed.The results are in Table 6 and compare Consensus Merging with baseline methods on NLP benchmarks.Source: §D.3 · p. 16, Table 6 · p. 17
- Removing selfish weights degrades performance, possibly because (IA)3 and full fine-tuning tune different weight profiles.The authors report consistent gains over Task Arithmetic and TIES when removing only catastrophic weights.Source: §D.3 · p. 16
D.4. Distribution of mask agreements with more tasks
- Figure 8 extends the mask-agreement histograms to 14 and 20 vision tasks for Task Arithmetic and TIES.The figure is described as complementing Figure 3.Source: §D.4 · p. 17
D.5. Result for ViT-B/16
- ViT-B/16 results in Table 7 show findings similar to those reported for ViT-B/32 and ViT-L/14.Table 7 complements Table 3.Source: §D.5 · p. 17
D.6.1. FULL RESULTS FOR VISION
- The appendix provides radar plots of individual-task results for ViT-B/16 (Figure 9) and ViT-L/14 (Figure 10), reporting findings similar to those observed for ViT-B/32 in the main text.Source: §D.6.1 · p. 17
D.6.2. FULL RESULTS FOR NLP
The appendix provides full individual-task NLP results in Figure 11, comparing the proposed method with individual fine-tuned models, Task arithmetic, Ties-merging, and zero-shot.
D.6.2. FULL RESULTS FOR NLP
- Figure 11 reports absolute accuracy for individual tasks on three T5-large NLP benchmarks, comparing the proposed method with fine-tuned models and merging baselines.The comparison also includes zero-shot accuracy.Source: §D.6.2 · p. 17
E. Sample subset selection protocol
For Figure 5, the authors select eight representative task combinations by ordering the 20 tasks in several ways. For each task count, they take the first n tasks from each sequence to account for difficulty and randomness.
E. Sample subset selection protocol
- The eight representative combinations include four alphabetic task orderings: ascending, descending, wave, and zigzag.These orderings are among the sequences used to construct Figure 5.Source: E. Sample subset selection protocol · pp. 18–19
- The other four sequences order tasks by zero-shot performance: ascending, descending, wave, and zigzag.The sequences provide alternative task orders for generating Figure 5.Source: E. Sample subset selection protocol · p. 18
- For each task count n, the protocol takes the first n tasks from each sequence to represent both task difficulty and randomness.The difficulty ordering is based on zero-shot performance, while alphabetic orderings supply alternative sequences.Source: E. Sample subset selection protocol · p. 19