🤖 machine learning

Online Vector Quantized Attention

本文介绍了在线向量量化（OVQ）注意力机制，这是一种序列混合层，通过稀疏内存更新实现线性计算和恒定内存成本，同时显著提升长上下文性能，以远低于自注意力机制的内存占用提供具有竞争力的替代方案。

原作者： Nick Alonso, Tomas Figliolia, Beren Millidge

发布于 2026-05-18

📖 1 分钟阅读☕ 轻松阅读

原作者： Nick Alonso, Tomas Figliolia, Beren Millidge

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是用通俗易懂的语言和生动的类比对论文《在线向量量化注意力》（Online Vector Quantized Attention）的解释。

问题：大脑“太大”与“太小”的困境

想象一下，你正在阅读一本长达 10 万页的巨著，只为回答一个关于第 5 页提到的某个角色的具体问题。

旧方法（自注意力机制）： 这就像拥有一位超级智能的助手，每次你提问时，他都会重读整本书。他能完美记住每一个字。然而，要做到这一点，他需要一个像城市一样巨大的图书馆来存放所有笔记。随着书籍变长，图书馆变得更大，助手也变得更慢。虽然准确，但代价高昂且效率低下。
高效方法（线性注意力/状态空间模型 SSMs）： 这就像一位只携带小记事本的助手。他在阅读过程中总结全书，只保留最重要的统计数据。他们速度极快，只需要一个口袋大小的记事本。但是，由于记事本太小，他们往往记不住长而复杂故事中的具体细节。当书籍长达 10 万页时，他们很难找到第 5 页提到的那个角色。

目标： 作者希望构建一位助手，既像“口袋记事本”那样快速且节省内存，又像“城市图书馆”那样聪明且细节丰富。

解决方案：“智能文件柜”（OVQ-Attention）

作者创造了一种名为在线向量量化（OVQ）注意力的新方法。你可以将其想象为一个实时更新的智能文件柜。

以下是其工作原理，分步说明：

1. 字典（文件夹）

该系统不使用“城市图书馆”那样记住每一个字，也不像“口袋记事本”那样只记摘要，而是使用文件夹字典。

想象你有一个文件柜，里面有比如 4000 个空文件夹。
当助手阅读书籍时，他不会记下每一个字。相反，他会查看当前的句子，并问自己：“这句话最适合放入这 4000 个文件夹中的哪一个？”
他将这句话放入那个文件夹中。

2. “在线”魔法（更新文件夹）

在以往类似的想法中，文件夹是在助手开始阅读之前预先写好的。如果书中出现了文件夹未预料到的词汇，助手就会感到困惑。

在OVQ-Attention中，文件夹在开始时是空的。随着助手阅读书籍：

如果看到独特的内容，他们会即时创建新文件夹。
他们会更新现有文件夹，以更好地匹配当前所见的内容。
关键在于： 他们只更新当前句子所属的特定文件夹。他们不会重写整个文件柜。这使得工作保持快速（线性）且内存占用低（恒定）。

3. “稀疏”更新（高效技巧）

通常，如果你想记住更多细节，就需要更大的文件柜，这会占用更多空间。

论文的技巧： 作者意识到，即使你有一个拥有 10 万个文件夹的文件柜，你一次也只会触碰其中几个。
因为更新是稀疏的（你只触碰与当前句子相关的特定文件夹），所以更新文件柜所需的精力并不会因为文件柜巨大而增加。
结果： 你可以拥有一个巨大的文件柜（高内存容量），而无需支付巨大文件柜带来的“精力税”。你得到了两全其美的效果：巨大的存储空间，极低的成本。

结果：效果如何？

作者在多个挑战中测试了这个“智能文件柜”：

“大海捞针”测试（上下文回忆）：
- 任务： 将特定的键值对（如电话号码）隐藏在一块巨大的文本中，然后要求模型稍后找到它。
- 结果： 旧的“口袋记事本”模型在超过一定长度后彻底失败。“城市图书馆”模型表现完美但速度缓慢。OVQ-Attention模型的表现几乎与城市图书馆一样完美，即使其内存占用小得多，它也能处理长达64,000 个词的文本。
“边读边学”测试（上下文学习）：
- 任务： 给模型提供一堆新规则的例子（例如“如果你看到 X，就执行 Y"），然后要求它在文本后面的新句子中应用该规则。
- 结果： OVQ 模型学习规则的能力与那些笨重缓慢的模型一样好，而那些高效但“愚蠢”的模型则无法学会复杂的模式。
阅读长故事（语言建模）：
- 当被要求预测长故事中的下一个词（使用 PG19 数据集）时，尽管只使用了极少量的内存，OVQ 模型的表现仍与最佳标准模型具有竞争力。

