[论文解读] Towards General Text Embeddings with Multi-stage Contrastive Learning

Zehan Li, Xin Zhang+4 more|arXiv (Cornell University)|Aug 7, 2023

Topic Modeling被引用 57

一句话总结

GTE 通过在多阶段对比学习上，在多样化的开放数据混合上训练一个通用文本嵌入模型，在文本、代码和检索任务上实现强大的零样本与有监督性能，且模型规模相对较小。

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$_ ext{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.

研究动机与目标

• 开发一个在多样化数据源上通过多阶段对比学习训练的统一文本嵌入模型。
• 利用大规模无监督预训练，随后进行有监督微调以提升泛化能力。
• 在文本与代码任务上，展示对更大尺寸或特定任务的嵌入模型具有竞争力甚至优越性。

提出的方法

• 使用带均值池化的双编码器 Transformer 主干网络来获取文本嵌入。
• 在来自多样开放来源的约 8e8 对文本对上进行无监督对比学习预训练。
• 在来自多个任务的约 300 万条注释文本三元组上进行有监督对比学习微调，其中包括难负样本。
• 通过在同批次中扩增查询/文档负样本以及双向项来改进对比损失。
• 以多项分布从多个来源抽样数据批次以平衡来源（α = 0.5）。
• 在零样本检索、MTEB 和代码检索基准上评估，而不使用特定任务提示。

实验结果

研究问题

• RQ1 是否一个通过在多阶段对比学习的多样开放数据上学习的一般文本嵌入模型，能够在标准基准上超越面向特定任务或更大尺寸的嵌入模型？
• RQ2 数据多样性、批量大小和模型规模如何影响无监督预训练和有监督微调的性能？
• RQ3 在不进行语言特定微调的情况下，作为文本嵌入进行训练的代码检索是否能得到有效改进？

主要发现

• GTEbase（110M）能够在大规模文本嵌入基准上超越 OpenAI 的嵌入 API，并超过一些大小约为其十倍以上的模型。
• 在开放数据上的无监督预训练加上有监督微调在 MTEB 任务、BEIR 和代码检索上表现强劲。
• GTEbase 在零样本文本检索和文本分类任务中的表现达到或超过更大模型，且无需任务特定提示。
• GTEbase 在 CodeSearchNet 上显著超越同等规模的最先进代码检索器，且无需语言特定微调。
• 扩大数据多样性和模型规模带来线性收益，批次大小在大约一万个负样本处达到饱和，而多阶段训练优于单阶段训练。