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

TL;DR

A novel method, Self Positioning, is introduced, which efficiently searches for optimal model combinations within the interpolation space of task vectors using stochastic gradient descent, which outperforms traditional resampling methods while reducing computational costs.

Abstract

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.

Figures (9)

• Figure 1: During joint training across multiple tasks, model weights fluctuate towards local minima for each task, which leads to two primary issues: 1) Gradient conflicts reduce performance; 2) Data imbalances result in unequal optimization frequencies, biasing the final weights towards more prevalent tasks. Model merging mitigates these challenges by isolating task training and achieving a more balanced weight position.
• Figure 2: Experimental Results on Task Conflict and Data Imbalance Problems.
• Figure 3: Performance of models under different languages, merging settings, and model merging methods during hyperparameter tuning.
• Figure 4: The relationship between the direction of task vectors and the average performance on STS and Retrieval tasks in different settings.
• Figure 5: The loss landscape on the interpolation space of Retrieval and STS task vectors, under English Separate Merging. Other settings yield similar results.