PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards General Text Embeddings with Multi-stage Contrastive Learning

Zehan Li, Xin Zhang, Yanzhao Zhang, Dingkun Long, Pengjun Xie, Meishan Zhang

arXiv:2308.03281v1cs.CL

TL;DR

Existing text embedding models often specialize in individual tasks, while pre-training data and task-specific prompts can limit generality and reuse. GTE addresses this with multi-stage contrastive learning over diverse open-source data, combining unsupervised pre-training with supervised fine-tuning. The resulting model shows strong generalization across text and code embedding benchmarks, although its current implementation is limited to English inputs shorter than 512 tokens and faces data-contamination concerns.

  • Problem

    Previous text embedding models primarily target specific tasks, while in-house pre-training data and task-specific prompts limit broad reuse and implementation simplicity.

  • Method

    GTE applies multi-stage contrastive learning, using diverse open-source text pairs for unsupervised pre-training and varied human-annotated data for supervised fine-tuning.

  • Results

    GTE demonstrates strong and broadly comparable or superior performance across multiple text embedding benchmarks and code search evaluations, including results competitive with larger models.

  • Takeaways & Limitations

    The approach provides a general text-and-code embedding model intended as a robust baseline across varied NLP and code-related tasks.

  • Takeaways & Limitations

    The current model handles only English text shorter than 512 tokens and uses exact-match deduplication, leaving multilingual, long-context, and contamination concerns unresolved.

Abstract

from arXiv · show

We present GTE, a general-purpose text embedding model trained with multi-stage contrastive learning. In line with recent advancements in unifying various NLP tasks into a single format, we train a unified text embedding model by employing contrastive learning over a diverse mixture of datasets from multiple sources. By significantly increasing the number of training data during both unsupervised pre-training and supervised fine-tuning stages, we achieve substantial performance gains over existing embedding models. Notably, even with a relatively modest parameter count of 110M, GTE$_\text{base}$ outperforms the black-box embedding API provided by OpenAI and even surpasses 10x larger text embedding models on the massive text embedding benchmark. Furthermore, without additional fine-tuning on each programming language individually, our model outperforms previous best code retrievers of similar size by treating code as text. In summary, our model achieves impressive results by effectively harnessing multi-stage contrastive learning, offering a powerful and efficient text embedding model with broad applicability across various NLP and code-related tasks.

1 Introduction

Text embeddings support many NLP applications, but task-specific training and in-house data limit their generality and accessibility. GTE addresses this with multi-stage contrastive learning over open-source data and reports strong performance across text and code embedding evaluations.

  • Text embeddings represent texts as low-dimensional vectors for tasks including classification, retrieval, question answering, and dialogue systems.
  • Previous embedding models often target specific tasks, so their training strategies or architectures may not transfer optimally across contexts.SimCSE, for example, uses symmetric text pairs and has limitations in text retrieval.
  • Existing research relies on in-house pre-training data and task-specific prompts, limiting reuse of model weights or APIs and increasing implementation effort.
  • GTE uses contrastive pre-training on large-scale unsupervised text pairs from diverse open-source sources, followed by supervised fine-tuning on high-quality human-labeled pairs.
  • A 110M-parameter BERT-based GTE model outperforms OpenAI's commercial embedding API and ranks highly on the MTEB benchmark after supervised fine-tuning.
  • GTE also evaluates code search across six programming languages without language-specific fine-tuning, supporting its broad text-and-code embedding scope.

2 Related Work

Text representation research spans sentence embeddings, dense retrieval, unified contrastive models, and cross-task benchmarks. GTE combines diverse open-source weakly supervised pre-training data with varied supervised fine-tuning and avoids task-specific prompts.

  • Text embeddings address lexical mismatch and improve retrieval efficiency, while pretrained language models require additional methods to produce high-quality sentence embeddings.
  • Dense retrieval methods train dual encoders with positive and negative query-document pairs, using self-supervised or web-structure-based strategies to construct training examples.
  • Recent work seeks unified text representation models through large-scale contrastive or prompt-based learning and evaluates stability across tasks and domains.
  • GTE pre-trains on naturally occurring text pairs from diverse public sources to improve domain generalization and better align language-model training with contrastive representation learning.
  • GTE uses varied supervised data and no task-specific prompts, improving versatility, reproducibility, and ease of use.

3 Approach

