Source-linked AI summary

AdaMerging: Adaptive Model Merging for Multi-Task Learning

Enneng Yang, Zhenyi Wang, Li Shen, Shiwei Liu, Guibing Guo, Xingwei Wang, Dacheng Tao

arXiv:2310.02575v2cs.LGcs.CV

TL;DR

Task-vector MTL can merge independently fine-tuned models without original training data, but its performance is sensitive to merging coefficients. AdaMerging automatically learns task- or layer-specific coefficients by minimizing entropy on unlabeled test data, and experiments report improved multi-task performance, generalization, and robustness.

  • Problem

    Task-vector MTL avoids original training data but is highly sensitive to merging coefficients, creating a challenge for effective model integration.

  • Method

    AdaMerging uses entropy minimization on unlabeled multi-task test data to automatically learn coefficients for each task vector or layer.

  • Results

    AdaMerging is superior to current state-of-the-art model-merging methods in multi-task performance, generalization, and robustness.

  • Takeaways & Limitations

    Adaptive coefficient learning improves task-vector model merging across multi-task performance, unseen-task generalization, and test-distribution-shift robustness.

Abstract

from arXiv · show

Multi-task learning (MTL) aims to empower a model to tackle multiple tasks simultaneously. A recent development known as task arithmetic has revealed that several models, each fine-tuned for distinct tasks, can be directly merged into a single model to execute MTL without necessitating a retraining process using the initial training data. Nevertheless, this direct addition of models often leads to a significant deterioration in the overall performance of the merged model. This decline occurs due to potential conflicts and intricate correlations among the multiple tasks. Consequently, the challenge emerges of how to merge pre-trained models more effectively without using their original training data. This paper introduces an innovative technique called Adaptive Model Merging (AdaMerging). This approach aims to autonomously learn the coefficients for model merging, either in a task-wise or layer-wise manner, without relying on the original training data. Specifically, our AdaMerging method operates as an automatic, unsupervised task arithmetic scheme. It leverages entropy minimization on unlabeled test samples from the multi-task setup as a surrogate objective function to iteratively refine the merging coefficients of the multiple models. Our experimental findings across eight tasks demonstrate the efficacy of the AdaMerging scheme we put forth. Compared to the current state-of-the-art task arithmetic merging scheme, AdaMerging showcases a remarkable 11\% improvement in performance. Notably, AdaMerging also exhibits superior generalization capabilities when applied to unseen downstream tasks. Furthermore, it displays a significantly enhanced robustness to data distribution shifts that may occur during the testing phase.

1 INTRODUCTION

Foundation-model MTL can reuse independently fine-tuned models, but task-vector merging is highly sensitive to coefficient choice and remains below traditional MTL. AdaMerging addresses this by learning task-wise or layer-wise coefficients from unlabeled test data through entropy minimization, improving performance, generalization, and robustness.

  • Motivation: Foundation-model MTL avoids the labeling and computation costs of traditional joint training by reusing independently fine-tuned models.The same pre-trained model may be fine-tuned separately for downstream tasks and later merged.
  • Task-vector merging: Task arithmetic merges task vectors into a pre-trained model to enable MTL without the original training data, but performance remains below traditional MTL.A task vector is formed by subtracting pre-trained weights from task-fine-tuned weights.
  • Challenge: Merging-coefficient choice strongly affects average accuracy because task conflicts and relationships make task-vector integration sensitive to λ.An ill-suited coefficient can make the merged model perform poorly across multiple tasks.
  • AdaMerging: AdaMerging learns distinct coefficients for tasks or layers by minimizing prediction entropy on unlabeled multi-task test data.The method uses entropy minimization as an unsupervised surrogate objective for coefficient adaptation.
  • Evaluation: Across eight tasks, AdaMerging improved average MTL accuracy by approximately 5.0% to 11.0% over Task Arithmetic and Ties-Merging.The evaluation also reports improved generalization on unseen downstream tasks and an average 8.45% robustness improvement across seven distribution drifts versus Task Arithmetic.
  • Contributions: Entropy minimization, adaptive coefficient learning, and comprehensive experiments support improvements in performance, generalization, and robustness over task-vector merging methods.The paper identifies these as central contributions of AdaMerging.

2 RELATED WORK

