PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving General Text Embedding Model: Tackling Task Conflict and Data Imbalance through Model Merging

Mingxin Li, Zhijie Nie, Yanzhao Zhang, Dingkun Long, Richong Zhang, Pengjun Xie

arXiv:2410.15035v1cs.CL

TL;DR

Joint training for general text embeddings faces task conflict and data imbalance, motivating alternatives that preserve performance across tasks. This paper explores model merging and introduces Self Positioning to search task-vector interpolation space for merge positions. The approach improves MTEB average performance by 0.7 points and outperforms resampling methods with less time cost.

  • Problem

    Joint training can suffer from conflicting task gradients and imbalanced training data, which can hurt performance across tasks.

    Source: §1 · p. 2

  • Method

    The paper combines independently trained models and uses Self Positioning with stochastic gradient descent to search task-vector interpolation space for merging positions.

    Source: §1 · p. 2, §5.1 · p. 9

  • Results

    0.7 points: the model merging strategy improves average performance on the MTEB dataset and outperforms resampling methods with significantly less time cost.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The experiments support finding high-performing general text embedding merges within task-vector interpolation space.

    Source: §1 · pp. 2–3

Abstract

from arXiv · show

Text embeddings are vital for tasks such as text retrieval and semantic textual similarity (STS). Recently, the advent of pretrained language models, along with unified benchmarks like the Massive Text Embedding Benchmark (MTEB), has facilitated the development of versatile general-purpose text embedding models. Advanced embedding models are typically developed using large-scale multi-task data and joint training across multiple tasks. However, our experimental analysis reveals two significant drawbacks of joint training: 1) Task Conflict: Gradients from different tasks interfere with each other, leading to negative transfer. 2) Data Imbalance: Disproportionate data distribution introduces biases that negatively impact performance across tasks. To overcome these challenges, we explore model merging-a technique that combines independently trained models to mitigate gradient conflicts and balance data distribution. We introduce a novel method, Self Positioning, which efficiently searches for optimal model combinations within the interpolation space of task vectors using stochastic gradient descent. Our experiments demonstrate that Self Positioning significantly enhances multi-task performance on the MTEB dataset, achieving an absolute improvement of 0.7 points. It outperforms traditional resampling methods while reducing computational costs. This work offers a robust approach to building generalized text embedding models with superior performance across diverse embedding-related tasks.

1 Introduction

General text embeddings extend pretrained language models to tasks such as retrieval and semantic similarity, but joint training can suffer from task conflict and data imbalance. The paper examines model merging and Self Positioning as alternatives, reporting improved MTEB performance and efficient search for high-performing merges.

1 Introduction

  • Text embeddings represent text in high-dimensional spaces to capture semantic information for tasks including retrieval, semantic similarity, and classification.The introduction situates general text embeddings alongside pretrained language models and the MTEB benchmark.Source: §1 · p. 1
  • Joint training can suffer from gradient conflicts that cause negative transfer and data imbalance that biases performance across tasks.The introduction gives a 50:1 Retrieval-to-Classification data ratio as an example of imbalance.Source: §1 · p. 2
  • The authors' empirical analysis finds poorer performance for jointly trained models than for models trained separately on individual tasks.They also report that training-data quantity affects task-vector positioning and performance across tasks.Source: §1 · p. 2
  • Model merging combines independently trained models, and evaluation of five common strategies shows significant multi-task performance improvements.The authors report that high-performing merges can be found within the interpolation space of task vectors.Source: §1 · p. 2
  • Self Positioning searches task-vector interpolation space for merging positions using stochastic gradient descent.The method seeks low-loss hyperparameter combinations for a target dataset and is reported to outperform resampling with significantly less time cost.Source: §1 · p. 2
  • Experiments demonstrate that the model merging strategy significantly improved performance on the MTEB dataset, with a 0.7-point increase in average performance.Source: §1 · p. 3

2 Background

This section reviews general text embedding approaches, task vectors, and model merging, then outlines common merging strategies and a comparison of joint and separate training.

2.1 General Text Embeddings

  • General text embedding work uses multi-task data and contrastive learning, while this paper focuses on limitations in the training process.Source: §2.1 · p. 3

2.2 Task Vectors

  • A task vector is the weight difference between an initial model and one trained on a specific task; its arithmetic can affect task capabilities.Source: §2.2 · p. 3
  • Task-vector norm measures length in weight space, while direction compares how the joint-task vector aligns with each individual-task vector.Source: §2.2 · p. 4

2.3 Model Merging

  • Model merging integrates parameters from same-architecture models into one unified model intended to combine their capabilities.Source: §2.3 · p. 4
  • The five evaluated strategies are Average Merging, SLERP, TIES Merging, Fisher Merging, and RegMean Merging.Source: §2.3 · p. 4
  • Figure 2 includes a comparison of joint-training performance with the best separately trained result on STS and Retrieval.Source: §2.3 · p. 5, Figure 2 · p. 5

3 Problems Analysis

