RetroMAE

RetroMAE 出自论文《RetroMAE: Pre-Training Retrieval-oriented Language Models Via Masked Auto-Encoder》，是一种针对于检索的基于自动编码器的预训练模型方法。

作者认为对于基于自动编码的预训练如下两个因素很重要：

• 重建任务必须对于编码质量有要求
• 预训练数据必须得到充分应用

因此，RetroMAE为了满足这两个因素有三个主要设计：

• MAE工作流，预训练方法遵循掩码自动编码工作流。一个输入句子经过两次不同的掩码处理后生成两个掩码输入：一个掩码输入经过编码器encoder生成句子向量；另一个掩码输入和句子向量一起输入到解码器decoder后通过掩码语言模型MLM(masked language modeling)重建原始句子。
• 非对称结构，RetroMAE的模型结构是非对称的，encoder是完整的BERT模型，可用来生成输入句子的向量。decoder只是一层简单的transformer，它被来重建输入句子。
• 非对称的掩码率，encoder输入使用一个适中的掩码率：15_{30%（比BERT略高），而decoder输入的掩码率激进多了：50}70%。

RetroMAE详情

编码

设输入句子X经过掩码后的受污染输入记为${\tilde{X}}_{enc}$，其中掩码率为15~30%。Bert-like encoder记作${\Phi }_{enc}(\cdot )$，它被用来将${\tilde{X}}_{enc}$转化为向量$h_{\tilde{X}}$：
$$h_{\tilde{X}}\leftarrow {\Phi }_{enc}({\tilde{X}}_{enc})(1)$$
作者使用CLS token的隐向量状态作为句子向量。

解码

设解码阶段输入句子X经过掩码后的受污染输入记为${\tilde{X}}_{dec}$，其中掩码率为50~70%，${\tilde{X}}_{dec}$和编码器生成的向量$h_{\tilde{X}}$会组合成如下序列（$e_{x_i}$是$x_i$的embedding， $p_i$是位置向量）：
$${\mathbf{H}}_{{\tilde{X}}_{dec}}\leftarrow [h_{\tilde{X}},e_{x_1}+p_1,\cdots ,e_{x_N}+p_N](2)$$
使用仅一层transformer作为decoder，记作 ${\Phi }_{dec}$，它通过优化如下目标函数来重建原始句子X（CE是交叉熵损失）：
$${\mathcal{L}}_{dec}=\sum _{x_i\in masked}CE(x_i|{\Phi }_{dec}({\mathbf{H}}_{{\tilde{X}}_{dec}}))(3)$$
因为decoder结构很简单，且使用了激进的掩码率，所以就强迫encoder必须生成高质量的句子向量来成功重建原始输入。

增强解码

前述解码过程中的交叉熵损失函数是从掩码token得到，而每一个掩码token总是从相同的上下文${\mathbf{H}}_{{\tilde{X}}_{dec}}$重建，作者认为如果满足下列因素预训练效果能够进一步提升：1）从输入句子中得到更多训练信号(training signals)，2）重建任务能够基于多样化的上下文。所以作者们提出了增强解码。

增强解码为解码操作生成两个输入流：${\mathbf{H}}_1(query)$和${\mathbf{H}}_2(context)$：
$$\begin{array}{c}{\mathbf{H}}_1\leftarrow [h_{\tilde{X}}+p_0,\cdots ,h_{\tilde{X}}+p_N]\\ {\mathbf{H}}_2\leftarrow [h_{\tilde{X}},e_{x_1}+p_1,\cdots ,e_{x_N}+p_N]\end{array}(4)$$
上式中$h_{\tilde{X}}$是句子向量，$e_{x_i}$是token embedding（在此处没有token会被掩码），$p_i$是位置向量。