Related work spans joint MTL and model merging, including task-vector methods that consolidate models trained for different tasks. AdaMerging differs by learning flexible task- or layer-specific coefficients through entropy minimization rather than relying on shared coefficients and costly grid search.

  • Joint Training for Multi-Task Learning: Joint-training methods gather data from multiple tasks and address task conflicts through architectural or optimization strategies.Architectural approaches include sparsification, branching, and modularization.
  • Model Merging for Multi-task Learning: Model merging combines models either to improve generalization or federated learning, or to merge different-task models for MTL.Task-vector merging belongs to the latter category.
  • Task-vector methods: PEM Composition extends task arithmetic to LoRA models, while Ties-Merging addresses redundancy and sign conflicts in task vectors.Ties-Merging resets redundant parameters, resolves sign conflicts, and merges sign-consistent parameters.
  • AdaMerging’s distinction: AdaMerging uses task- or layer-specific coefficients, entropy minimization, and automatic optimization instead of shared coefficients and costly grid search.These differences target flexibility and scalability as the number of tasks grows.

3 METHODOLOGY

The methodology formulates multi-task model merging without retraining data, develops task-wise and layer-wise adaptive coefficients, and optimizes them using entropy minimization on unlabeled test samples. The approach is motivated by the sensitivity of task-vector merging to coefficient choices and by the observed relationship between prediction entropy and loss.

  • Task-vector preliminaries: Task arithmetic forms task vectors by subtracting pre-trained weights from task-finetuned weights, then adds them to the pre-trained model for multi-task learning.Ties-Merging additionally trims redundant parameters, elects parameter signs, and performs disjoint merging before combination.
  • Task-vector preliminaries: Merging coefficients strongly affect average multi-task accuracy, while ill-suited values can make task-vector merging worse than weight averaging.Grid-searching separate coefficients becomes expensive as the number of tasks grows.
  • Adaptive model merging: AdaMerging learns separate coefficients for task vectors, allowing beneficial vectors to contribute more and harmful vectors to contribute less.The task-wise formulation assigns one coefficient λk to each task vector Tk.
  • Adaptive model merging: Layer-wise AdaMerging assigns distinct coefficients to different layers because layers can encode general or task-specific information.AdaMerging++ applies the same adaptive strategy after Ties-Merging’s parameter trimming and sign-conflict handling.
  • Entropy optimization: Entropy minimization on unlabeled multi-task test samples serves as a surrogate objective for optimizing merging coefficients without the original training data.The method draws inspiration from test-time adaptation and minimizes Shannon entropy of model predictions.
  • Entropy optimization: Across eight tasks, smaller prediction entropy corresponded to smaller average loss, with a high positive Spearman correlation between entropy and prediction loss.The analysis grouped samples into eleven entropy intervals and compared each interval’s average prediction loss.

4 EXPERIMENT

The experiments evaluate AdaMerging on eight image-classification tasks, comparing performance, generalization to unseen tasks, robustness to distribution shifts, and learned coefficient behavior. AdaMerging improves task-vector merging through unsupervised coefficient learning and entropy-based adaptation.

  • Experimental setup: The experiments compare model-merging methods on eight image-classification tasks using ViT-B/32 and ViT-L/14 models.Baselines include non-merging methods and model-merging approaches such as Weight Averaging, Fisher Merging, RegMean, Task Arithmetic, and Ties-Merging.
  • Performance: 11% improvement over Task Arithmetic and 8.7% over Ties-Merging are achieved by Layer-wise AdaMerging++ on ViT-B/32.On ViT-L/14, the corresponding improvements are 6.3% and 5.0%.
  • Generalization: 8.3% and 9.1% average improvements on MNIST and EuroSAT show that AdaMerging and AdaMerging++ generalize better to unseen tasks.On RESISC45 and SVHN, the average accuracy improvements are 4.4% and 5.4%, respectively.
  • Robustness: 11.2% higher performance under Motion Blur and 10.1% under Contrast demonstrate AdaMerging’s robustness to test-data distribution shifts.Compared with Task Arithmetic, improvements on clean data and the other corruption types range from 5.8% to 9.8%.
  • AdaMerging analysis: Task-wise coefficients are inconsistent across task vectors, making grid search undesirable as the number of tasks grows; AdaMerging avoids this manual search.The coefficients are learned rather than manually searched.
  • AdaMerging analysis: Layer-wise coefficients differ across task vectors and layers, with shallow layers generally receiving smaller coefficients than deep layers.This pattern indicates different layer contributions during model merging and greater reliance of shallow layers on pretrained weights.

5 CONCLUSION AND FUTURE WORK

The paper concludes that AdaMerging addresses the sensitivity of task-vector multi-task learning to merging coefficients. It uses entropy minimization to learn task-wise or layer-wise coefficients and reports gains in multi-task performance, generalization, and robustness.

  • Conclusion: AdaMerging automatically learns task-vector or layer-specific merging coefficients by minimizing entropy as a surrogate objective.The method is designed for model merging without the original training data.
  • Conclusion: AdaMerging is reported to outperform current state-of-the-art model-merging methods in multi-task performance, generalization, and robustness.The paper identifies model-merging solutions for different architectures as future work.