秘诀：高斯混合回归

论文使用高斯混合回归这一概念来解释其背后的数学原理。

简单类比： 想象你正试图根据一朵云来猜测天气。
- 标准模型试图将当前的云与他们曾经见过的每一朵过去的云进行匹配。
- OVQ-Attention 将云归类为“类型”（例如“风暴云”、“蓬松云”）。
- 当新云出现时，系统不仅仅是死记硬背它们，而是调整“风暴云”的定义，以更好地拟合新数据。它在阅读过程中学习云类型的形状，从而在长时间内使其猜测更加准确。

总结

这篇论文介绍了OVQ-Attention，这是一种让 AI 处理长文本的新方法。

旧的高效模型： 快速且小巧，但容易遗忘。
旧的强大模型： 聪明且细节丰富，但缓慢且消耗大量内存。
OVQ-Attention： 使用动态、自我更新的文件系统。它保持少量且恒定的活跃内存，但通过将信息组织成簇，可以存储海量信息。它在阅读过程中学习这些簇，使其在极长的上下文（高达 64k 个 token）中既能保持快速，又能变得极其聪明。

作者得出结论，这是迈向开发既高效又能处理长而复杂任务（且无需超级计算机）的 AI 模型的重要一步。

技术摘要：在线向量量化注意力

1. 问题陈述

大型语言模型（LLM）中的标准序列混合层在计算效率与长上下文性能之间面临根本性的权衡。

自注意力（Self-Attention）： 虽然在长上下文任务（上下文学习、推理、检索）中非常有效，但其计算成本呈二次方增长（ $O(T^2)$ ），内存成本呈线性增长（ $O(T)$ ），使其在处理极长序列时代价高昂。
线性注意力与状态空间模型（SSMs）： 这些模型提供线性计算成本（ $O(T)$ ）和恒定内存（ $O(1)$ ）。然而，由于内存容量有限，它们在需要精确长程信息处理的任务中表现显著不足。
向量量化（VQ）注意力： 一种先前的方法（Lingle, 2023），通过将键（keys）量化为固定字典中的质心，实现了线性计算和恒定内存。然而，初步测试表明，原始的 VQ 注意力在长上下文任务中表现挣扎。主要瓶颈被确定为静态的、预训练的键字典（ $D_k$ ），它无法在前向传播过程中有效地建模键，导致量化误差并降低性能。增加字典大小带来的收益递减。

2. 方法论：在线向量量化（OVQ）注意力

作者提出了在线向量量化（OVQ）注意力，这是一种新颖的序列混合层，它在前向传播过程中在线学习键字典（ $D_k$ ）和值字典（ $D_v$ ）。

理论基础：高斯混合回归（GMR）

该论文建立了 OVQ 注意力与**高斯混合回归（GMR）**之间的理论联系。

自注意力作为 GKR： 在特定假设下，标准自注意力等价于高斯核回归（GKR）。
VQ 注意力作为 GMR： 作者证明，VQ 注意力（进而 OVQ）与 GMR 的关系更为密切。在 GMR 中，高斯混合模型（GMM）被拟合到输入 - 输出对（键和值）的数据集上。
联系： 在假设查询/键为单位范数且协方差各向同性的情况下，GMR 的预测公式简化为线性时间的 VQ 注意力方程：
$O_{ovq} = \text{softmax}(\beta Q D_k^\top + \log(c)) D_v$
其中 $c$ 代表分配给每个质心的键的数量。
启示： 与分别学习 $D_k$ 和 $D_v$ 的标准 VQ 注意力（预训练与在线）不同，GMR 同时学习它们以拟合同一数据集。这激发了 OVQ 方法，即在线学习两个字典以准确建模条件分布 $P(V|K)$ 。

在线学习公式

OVQ 注意力在测试时训练框架内运行，采用分块并行公式：

