如果说自回归(auto-regressive)模型的目标是让人能够更好地使用AI的话, 那么embedding模型就是要让机器更好地使用AI。 不同于自回归模型对输入做进一步展开,输出更多的自然语言内容,embedding模型输出的是对输入的一个向量化表达, 这个由一堆浮点数组成的向量蕴含了机器对输入句子的数值化理解。 由于语义相似的两句话在这个向量化表达下计算出的cosine距离也更相近, 那么它就可以用来进一步做分类,聚类,搜索等等一系列任务。这也是向量化数据库的基础。


模型结构

通常来说,encoder是这类任务最自然的选择,因为在transformer中,encoder天生就是生成一个向量表达给到decoder使用的。 由于模型的输出大小是[batch, seqlen, dim],需要有一个方法来收缩seqlen这一维。 最简单的就是直接取[CLS]这个特殊token的输出。也可以使用更复杂的方法例如取平均等等。

这个任务当然也可以由decoder来承担。从模型结构上,decoder只是多了一层masked-attention, 去掉自回归相关的结构和机制,decoder的输出和encoder没有本质的区别。 OpenAI的embedding模型使用了以next-token为目标预训练的decoder初始化, 在相似度数据集上做了进一步的contrastive预训练任务。


训练和数据

同自回归模型一样,embedding的训练也可以分成预训练(pre-training)和微调(fine-tuning) 两步。预训练通过大量数据,让模型学会分辨哪些句子是相似的,哪些是不相似的,从而产生较为通用的编码能力。 预训练的模型针对特定任务,比如搜索,或者特定的领域,比如医药,可以做进一步的微调, 使得模型的编码能够更贴合下游任务的需求。

预训练(pre-training)

预训练要求模型学会区分相似和不相似的语句,因此需要从数据中构造相似语句对。 这个构造方法有很多,比如OpenAI就简单地从文章中选取相邻的两句话作为相似对。 另外常见的做法例如(标题,正文)(正文,评论)等等。 微软的Text Embeddings by Weakly-Supervised Contrastive Pre-training一文详细 介绍了他们从多种数据源中提取相似语句对的做法。

由于预训练需要的数据量巨大,这一步通常只能来源于无标注数据。于是数据的清理也是非常重要的, 同样是上面这篇文章,提到了一种数据清理的思路:先通过1.3B的数据对训练一个模型, 然后用这个模型再去给原数据打分和清理,最后挑出270M个数据对再去做训练。 这个方法的原理在于大模型通常会首先记住干净数据,然后才缓慢地overfit脏数据, 因此让1.3B数据训练的模型从一大堆候选集里做排序,它通常还是能够排得比较准的。 这个数据清洗的思路理论上可以一直迭代,清理出更加优质的数据集。

同样是数据量的原因,不相似的数据对一般就直接从训练时同一个batch里找了。 具体来讲,预训练优化一个被称为in-batch negative contrastive loss的目标,

\[L = - \frac{1}{N} \sum_{i} \log \frac{e^{s_\theta (q_i, p_i^+)}}{e^{s_\theta (q_i, p_i^+)} + \sum_j e^{s_\theta (q_i, p_{i,j}^-)}}\]

其中\(s_\theta (q, p)\)表示查询\(q\)和结果\(p\)之间的相似度,通常定义为embedding向量 \(\mathbb{E}_q\)和\(\mathbb{E}_p\)的余弦相似度,

\[s_\theta (q, p) = \cos (\mathbb{E}_q, \mathbb{E}_p) / \tau\]

相应的,\(p_i^+\)是正例样本,\(p_{i,j}^-\)是负例样本,\(N\)是batch size, \(\tau\)是一个被称作”temperature”的超参数,可以通过训练得到,也可以预设一个例如0.01这样的经验值。

几乎所有embedding预训练的论文都提到,模型的效果会随着batch-size的增大而提升。 为什么呢?因为越大的batch-size意味着更多的负例样本,而这一点对于检索而言尤其重要, 在实际检索中,向量数据库面对的将是大得多的候选数据集。 让模型在预训练这一步学到从海量数据里区分正例样本的能力于是就显得尤其重要了。 在上面微软的文章中,他们的batch-size高达32K。

也有一些预训练工作使用了不同于contrastive的目标函数。例如BGE 用了MAE的思路,让模型尝试通过被污染文本的embedding \(\mathbb{E}_{\tilde{X}}\)来还原原始文本:

\[\min \sum_{x\in X} -\log Dec(x \| e_{\tilde{X}})\]

其中Dec是一个轻量的decoder,\(X\)和\(\tilde{X}\)分别是原始和被污染的文本。

微调(fine-tuning)

相较之下,微调需要的数据量就小多了,对数据质量的要求也就相对比较高一些, 通常会使用人工标注的数据。再者,虽然同样是contrastive learning, 微调阶段一般不会用到in-batch negative,也不会使用很大的batch size, 取而代之的是预先从数据集中挖掘一些负例样本。 挖掘的方式也很多,常见的比如使用一个已有的embedding模型, 对每个query找到一组相似度非常低的文本。

通常来说,微调是针对某个特定领域或者特定任务来做的。如果想让模型同时胜任多种不同的任务, 比如既要做检索,又要做分类,或者既做医药,又做金融,那就要想个办法让模型能够区分不同的任务, 从而在生成embedding的时候有所侧重。一个常用的方法是在文本输入中加上对这个任务的描述, 比如把给到模型的输入从普通的一句话X,改成“给出这句话的分类:X”,或者“从Wikipedia里搜索相关文章:X”,等等。 也就是给prompt加上了指令,被称为instruction-finetuning。


参考文献

[1] Xiao, Shitao, et al. “C-pack: Packaged resources to advance general chinese embedding.” arXiv preprint arXiv:2309.07597 (2023).

[2] Gao, Luyu, and Jamie Callan. “Condenser: a pre-training architecture for dense retrieval.” arXiv preprint arXiv:2104.08253 (2021).

[3] Giorgi, John, et al. “Declutr: Deep contrastive learning for unsupervised textual representations.” arXiv preprint arXiv:2006.03659 (2020).

[4] Su, Hongjin, et al. “One embedder, any task: Instruction-finetuned text embeddings.” arXiv preprint arXiv:2212.09741 (2022).

[5] Zhang, Peitian, et al. “Retrieve anything to augment large language models.” arXiv preprint arXiv:2310.07554 (2023).

[6] Qu, Yingqi, et al. “RocketQA: An optimized training approach to dense passage retrieval for open-domain question answering.” arXiv preprint arXiv:2010.08191 (2020).

[7] Gao, Tianyu, Xingcheng Yao, and Danqi Chen. “Simcse: Simple contrastive learning of sentence embeddings.” arXiv preprint arXiv:2104.08821 (2021).

[8] Asai, Akari, et al. “Task-aware retrieval with instructions.” arXiv preprint arXiv:2211.09260 (2022).

[9] Neelakantan, Arvind, et al. “Text and code embeddings by contrastive pre-training.” arXiv preprint arXiv:2201.10005 (2022).

[10] Wang, Liang, et al. “Text embeddings by weakly-supervised contrastive pre-training.” arXiv preprint arXiv:2212.03533 (2022).

[11] Izacard, Gautier, et al. “Unsupervised dense information retrieval with contrastive learning.” arXiv preprint arXiv:2112.09118 (2021).