Int4:Lucene 中的更多标量量化

时间:2024-04-28 07:21:12

作者:来自 Elastic Benjamin Trent, Thomas Veasey

在 Lucene 中引入 Int4 量化

在之前的博客中,我们全面介绍了 Lucene 中标量量化的实现。 我们还探索了两种具体的量化优化。 现在我们遇到了一个问题:int4 量化在 Lucene 中是如何工作的以及它是如何排列的?

存储量化向量并对其进行评分

Lucene 将所有向量存储在一个平面文件中,从而可以根据给定的序数检索每个向量。 你可以在我们之前的标量量化博客中阅读对此的简要概述。

现在 int4 为我们提供了比以前更多的压缩选项。 它将量化空间减少到只有 16 个可能的值(0 到 15)。 为了更紧凑的存储,Lucene 使用一些简单的位移操作将这些较小的值打包到单个字节中,在 int8 已经节省 4 倍空间的基础上,还可以节省 2 倍空间。 总之,使用位压缩存储 int4 比 float32 小 8 倍。

图 1:这显示了 int4 所需的字节数减少,压缩后大小比 float32 减少了 8 倍。

int4 在计算延迟方面也有一些好处。 由于已知值在 0-15 之间,因此我们可以利用准确了解何时担心值溢出并优化点积计算。 点积的最大值为 15*15=225,可以容纳在一个字节中。 ARM 处理器(如我的 MacBook)的 SIMD 指令长度为 128 位(16 字节)。 这意味着对于 Java Short,我们可以分配 8 个值来填充通道。 对于 1024 维,每个通道最终将累积总共 1024/8=128 次乘法,最大值为 225。所得的最大和 28800 完全符合 Java 短值的限制,我们可以迭代更多值​​一次比。 下面是 ARM 的一些简化代码。

// snip preamble handling vectors longer than 1024
// 8 lanes of 2 bytes 
ShortVector acc = ShortVector.zero(ShortVector.SPECIES_128);
for (int i = 0; i < length; i += ByteVector.SPECIES_64.length()) {
  // Get 8 bytes from vector a
  ByteVector va8 = ByteVector.fromArray(ByteVector.SPECIES_64, a, i);
  // Get 8 bytes from vector b
  ByteVector vb8 = ByteVector.fromArray(ByteVector.SPECIES_64, b, i);
  // Multiply together, potentially saturating signed byte with a max of 225
  ByteVector prod8 = va8.mul(vb8);
  // Now convert the product to accumulate into the short
  ShortVector prod16 = prod8.convertShape(B2S, ShortVector.SPECIES_128, 0).reinterpretAsShorts();
  // Ensure to handle potential byte saturation
  acc = acc.add(prod16.and((short) 0xFF));
}
// snip, tail handling

计算误差修正

有关误差校正计算及其推导的更详细说明,请参阅误差校正标量点积

这是一个简短的总结,遗憾的是(或高兴的是)没有复杂的数学。

对于存储的每个量化向量,我们还跟踪量化误差校正。 回到标量量化 101 博客中,提到了一个特定的常数:

该常数是从基本代数导出的简单常数。 但是,我们现在在存储的浮点数中包含与舍入损失相关的附加信息。

在这里,i 是每个浮点向量维度,i^{^{'}} 是每个浮点向量维度,且 \alpha = \frac{max-min}{(1<<bits)-1}

这有两个后果。 第一个是直观的,因为这意味着对于给定的一组量化桶,我们会稍微更准确,因为我们考虑了量化的一些损耗。 第二个结果有点微妙。 现在这意味着我们有一个受量化分桶影响的纠错措施。 这意味着它可以被优化。

寻找最佳分桶

进行标量量化的简单而简单的方法可以让你走得很远。 通常,你选择一个置信区间,从中计算向量值允许的极端边界。 Lucene 和 Elasticsearch 中的默认值是

1−1/(dimensions+1) 。图 2 显示了一些 CohereV3 嵌入样本的置信区间。 图 3 显示了相同的向量,但使用静态设置的置信区间进行标量量化。

图 2:CohereV3 维度值的样本。

图3:CohereV3 维度值量化为 int7 值。 最后的那些尖峰是什么? 嗯,这就是量化过程中截断极值的结果。

但是,我们留下了一些不错的优化。 如果我们可以调整置信区间来移动存储桶,从而允许更重要的维度值具有更高的保真度,该怎么办? 为了优化,Lucene 执行以下操作:

  • 从数据集中采样大约 1,000 个向量并计算它们真正的最近 10 个邻居。
  • 计算一组候选上分位数和下分位数。 该集合是通过使用两个不同的置信区间来计算的:1−1/(dimensions+1) 及 1−(dimensions/10)/(dimensions+1)。这些间隔处于相反的极端。 例如,具有 1024 个维度的向量将搜索置信区间 0.99902 和 0.90009 之间的候选分位数。
  • 对这两个置信区间之间存在的分位数子集进行网格搜索。 网格搜索找到使量化得分误差的确定系数与之前计算的真实 10 个最近邻相比最大的分位数。
图 3:Lucene 搜索置信区间空间并测试 int4 量化的各种存储桶。

图 4:为此 CohereV3 样本集找到的最佳 int4 量化桶。

有关优化过程和优化背后的数学原理的更完整说明,请参阅优化截断间隔

量化速度与大小

正如我之前提到的,int4 为你提供了性能和空间之间有趣的权衡。 为了说明这一点,以下是 CohereV3 500k 向量的一些内存要求。

图 5:CohereV3 500k 向量的内存要求。

当然,我们看到常规标量量化典型地减少了 4 倍,但使用 int4 又减少了 2 倍。 将所需内存从 2GB 减少到 300MB 以下。 请记住,这是启用压缩的情况。 解压缩和压缩字节在搜索时确实会产生开销。 对于每个字节向量,我们必须在进行 int4 比较之前对其进行解压缩。 因此,当在 Elasticsearch 中引入这一功能时,我们希望让用户能够选择压缩或不压缩。 对于某些用户来说,更便宜的内存要求是不容错过的,而对于其他用户来说,他们关注的是速度。 Int4 提供了调整设置以适合你的用例的机会。

图 6:CohereV3 500k 向量的速度比较。

结束 ?

在过去的两篇大型技术博客文章中,我们回顾了有关优化的数学和直觉以及它们给 Lucene 带来的影响。 这是一段漫长的旅程,我们还远远没有完成。 请在未来的 Elasticsearch 版本中关注这些功能!

准备好将 RAG 构建到你的应用程序中了吗? 想要尝试使用向量数据库的不同 LLMs?
Github 上查看我们的 LangChain、Cohere 等示例笔记本,并参加即将开始的 Elasticsearch 工程师培训

原文:Int4: More Scalar Quantization in Lucene — Elastic Search Labs