PaperlayerView exact revision on arXiv

Source-linked AI summary

MetaGPT: Merging Large Language Models Using Model Exclusive Task Arithmetic

Yuyan Zhou, Liang Song, Bingning Wang, Weipeng Chen

arXiv:2406.11385v2cs.CL

TL;DR

Existing task arithmetic methods do not jointly provide optimal performance, computational efficiency, and data privacy for LLM-scale multi-task learning. MetaGPT formalizes a loss-based merging objective and derives data-free closed-form scaling coefficients; experiments report that it outperforms existing state-of-the-art model-exclusive merging methods and can be combined with task-vector refinements. The method relies on common initialization and architecture, and its NTK-linearization basis may limit performance on small models.

  • Problem

    Existing task arithmetic approaches trade off computational cost, performance, and privacy, limiting the availability of suitable methods for billion-scale models.

    Source: §1 · p. 1

  • Method

    MetaGPT formalizes merging as minimizing average loss difference across task models and separates data terms from scaling coefficients to derive a model-exclusive closed-form solution.

    Source: §1 · p. 2

  • Results

    MetaGPT outperforms the existing state-of-the-art model-exclusive merging method and can be integrated with Ties-Merging and DARE.

    Source: §7 · p. 10

  • Takeaways & Limitations

    The authors present MetaGPT as an efficient avenue for implementing task arithmetic in large-scale multi-task learning.

    Source: §7 · p. 10

    MetaGPT relies on common initialization and model architecture, and its NTK-linearization basis may limit performance on small models.

    Source: §8 · p. 10

Abstract

from arXiv · show

The advent of large language models (LLMs) like GPT-4 has catalyzed the exploration of multi-task learning (MTL), in which a single model demonstrates proficiency across diverse tasks. Task arithmetic has emerged as a cost-effective approach for MTL. It enables performance enhancement across multiple tasks by adding their corresponding task vectors to a pre-trained model. However, the current lack of a method that can simultaneously achieve optimal performance, computational efficiency, and data privacy limits their application to LLMs. In this paper, we propose \textbf{M}odel \textbf{E}xclusive \textbf{T}ask \textbf{A}rithmetic for merging \textbf{GPT}-scale models, which formalizes the objective of model merging into a multi-task learning framework, aiming to minimize the average loss difference between the merged model and each individual task model. Since data privacy limits the use of multi-task training data, we leverage LLMs' local linearity and task vectors' orthogonality to separate the data term and scaling coefficients term and derive a model-exclusive task arithmetic method. Our proposed MetaGPT is data-agnostic and bypasses the heavy search process, making it cost-effective and easy to implement for LLMs.Extensive experiments demonstrate that MetaGPT leads to improvements in task arithmetic and achieves state-of-the-art performance on multiple tasks.

1 Introduction

The introduction motivates task arithmetic as a way to merge task-specific models for multi-task learning, but existing approaches trade off performance, computational cost, and data privacy. MetaGPT separates data terms from scaling coefficients to derive a data-free closed-form solution, and the authors report state-of-the-art performance on multiple tasks.

1 Introduction

  • MTL aims to let one model handle multiple tasks, avoiding the need to fine-tune and deploy a separate model for each new application.The introduction notes that conventional MTL jointly trains on raw data across tasks, while LLM scale makes this process computationally intensive.Source: §1 · p. 1
  • Existing task arithmetic methods trade off performance, compute, and privacy: optimal tuning can be expensive, fixed coefficients can be sub-optimal, and data-based search raises privacy concerns.The authors conclude that no existing method is suitable for billion-scale models while performing satisfactorily in practice.Source: §1 · pp. 1–2
  • MetaGPT formalizes merging as minimizing average loss difference across task models, then separates data terms from scaling coefficients to obtain a closed-form solution.The proposed method is intended to work without multi-task training data.Source: §1 · p. 2
  • MetaGPT outperforms previous merging methods on several tasks across the LLaMA-2 and Mistral series and achieves state-of-the-art performance in task arithmetic.The authors also report that it can be integrated with task-vector-improving methods such as Ties-Merging and DARE.Source: §1 · p. 2
  • The paper contributes a mathematical optimization objective and a theoretical performance bound for task arithmetic.The bound is used alongside coefficient separation in the proposed method.Source: §1 · p. 2
  • MetaGPT is presented as orthogonal to task-vector-improving methods, allowing integration with them for higher performance.The authors identify Ties-Merging and DARE as examples of methods with which MetaGPT can be integrated.Source: §1 · p. 2

