ATM: Improving Model Merging by Alternating Tuning and Merging

Abstract: Model merging has recently emerged as a cost-efficient paradigm for multi-task learning. Among current approaches, task arithmetic stands out for its simplicity and effectiveness. In this paper, we motivate the effectiveness of task vectors by linking them to multi-task gradients. We show that in a single-epoch scenario, if the optimization is performed via gradient descent, task vectors are after one step mathematically equivalent to the gradients obtained via gradient descent in a multi-task setting, and still approximate these gradients in subsequent epochs. Furthermore, we show that the effectiveness of task vectors is largely driven by the first epoch's gradient. Given this parallel between task vectors and gradients, we propose viewing model merging as a single step in an iterative process that alternates between tuning and merging (ATM). We then propose two ways to utilize ATM. The first is to replace multi-task learning with ATM in scenarios where data sharing is prohibited, such as federated learning. The second is to improve the outcome of any model merging algorithm by applying a few post-hoc iterations of ATM on a small validation dataset, which is commonly available for hyperparameter tuning. Finally, we provide both empirical and theoretical support for the effectiveness of ATM, demonstrating that it minimizes an upper bound on the loss obtained by jointly finetuning all tasks.

• The paper proposes ATM as an iterative method that alternates fine-tuning and merging to prevent overshooting multi-task optima.
• It demonstrates that task vectors act like gradients, achieving up to 20% accuracy improvements across diverse datasets.
• ATM’s framework balances task-specific and collective performance, integrating with existing methods without extra computational cost.

Overview of Alternating Tuning and Merging (ATM) for Model Merging

The research paper titled "ATM: Improving Model Merging by Alternating Tuning and Merging" presents a novel approach to tackle the challenges inherent in model merging methodologies, particularly in multi-task settings. The ATM framework underscores the intricate relationship between task arithmetic—an established model merging technique—and gradient descent processes. This paper ascertains that task vectors are conceptually equivalent to gradients computed over multi-task datasets, thus proposing Alternating Tuning and Merging (ATM) as an iterative alternative to conventional model aggregation methods.

Theoretical Insights and Contributions

The paper illuminates the theoretical oversight in standard one-step model merging methods that utilize task arithmetic, highlighting their propensity to overshoot the multi-task optimum. The authors show that in a single epoch, task vectors act as the additive inverse of gradient steps. This revelation extends to the fact that task vectors' efficacy can largely stem from the gradient direction determined in the initial finetuning epoch.

ATM is introduced as a framework that iteratively fine-tunes the model on individual tasks before merging. This iterative alternation reduces the likelihood of overshooting while incorporating interference-resolution strategies to enhance the final model's performance. The framework's flexibility allows it to integrate seamlessly with existing task-vector methods, circumventing additional computational burdens typically associated with task-vector pruning or elaborate weight adjustments.

Key contributions from the research include:

• Demonstrating that task vectors, under specific conditions, either equate to or closely approximate gradients of task losses.
• Highlighting that prevalent one-shot merging frameworks can overshoot multi-task optima, especially when task vectors possess large norms.
• Introducing ATM as a generalized, iterative merging framework with empirical demonstrations of increased task vector orthogonality and theoretical validations on minimizing multi-task loss.

Empirical Evaluation and Results

Extensive experiments validate ATM's superiority over established baselines across diverse datasets in computer vision (ViT-B-16 backbone) and NLP (RoBERTa-base and BERT-base-uncased) tasks. Notably, ATM achieves up to 20% greater accuracy than current baseline methods. These results are consistent irrespective of compute budgets, with ATM showing pronounced improvements as more computational resources are apportioned.

An essential characteristic observed is ATM's ability to balance specialist performance with collective multi-task effectiveness, effectively merging the benefits of task-specific learning trajectories without succumbing to the limitations of prior one-step assumptions.

Implications and Future Directions

Practically, the ATM framework offers a robust method for multi-task learning scenarios where model storage constraints necessitate a single strong performance model. Theoretically, the insights on gradient alignment provide a foundation for further investigation into model merging dynamics. Future research might explore leveraging ATM's iterative framework with advanced gradient-descent techniques or interference-mitigation strategies to further enhance performance without detracting from computational efficiency.

Conclusion

ATM represents a significant advancement in the model merging field by addressing the inherent deficiencies of one-shot task-vector methods. Its design and empirical success open pathways for more nuanced, data-privacy-preserving multi-task models that can seamlessly integrate into diverse machine learning pipelines. The research underscores the value of iteratively tuning and merging models to achieve optimal balance and performance in multi-task settings.