Meilisearch 使用二值量化将嵌入索引速度提高7倍

这篇博客文章最初发表在 kero 的个人博客上。

在过去的一年里，Meilisearch 一直在努力开发混合搜索，它将关键字搜索与语义搜索结合起来，利用 Arroy：我们的向量存储，基于 Spotify 技术。向量存储是一种数据结构，可以高效地存储嵌入（向量），以便快速且相关地检索，这被称为近似最近邻 (ANN) 搜索。

当我们的客户开始大量使用 Arroy 时，他们中的一些人开始达到机器的极限。例如，对于使用 768 维模型的客户，我们发现一台拥有 64GiB RAM 的机器无法索引超过 15M 的嵌入。

二值量化来救援

这个复杂术语背后的概念是我们将要对嵌入进行*量化*。量化一个数字涉及将其分解为一组定义的数值；对于*二值*量化，两个数值可以表示在一个比特上，例如，0.2561 变为 1，-0.568 变为 -1。

这意味着我们可以将一个 32 位浮点数转换为一个 1 位数字，这将使磁盘和 RAM 使用量减少 32 倍。这意味着我们可以使用 64GiB RAM 索引多达 4.8 亿个嵌入，而不是仅限于 15M 个嵌入。

我们面临的一个问题是，比特只能表示 0 和 1，而不能表示 -1 和 1。我们必须使用一些技巧才能使其工作，并在计算中将 0 伪装成 -1。例如，我们不得不大量优化的一项功能是将二值量化向量转换为 32 位浮点向量的功能。

我们不公开原始量化嵌入，而是利用此函数向用户提供真实的 32 位浮点嵌入。然而，将 1 位量化嵌入转换为 32 位浮点的主要原因是为了在生成超平面以将嵌入分组和保持相关 ANN 时确保高精度。

在 1 分 53 秒的运行中，将量化嵌入转换回 32 位浮点数需要 24.21 秒。

下面是此函数的 SIMD Neon 专用版本。它处理对应于 1 位量化嵌入的原始字节，将它们转换为正确对齐的 32 位浮点嵌入。使用掩码，该函数分离低位和高位，以四批次的形式在 128 位寄存器中处理它们，然后将其转换为四个 32 位浮点数。这些值被整体写入输出向量。

unsafe fn to_vec_neon(bytes: &[u8]) -> Vec<f32> {
    let mut output: Vec<f32> = vec![0.0; bytes.len() / u32::BITS];
    let output_ptr = output.as_mut_ptr();
    let low_mask = [0b_0000_0001, 0b_0000_0010, 0b_0000_0100, 0b_0000_1000];
    let high_mask = [0b_0001_0000, 0b_0010_0000, 0b_0100_0000, 0b_1000_0000];
    let ones = unsafe { vld1q_dup_f32(&1.0) };
    let minus = unsafe { vld1q_dup_f32(&-1.0) };

    for (current_byte, base) in bytes.iter().enumerate() {
        unsafe {
            let base = *base as u32;
            let base = vld1q_dup_u32(&base);
            for (i, mask) in [low_mask, high_mask].iter().enumerate() {
                let mask = vld1q_u32(mask.as_ptr());
                let mask = vandq_u32(base, mask);
                // 0xffffffff if equal to zero and 0x00000000 otherwise
                let mask = vceqzq_u32(mask);
                let lane = vbslq_f32(mask, minus, ones);
                let offset = output_ptr.add(current_byte * 8 + i * 4);
                vst1q_f32(offset, lane);
            }
        }
    }

    output
}

Daniel Lemire 似乎没有找到更好的算法 🤭

结果

我们前面讨论了这种方法的理论优势，但在实际生活中效果如何呢？以下是为您总结结果的不同表格。

打开查看原始基准

1024 维嵌入

在检查小型嵌入时，我们可以观察到二值量化对相关性有显著影响。这种方法的主要优点在于更快的插入时间和更少的存储需求。不建议为小型嵌入启用此功能，因为它们已经占用最少的磁盘空间。存储超过四百万个 1024 维嵌入仅需要 8GiB 的磁盘空间。

1536 维嵌入

二值量化在管理大型嵌入方面显示出优势。它显著缩短了搜索时间，并且仅使用原始磁盘空间的 15%。索引约 50 万个 1536 维嵌入只需四分之一的时间。尽管如此，对相关性仍有明显影响……

3072 维嵌入

二值量化对大型嵌入的影响显而易见。虽然相关性可能会略受影响，但磁盘使用量、搜索时间和索引时间的减少是显著的。能够以仅 15% 的空间存储相同数量的嵌入，并将索引速度提高 6 倍，这是一个重大成就！存储一百万个 3072 维二值量化嵌入所需的磁盘空间是存储四百万个 1024 维嵌入的一半。

结论

考虑到算法中信息的减少，二值量化对召回率的影响预计会降低。然而，对于超过 1500 维的数据集，与磁盘空间和索引速度的显著优势相比，对召回率的影响相对较小。

实际上，我们观察到空间利用率并没有减少 32 倍，而是仅减少到原始空间的 15%，这是一个显著的成就。这主要是因为 Arroy 不仅存储用户向量，还存储内部节点以进行高效搜索。分割节点（超平面）以完整的 32 位浮点数形式保留，以保持精度。

• 在 1536 维度下，索引速度提高了近 5 倍；在 3072 维度下，索引速度提高了约 7 倍。
• 磁盘大小缩小了约 6.5 倍。
• 从 1536 维开始的嵌入搜索时间略有改善。

我们建议 OpenAI 用户使用 Meilisearch/Arroy 的二值量化功能处理大型嵌入。在即将发布的一篇文章中，我们将比较 Meilisearch/Arroy 与竞争对手，重点介绍我们的过滤式磁盘 ANN 功能。敬请期待！