Experiments on English and Chinese STS and Retrieval identify task conflict and data imbalance as challenges in joint training. Dataset size is associated with changes in task-vector norms and directions, while joint training can favor the larger dataset's task vector.

3.1 Training and Evaluation Setup

  • The controlled setup uses ten datasets across STS and Retrieval in English and Chinese, evaluated with MTEB tasks.STS and Retrieval were chosen for their widespread application and differing data requirements.Source: §3.1 · p. 5

3.2 Task Conflict

  • Joint training underperforms task-specific models on 13 of 16 Chinese evaluation tasks and 12 of 25 English tasks.The authors note that detrimental gradient interference may explain the declines.Source: §3.2 · p. 6

3.3 Data Imbalance

  • Dataset sizes vary by nearly a factor of 10, while current methods combine the datasets in their entirety for training.The section examines how these training-data differences affect embedding-model performance.Source: §3.3 · p. 6

3.3.1 Impact on Norm

  • Task-vector norm increases proportionally with the amount of training data across individual-dataset, single-task, and combined-data training.The relationship is reported from Figure 2b.Source: §3.3.1 · p. 6

3.3.2 Impact on Direction

  • Jointly trained task vectors become more aligned with the larger dataset's vector as data-size ratios rise, while dataset size biases vector properties without directly tracking performance.The direction analysis compares jointly trained and separately trained task vectors across dataset pairs.Source: §3.3.2 · p. 6

4 Solutions via Model Merging

The paper uses model merging to address gradient conflict and data imbalance in joint training. Across English and Chinese settings and both training pipelines, merging outperforms joint training on average; direction and norm both matter when selecting task vectors.

4 Solutions via Model Merging

  • Model merging trains tasks independently to avoid gradient conflicts, while scaling and weighting task vectors can mitigate data-imbalance effects.The section presents merging as a response to negative transfer and the influence of training-data amounts on task-vector positions.Source: §4 · p. 6

4.1 Model Merging Pipelines

  • Separate Merging trains distinct models from a common backbone, while Iterative Merging trains sequentially on datasets before merging models.Separate Merging can transform an existing backbone into an embedding model; Iterative Merging is described as useful for improving existing embedding models on specific tasks.Source: §4.1 · pp. 6–7

4.2 Effectiveness of Model Merging

  • Across English and Chinese settings and both training pipelines, model merging methods significantly outperform joint training in average performance.The comparison includes results from joint training and model-merging methods under tuned hyperparameters.Source: §4.2 · p. 7, Figure 3 · p. 7

4.3 Key Factors in Model Merging

  • SLERP has the highest upper-performance limit in three of four settings and performs well across all settings.The section compares SLERP with Averaging because both produce merging results in the interpolation space of task vectors.Source: §4.3 · p. 8

4.3.1 Direction of Vm

  • Task-vector direction significantly affects performance, making direction a key factor in seeking an optimal task vector.The analysis compares Average and SLERP results.Source: §4.3.1 · p. 8, Figure 4 · p. 8

4.3.2 Norm of Vm

  • SLERP task vectors are consistently longer than Average vectors and remain closer to lower-loss regions in the training-data loss landscape.The comparison is reported for the interpolation space under English Separate Merging; other settings yield similar landscapes.Source: §4.3.2 · p. 8, Figure 5 · p. 8

5 Self Positioning in the Interpolation Space

Self Positioning searches the interpolation space of task vectors for low-loss merging positions. Experiments show strong search results and improved performance and time efficiency when addressing data imbalance.

5 Self Positioning in the Interpolation Space

  • The high-performing region depends on task-vector direction and norm, while exhaustive grid search becomes impractical as the number of models grows.The number of direction parameters grows linearly with model count, making grid-search time complexity grow exponentially.Source: §5 · p. 9

5.1 Methodology

  • The method controls merged-vector direction through weights and its norm through a scaling factor, and uses SLERP in the equations and experiments.The approach is framed as a probe-dataset hyperparameter search applicable to any merging method; SLERP is chosen for its superior performance.Source: §5.1 · p. 9
  • Self Positioning uses stochastic gradient descent on a target dataset to search task-vector interpolation positions for low loss.The target dataset can be sampled from training data or constructed from another source, and requires less data than model training; µ is used to prevent overfitting.Source: §5.1 · p. 9

5.2 Experiments

  • Self Positioning achieves the best results in three of four merging settings and performs comparably to joint training in English iterative merging.Table 2 reports average performance across STS and Retrieval tasks; competing merging methods use their best hyperparameter-tuned results.Source: §5.2 · p. 10, Table 2 · p. 10

5.3.1 Effectiveness of the Search Results

  • Search results lie in high-performing regions in most scenarios, and English Separate Merging remains close to its optimum when it misses it.The analysis varies the task-vector weight ratio at fixed λ and varies λ with weights fixed.Source: §5.3.1 · p. 10
  • Low µ values, typically µ ≤0.05, yield optimal performance; µ changes search-result magnitude substantially but direction only slightly.Equation 11 constrains λ to prevent overfitting during optimization.Source: §5.3.1 · p. 11

