论文分享|NeurIPS2022‘华盛顿大学|俄罗斯套娃表示学习（OpenAI使用的文本表示学习技术）

论文题目：Matryoshka Representation Learning

来源：NeurIPS2022/华盛顿大学+谷歌

方向：表示学习

开源地址：https://github.com/RAIVNLab/MRL

摘要

学习表征对于现代机器学习很重要，广泛用于很多下游任务。大多数情况下，每个下游任务的计算资源和需求都是未知的，因此对于某个下游任务，一个固定长度的表征可能会过大或过小。这引出一个问题：我们能否设计出一个灵活的表示方式，以适应具有不同计算资源的多个下游任务？

本文提出了俄罗斯套娃表示学习Matryoshka Representation Learning (MRL)，编码不同粒度的信息，并让一个编码能够适应不同计算资源的下游任务。MRL最小限度地修改了当前的表示学习流程，没有引入额外的推理和部署开销。MRL学习到了至少和单独训练的低维表征一样准确的从粗到细的表征。MRL具有以下出色性能：(a)在保证ImageNet-1K分类任务精确率的前提下，将表征大小缩小为原来的14倍。(b)将ImageNet-1K和4K上的大规模检索速度加快到原来的14倍。©将长尾少样本分类的准确率提升了2%。(d)提升的同时也能保证鲁棒性。最后本文发现MRL可以很容易地扩展到各种模态的模型，比如视觉模型ViT/ResNet，视觉语言模型ALIGN，语言模型BERT等。

介绍

深度表征进行部署时一般分为两步：(1)昂贵但是固定的前向推理来计算表征 (2)在下游应用中使用表征。使用成本一般需要使用编码维度，数据量(N)，标签空间(L)三方面进行计算。对于大规模数据的网络应用，使用成本一般会超过计算成本。固定的表征大小使得不同任务都需要使用高维编码向量，尽管许多应用的资源和精度需求不同。

人类对于自然世界的理解有很自然的从粗到细的粒度，然而由于梯度训练的归纳偏置，深度学习模型倾向于向整个表征向量扩散信息。之前的方法一般使用三种方式来得到有弹性的表示模型：训练多个低维度模型（ResNet）；联合优化不同容量的子网络[1,2]；向量压缩（SVD）。但每种方法都因为训练和维护费用**，大量昂贵的对所有的数据的前向推理，存储和内存成本，昂贵的实时特征选择或准确率的显著下降**很难满足大规模部署。通过编码不同粗细粒度的与独立训练准确率相当的表征，本文以最小的额外开销学习一种可以在推理过程中无需额外成本自适应部署的表征。

MRL通过嵌套的形式在同一个高维向量中对 $O(\log (d))$ 维度的低维向量做显式优化。下图展示了核心技术，随着维度的增加，表征的粒度也越来越细。

本文主要关注最重要的两项任务，大规模分类和检索。对于分类任务，使用MRL训练的模型中的可变大小表征进行自适应级联（小于等于某个阈值则使用更高维度的向量进行分类**），显著降低了实现特定精度所需的编码的平均维数。对于检索任务，先使用最前面的维度进行粗排**，再使用更多的维度对粗排结果进行精排。同样还可以用在长尾持续学习以及判断样本间分类的困难程度和信息瓶颈。

方法

MRL的目的是学习许多个小于等于 $\lfloor \log d\rfloor$ 的前 $M$ 维表征，即总维度的前 $\frac{1}{2^k}$ 维。以分类任务为例，其实就是将每个维度得到的表征通过各自的分类头计算多分类交叉熵损失再相加。对于ImageNet来说，本文选取了 M = { 8 , 16 , . . . , 1024 , 2048 } M=\{8,16,...,1024,2048\} M={8,16,...,1024,2048} 。 $c_m$ 是每个维数的重要性，本文都设置为1，即看成同等重要。尽管仅对这些前 $\lfloor \log d\rfloor$ 维度向量进行了优化，对于这些维度之间的维度也能取得比较好的效果。

一种比较高效的做法是将每个投射头 $W^{(m)}\in {\mathbb{R}}^{L\times m}$ 看成是一个大投射头的 $W\in {\mathbb{R}}^{L\times d}$ 的一部分，即 $W^{(m)}=W_{1:m}$ ，这种做法在大输出空间时尤其重要，本文称之为Efficient Matryoshka Representation Learning (MRL–E)

对于视觉模型、视觉语言模型中的对比学习，以及语言模型中的掩蔽语言模型，下一词预测等任务都可以转换为类似分类的方式进行处理。(毕竟万物皆是分类

应用

本文将MRL/MRL-E模型与单独训练的低维表征（FF），SVD分解，子网络[2]方法进行了比较

首先是分类任务。对于在ImageNet上训练的模型，线性分类准确率基本和FF保持一致，1-NN准确率甚至在低维时高于FF

对于大规模数据集上训练的模型也取得了很好的精度与速度间的平衡

对于适应性分类，期望的表征大小相比FF减小了14倍。

图像检索的结果也超越了baseline，最高超过了FF 3%。适应性图像检索也达到了效率和精度的权衡，16维度做粗排，2048维度做精排的准确率已经和直接使用2048维度做排序的精度还高，但计算量大幅减小。值得一提的是本文提出了一个漏斗检索方法，即使用逐渐增大的维度16-32-64-128-256-2048 对前200-100-50-25-10个样本的逐步重排，这种方法可以省去调参，应用比较方便。

分析与消融

鲁棒性：在分类和检索任务都具有一定鲁棒性。鲁棒性研究最好选用最近邻分类或检索而不是线性探针

少样本和长尾学习：长尾分布中的新类别可以有2%的提升。且高维度的表征对于比较难的类别分类更准确

不同维度间的分歧：某些类别和实例低维效果可能超过高维，MRL可能可以被用作分析信息瓶颈的工具

超类准确率：维度越高，超类准确率越高

结论

可能的优化方向：(1)优化嵌套损失的权重 (2)在不同的保真度下使用不同的损失，旨在解决自适应部署的特定需求，例如，8维的高召回率和2048维的鲁棒性。(3)在矩阵数据表示之上学习搜索数据结构，如可微k-d树，以实现数据集和表征感知检索。(4)最后，结合多目标MRL与端到端可学习搜索数据结构的联合优化，实现数据驱动的面向大规模搜索应用的自适应大规模检索。

笔者评价：这篇文章主要是对检索过程中KNN相似度匹配部分的速度优化，且能够保持不同维度表征的性能和一致性，可以让大公司为不同用户（个人，开发者）的各种编码需求提供一个相对通用的服务

---

大家好，我是NLP研究者BrownSearch，如果你觉得本文对你有帮助的话，不妨点赞或收藏支持我的创作，您的正反馈是我持续更新的动力！如果想了解更多LLM/检索的知识，记得关注我！

引用

[1]Cai H, Gan C, Wang T, et al. Once-for-all: Train one network and specialize it for efficient deployment[J]. arXiv preprint arXiv:1908.09791, 2019.

[2]Yu J, Yang L, Xu N, et al. Slimmable neural networks[J]. arXiv preprint arXiv:1812.08928, 2018.