[论文解读] Text and Code Embeddings by Contrastive Pre-Training

Arvind Neelakantan, Tao Xu+23 more|arXiv (Cornell University)|Jan 24, 2022

Topic Modeling被引用 149

一句话总结

该论文表明，在大规模对比预训练上，针对未标注文本和代码，使用大批量可以得到高质量嵌入，在分类、检索和代码搜索等任务中取得强劲结果，常常接近或超过有监督微调模型。

ABSTRACT

Text embeddings are useful features in many applications such as semantic search and computing text similarity. Previous work typically trains models customized for different use cases, varying in dataset choice, training objective and model architecture. In this work, we show that contrastive pre-training on unsupervised data at scale leads to high quality vector representations of text and code. The same unsupervised text embeddings that achieve new state-of-the-art results in linear-probe classification also display impressive semantic search capabilities and sometimes even perform competitively with fine-tuned models. On linear-probe classification accuracy averaging over 7 tasks, our best unsupervised model achieves a relative improvement of 4% and 1.8% over previous best unsupervised and supervised text embedding models respectively. The same text embeddings when evaluated on large-scale semantic search attains a relative improvement of 23.4%, 14.7%, and 10.6% over previous best unsupervised methods on MSMARCO, Natural Questions and TriviaQA benchmarks, respectively. Similarly to text embeddings, we train code embedding models on (text, code) pairs, obtaining a 20.8% relative improvement over prior best work on code search.

研究动机与目标

• 使用基于成对数据的对比预训练来开发高质量的无监督文本和代码嵌入。
• 证明大批量训练和规模扩大在分类、检索和代码检索等任务中提升性能。
• 在标准基准上评估迁移学习和零-shot 能力。
• 分析训练动态以及模型规模和数据量的影响，尽管在不同任务上性能有所差异。
• 讨论大规模嵌入模型的更广泛影响及计算方面的考量。

提出的方法

• 使用 Transformer 编码器通过特殊 EOS 标记的最后隐藏状态将输入映射到嵌入。
• 使用带有批内负样本的对比目标进行训练，并使用可训练的温度参数 tau。
• 从自然成对数据中构造正样本对（文本的相邻文本片段；代码的 (文本、代码) 对）。
• 将 CPT-text 初始化为来自 GPT 系列模型，将 CPT-code 初始化为来自 Codex 模型，以引导学习。
• 利用极大的批量大小以实现有效的难负样本挖掘并提升嵌入表征。
• 在线性探针分类、句子相似性、大规模语义搜索以及代码搜索等任务上评估嵌入。

实验结果

研究问题

• RQ1 大规模的无监督对比预训练是否能产生高质量的文本和代码嵌入？
• RQ2 这些嵌入在线性探针分类、零-shot、k-NN 和语义/检索任务上的表现如何？
• RQ3 模型规模和批量大小对不同任务的嵌入质量有何影响？
• RQ4 在微调和 BEIR 风格的零-shot 评估中，学习到的嵌入的迁移能力如何？

主要发现

• 最大的无监督文本模型在 7 个线性探针任务上，相对于之前最佳无监督文本嵌入实现了 4% 的相对提升，相对于有监督文本嵌入提升了 1.8%。
• 在大规模文本检索基准测试中，CPT-text 相对于先前的无监督方法，在 MSMARCO 提升 23.4%，Natural Questions 提升 14.7%，TriviaQA 提升 10.6%。
• CPT-code 在 CodeSearchNet 上相对于之前最佳的代码检索结果实现了 20.8% 的相对提升。
• 6B CPT-text 的零-shot SST-2 结果达到 88.1% 的准确率，带提示 89.1%，使用 256-NN 为 93.3%，线性探针 95.7%，在监督设置下接近最先进的 97.5%。
• 在 BEIR 中，无监督 CPT-text 的迁移相对于先前的嵌入方法实现了 5.2% 的相对提升，并与更重的测试时推理基线竞争。
• 代码搜索结果显示 CPT-code 超越跨语言的先前方法，CPT-code S/M 在 CodeSearchNet 基准上达到最先进水平。