Meilisearch 通过二进制量化实现嵌入索引速度提升7倍

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

在 Meilisearch 过去的一年里，我们一直致力于混合搜索，它将关键词搜索与语义搜索相结合，使用Arroy：我们基于 Spotify 技术的向量存储。向量存储是一种数据结构，它能有效地存储嵌入（向量），以实现快速且相关的检索，这被称为近似最近邻 (ANN) 搜索。

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

二进制量化来解救

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

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

我们面临的一个问题是，比特无法表示 -1 和 1，而只能表示 0 和 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 维嵌入仅需 8 GiB 的磁盘空间。

1536 维嵌入

二进制量化被证明对管理大型嵌入具有优势。它显著减少了搜索时间，并且仅使用原始磁盘空间的 15%。索引约 500,000 个 1536 维嵌入所需时间仅为原先的四分之一。尽管如此，对相关性仍然存在明显影响...

3072 维嵌入

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

结论

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

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

• 对于 1536 维数据，索引速度提升近 5 倍；对于 3072 维数据，索引速度提升约 7 倍。
• 磁盘大小缩小到约 6.5 倍。
• 对于 1536 维及以上的嵌入，搜索时间略有改善。

我们建议 OpenAI 用户使用 Meilisearch/Arroy 的二进制量化功能处理大型嵌入。在即将发布的文章中，我们将比较 Meilisearch/Arroy 与竞争对手，重点关注我们的过滤 DiskANN 功能。敬请期待！