- 论文：https://proceedings.neurips.cc/paper_files/paper/2022/file/c32319f4868da7613d78af9993100e42-Paper-Conference.pdf
- 代码：GitHub - RAIVNLab/MRL: Code repository for the paper - "Matryoshka Representation Learning"
- 机构：McGill University, Mila ServiceNow Research ，Facebook CIFAR AI Chair
- 领域：embedding model
- 发表：NeurIPS 2022
研究背景
- 研究问题：这篇文章要解决的问题是如何设计一种灵活的表示学习方法，使其能够适应多个下游任务，并且能够根据任务的计算资源需求进行调整。
- 研究难点：该问题的研究难点包括：现有固定容量的表示在学习新任务时可能过度或不足；如何在保持准确性的前提下，显著减少表示的大小和计算成本；如何扩展表示学习方法以适应不同模态（如视觉、语言）和数据规模（如网页规模）。
- 相关工作：该问题的研究相关工作包括大规模数据集上的通用表示学习（如ImageNet和JFT），对比学习（如Contrastive Learning），以及自然语言处理中的预训练模型（如BERT）。这些工作通常依赖于独立的低维模型、子网络优化或后处理压缩来实现表示的灵活性，但这些方法在训练/维护开销、多次前向传播、存储和内存成本等方面存在不足。
研究方法
这篇论文提出了Matryoshka Representation Learning（MRL）用于解决表示学习中的灵活性问题。具体来说，
-
多粒度表示：MRL通过显式优化嵌套的O(log(d))个低维向量，在高维向量中捕获多粒度信息。每个嵌入的前几个维度是一个信息丰富的低维向量，随着维度的增加，表示逐渐变得粗糙。
优化目标：MRL的目标是学习一个d维表示向量z∈Rd，使得每个嵌套维度m∈M都能独立地作为数据点x的可迁移通用表示。优化目标是使用标准经验风险最小化方法，通过单独的线性分类器来优化每个嵌套维度的多类分类损失。
其中，L是多类softmax交叉熵损失函数，cm是相对重要性权重。
3. 高效实现：为了提高效率，MRL采用了权重绑定技术，即所有线性分类器的权重相同，从而减少内存成本。这种变体称为Efficient Matryoshka Representation Learning（MRL-E）。
实现代码为：
class MRL_Linear_Layer(nn.Module):
def __init__(self, nesting_list: List, num_classes=1000, efficient=False, **kwargs):
super(MRL_Linear_Layer, self).__init__()
self.nesting_list = nesting_list
self.num_classes = num_classes # Number of classes for classification
self.efficient = efficient
if self.efficient:
setattr(self, f"nesting_classifier_{0}", nn.Linear(nesting_list[-1], self.num_classes, **kwargs))
else:
for i, num_feat in enumerate(self.nesting_list):
setattr(self, f"nesting_classifier_{i}", nn.Linear(num_feat, self.num_classes, **kwargs))
def reset_parameters(self):
if self.efficient:
self.nesting_classifier_0.reset_parameters()
else:
for i in range(len(self.nesting_list)):
getattr(self, f"nesting_classifier_{i}").reset_parameters()
def forward(self, x):
nesting_logits = ()
for i, num_feat in enumerate(self.nesting_list):
if self.efficient:
if self.nesting_classifier_0.bias is None:
nesting_logits += (torch.matmul(x[:, :num_feat], (self.nesting_classifier_0.weight[:, :num_feat]).t()), )
else:
nesting_logits += (torch.matmul(x[:, :num_feat], (self.nesting_classifier_0.weight[:, :num_feat]).t()) + self.nesting_classifier_0.bias, )
else:
nesting_logits += (getattr(self, f"nesting_classifier_{i}")(x[:, :num_feat]),)
return nesting_logits借用一张图，很直观：
实验设计
- 数据集：实验使用了多个大规模数据集，包括ImageNet-1K、JFT-300M和ALIGN数据集。对于视觉任务，使用了ResNet50和ViT-B/16模型；对于视觉+语言任务，使用了ALIGN模型；对于语言任务，使用了BERT模型。
- 实验设置：实验中，MRL和MRL-E模型与独立训练的低维表示（FF）、降维（SVD）、子网络方法（slimmable networks）和随机选择的高容量特征进行比较。实验评估了线性分类/探测（LP）和1-最近邻（1-NN）准确性。
- 参数配置：实验中使用的超参数与独立训练的基线模型相同。例如，ResNet50输出2048维表示，ViT-B/16和BERT-Base输出768维嵌入。
本文将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个样本的逐步重排，这种方法可以省去调参，应用比较方便。
不足与反思
- 嵌套损失权重的优化：未来的工作可以探索自适应损失平衡方法，以实现更优的准确性-效率权衡。
- 不同保真度的损失函数：可以考虑使用针对不同保真度的损失函数，以解决特定方面的自适应部署问题，例如高召回率的8维表示和鲁棒的2048维表示。
- 搜索数据结构的集成：可以在MRL上学习一个可微分的k-d树，以实现数据集和表示感知的检索。
- 多目标MRL的联合优化：结合端到端可学习的搜索数据结构，进行数据驱动的自适应大规模检索，适用于Web规模的搜索应用。