2 Related Work

Related work covers model-merging approaches for single-task performance, generalization, and multi-task learning, alongside methods for representation sharing and task balancing. The paper notes that conventional data-sharing MTL is poorly suited to LLMs because of computational costs and data-protection concerns.

Model Merging

  • Model merging has been studied for single-task performance, out-of-domain generalization, and multi-task learning, the paper’s central focus.Methods beyond parameter averaging include Fisher merging, RegMean, Task Arithmetic, Ties Merging, DARE, and tangent-space fine-tuning.Source: §2 · p. 3

Multi-Task Learning

  • MTL methods commonly learn shared representations through architecture design or balance tasks through optimization methods.The cited approaches include representation-sharing modules, branching, task-weight weighting, gradient-dominance handling, and gradient-conflict resolution.Source: §2 · p. 3
  • Conventional MTL that jointly trains on raw data across tasks is unsuitable for LLMs because model updates are costly and data proprietors may not disclose sensitive data.Source: §2 · p. 3

3 Preliminaries

The preliminaries define the task-model setup and task arithmetic, then review coefficient-selection strategies and their scalability limits for LLMs.

3.1 Notation

  • Each task is represented by training, validation, and test data, a loss function, and parameters fine-tuned from a shared pretrained model.Source: §3.1 · p. 4

3.2 Task Arithmetic

  • Task arithmetic merges task models by adding scaled task vectors—the differences between fine-tuned and pretrained weights—to the pretrained weights.Each task vector has a scaling coefficient, and coefficient choice significantly affects merged-model performance.Source: §3.2 · p. 4

3.3 Existing Methods

  • Prior coefficient-selection methods include validation-set grid search, shared fixed coefficients such as 0.3, and entropy minimization on unlabeled test samples.Grid search can face a curse of dimensionality as the number of tasks increases.Source: §3.3 · p. 4

3.4 Scalability Challenges for LLMs

  • For LLMs, grid search requires extra data and grows costly with more coefficients, while fixed coefficients are efficient but can be sub-optimal.Unsupervised test-sample optimization requires extra data and simultaneous loading of multiple models, making it costly in time and memory.Source: §3.4 · p. 4

4 Our Proposed MetaGPT

MetaGPT frames model-exclusive task arithmetic as minimizing average loss difference across tasks. It separates data-related terms from coefficient/model terms under linearity and task-vector orthogonality, yielding closed-form scaling coefficients.

4.1 Overview

  • MetaGPT is presented as an efficient, data-free task-arithmetic method intended to scale to large models while protecting data privacy.The method’s overview describes a closed-form solution based on careful approximations.Source: §4.1 · p. 5

4.2 MetaGPT Optimization Objective

  • MetaGPT minimizes average loss difference across tasks, where each task loss difference compares merged-model loss with the fine-tuned task model’s loss.A smaller task loss difference means the merged model’s loss is closer to or below the fine-tuned model’s loss on that task.Source: §4.2 · p. 5

4.3 Separating Data and Coefficients

  • Taylor expansion reformulates each task loss difference as a quadratic form in a linear combination of scaling coefficients and model weights.The resulting task loss difference depends on data, models, and scaling coefficients.Source: §4.3 · p. 5
  • NTK linearity near initialization and orthogonal task vectors let MetaGPT upper-bound task loss differences and separate data terms from coefficient/model terms.The linear approximation becomes exact as network width approaches infinity.Source: §4.3 · pp. 5–6

4.4 The Optimal Solution

  • After separating these terms, MetaGPT decomposes the average-loss objective and derives a model-exclusive closed-form solution for each scaling coefficient.The solution follows from the decomposition of the objective in Theorem 9.Source: §4.4 · p. 6

5 Property Verification

The section checks two properties used in the method: near-linear outputs under interpolation and near-orthogonality among LLM task vectors.

5.2 Task Vector Orthogonality

  • Cosine similarities among different LLM task vectors are nearly zero, supporting the orthogonality property used in the proof.The figure reports cosine similarities for six task vectors.Source: Figure 4 · p. 7, §5.2 · p. 7

5.1 NTK Linearization

  • Llama-2-7b-chat-hf outputs on AGIEval are almost linear as α varies from 0 to 1, supporting kernel-like behavior during fine-tuning.The authors sampled three outputs at α values from 0 to 1.Source: §5.1 · p. 7

