基于自动编码器的预训练模型方法模型预训练方法RetroMAE和RetroMAE-2

文章目录

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经过掩码后的受污染输入记为 X ~ e n c \tilde{X}{enc} X~enc，其中掩码率为15~30%。Bert-like encoder记作 Φ e n c ( ⋅ ) \Phi{enc}(\cdot) Φenc(⋅)，它被用来将 X ~ e n c \tilde{X}{enc} X~enc转化为向量 h X ~ h{\tilde{X}} hX~：
h X ~ ← Φ e n c ( X ~ e n c ) ( 1 ) h_{\tilde{X}} \leftarrow \Phi_{enc}(\tilde{X}_{enc}) \qquad \qquad (1) hX~←Φenc(X~enc)(1)

作者使用CLS token的隐向量状态作为句子向量。

解码

设解码阶段输入句子X经过掩码后的受污染输入记为 X ~ d e c \tilde{X}{dec} X~dec，其中掩码率为50~70%， X ~ d e c \tilde{X}{dec} X~dec和编码器生成的向量 h X ~ h_{\tilde{X}} hX~会组合成如下序列（ e x i e_{x_i} exi是 x i x_i xi的embedding， p i p_i pi是位置向量）：
H X ~ d e c ← [ h X ~ , e x 1 + p 1 , ⋯ , e x N + p N ] ( 2 ) \mathbf{H}{\tilde{X}{dec}} \leftarrow [h_{\tilde{X}}, e_{x_1}+ p_1, \cdots, e_{x_N} + p_N] \qquad \qquad (2) HX~dec←[hX~,ex1+p1,⋯,exN+pN](2)

使用仅一层transformer作为decoder，记作 Φ d e c \Phi_{dec} Φdec，它通过优化如下目标函数来重建原始句子X（CE是交叉熵损失）：
L d e c = ∑ x i ∈ m a s k e d C E ( x i ∣ Φ d e c ( H X ~ d e c ) ) ( 3 ) \mathcal{L}{dec} = \sum{x_i \in masked} CE(x_i |\Phi_{dec}(\mathbf{H}{\tilde{X}{dec}})) \qquad \qquad (3) Ldec=xi∈masked∑CE(xi∣Φdec(HX~dec))(3)

因为decoder结构很简单，且使用了激进的掩码率，所以就强迫encoder必须生成高质量的句子向量来成功重建原始输入。

增强解码

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

增强解码为解码操作生成两个输入流： H 1 ( q u e r y ) \mathbf{H}1(query) H1(query)和 H 2 ( c o n t e x t ) \mathbf{H}2(context) H2(context)：
H 1 ← [ h X ~ + p 0 , ⋯ , h X ~ + p N ] H 2 ← [ h X ~ , e x 1 + p 1 , ⋯ , e x N + p N ] ( 4 ) \begin{gathered} \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{gathered} \qquad \qquad (4) H1←[hX~+p0,⋯,hX~+pN]H2←[hX~,ex1+p1,⋯,exN+pN](4)

上式中 h X ~ h_{\tilde{X}} hX~是句子向量， e x i e_{x_i} exi是token embedding（在此处没有token会被掩码）， p i p_i pi是位置向量。

然后引入位置相关注意力掩码矩阵 M ∈ R L × L \mathbf{M} \in \mathbb{R}^{L \times L} M∈RL×L，则自注意力过程可表示为下列式子：
Q = H 1 W Q , K = H 2 W K , V = H 2 W V ; M i j = { 0 , can be attended, − ∞ , masked; A = softmax ⁡ ( Q T K d + M ) V . ( 5 ) \begin{gathered} \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}{i j}= \begin{cases}0, & \text { can be attended, } \\ -\infty, & \text { masked; }\end{cases} \\ \mathbf{A}=\operatorname{softmax}\left(\frac{\mathbf{Q}^T \mathbf{K}}{\sqrt{d}}+\mathbf{M}\right) \mathbf{V} . \end{gathered} \qquad \qquad (5) Q=H1WQ,K=H2WK,V=H2WV;Mij={0,−∞, can be attended, masked; A=softmax(d QTK+M)V.(5)

输出 A \mathbf{A} A和 H 1 \mathbf{H}1 H1（残差连接）一起用来重建输入，目标函数如下：
L d e c = ∑ x i ∈ X C E ( x i ∣ A , H 1 ) ) ( 6 ) \mathcal{L}{dec} = \sum_{x_i \in X} CE(x_i |\mathbf{A}, \mathbf{H_1})) \qquad \qquad (6) Ldec=xi∈X∑CE(xi∣A,H1))(6)

