TurboQuant:通过极端压缩重新定义AI效率

这篇文章介绍了Google Research开发的TurboQuant、QJL和PolarQuant等先进量化算法,旨在通过对大型语言模型和向量搜索引擎进行极端压缩来提高AI效率。这些算法能够显著减少内存占用和加速运行,同时保持模型性能,解决了关键值缓存瓶颈和向量搜索的效率问题。

我们引入了一套先进的、有理论基础的量化算法,能够对大型语言模型和向量搜索引擎进行大规模压缩。

向量是AI模型理解和处理信息的基本方式。小向量描述简单的属性,例如图中的一个点,而“高维”向量则捕获复杂信息,例如图像的特征、词语的含义或数据集的属性。高维向量功能强大,但它们也消耗大量的内存,导致键值缓存(key-value cache)出现瓶颈。键值缓存是一种高速的“数字备忘录”,它将频繁使用的信息存储在简单的标签下,以便计算机可以立即检索,而无需搜索缓慢、庞大的数据库。

向量量化是一种强大的经典数据压缩技术,可以减小高维向量的大小。这种优化解决了AI的两个关键方面:它通过实现更快的相似性查找来增强向量搜索(为大规模AI和搜索引擎提供动力的高速技术);并通过减小键值对的大小来帮助疏通键值缓存瓶颈,从而实现更快的相似性搜索并降低内存成本。然而,传统的向量量化通常会引入其自身的“内存开销”,因为大多数方法需要为每个小数据块计算和存储(以全精度)量化常数。这种开销可能会为每个数字增加1或2个额外的比特,部分抵消了向量量化的目的。

今天,我们介绍了 TurboQuant(将在 ICLR 2026 上发表),这是一种压缩算法,它以最佳方式解决了向量量化中的内存开销挑战。我们还介绍了量化 Johnson-Lindenstrauss (QJL) 和 PolarQuant(将在 AISTATS 2026 上发表),TurboQuant 使用它们来实现其结果。在测试中,所有这三种技术都显示出在不牺牲AI模型性能的情况下减少键值瓶颈的巨大潜力。这对于所有依赖压缩的用例,特别是搜索和AI领域,都具有潜在的深远影响。

TurboQuant 的工作原理

TurboQuant 是一种压缩方法,它在不损失准确性的情况下实现了模型尺寸的大幅缩小,使其成为支持键值 (KV) 缓存压缩和向量搜索的理想选择。它通过以下两个关键步骤实现:

  1. 高质量压缩(PolarQuant 方法):TurboQuant 首先随机旋转数据向量。这一巧妙的步骤简化了数据的几何结构,使得可以轻松地将标准的、高质量的量化器(一种将大量连续值(如精确小数)映射到更小、离散的符号或数字(如整数)的工具:例如音频量化和 JPEG 压缩)应用于向量的每个部分。第一阶段使用大部分压缩能力(大部分比特)来捕获原始向量的主要概念和强度。
  2. 消除隐藏错误:TurboQuant 使用少量剩余的压缩能力(仅1比特)将 QJL 算法应用于第一阶段剩余的微小误差。QJL 阶段充当数学误差检查器,消除偏差,从而获得更准确的注意力分数。

为了充分理解 TurboQuant 如何实现这种效率,我们更深入地了解 QJL 和 PolarQuant 算法的工作原理。

QJL:零开销的 1 比特技巧

QJL 使用一种称为 Johnson-Lindenstrauss 变换的数学技术来缩小复杂的、高维数据,同时保留数据点之间的基本距离和关系。它将每个结果向量数减少到单个符号位(+1 或 -1)。该算法本质上创建了一种高速速记,不需要任何内存开销。为了保持准确性,QJL 使用一种特殊的估计器,策略性地平衡高精度查询与低精度、简化数据。这使得模型能够准确计算注意力分数(用于决定其输入哪些部分重要以及哪些部分可以安全忽略的过程)。

PolarQuant:压缩的新“角度”

PolarQuant 使用一种完全不同的方法来解决内存开销问题。它不是使用标准坐标(即 X、Y、Z)来表示沿每个轴的距离来查看内存向量,而是使用笛卡尔坐标系将向量转换为极坐标。这类似于将“向东走 3 个街区,向北走 4 个街区”替换为“以 37 度角总共走 5 个街区”。这会产生两部分信息:半径(表示核心数据的强度)和角度(表示数据的方向或含义)。由于角度的模式是已知的且高度集中,模型不再需要执行昂贵的数据归一化步骤,因为它将数据映射到一个固定的、可预测的“圆形”网格,其中边界是已知的,而不是一个边界不断变化的“方形”网格。这使得 PolarQuant 能够消除传统方法必须承担的内存开销。

