Matryoshka Representation Learning (MRL)-俄罗斯套娃向量表征学习

前言

在2024年1月底OpenAI发布新的向量模型，并提到新的向量模型支持将向量维度缩短。向量模型支持缩短维度而又不会威胁到向量的表示能力的原因在于使用了Matryoshka Representation Learning。

Matryoshka Representation Learning (MRL)是2022年发表的论文，由于OpenAI的使用得到了很多关注，论文共同一作甚至写了一篇博客来解释MRL的原理。开源文本向量nomic-embed-text-v1.5 也应用了MRL支持使用时调整向量维度。

MRL原理

MRL名字里的Matryoshka是指俄罗斯套娃，俄罗斯套娃是一系列可以一个套一个的娃娃，所以MRL希望训练好的向量在去掉向量后面的一些维度后仍有表征含义。

MRL learns representations of varying capacities within the same high-dimensional vector through explicit optimization of O(log(d)) lower-dimensional vectors in a nested fashion, hence the name Matryoshka.

MRL的架构图如下

假设要训练一个$d$维向量，定义一个嵌套维度集合$\mathcal{M}$，希望对于$m\in \mathcal{M}$，MRL可以使得$m$向量$z_{1:m}\in {\mathbb{R}}^m$能够作为可独立表示一个输入x的向量。

$\mathcal{M}$的大小需小于$log(d)$，即$|\mathcal{M}|\le \lfloor log(d)\rfloor$，一般$\mathcal{M}$取值为最大向量大小不断打对折直到可包含最低信息含量的大小，比如论文中训练MRL选取的嵌套维度集$\mathcal{M}$为 M = { 8 , 16 , 32 , ⋯   , 1024 , 2048 } \mathcal{M}=\{8, 16, 32, \cdots,1024,2048\} M={8,16,32,⋯,1024,2048}。

训练的损失被定义为嵌套维度集合$\mathcal{M}$中各个向量对应的损失之和（参考上面MRL架构图的右半部分）。论文以多标签分类为例定义了如下损失函数示例，$x_i\in \mathcal{X}$是输入，$y_i\in [L]$是其对应的标签，N是训练集大小，$(c_m\ge 0{)}_{m\in \mathcal{M}}$表示每个向量维度的相对重要度（论文作者取的是1），$W^{(m)}\in {\mathbb{R}}^{L\times m}$是每一个m维向量对应的线性层参数；$F\left(x_i;{\theta }_F\right)$是参数为${\theta }_F$的深度学习网络，它会得到一个m为的嵌入向量$z_{1:m}\in {\mathcal{R}}^m$；$\mathcal{L}$是多标签交叉熵损失函数。
min ⁡ { W ( m ) } m ∈ M , θ F 1 N ∑ i ∈ [ N ] ∑ m ∈ M c m ⋅ L ( W ( m ) ⋅ F ( x i ; θ F ) 1 : m ; y i ) \min _{\left\{\mathbf{W}^{(m)}\right\}_{m \in \mathcal{M}}, \theta_F} \frac{1}{N} \sum_{i \in[N]} \sum_{m \in \mathcal{M}} c_m \cdot \mathcal{L}\left(\mathbf{W}^{(m)} \cdot F\left(x_i ; \theta_F\right)_{1: m} ; y_i\right) {W(m)}m∈M​,θF​min​N1​i∈[N]∑​m∈M∑​cm​⋅L(W(m)⋅F(xi​;θF​)1:m​;yi​)
上式数学公式看起来很复杂，简单理解MRL就是在训练时使模型同时输出多个指定维度的向量，累加应用这些向量后的损失之和作为目标损失函数，通过大量实验表明这样做达到了预期效果。但是作者们并不知道为什么有效，只是假设MRL学习了嵌套子空间的从粗到细的层次结构，并有效地将信息打包在所有向量空间。

特别的是，MRL训练好的向量除了在训练时的指定维度表现很好外，在其他维度上效果也很好（如下图），有插值效果。所以OpenAI的新向量模型才说用户可以指定比最大维度小的任意维度（相当于把向量序列末尾部分去掉）。

注：论文作者写的博客中写的是MRL的损失是各个子向量损失的平均，对此邮件咨询了作者，作者说求和或求平均对于结果没有影响，关键是在训练时要学习多个向量维度。

Whether you just sum or average doesn’t really matter - the key is that the model needs explicitly supervision at log(d) scales to learn representations (dolls) at each scale. Averaging keeps the loss scale the same as a vanilla (“fixed feature” from our paper), and not averaging will just have a larger scale than the vanilla model. Practically this will not affect training convergence

参考资料：

1. Matryoshka Representation Learning

2. 论文作者写的blog

3. 知乎文章、HN讨论

4. Nomic 开源的contrastors 实现了对比学习版的MRL