Meilisearch 通过二值量化将 embeddings 索引速度提升 7 倍

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

在 Meilisearch 的过去一年里，我们一直在努力进行混合搜索，它将关键词搜索与使用 Arroy：我们基于 Spotify 技术的向量存储的语义搜索相结合。向量存储是一种数据结构，可以高效地存储 embeddings（向量），以便快速和相关地检索，称为近似最近邻 (ANN) 搜索。

当我们的客户开始大量使用 arroy 时，他们中的一些人开始达到机器的极限。例如，对于使用 768 维模型的客户，我们发现在一台拥有 64GiB 内存的机器上，我们无法索引超过 1500 万个 embeddings。

二值量化来救援

这个复杂术语背后的概念是我们将要量化 embeddings。量化数字涉及将其分散到定义数量的值中；对于二值量化，两个值可以用一个比特表示，例如，0.2561 变为 1，而 -0.568 变为 -1。

这意味着我们可以将 32 位浮点数转换为 1 位数，这将使磁盘和内存使用量减少 32 倍。这意味着我们可以在 64GiB 内存中索引多达 4.8 亿个 embeddings，而不是仅限于 1500 万个 embeddings。

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

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

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

以下是此函数的 SIMD neon 特定版本。它处理对应于 1 位量化 embedding 的原始字节，将其转换为正确对齐的 32 位浮点 embedding。使用掩码，该函数分离低位和高位，在 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
}

丹尼尔·莱米尔似乎没有找到更好的算法 🤭

结果

我们之前讨论了这种方法的理论优势，但它在现实生活中效果如何？以下是总结结果的不同表格供您参考。

展开查看原始基准

1024 维度的 Embeddings

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

1536 维度的 Embeddings

二值量化证明对于管理更大的 embeddings 很有优势。它显著减少了搜索时间，仅使用了原始磁盘空间的 15%。大约 50 万个 1536 维 embeddings 的索引在四分之一的时间内完成。尽管如此，对相关性的影响仍然很明显...

3072 维度的 Embeddings

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

结论

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

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

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

我们建议 OpenAI 用户对大型 embeddings 使用 Meilisearch/arroy 的二值量化功能。在即将发布的文章中，我们将把 Meilisearch/arroy 与竞争对手进行比较，重点是我们的过滤 DiskANN 功能。敬请关注！