PolarQuant 充当高效压缩桥梁,将笛卡尔输入转换为紧凑的极坐标“速记”以进行存储和处理。该机制首先将 d 维向量中的坐标对分组,并将它们映射到极坐标系。然后将半径成对收集,进行递归极坐标变换——这个过程重复进行,直到数据被提炼成一个最终半径和一组描述性角度。

实验与结果

我们使用开源 LLM(Gemma 和 Mistral)在标准长上下文基准测试(包括 LongBench、Needle In A Haystack、ZeroSCROLLS、RULER 和 L-Eval)中严格评估了所有三种算法。实验数据表明,TurboQuant 在点积失真和召回率方面均实现了最佳评分性能,同时最大限度地减少了键值 (KV) 内存占用。下表显示了 TurboQuant、PolarQuant 和 KIVI 基线在各种任务(包括问答、代码生成和摘要)中的聚合性能得分。

Quantization-2

TurboQuant 在 LongBench 基准测试中,相对于 Llama-3.1-8B-Instruct 模型(比特宽度在括号中表示)的各种压缩方法,展示了强大的 KV 缓存压缩性能。

长上下文“大海捞针”任务(即旨在测试模型是否能从大量文本中找到一个特定的微小信息)的结果如下所示。同样,TurboQuant 在所有基准测试中都取得了完美的下游结果,同时将键值内存大小至少减少了 6 倍。PolarQuant 在此任务中也几乎是无损的。

TurboQuant 证明它可以将键值缓存量化到仅 3 比特,而无需训练或微调,也不会损害模型准确性,同时实现了比原始 LLM(Gemma 和 Mistral)更快的运行时间。它实现起来非常高效,并且运行时开销可以忽略不计。下图说明了使用 TurboQuant 计算注意力对数的速度提升:具体来说,4 比特 TurboQuant 在 H100 GPU 加速器上比 32 比特未量化键实现了高达 8 倍的性能提升。

Quantization-3

TurboQuant 展示了在键值缓存中计算注意力对数的显著性能提升,跨越各种比特宽度级别,相对于高度优化的 JAX 基线进行测量。

这使其成为支持向量搜索等用例的理想选择,在这些用例中,它极大地加快了索引构建过程。我们使用 1@k 召回率(衡量算法在 k 个近似值中捕获真实顶部内积结果的频率)评估了 TurboQuant 在高维向量搜索中对抗最先进方法(PQ 和 RabbiQ)的有效性。尽管这些基线使用了低效的大码本和特定数据集的调整,但 TurboQuant 始终比基线方法实现更高的召回率(下图)。这证实了 TurboQuant 在高维搜索任务中的鲁棒性和效率。

Quantization-4

TurboQuant 展示了强大的检索性能,在 GloVe 数据集(d=200)上相对于各种最先进的量化基线实现了最佳的 1@k 召回率。

TurboQuant 证明了高维搜索的变革性转变。通过为可实现的速度设定新的基准,它以数据无关的方式提供了接近最佳的失真率。这使得我们的最近邻引擎能够以 3 比特系统的效率运行,同时保持比更重模型更高的精度。更多详情请参阅论文。

展望

TurboQuant、QJL 和 PolarQuant 不仅仅是实用的工程解决方案;它们是具有强大理论证明的基础算法贡献。这些方法不仅在实际应用中表现良好;它们被证明是高效的,并且在理论下限附近运行。这种严谨的基础使其对于关键的大规模系统来说是稳健和值得信赖的。

虽然一个主要应用是解决 Gemini 等模型中的键值缓存瓶颈,但高效的在线向量量化的影响甚至更广。例如,现代搜索正在超越关键词,以理解意图和含义。这需要向量搜索——在包含数十亿向量的数据库中找到“最近”或语义最相似的项目。

像 TurboQuant 这样的技术对于这项任务至关重要。它们允许以最小的内存、接近零的预处理时间和最先进的准确性来构建和查询大型向量索引。这使得 Google 规模的语义搜索更快、更高效。随着 AI 越来越集成到所有产品中,从 LLM 到语义搜索,这项基础向量量化工作将比以往任何时候都更加关键。

致谢

这项研究是与 Google 研究员 Praneeth Kacham、Google DeepMind 首席工程师 Majid Hadian、KAIST 助理教授 Insu Han、NYU 博士生 Majid Daliri、Google 研究员 Lars Gottesbüren 和 Google 研究员 Rajesh Jayaram 合作进行的。

  • 原文链接: research.google/blog/tur...
  • 登链社区 AI 助手,为大家转译优秀英文文章,如有翻译不通的地方,还请包涵~
点赞 0
收藏 0
分享
本文参与登链社区写作激励计划 ,好文好收益,欢迎正在阅读的你也加入。

0 条评论

请先 登录 后评论
research
research
江湖只有他的大名,没有他的介绍。