6 Experiments

Experiments compare MetaGPT across model sizes and architectures, test integration with task-vector improvement methods, and assess out-of-distribution generalization.

6.1 Merging Models Using MetaGPT

  • 31.51 absolute average performance and 1.31 normalized average performance are MetaGPT’s rank-1 scores on LLaMA-2-7B.MetaGPT leads on WinoGrande, AGIEval, GSM8K, and MATH, while DARE leads on HumanEval and Weight Average leads on MBPP.Source: §6.1 · p. 9
  • On Mistral-7B, MetaGPT leads on WinoGrande, AGIEval, GSM8K, and MATH, while other methods lead on HumanEval and MBPP.Weight Average, Task Arithmetic, and Ties-Merging lead on HumanEval; DARE leads on MBPP.Source: §6.1 · p. 9
  • On LLaMA-2-13B, MetaGPT outperforms the other methods on AGIEval, Math, and MBPP and leads under both average measures.DARE, Weight Average, and Ties-Merging lead on WinoGrande, GSM8K, and HumanEval, respectively.Source: Table 3 · p. 8, §6.1 · p. 9
  • MetaGPT integrates with DARE and Ties-Merging, with reported DARE averages improving from 30.72 to 31.57 absolute and 1.26 to 1.3 normalized.The passage also reports improvements for Ties-Merging, giving absolute values 30.20 to 31.57 and normalized values 1.26 to 1.33.Source: Table 4 · p. 8, §6.1 · p. 9

6.2 Out of Distribution Generalization

  • On unseen JEC-QA, FinanceIQ, and MedQA datasets, MetaGPT outperforms the compared methods in 5-shot accuracy.The authors report this as evidence of greater robustness to test-distribution shifts.Source: §6.2 · p. 10, Table 5 · p. 10

7 Conclusion

The conclusion presents MetaGPT as a data-free task-arithmetic method for LLMs, based on a formalized objective and a closed-form solution.

7 Conclusion

  • MetaGPT separates data and scaling-coefficient terms to obtain a closed-form, data-free solution and can integrate with Ties-Merging and DARE.The paper reports that it outperforms existing state-of-the-art model-exclusive merging methods.Source: §7 · p. 10

8 Limitations

The authors identify shared initialization and architecture as conditions for their task-vector orthogonality assumption, and note a small-model boundary.

8 Limitations

  • MetaGPT relies on common initialization and model architecture for task-vector orthogonality, and may perform poorly on small models.The authors connect the small-model caveat to the method’s reliance on NTK linearization and its LLM-specific design.Source: §8 · p. 10

A Proof

The proof excerpts use gradient and linear-regime relations to develop task-loss bounds and derive scaling coefficients.

A.1 Proof of Lemma 4

  • Lemma 4’s proof substitutes the interpolation γ_t(β) and h definitions, then uses the zero task-loss gradient at θ_t for the first-order expansion.The passage states that the gradient of L_t at the fine-tuned θ_t is zero.Source: §A.1 · p. 15

A.2 Proof of Theorem 7

  • Under Property 5’s linear regime, the proof represents loss, Hessian, and TLD terms through initialization gradients, then applies Property 6 to bound the h term.The cited steps include a squared-loss expression, a Hessian representation, and a trace bound.Source: §A.2 · pp. 15–16

A.3 Proof of Theorem 8

A.4 Proof of Theorem 9

  • For Eq. 37, δ0 is defined as max{δi}; for Eq. 38, terms containing λt are represented, and the ALD is consequently upper-bounded.Source: §A.4 · p. 16

A.5 Proof of Theorem 10

  • Theorem 10’s proof treats each ALD_λt separately; its quadratic dependence on λ_t yields an optimal scaling coefficient.Source: §A.5 · p. 17

B Details of Models and Datasets

The appendix lists model and dataset references, evaluation hardware, and settings for DARE and TIES-Merging.

B Details of Models and Datasets

C Infra and hardware details

D Hyper-parameter Setting

E Details of different Methods

The appendix compares merging methods by data requirements and complexity, then reports MetaGPT results for merging intermediate Baichuan2-7B checkpoints.

E Details of different Methods

F Merging Checkpoints of the Pre-trained Model

Related papers

Loading 2406.11385v2…