Matryoshka Embeddings：多尺度的细节

Engineering

October 30, 2024

Stefan Webb, David Wang

什么是 Matryoshka Embeddings？

在构建高效的向量搜索系统时，一个关键的挑战是管理存储成本，同时保持可接受的延迟和召回率。现代Embeddings 模型输出的向量有成百上千个维度，这给原始向量和索引带来了巨大的存储和计算开销。

传统方法是在建立索引前应用量化或降维方法来降低存储需求。例如，我们可以使用乘积量化（PQ）降低精度，或使用主成分分析（PCA）降低维数，从而节省存储空间。这些方法会对整个向量集进行分析，以找到一个能保持向量间语义关系的更紧凑的向量集。

这些标准方法虽然有效，但只能在单一尺度上降低一次精度或维度。但是，如果我们能同时保持多层细节，就像一座精确度越来越高的表征金字塔呢？

这就是Matryoshka Embeddings。这些巧妙的结构以俄罗斯嵌套娃娃命名（见插图），将多级表示嵌入到单个向量中。与传统的后处理方法不同，Matryoshka 嵌入在初始训练过程中就能学习这种多尺度结构。结果非常显著：不仅完整的 Embeddings 能够捕捉输入语义，而且每个嵌套的子集前缀（前半部分、前四分之一等）都能提供连贯的表示，尽管不那么详细。

Figure: Visualization of Matryoshka embeddings with multiple layers of detail 图具有多层细节的可视化 Matryoshka 嵌入

图具有多层细节的可视化马特罗什卡嵌入法

这种方法与传统的嵌入式方法形成了鲜明的对比，在传统的嵌入式方法中，使用向量维度的任意子集通常会破坏语义。通过 Matryoshka 嵌入，您可以选择最能平衡特定任务的精度和计算成本的粒度。

需要快速近似搜索？使用最小的 "娃娃"。需要最高精度？使用完整的 Embeddings。这种灵活性使它们对于适应不同性能要求或资源限制的系统尤为重要。

推理

Matryoshka 嵌入的一个重要应用是在不影响召回率的情况下加速相似性搜索。通过利用查询和数据库嵌入的较小子集，例如其维度的前 1/32，我们可以在这个缩小的空间上建立一个索引，该索引仍然保留了大量的相似性信息。这种较小嵌入空间的初始结果可以直接使用。不过，还有一种技术可以提高召回率，并考虑到维度缩减带来的任何微小缩减，从而使这种方法对相似性搜索任务而言既高效又有效。

图漏斗搜索如何与 Matryoshka 嵌入配合使用

图漏斗搜索如何与 Matryoshka 嵌入一起工作

为了在保持准确性的同时有效加快相似性搜索速度，我们可以使用 "漏斗搜索 "方法。首先，我们仅使用嵌入维度的前 1/32 执行初始相似性搜索，生成一个广泛的候选项库。然后，我们根据这些候选项与查询的相似度，使用前 1/16 个维度对其进行 Rerankers，从而剪切出一部分列表。这个过程会反复进行，使用嵌入维度中越来越大的子集--1/8、1/4 等进行重新排序和剪枝。重要的是，我们只在这个低维空间中执行一次初始相似性搜索，嵌入模型的一次传递就能计算出查询嵌入。这种漏斗过程在每一步都会缩小候选范围，比直接在全维空间中搜索更快、更高效。从 1/32 维空间中提取许多匹配项，然后通过漏斗搜索对其进行提炼，可以大大加快相似性搜索的速度，同时保持较高的召回率。

培训

让我们来了解一些技术细节。该方法的应用非常简单。考虑微调用于句子嵌入的BERT 模型的上下文。为了将预先训练了掩码标记损失的 BERT 模型转换为句子嵌入模型，我们将句子嵌入形成为最后一层的平均值，即每个标记上下文化嵌入的平均值。

训练目标的一种选择是余弦句子（CoSENT）损失，即 $L(u, v; s)$ L $(u,$ v $;$ s $)$ 。它输入一对句子嵌入， $,v$ u $,$ v，以及它们期望的相似度得分， $ss$ （公式见上面的链接）。现在，为了学习 Matryoshka 嵌入，我们对训练目标稍作修改：

$(u, v) = w_0L(u_{1:d}, v_{1:d}) + w_1L(u_{1:d/2}, v_{1：$ d $/2}) + w_2L(u_{1:d/4}, v_{1:d/4}) + \cdots$ $_{:}$ $, v 1_{:}$ $) + w 1 L(u$ $1_{:}$ $, v 1_{:}$ $) + w L (u$ $1_{:}$ $, v xml-ph$

其中，通过计算前一项输入的一半损失来继续求和，直到达到信息瓶颈。作者建议设置

w $0=w1=⋯=1w_0=w_1=\cdots=1$ w $=$ $w$ = $\dots$ $=$ 。

简单地说，Matryoshka 损失是原始损失在输入递归子集上的加权和。

从上式中可以得到的一个重要启示是，Matryoshka 损失通过在嵌入模型间共享权重，实现了多尺度表征的高效学习（例如，同一模型用于编码 $du_{1:$ u $_{1}$ $_{:}$ 和 $:$ $d/2}$ u $_{1}$ $_{:}$ $d/2$ ），并跨尺度共享维度 $（$ $u1$ $:$ $d/2u_{$ $1:$ $d / 2$ } $u$ $_{1}$ $_{:}$ 是 $uu$ u 的子集）。