GTE uses a two-stage contrastive-learning pipeline that combines diverse unsupervised text pairs with supervised relevance data to train general text embeddings. Its architecture, sampling, loss, and training choices are designed to support broad task coverage and efficient optimization.

  • Training pipeline: Training proceeds through unsupervised contrastive pre-training followed by supervised fine-tuning with human-annotated relevance data and hard negatives.The fine-tuning mixture covers symmetric and asymmetric tasks across web search, open-domain QA, NLI, fact verification, and paraphrase data.
  • Model architecture: The model uses a Transformer dual encoder with mean pooling to convert contextualized token representations into dense text vectors.A language model produces token representations, and lightweight mean pooling produces the text representation used by the encoder.
  • Contrastive objective: Contrastive training distinguishes relevant query-document pairs from irrelevant documents using similarity scores and an InfoNCE-style objective.The objective uses positive and negative pairs, with similarity estimated from the resulting text embeddings.
  • Training data: Unsupervised pre-training draws text pairs from diverse public sources, including weakly supervised web relationships, while supervised fine-tuning uses approximately 3M annotated pairs.The data mixture spans multiple formats, domains, and tasks to support versatile text representations.
  • Data sampling: Data sampling uses a multinomial distribution over pre-training subsets to account for differences in source sizes.The sampling probability depends on each subset’s size, and the method sets α = 0.5 while keeping batch instances within the same task.
  • Improved contrastive loss: The improved contrastive loss is bidirectional, using both in-batch queries and documents as expanded negative samples.The first terms contrast queries with documents, while the remaining terms perform the inverse contrast; cosine similarity is used with temperature τ = 0.01.
  • Training details: Pre-training uses large batches, distributed negatives, a maximum sequence length of 128, and 50,000 optimization steps, while fine-tuning uses batch size 128 and group size 16.The implementation supports large-scale training across multiple GPUs and uses mixed-precision computation.

4 Experiments

The experiments evaluate GTE across zero-shot classification, unsupervised retrieval, MTEB, and code search, using benchmark comparisons across supervised and unsupervised settings. GTE generally matches or outperforms larger or task-specific baselines, including without task-specific prompts or language-specific code fine-tuning.

  • Zero-shot Text Classification: 110M-parameter GTE matches prompted E5large on SST-2 zero-shot classification in the vanilla setting.Prompting further improves GTE and narrows the gap with larger models.
  • Unsupervised Text Retrieval: GTEbase significantly outperforms comparable-size SimCSE, Contriever, and E5 on unsupervised BEIR retrieval, while remaining comparable to E5large without human supervision.Figure 2 reports Recall@100, while Table 5 reports nDCG@10; the comparisons use unsupervised checkpoints without annotated data.
  • Evaluation setup: GTE is evaluated on 56 English MTEB datasets spanning seven tasks, alongside zero-shot classification, BEIR retrieval, and CodeSearchNet code search.MTEB covers classification, clustering, pairwise classification, reranking, retrieval, semantic textual similarity, and summarization.
  • Massive Text Embedding Benchmark: In unsupervised MTEB evaluation, GTE outperforms E5 across all considered tasks without task-specific prompts and narrows the gap with larger supervised baselines.The reported improvement is attributed to more training data formats and diverse self-supervision sources.
  • Code Search: On CodeSearchNet, GTE surpasses code-pre-trained models despite treating programming languages as text and avoiding separate fine-tuning for each language.The evaluation uses all code candidates from the development and test sets rather than 1,000 randomly sampled candidates.

5 Analysis

The analysis examines how data, batch size, model scale, training duration, training strategy, sampling ratio, and contrastive objective affect MTEB performance. Results support diverse data, multi-stage training, enlarged negative pools, and appropriate scaling choices.

  • Number of Training Datasets: 33 pre-training datasets were compared with smaller mixtures of 5 and 15 datasets to assess the effect of training-data diversity.The ablation also varied the datasets used for fine-tuning.
  • Pre-training Batch Size: Performance saturates at around a batch size of 10,000, with no further gain from larger batches.Training steps were held fixed while batch size was doubled progressively.
  • Number of Model Parameters: As model size grows from 30M to 330M parameters, performance improves linearly as size increases exponentially.The compared sizes correspond to the small, base, and large BERT scales.
  • Training Steps: 20k training steps approximately mark performance saturation and training convergence.The comparison evaluated model performance at different training steps during unsupervised contrastive pre-training.
  • Training Strategies: Unsupervised pre-training on web-scale text pairs outperforms supervised fine-tuning alone, while sequential pre-training and fine-tuning further refines the embeddings.The three compared strategies were pre-training only, fine-tuning only, and pre-training followed by fine-tuning.
  • Training Data Mixture: α = 0.5 improves retrieval, STS, and average MTEB performance, whereas uniform sampling and direct combination of all sources are not optimal.The study reports task-category and average MTEB results for different pre-training sampling ratios.
  • Contrastive Objective: The improved contrastive loss consistently outperforms vanilla in-batch-negative loss during both pre-training and fine-tuning.The ablation compares an enlarged negative pool with only in-batch negatives and reports average MTEB scores.

6 Discussion

