Towards General Text Embeddings with Multi-stage Contrastive Learning
Proposes GTE, a general text embedding model trained via multi-stage contrastive learning, with only 110M parameters, surpassing large models in multiple benchmarks.
Key Findings
Methodology
The GTE model employs a deep Transformer encoder within a dual-encoder architecture, utilizing mean pooling to derive dense text vectors. The training process involves two stages: unsupervised pre-training on approximately 800 million text pairs collected from diverse open-source datasets such as CommonCrawl, arXiv, StackExchange, Reddit, Wikipedia, and GitHub, and supervised fine-tuning on around 3 million annotated pairs. During pre-training, an improved contrastive loss (variants of InfoNCE) is used to enhance semantic discrimination, with multi-source data sampling (α=0.5) to balance domain diversity. The fine-tuning stage incorporates hard negatives and task-specific data, optimizing the model further. The training leverages large batch sizes, mixed precision, and distributed GPU setups to ensure efficiency. The model parameters range from 30M to 330M, with maximum sequence lengths from 128 to 512 tokens, tailored for different tasks and scales.
Key Results
- In zero-shot text classification tasks such as SST-2, GTE_base (110M) achieved 87.2% accuracy, rivaling larger models like prompt-tuned E5-6B. On the Massive Text Embedding Benchmark (MTEB), GTE outperformed models like Contriever and E5, demonstrating superior generalization across 56 datasets. In unsupervised text retrieval benchmarks (BEIR), GTE_base significantly outperformed comparable models, with Recall@100 surpassing SimCSE, Contriever, and BM25. Notably, in code retrieval tasks on CodeSearchNet, the model outperformed specialized code retrievers without language-specific fine-tuning, indicating strong cross-task transferability. These results highlight the effectiveness of multi-source, multi-stage contrastive training in producing versatile embeddings.
- The experimental data confirms that multi-source data integration and contrastive learning substantially improve semantic understanding and robustness. The model's ability to perform well across diverse tasks with limited parameters underscores its efficiency. Ablation studies reveal that data diversity and loss function design are critical factors for performance gains. The model's competitive results on both retrieval and classification tasks validate its potential as a universal embedding solution.
Significance
This research addresses longstanding challenges in creating universal text representations capable of handling multiple NLP and code-related tasks. By leveraging large-scale open-source data and multi-stage contrastive learning, the GTE model achieves a remarkable balance between efficiency and performance. Its ability to outperform larger models in zero-shot and few-shot settings signifies a breakthrough in reducing computational costs while maintaining high accuracy. The model's versatility across tasks—from text classification and retrieval to code search—demonstrates its potential to serve as a foundational component in diverse AI applications, lowering barriers for deployment and fostering innovation in industry and academia.
Technical Contribution
The primary technical contribution lies in designing a multi-stage contrastive learning framework that effectively integrates vast unlabeled web data with limited high-quality annotations. The approach introduces a novel data sampling strategy to balance sources, employs an improved contrastive loss (including bidirectional variants with in-batch negatives and hard negatives), and adapts a deep Transformer encoder with mean pooling for robust semantic embedding. This framework enables training a parameter-efficient model (110M) that surpasses larger models in multiple benchmarks. The combination of large-scale open data, multi-task training, and loss function innovations provides a new paradigm for scalable, universal text embedding models.
Novelty
This work is the first to systematically combine multi-source open data with multi-stage contrastive learning to produce a compact yet highly effective universal text embedding model. Unlike prior models focusing on single-task or domain-specific training, GTE leverages diverse datasets and a multi-phase training process to achieve broad applicability. Its ability to outperform models with ten times more parameters, without task-specific prompts or extensive fine-tuning, marks a significant advancement in the field. The integration of code data as text further extends its versatility, making it a pioneering effort in unified text and code representation.
Limitations
- The model's performance may decline in highly specialized or niche domains not well represented in the training data, such as certain scientific or technical fields.
- Despite parameter efficiency, large-scale training still requires significant computational resources, limiting rapid iteration and deployment in resource-constrained environments.
- While the model generalizes well across many tasks, some specific tasks or languages might benefit from targeted fine-tuning, which is not explored in this work.
Future Work
Future research will focus on enhancing domain adaptation capabilities, reducing training costs through model compression, and extending multi-modal learning to incorporate images and videos. Additionally, exploring few-shot and zero-shot learning further, as well as developing more efficient fine-tuning strategies, will be key to deploying such models in real-world applications. The authors also aim to investigate the integration of this embedding framework into larger reasoning systems and multi-modal AI architectures.
AI Executive Summary
The rapid evolution of natural language processing has underscored the importance of high-quality text embeddings as foundational tools for understanding and manipulating textual data. Traditional models, often tailored for specific tasks like sentence similarity or information retrieval, struggle to generalize across diverse NLP applications. Recent advances in large-scale pretraining, exemplified by models like BERT and GPT, have opened new avenues for creating more universal representations. However, these models often require task-specific fine-tuning, and their embeddings can suffer from anisotropic spaces that hinder similarity computations.
Addressing these challenges, this study introduces GTE (General Text Embedding), a novel model trained through a multi-stage contrastive learning framework. The core idea is to leverage vast amounts of openly available web and domain data, combined with a carefully designed training strategy, to produce a compact yet powerful text embedding model. The training pipeline involves two main phases: an unsupervised pretraining stage utilizing approximately 800 million text pairs from diverse sources, and a supervised fine-tuning stage with around 3 million annotated pairs, including hard negatives. This approach ensures the model learns robust semantic representations capable of generalizing across tasks.
The architecture of GTE is based on a deep Transformer encoder, such as BERT or MiniLM, with a dual-encoder setup and mean pooling to generate dense vectors. During training, the model employs an improved contrastive loss, incorporating bidirectional negatives and multi-source sampling, to maximize semantic discrimination. Large batch sizes, mixed-precision training, and distributed GPU setups enable efficient scaling. The model parameters range from 30M to 330M, with maximum sequence lengths adapted for different tasks.
Extensive evaluation demonstrates that GTE outperforms state-of-the-art models across multiple benchmarks. In the Massive Text Embedding Benchmark (MTEB), GTE surpasses models like Contriever and E5, achieving higher accuracy and retrieval metrics across 56 datasets. On zero-shot classification tasks like SST-2, GTE_base reaches 87.2% accuracy, rivaling larger models. In unsupervised retrieval benchmarks such as BEIR, it significantly outperforms comparable models, confirming its robustness. Notably, in code retrieval scenarios, GTE effectively handles multiple programming languages without language-specific fine-tuning, outperforming specialized models.
These results highlight the potential of multi-source, multi-stage contrastive learning to produce versatile, efficient embeddings. The model's ability to perform well across diverse NLP and code tasks with limited parameters suggests a promising direction for scalable AI systems. Its broad applicability, combined with low inference costs, makes GTE a valuable tool for industry applications like search engines, question-answering systems, and software development tools. Looking ahead, future work will focus on further reducing training costs, extending multi-modal capabilities, and enhancing domain adaptation, aiming to make such models more accessible and adaptable for real-world deployment.
Deep Analysis
Background
The field of text embedding has evolved from simple word vectors like Word2Vec and GloVe to contextualized models such as BERT and GPT, which capture richer semantic information. Early models focused on static embeddings, but their inability to handle polysemy limited their effectiveness. The advent of deep Transformer-based models addressed this, enabling context-aware representations. However, these models often produce embeddings with anisotropic spaces, reducing their utility in similarity tasks. To mitigate this, techniques like SimCSE introduced contrastive learning to produce more isotropic embeddings. Simultaneously, large-scale pretraining on web data, combined with task-specific fine-tuning, became the dominant paradigm. Benchmarks like MTEB and datasets like BEIR have provided standardized evaluation frameworks, revealing the strengths and weaknesses of existing models. Despite progress, challenges remain in creating truly universal, parameter-efficient embeddings that generalize across tasks, domains, and modalities. This research aims to fill this gap by leveraging multi-source data and multi-stage contrastive learning to develop a more robust, versatile embedding model.
Core Problem
Existing text embedding models often excel in specific tasks but lack the ability to generalize across multiple NLP and code-related applications. Many rely heavily on task-specific fine-tuning, which limits their flexibility and increases deployment complexity. Additionally, models with large parameters, such as 10B+ models, demand substantial computational resources, making them impractical for many real-world scenarios. Conversely, smaller models tend to underperform in complex semantic tasks. The core challenge is to develop a parameter-efficient, universally applicable embedding model that can handle diverse tasks—from sentence similarity and retrieval to code understanding—without extensive fine-tuning or task-specific prompts. Achieving this requires innovative training strategies that can extract maximum semantic information from large-scale open-source data while maintaining computational efficiency.
Innovation
The key innovation of this work is the integration of multi-source open data with a multi-stage contrastive learning framework. Unlike prior models that focus on single-task or domain-specific training, GTE leverages a diverse dataset comprising web pages, scientific papers, QA forums, social media, knowledge bases, and code repositories. The multi-stage approach involves initial unsupervised pretraining with weakly supervised data, followed by supervised fine-tuning with high-quality annotated pairs, including hard negatives. This strategy enhances the semantic richness and robustness of the embeddings. Additionally, the model employs an improved contrastive loss that incorporates bidirectional negatives and multi-source sampling, ensuring better discrimination of semantically similar and dissimilar pairs. The architecture, based on a deep Transformer encoder with mean pooling, balances parameter efficiency with high performance, enabling the model to outperform larger counterparts in multiple benchmarks.
Methodology
- �� Data collection: Gathered approximately 800 million text pairs from sources such as CommonCrawl, arXiv, StackExchange, Reddit, Wikipedia, DBPedia, and GitHub, ensuring diverse domain coverage.
- �� Unsupervised pretraining: Utilized an improved contrastive loss (variants of InfoNCE) to train the model on large-scale web data, employing multi-source sampling (α=0.5) to balance domain influence. Distributed training with large batch sizes, mixed precision, and gradient accumulation ensured efficiency.
- �� Supervised fine-tuning: Fine-tuned on about 3 million annotated text pairs, incorporating hard negatives mined via an auxiliary retriever. The maximum sequence length increased from 128 to 512 tokens, optimizing for longer texts.
- �� Model architecture: Based on deep Transformer encoders (e.g., BERT, MiniLM), with mean pooling to generate dense vectors. Cosine similarity measures semantic relatedness.
- �� Optimization: Employed AdamW optimizer with linear learning rate decay, warm-up steps, and distributed GPU training across multiple nodes, ensuring convergence and stability.
- �� Evaluation: Conducted extensive testing on benchmarks like MTEB, BEIR, and CodeSearchNet, analyzing metrics such as accuracy, Recall@100, nDCG@10, and average precision. Ablation studies validated the impact of data diversity and loss design.
Experiments
The experimental setup involved evaluating GTE across diverse NLP and code tasks. Datasets included SST-2 for classification, MS MARCO and BEIR for retrieval, STS for semantic similarity, and CodeSearchNet for code retrieval. The models were compared against baselines like SimCSE, Contriever, and E5, with parameters ranging from 30M to 330M. Hyperparameters such as learning rate, batch size, and sequence length were tuned for optimal performance. The training process emphasized large batch sizes and mixed precision to handle computational demands. Ablation experiments assessed the contribution of multi-source data, loss functions, and model size. Performance metrics included accuracy, Recall@100, nDCG@10, and average precision, providing a comprehensive evaluation of the model’s versatility and robustness.
Results
GTE_base (110M) achieved 87.2% accuracy on SST-2 in zero-shot classification, matching larger models. In MTEB, it outperformed models like Contriever and E5 across multiple tasks, with significant gains in retrieval and semantic similarity metrics. On BEIR, Recall@100 exceeded 64% on average, surpassing SimCSE and BM25. In code retrieval, GTE outperformed specialized models without language-specific fine-tuning, demonstrating remarkable cross-task transferability. The ablation studies confirmed that multi-source data and contrastive loss improvements directly contributed to performance boosts. These results collectively validate the effectiveness of the multi-stage contrastive learning framework in producing versatile, high-quality embeddings.
Applications
- �� Search engines and enterprise document retrieval: Enhancing relevance and efficiency across diverse domains.
- �� Question-answering systems: Providing robust semantic representations for multi-turn dialogues and knowledge bases.
- �� Code understanding and search: Supporting multi-language code retrieval without language-specific tuning.
- �� Content recommendation: Improving personalized content filtering and ranking.
- �� Multi-modal integration: Extending to include images and videos for comprehensive AI understanding.
- �� Future AI systems: Serving as a foundational component for large reasoning and multi-task AI architectures.
Limitations & Outlook
- �� Domain-specific performance may decline in highly specialized fields not well represented in training data.
- �� Large-scale training incurs high computational costs, limiting rapid updates.
- �� Some tasks or languages might still require fine-tuning for optimal results.
- �� Future work should focus on model compression, domain adaptation, and multi-modal extension to address these issues.
Plain Language Accessible to non-experts
想象你在一个巨大的厨房里,有许多不同的厨师负责不同的菜系,比如中餐、西餐、甜点等。每个厨师都用不同的食材和方法做菜,但他们都在追求一道目标:做出美味的菜肴。现在,你想让这个厨房变得更聪明,让厨师们能更快找到合适的食材,做出符合客户口味的菜。
为了实现这个目标,你给每个厨师配备了一个智能助手,这个助手可以观察所有厨师的工作,学习他们用的食材、做菜的步骤,然后用一种统一的“厨艺语言”表达出来。这个助手不断看、学、比较,学会了哪些食材和步骤是相似的,哪些是不一样的。
有了这个助手,厨师们只要告诉它“我想做一道辣的川菜”,它就能快速找到用到辣椒、花椒的菜谱,帮厨师们节省了很多时间。这个助手还可以帮厨师们发现不同菜系之间的相似点,比如川菜和湘菜都喜欢用辣椒,帮助厨师们互相借鉴。
这个系统的厉害之处在于,它不用每次都重新学习,而是通过观察大量的菜谱和厨艺资料,建立了一套通用的理解方式。无论是新加入的厨师还是不同的菜系,它都能快速适应,帮你做出更好、更符合需求的菜肴。就像一个聪明的厨房助手,让整个厨房变得更高效、更有创意。
ELI14 Explained like you're 14
想象你有一个超级聪明的朋友,他可以帮你理解任何事情。比如你问他:“这个问题的答案是什么?”他不用看书,只是根据他平时听到的很多话和看到的很多例子,快速猜出答案。这个朋友就像一个特别聪明的机器人,学会了用一种特别的方式,把各种信息变成数字,然后用这些数字来判断两个东西是不是一样。
比如,你给他两个句子,他会把它们变成两个数字的“密码”,然后比较这两个密码的相似度。如果密码很相似,说明两个句子意思差不多;如果差别很大,说明它们的意思不同。这就像你用尺子量两个东西的长度,如果很接近,就说明它们差不多一样。
这个机器人是怎么学会的呢?它看了很多很多的例子,比如网页上的文章、问答、社交媒体的帖子,还有代码片段。它不断学习这些例子,慢慢地明白了什么样的句子意思相似,什么样的句子意思不同。通过这种学习,它变得非常聪明,可以帮你找到相关的内容,甚至帮你写作文或解答问题。
所以,这个研究就是让机器变得像这个聪明的朋友一样,能理解各种不同的文字和代码,不管是问问题、找资料,还是帮忙写东西,都能做到又快又准。它用了一种叫“对比学习”的方法,让机器学会区分相似和不同,从而变得更聪明、更有用。
Abstract
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.
References (20)
Unsupervised Dense Information Retrieval with Contrastive Learning
Gautier Izacard, Mathilde Caron, Lucas Hosseini et al.
Sentence-T5: Scalable Sentence Encoders from Pre-trained Text-to-Text Models
Jianmo Ni, Gustavo Hernández Abrego, Noah Constant et al.
MTEB: Massive Text Embedding Benchmark
Niklas Muennighoff, Nouamane Tazi, L. Magne et al.
Task-aware Retrieval with Instructions
Akari Asai, Timo Schick, Patrick Lewis et al.
Text Embeddings by Weakly-Supervised Contrastive Pre-training
Liang Wang, Nan Yang, Xiaolong Huang et al.
One Embedder, Any Task: Instruction-Finetuned Text Embeddings
Hongjin Su, Weijia Shi, Jungo Kasai et al.
GraphCodeBERT: Pre-training Code Representations with Data Flow
Daya Guo, Shuo Ren, Shuai Lu et al.
SimCSE: Simple Contrastive Learning of Sentence Embeddings
Tianyu Gao, Xingcheng Yao, Danqi Chen
BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
Jacob Devlin, Ming-Wei Chang, Kenton Lee et al.
Challenging Decoder helps in Masked Auto-Encoder Pre-training for Dense Passage Retrieval
Zehan Li, Yanzhao Zhang, Dingkun Long et al.
MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers
Wenhui Wang, Furu Wei, Li Dong et al.
Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval
Lee Xiong, Chenyan Xiong, Ye Li et al.
CodeBERT: A Pre-Trained Model for Programming and Natural Languages
Zhangyin Feng, Daya Guo, Duyu Tang et al.
Language Models are Few-Shot Learners
Tom B. Brown, Benjamin Mann, Nick Ryder et al.
Dense Passage Retrieval for Open-Domain Question Answering
Vladimir Karpukhin, Barlas Oğuz, Sewon Min et al.
MIND: A Large-scale Dataset for News Recommendation
Fangzhao Wu, Ying Qiao, Jiun-Hung Chen et al.
REALM: Retrieval-Augmented Language Model Pre-Training
Kelvin Guu, Kenton Lee, Zora Tung et al.
Pre-training Tasks for Embedding-based Large-scale Retrieval
Wei-Cheng Chang, Felix X. Yu, Yin-Wen Chang et al.
Latent Retrieval for Weakly Supervised Open Domain Question Answering
Kenton Lee, Ming-Wei Chang, Kristina Toutanova
CCNet: Extracting High Quality Monolingual Datasets from Web Crawl Data
Guillaume Wenzek, M. Lachaux, Alexis Conneau et al.
Cited By (20)
Denoising-Aware Inversion: Revealing Privacy Risks in Noise-Protected Text Embeddings
EXCISE: Query-Side Exclusion for Late-Interaction Retrieval
Bekko Embedding: Parameter-Efficient Multilingual Retrieval with Ultra-Compact Encoders
From Queries to Playlists: An LLM-Driven Architecture for Semantic Music Search at Scale
Scaling Dense Retrieval with LLM-Annotated Training Data: Structured Mining and Progressive Curriculum for E-Commerce Sponsored Search
Query-Focused Event Summarization: A Dataset and Benchmark
Agentic table talk
PETRA: Transforming Web Text for Petroleum-Engineering Domain Adaptation
FlowPipe: LLM-Enhanced Conditional Generative Flow Networks for Data Preparation Pipeline Construction
KaLM-Reranker-V1: Fast but Not Late Interaction for Compressed Document Reranking
Scalable text clustering based on word embeddings and noise analysis
STEB: Style Text Embedding Benchmark
Separating Representation from Reconstruction Enables Scalable Text Encoders
MTEB-BR: A Text Embedding Benchmark for Brazilian Portuguese
Uncertainty-Aware Cross-Modal Remote Sensing Image-Text Retrieval via Evidential Learning
Inject or Navigate? Token-Efficient Retrieval for LLM Analysis of Transactional Legal Documents
Large language models create an uneven informational layer over cities
Scaling and Stabilizing Large-Scale Embedding-Based Retrieval
ProgramTab: Boosting Table Reasoning of LLMs via Programmatic Paradigm
PCTD: Preference-Guided Counterfactual Task Decomposition for Agent Tool Retrieval