【LLM 论文】OpenAI 基于对比学习微调 LLM 得到嵌入模型

论文：Text and Code Embeddings by Contrastive Pre-Training
⭐⭐⭐⭐
OpenAI

一、论文速读

这篇论文基于大型生成式 LLM 通过对比学习来微调得到一个高质量的 text 和 code 的 embedding 模型。

训练数据的格式：是一堆 $(x_i,y_i)$ 形式的 positive pairs，这些 pairs 的 $x_i$ 与 $y_i$ 存在着语义相似性或者上下文相关性。在具体训练时，使用 in-batch negatives 的方法来构造出 negative pairs。

在训练 text embedding model 和训练 code embedding model 时，所使用的数据不一样：

• 训练 text embedding model 时，互联网上相邻 piece 的文本就作为 positive pair
• 训练 code embedding model 时，function 上方的 docstring 与其实现作为 positive pair $(text,code)$

Embedding Model：这个 encoder 将一个 sentence 转为一个 vector representation：

具体做法是：对于输入 sentence input $x$，将其前后拼接上 [SOS] 和 [EOS] 两个特殊 token，之后输入给 LLM，然后将 [EOS] token 的 hidden state 作为编码后的 embedding 结果。

相似度计算方式：将 x 和 y 分别输入给 encoder 得到 representation，然后计算 cosine 相似度即可：

该工作训练得到的模型，可以为 text 和 code 编码出高质量的 representation，其 embedding 效果在文本 classification、sentence similarity、semantic search 等实验中都表现出不错的效果。

二、模型的训练

模型的推理过程较为简单直接，这里重点详细介绍一下他的训练过程。

前面说了训练数据的格式，按照上面所述，我们可以得到一批用于训练模型的数据，然后将其划分为多个 mini-batch 对模型迭代训练微调。

每一个 mini-batch 由 $M$ 个 positive pairs，在使用它们的训练过程中，首先使用 in-batch negatives 的方法可以得到这个 batch 的 negative pairs，然后针对每一个 pair，将 x 与 y 输入给 encoder 得到其向量表示，然后计算这个 pair 的 logit 值：

其中 $\tau$ 是一个可训练的温度参数，用于调整相似度的缩放。

这样，可以计算出 $M\times M$ 个 logits 组成一个 logit matrix，其中对角线上的元素是 positive pairs 的 logits，其余是负样本的 logits，然后针对 matrix 的每一行每一列计算 cross-entropy 损失，这样的损失是希望模型让正样本的 logit 最大化，让负样本的 logit 最小化，这里计算损失的伪代码如下：

可以看到，$l_r$ 和 $l_c$ 分别是行方向和列方向上计算的交叉熵损失，然后取它们的平均值作为最终的损失函数值，它衡量了区分 positive pairs 和 negative pairs 方面的表现。

通过最小化这个损失函数，模型学习到的嵌入表示能够捕捉输入样本之间的语义相似性，从而在不同的下游任务中发挥作用。

关于 cross entropy 损失，如果忘了可以参考相关资料，从而理解这里为什么用这个函数。

三、实验

论文使用 GPT-3 来初始化 text embedding model 的基座模型，使用 Codex 来初始化 code embedding model 的基座模型，分别对两个模型进行了训练。

论文做了很多实验，这里选取了在文本检索任务上做的实验结果：

其中，cpt-text 就是本论文基于 GPT-3 微调得到的 text embedding model，可以看到：

• 模型越大（参数越多），其效果可以越好
• 该模型的效果远远超过了无监督的检索模型，其中 cpt-text XL 版本的表现已经可以与微调的 SOTA 模型相比较了

四、分析 ⭐

4.1 Batch Size 的影响

论文的消融实验研究了 batch size 对模型训练的影响。

已经公认的是，由于对比学习的机制，当 batch size 越大时，学习难度就越高，使得模型可以学得更好。论文的实验也证明了这一点：

4.2 训练时长的影响

一般认为，训练越久，效果越好，后面会趋于稳定。但是该论文的实验显示，训练越久，在检索和分类任务上效果越好，但是对于句子相似度任务来说，却恰好相反。

这其实涉及到一个很重要的讨论点：训练 embedding 时，检索任务和句子相似度任务是存在冲突的。

这两个任务看似类似，都可以基于 sentence embedding 完成，但两种却是属于两个研究方向，各有自己侧重的研究 topic。从本质上，这两个任务就是存在冲突的，比如一个句子及其否定可以被相似度任务认为是相关的，但在句子相似度任务中则却不应该被认为是“相似-相似”的。

在本论文的 model checkpoint 评估时，也是给检索任务和分类任务赋予了更高的重要性，因为它们往往与现实世界的应用相关联，而句子相似度任务则没那么重要。

总结

这篇论文表明，基于 GPT 系列的 LLM，在足够大的批处理的无监督数据上进行对比预训练，可以获得高质量的文本和代码的 embedding 表示，并可以应用于其他的下游任务中得到不错的表现。