5.3.3 Efficiency in Addressing Data Imbalance

  • Self Positioning significantly improves performance and time efficiency compared with resampling across the tested data ratios and merging settings.Its target dataset remains fixed at 32,000 instances while the STS-to-Retrieval sampling ratios vary.Source: §5.3.3 · pp. 11–12

6 Applications on Instructor

The paper applies Self Positioning to Instructor, an instruction-following embedding model, using separate and iterative merging. The reported results improve Instructor’s average performance and show that iterative merging can raise classification performance while maintaining average performance.

6 Applications on Instructor

  • The application tests whether Self Positioning can extend merging to a model already jointly trained on 330 tasks.The paper applies Self Positioning and the two merging pipelines to Instructor.Source: §6 · p. 12

6.1 Separate Merging

  • Separate Merging consistently improves on the original Instructor across settings, while Self Positioning remains strong across different target-data sources.The evaluation covers 56 MTEB tasks; Fisher and RegMean show substantial drops when using data from sources other than training data.Source: §6.1 · p. 12

6.2 Iterative Merging

  • Adding classification data and Iterative Merging significantly improves average and classification performance while maintaining average performance; Self Positioning has the highest average among merging methods.The classification training set contains 209,375 instances, and the target dataset combines Instructor and classification training data at a 2:1 ratio.Source: §6.2 · p. 13

7 Conclusion

The conclusion identifies task conflict and data imbalance as challenges for joint training of general text embeddings. It presents model merging and Self Positioning as an effective approach, with Instructor results illustrating applicability to mainstream methods.

7 Conclusion

  • Self Positioning searches the task-vector interpolation space efficiently for high-performing merges, addressing the search cost of model merging and data imbalance.The paper reports that applying the method to Instructor significantly enhances its performance.Source: §7 · p. 13

A.1 Experimental Environment

The appendix specifies the experimental environment, MTEB evaluation coverage, and a RegMean hyperparameter search range.

A.1 Experimental Environment

  • BERTbase experiments use one A100 GPU, while Instructor experiments use two A100 GPUs.The environment uses Python 3.10.14, PyTorch 2.3.0, Transformers 4.37.2, and MTEB 1.2.0.Source: §A.1 · p. 18

A.2 Details of Evaluation

  • MTEB evaluation covers seven task types across 56 datasets, with a designated main metric for each task.The listed task types include classification, clustering, pair classification, reranking, retrieval, STS, and summarization.Source: §A.2 · p. 18

A.3 Hyperparameter Tuning for Model Merging Methods

  • In Section 4.2’s hyperparameter search, RegMean varies the ratio for reducing non-diagonal elements σ over the values 0.0, 0.1, 0.2, 0.3, and 0.4.Source: §A.3 · p. 18

B Task Conflict

English joint training still shows task conflict, while the loss landscape analysis evaluates merged models across interpolations of the STS and Retrieval task vectors.

B Task Conflict

  • 12 of 25 English tasks show a performance decline under joint training, although the negative-transfer effect is less pronounced than in Chinese.Figure 9 compares joint-training performance with the best separately trained performance on STS and Retrieval.Source: B Task Conflict · p. 18, Figure 9 · p. 19

C Loss Landscape

  • The loss-landscape setup positions the STS and Retrieval task vectors using their angle and relative magnitudes, then samples a 20x20 grid of merging weights.The backbone parameters anchor the lower-left corner, and the longer task vector is scaled to 0.7 on an axis.Source: C Loss Landscape · pp. 18–19
  • Table 4 presents clustering statistics for 330 tasks.Source: Table 4 · p. 19
  • Each grid cell’s merged model is evaluated by multi-task loss using proportionally sampled STS and Retrieval data, batch size 32, and an average over 200 steps.Source: C Loss Landscape · p. 19

D Implementation of Clustering

The clustering procedure represents each of 330 tasks with pooled embeddings of sampled query and passage instances before applying KMeans.

D Implementation of Clustering

E Extra Classification Training Data

The extra classification training data comprises English MTEB datasets, with each dataset capped at 25,000 instances; Table 5 also reports broader MTEB method results.

E Extra Classification Training Data

  • The English training data includes eight classification datasets from MTEB, including AmazonReviews, Banking77, Emotion, IMDB, and MTOPIntent.The other listed datasets are AmazonCounterfactual, ToxicConversations, and TweetSentimentExtraction.Source: E Extra Classification Training Data · p. 19
  • Each classification dataset is limited to a maximum of 25,000 instances.Source: E Extra Classification Training Data · p. 19
  • Table 5 reports MTEB results across retrieval, pair classification, classification, and summarization, using T5-large encoders.For Fisher, RegMean, and Self Positioning, the first result in each group uses training data and the second uses alternative data sources.Source: Table 5 · p. 20

Related papers

Loading 2410.15035v1…