每一个token x i x_i xi基于掩码矩阵 M \mathbf{M} M的第i行重建，所以注意力掩码矩阵 M \mathbf{M} M由如下规则生成（可参考上面图片右侧矩阵来理解），第一个位置(除第一行在对角线上的元素外)和 s ( X ≠ i ) s(X_{\neq i}) s(X=i)对于重建 x i x_i xi是可见的，而对角线上的元素总是会被掩码的（也就是说每一个token在重建过程中不能将自己包括进来）：
M i j = { 0 , x j ∈ s ( X ≠ i ) , or j ∣ i ≠ 0 = 0 − ∞ , otherwise. ( 7 ) \mathbf{M}{ij} = \begin{cases}0, & x_j \in s(X{\neq i}), \text{or } j_{|i \neq 0} =0 \\ -\infty, & \text { otherwise. }\end{cases} \qquad \qquad (7) Mij={0,−∞,xj∈s(X=i),or j∣i=0=0 otherwise. (7)

使用了增强解码的预训练流程算法如下图，encoder使用与BERT一致的掩码语言模型任务(MLM)，其损失记作 L e n c \mathcal{L}_{enc} Lenc，它与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经过掩码后的受污染输入记为 X ~ e n c \tilde{X}{enc} X~enc，其中掩码率为30%。Bert-like encoder记作 Φ e n c ( ⋅ ) \Phi{enc}(\cdot) Φenc(⋅)，它被用来将 X ~ e n c \tilde{X}{enc} X~enc转化为向量[CLS]向量 h X ~ h{\tilde{X}} hX~和普通token向量 H X ~ e n c \mathbf{H}{\tilde{X}{enc}} HX~enc：
h X ~ , H X ~ e n c ← Φ e n c ( X ~ e n c ) ( 1 ) h_{\tilde{X}},\ \mathbf{H}{\tilde{X}{enc}} \leftarrow \Phi_{enc}(\tilde{X}_{enc}) \qquad \qquad (1) hX~, HX~enc←Φenc(X~enc)(1)

掩码token由常规MLM预测得到，MLM的损失函数记为 L m l m \mathcal{L}_{mlm} Lmlm

[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 ∈ R d × ∣ V ∣ W^O \in \mathbb{R}^{d \times |V|} WO∈Rd×∣V∣，d是向量尺寸， ∣ V ∣ |V| ∣V∣是词汇表尺寸）：
μ x i ← h x i T W O , x i ∈ X ~ e n c \mu_{x_i} \leftarrow h^T_{x_i} W^O,\ x_i \in \tilde{X}_{enc} μxi←hxiTWO, xi∈X~enc

上述结果接着经过token-wise max-pooling后聚合得到（对于每个词汇在 X ~ e n c \tilde{X}{enc} 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是整个词汇表)：
m i n . − ∑ x ∈ s e t ( X ) l o g e x p ( μ X ~ e n c [ x ] ) ∑ x ′ ∈ V e x p ( μ X ~ e n c [ x ′ ] ) 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}]) } min.−x∈set(X)∑log∑x′∈Vexp(μX~enc[x′])exp(μX~enc[x])

整个训练任务的目标函数为encoder损失、[CLS]解码损失和上式损失之和：
m i n . L m l m + L d e c + L B o W min. \mathcal{L}{mlm} + \mathcal{L}{dec} + \mathcal{L}_{BoW} min.Lmlm+Ldec+LBoW

向量表征

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

将[CLS]向量 h X h_X hX经过线性转换到更低维度 d ′ d^{\prime} d′： h ^ X ← h X T W c l s , W c l s ∈ R d × d ′ \hat{h}_X \leftarrow h^T_X W^{cls},\ W^{cls} \in \mathbb{R}^{d \times d^{\prime}} h^X←hXTWcls, Wcls∈Rd×d′
使用稀疏化的方式将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 I_X IX是 μ X [ i ] ∈ T o p − k ( μ X ) \mu_X[i] \in Top-k(\mu_X) μX[i]∈Top−k(μX)，k是 μ X \mu_X μX里要被保留的元素个数。
对于每一个文档，将前面两项拼起来作为其语义表征： [ h ^ X ; μ ^ x ] [\hat{h}_X; \hat{\mu}_x] [h^X;μ^x]。

对于每一个query，其与文档的相关性基于如下形式的内积得到：
⟨ q , d ⟩ = h ^ q T h ^ d + ∑ I d μ q [ i ] μ d [ i ] \langle q, d \rangle = \hat{h}^T_q \hat{h}d + \sum{I_d} \mu_q[i] \mu_d[i] ⟨q,d⟩=h^qTh^d+Id∑μq[i]μd[i]

总结

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

参考资料

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