预测： 对于每个查询分块，模型使用字典（ $D_k, D_v$ ）和计数向量（ $c$ ）的当前状态计算注意力。与原始 VQ 注意力不同，当前分块的键和值在用于注意力计算前不被量化；它们被直接拼接到字典中。
状态更新： 生成预测后，模型使用当前分块的键和值更新字典。
- 字典增长： 质心数量（ $N_t$ ）通过受狄利克雷过程启发的平台化增长函数，渐近增长至最大值 $N$ 。这允许模型动态调整其内存容量。
- 质心初始化： 新质心使用“最大化分散”方案（类似于 k-means++）进行初始化，选择与现有质心点积最小的数据点，以确保多样性。
- 更新规则： 现有质心通过批处理 K-means 更新的在线版本进行更新。作者证明，在他们的假设下，GMM 上的第一批 EM 更新等价于负对数似然（NLL）上的牛顿更新，为最小化损失提供了二阶近似。

计算特性

内存： 无论序列长度如何，OVQ 保持恒定内存（ $O(N)$ ），因为状态由字典和计数定义，而非完整序列。
计算： 该层产生线性计算成本（ $O(T)$ ）。更新操作是稀疏的；分块中的每个令牌仅更新字典的单行。这使得状态大小（ $N$ ）可以显著增加，而不会增加更新张量的内存占用，这与状态更新随状态维度缩放的线性注意力不同。

3. 主要贡献

提出 OVQ 注意力： 一种新的序列混合层，在线学习键和值字典，克服了先前 VQ 注意力中静态字典的局限性。
理论框架： 一种新颖的推导，将 OVQ 注意力与高斯混合回归联系起来，为在线学习算法（包括字典增长和稀疏更新规则）提供了原则性基础。
实证验证： 在合成任务（上下文回忆、上下文学习）和长上下文语言建模（PG19）上进行了全面测试。

4. 实验结果

作者将 OVQ 注意力与强大的基线进行了评估，包括完整自注意力、线性注意力和 SSM（例如 Mamba、Gated Delta Net）。

合成上下文回忆（ICR）：
- 在“基本 ICR"和“位置 ICR"任务上，OVQ 注意力显著优于原始 VQ 注意力和线性注意力基线。
- 尽管使用的状态大小仅为完整自注意力所用大小的 10-25%，但在高达64k 序列长度的情况下，其表现与完整自注意力基线相当或略有偏差。
- 随着测试时字典大小的增加，性能得到提升，证明了有效的长度外推能力。
合成上下文学习（ICL）：
- 在具有多个函数的长上下文线性回归任务中，OVQ 注意力在多达 128 个不同函数的场景下成功学习了所有函数，其表现与完整自注意力相当。
- 原始 VQ 注意力和线性注意力基线在这些困难的长上下文场景中甚至无法学习单个函数。
长上下文语言建模（PG19）：
- 将 OVQ 与滑动窗口层交错混合的混合架构，其性能显著优于纯滑动窗口或 Gated Delta Net 模型。
- 在 PG19 数据集上，OVQ 注意力在 16k 上下文长度下与标准自注意力基线相当或略胜一筹。
短上下文基准：
- 在标准短上下文基准（PIQA、HellaSwag 等）上，OVQ 注意力的表现与标准自注意力相当，证实了它在不需要压缩的任务中不会降低性能。

5. 意义与主张

该论文声称，OVQ 注意力代表了在长上下文处理中平衡效率与性能的重要一步。

效率与容量： 与线性注意力和 SSM 不同，OVQ 注意力使用稀疏内存更新，这些更新独立于总状态大小。这使得模型能够利用非常大的内存状态（高容量），同时保持恒定的内存占用和线性计算成本。
持续学习潜力： 作者指出，OVQ 注意力的稀疏、在线聚类特性使其比线性注意力/SSM 中使用的密集梯度更新更能抵抗“灾难性遗忘”。这将 OVQ 定位为数据分布随时间变化的持续学习场景中有前途的候选者。
可扩展性： 虽然当前实验规模相对较小（<5 亿参数），但该架构的理论特性表明，它是通往高效长上下文 LLM 的可行路径，无需承担自注意力的二次方成本。

该论文得出结论，OVQ 注意力提供了一种引人注目的自注意力替代方案，以一小部分内存需求实现了长上下文任务上的竞争性性能，标志着高效序列建模的一个重要研究方向。