然后引入位置相关注意力掩码矩阵$\mathbf{M}\in {\mathbb{R}}^{L\times L}$，则自注意力过程可表示为下列式子：
$$\begin{array}{c}\mathbf{Q}={\mathbf{H}}_1{\mathbf{W}}^Q,\mathbf{K}={\mathbf{H}}_2{\mathbf{W}}^K,\mathbf{V}={\mathbf{H}}_2{\mathbf{W}}^V;\\ {\mathbf{M}}_{ij}=\{\begin{array}{ll}0,& \text{ can be attended, }\\ -\infty ,& \text{ masked; }\end{array}\\ \mathbf{A}=\mathrm{softmax}\left(\frac{{\mathbf{Q}}^T\mathbf{K}}{\sqrt{d}}+\mathbf{M}\right)\mathbf{V}.\end{array}(5)$$
输出$\mathbf{A}$和${\mathbf{H}}_1$（残差连接）一起用来重建输入，目标函数如下：
$${\mathcal{L}}_{dec}=\sum _{x_i\in X}CE(x_i|\mathbf{A},{\mathbf{H}}_1))(6)$$
每一个token $x_i$基于掩码矩阵$\mathbf{M}$的第i行重建，所以注意力掩码矩阵$\mathbf{M}$由如下规则生成（可参考上面图片右侧矩阵来理解），第一个位置(除第一行在对角线上的元素外)和$s(X_{\ne i})$对于重建$x_i$是可见的，而对角线上的元素总是会被掩码的（也就是说每一个token在重建过程中不能将自己包括进来）：
$${\mathbf{M}}_{ij}=\{\begin{array}{ll}0,& x_j\in s(X_{\ne i}),\text{or }{j}_{|i\ne 0}=0\\ -\infty ,& \text{ otherwise. }\end{array}(7)$$

使用了增强解码的预训练流程算法如下图，encoder使用与BERT一致的掩码语言模型任务(MLM)，其损失记作${\mathcal{L}}_{enc}$，它与decoder的损失求和作为预训练模型的最终损失函数。

RetroMAE-2

RetroMAE-2出自论文《RetroMAE-2: Duplex Masked Auto-Encoder For Pre-Training Retrieval-Oriented Language Models》，是RetroMAE团队提出的改进版本。

RetroMAE只使用了CLS的token的隐状态向量作为语义表示，RetroMAE-2希望将其他token的embedding也利用起来，所以提出了名为DupMAE（Duplex Masked Auto-Encoder）的自动编码框架。

RetroMAE-2详情

编码

设输入句子X经过掩码后的受污染输入记为${\tilde{X}}_{enc}$，其中掩码率为30%。Bert-like encoder记作${\Phi }_{enc}(\cdot )$，它被用来将${\tilde{X}}_{enc}$转化为向量[CLS]向量$h_{\tilde{X}}$和普通token向量${\mathbf{H}}_{{\tilde{X}}_{enc}}$：
$$h_{\tilde{X}},{\mathbf{H}}_{{\tilde{X}}_{enc}}\leftarrow {\Phi }_{enc}({\tilde{X}}_{enc})(1)$$
掩码token由常规MLM预测得到，MLM的损失函数记为${\mathcal{L}}_{mlm}$

[CLS]解码

与RetroMAE的增强解码实现方式一模一样，参见前面记录。

OT解码和训练目标

OT向量的解码任务基于两个考虑：

• 与[CLS]解码一样，解码网络尽可能简单化
• 与[CLS]解码任务的目标函数不一样，因此两个类型的向量可以捕捉互补信息

OT向量(掩码token除外) H X ~ e n c : { h x 1 , . . . , h x N } \mathbf{H}_{\tilde{X}_{enc}}:\{h_{x1}, ..., h_{x_N} \} HX~enc​​:{hx1​,...,hxN​​}经线性转换到词汇空间（$W^O\in {\mathbb{R}}^{d\times |V|}$，d是向量尺寸，$|V|$是词汇表尺寸）：
$${\mu }_{x_i}\leftarrow h_{x_i}^T{W}^O,x_i\in {\tilde{X}}_{enc}$$
上述结果接着经过token-wise max-pooling后聚合得到（对于每个词汇在${\tilde{X}}_{enc}$里所有token中的最大激活值将被保留）：
μ X ~ e n c ← t o k e n . M a x ( { μ x i ∣ X ~ e n c } ) \mu_{\tilde{X}_{enc}} \leftarrow token.Max(\{\mu_{x_i} | \tilde{X}_{enc} \}) μX~enc​​←token.Max({μxi​​∣X~enc​})
尝试去恢复输入的BoW特征的目标函数如下，目的是使OT向量能够更好编码词汇信息(式中，$x \in set(X) $是输入X中的去重token，V是整个词汇表)：
$$min.-\sum _{x\in set(X)}log\frac{exp({\mu }_{{\tilde{X}}_{enc}}[x])}{{\sum }_{x^{\prime }\in V}exp({\mu }_{{\tilde{X}}_{enc}}[x^{\prime }])}$$

整个训练任务的目标函数为encoder损失、[CLS]解码损失和上式损失之和：
$$min.{\mathcal{L}}_{mlm}+{\mathcal{L}}_{dec}+{\mathcal{L}}_{BoW}$$

向量表征

使用如下的聚合方法将[CLS]向量和OT向量统一起来。

1. 将[CLS]向量$h_X$经过线性转换到更低维度$d^{\prime }$：${\hat{h}}_X\leftarrow h_X^T{W}^{cls},W^{cls}\in {\mathbb{R}}^{d\times d^{\prime }}$
2. 使用稀疏化的方式将OT embedding减少维度： μ ^ X ← { i : μ X [ i ] ∣ i ∈ I X } \hat {\mu}_X \leftarrow \{i: \mu_X[i] \ |\ i \in I_X \} μ^​X​←{i:μX​[i] ∣ i∈IX​}，式中的$I_X$是${\mu }_X[i]\in Top-k({\mu }_X)$，k是${\mu }_X$里要被保留的元素个数。
3. 对于每一个文档，将前面两项拼起来作为其语义表征：$[{\hat{h}}_X;{\hat{\mu }}_x]$。

对于每一个query，其与文档的相关性基于如下形式的内积得到：
$$⟨q,d⟩={\hat{h}}_q^T{\hat{h}}_d+\sum _{I_d}{\mu }_q[i]{\mu }_d[i]$$

总结

本文记录了RetroMAE和RetroMAE-2的原理，RetroMAE是针对文本检索专门优化的预训练方法，BGE文本向量使用了RetroMAE来预训练基础模型。RetroMAE-2进一步在RetroMAE基础上考虑利用词性信息，看到网上有人将其称为RetroMAE+keyword，而2024年的BGE-M3向量模型虽然仍基于RetroMAE来预训练模型，但是向量本身也是考虑了稀疏向量的，与RetroMAE-2有相似之处。

参考资料

1. Liu, Zheng, and Yingxia Shao. 2022. “RetroMAE: Pre-Training Retrieval-Oriented Transformers via Masked Auto-Encoder,” May.
2. Xiao, Shitao, and Zheng Liu. 2022. “RetroMAE v2: Duplex Masked Auto-Encoder For Pre-Training Retrieval-Oriented Language Models,” November.
3. RetroMAE github
4. 知乎文章：RetroMAE+key word=RetroMAE-2