A EXPERIMENT SETTINGS

The experiments evaluate model merging on eight image-classification tasks using established baselines and AdaMerging variants. Coefficients are initialized uniformly and updated unsupervised with Adam for 500 iterations.

  • Datasets: Experiments study multi-task model merging on eight image-classification datasets.The setup follows Task Arithmetic and Ties-Merging.
  • Baselines: Seven comparison methods and four AdaMerging variations are evaluated.The baselines include Individual, Traditional MTL, Weight Averaging, Fisher Merging, RegMean, Task Arithmetic, and Ties-Merging.
  • Baselines: Individual models avoid task interference but cannot perform multiple tasks simultaneously, whereas Traditional MTL trains jointly on all original task data.These methods provide reference endpoints for evaluating model merging.
  • Implementation: Task Arithmetic and Ties-Merging use λ=0.3 by default, while AdaMerging coefficients are initialized at 0.3 before unsupervised updates.The coefficient updates use Adam with learning rate 0.001, momentum (0.9, 0.999), and batch size 16.
  • Implementation: AdaMerging updates coefficients for 500 iterations to limit additional training costs.The experiments use pretrained ViT-B/32, ViT-B/16, and ViT-L/14 models from CLIP.

B.1 PERFORMANCE, GENERALIZATION AND ROBUSTNESS

AdaMerging improves multi-task performance over task-vector baselines, including on unseen tasks and corruption-shifted data. The reported gains cover average accuracy across eight tasks, generalization to held-out tasks, and robustness under seven distribution shifts.

  • Performance: 3.2% performance improvement over Task Arithmetic is achieved by Ties-Merging on ViT-B/16.The comparison concerns average accuracy when merging models on eight tasks.
  • Performance: 2.2% and 1.0% performance improvements are brought by Task-wise AdaMerging and AdaMerging++, respectively.These methods learn a merging coefficient for each task vector in Task Arithmetic and Ties-Merging.
  • Generalization: 2.3% higher average accuracy is achieved on unseen EuroSAT and MNIST tasks, compared to Task Arithmetic.Their corresponding task vectors are not merged.
  • Generalization: 1.1% higher average accuracy is achieved on unseen RESISC45 and SVHN tasks, compared to Task Arithmetic.The result is reported as evidence of improved generalization properties.
  • Robustness: AdaMerging exceeds Task Arithmetic by 9.9% to 12.4% across seven corruption datasets.The reported gains are 9.9%, 8.2%, 7.8%, 6.8%, 12.4%, 9.5%, and 9.7% for Motion Blur, Impulse Noise, Gaussian Noise, Pixelate, Spatter, Contrast, and JPEG Compression, respectively.

B.2 ANALYSIS EXPERIMENT

The analysis examines task relationships, entropy as a surrogate objective, data availability, supervised comparison, computational cost, and learned coefficient patterns. Across these analyses, the passages report near-orthogonal task vectors, strong entropy–loss correlation, and measurable gains under constrained data and time budgets.

  • Task Relationship Analysis: Task vectors are almost orthogonal on ViT-B/32 and ViT-L/14.The analysis also notes higher similarity for related tasks such as SVHN and MNIST.
  • Test Data Availability: 5% of available unlabeled test data can nearly match the performance obtained with 100%.The study evaluates 0.1%, 1%, 5%, and 100% data availability; the method requires unlabeled test data.
  • Supervised Comparison: Unsupervised AdaMerging performs very close to supervised AdaMerging despite using entropy minimization as a proxy objective.The supervised version is treated as an upper bound for the unsupervised version.
  • Parameter Cost: Task-wise AdaMerging adds 8 parameters, while Layer-wise AdaMerging adds 1,248 to eight task vectors totaling 907,589,640 parameters.The analysis reports that AdaMerging introduces very few coefficients to update.
  • Time Cost: 8% performance improvement is reported after 50 minutes of training, while 2% is reported with 7.5 additional minutes over Task Arithmetic.These time-cost measurements use a single GeForce RTX 3090.
  • Coefficient Visualization: Layer-wise coefficient visualizations show learned coefficients across task vectors and layers for ViT-B/16 and ViT-L/14.Rows represent task vectors, columns represent layers, and intersections represent λ^l_k.
  • Entropy–Loss Correlation: Entropy minimization and supervised cross-entropy loss maintain a high Spearman correlation throughout coefficient learning.The comparison covers iterations 0, 100, 200, 300, 400, and 500.
  • Entropy–Loss Correlation: Smaller prediction entropy generally corresponds to smaller average prediction loss across eight tasks.Entropy groups are formed by sorting test samples into eleven groups from small to large entropy.
Loading 2310.02575v2…