The discussion identifies scope constraints in the current model: English-only, sub-512-token inputs, possible Internet-data contamination, and a non-causal architecture. It also points to extensions for longer, multilingual, and causal or prefix models.

  • Scope: The model handles English text shorter than 512 tokens because it is initialized from BERT and lacks multilingual capabilities.Longer texts must be truncated or split for encoding.
  • Data Contamination: Large-scale Internet pre-training creates a data-contamination risk, and exact text-pair deduplication is described as overly strict.The authors note that quantifying contamination is challenging without detailed training-source information.
  • Architecture: The models use non-causal bidirectional attention, leaving causal or prefix language-model variants for future exploration.The proposed direction could jointly optimize generation and retrieval within one model.

7 Conclusion

The paper concludes that multi-stage contrastive learning and diverse training mixtures produce a general text embedding model with strong cross-task generalization. Future work targets longer contexts, multilingual and multimodal applications, and prompts or instructions.

  • Conclusion: Multi-stage contrastive learning with diverse training data yields a text embedding model applicable across various tasks.The conclusion links the diverse mixture to good generalization for single-vector embeddings.
  • Future Work: Future work will extend the model to longer contexts, multilingual and multimodal applications, and prompts or instructions.These directions are presented as planned areas for further exploration.

A.1 Pre-training Data

The pre-training data combines naturally paired text from web pages, academic papers, hyperlinks, community QA, social media, news, knowledge bases, and code.

  • Web Page: Web-page pairs use titles as queries and body text as documents, drawing from Common Crawl, ClueWeb, and MS MARCO documents.The task retrieves relevant body texts for a short title.
  • Academic Paper: Academic-paper pairs use titles as queries and abstracts as documents from arXiv, bioRxiv, medRxiv, PubMed, and Semantic Scholar.The sources cover a wide range of topics.
  • Hyperlink: Hyperlink pairs connect citation arguments with referenced text from ClueWeb, Wikipedia, and Semantic Scholar citations.The paper characterizes this data as involving more challenging multi-hop reasoning.
  • Community QA: Community-QA data uses title-body consistency and question-answer pairs from StackExchange, Yahoo Answers, WikiHow, and Amazon QA.Text lengths and voting numbers provide simple quality filters.
  • Social Media: Social-media pairs treat posts and comments as positive pairs, using Reddit as a source.Posts are structured into title-body pairs when available.
  • News: News data uses title-body pairs and highlighted sentences from CCNews, MicrosoftNews, and NPR.These elements are used to construct query-document pairs.
  • Knowledge Base: Knowledge-base data mines entity-description pairs from Wikipedia and DBPedia.The descriptions encode textual knowledge about entities or events.
  • Code: Code is treated as text through naturally paired text-code examples from GitHub and StackOverflow, including CodeSearchNet training data.This repurposes text-code relationships as positive pairs.

A.2 Fine-tuning Data

Fine-tuning combines contrastive training data from retrieval, question answering, inference, fact verification, paraphrase, and miscellaneous NLP sources.

  • MS MARCO uses query passages with hard negatives mined from highly ranked documents that exclude positives.
  • Open-domain QA pairs questions with supporting evidence passages and treats highly ranked passages lacking the answer as hard negatives.
  • NLI training uses entailment as positive pairs and contradiction as negative pairs, combining MNLI and SNLI.
  • FEVER supplies argument–Wikipedia source pairs for fact verification, while Quora and StackExchangeDupquestion provide similar-meaning sentence pairs.
  • Additional MEDI and BERRI datasets broaden task and domain coverage, with subsampled pre-training data included during fine-tuning to avoid catastrophic forgetting.
  • The pre-training corpus draws on CommonCrawl, Reddit, hyperlink-mined pairs, and S2ORC citation pairs, while fine-tuning reuses prior research data and training triples.

B Massive Text Embedding Benchmark

The benchmark evaluates frozen text embeddings across classification, clustering, reranking, retrieval, pair classification, semantic similarity, and summarization using task-specific downstream procedures.

  • Classification: Classification freezes the embedding model, trains logistic regression on train-set embeddings, and reports test accuracy.
  • Clustering: Clustering applies minibatch k-means to sentence embeddings, using the number of labels as k and v-measure for evaluation.
  • Reranking: Reranking scores query–reference pairs with cosine similarity and reports mean average precision over smaller, harder-to-distinguish reference sets.
  • Retrieval: Retrieval evaluation is omitted because it is similar to the procedure introduced in the previous section.
  • Pair Classification: Pair classification uses cosine similarity for text pairs and reports average precision at the best binary threshold.
  • Semantic Textual Similarity / Summarization: Semantic textual similarity and summarization compare cosine-similarity rankings against human or reference scores using Spearman correlation.

C Original CodeSearchNet Results

On CodeSearchNet, the model retrieves relevant code from 1K candidates and generally outperforms similarly sized open-source code models across programming languages.

  • 1K candidates form the CodeSearchNet retrieval corpus for each natural-language query.
  • The model is superior in most programming languages to CodeBERT and GraphCodeBERT, which have similar architecture and size.
  • A performance gap remains versus cpt-code, trained with Codex on large-scale code–text pairs, motivating further work to close it.
Loading 2308.03281v1…