AVQ-Attention: Adaptive Vector-Quantized Attention¶

会议: ECCV2026
Paper: https://eccv.ecva.net/virtual/2026/poster/4988
PDF: https://media.eventhosts.cc/Conferences/ECCV2026/pdfs/8320.pdf
领域: 模型压缩
关键词: 向量量化注意力、自适应码本、层次码本、算子融合、注意力加速

一句话总结¶

AVQ-Attention 先用粗粒度码本计算注意力，再按每个查询块的注意力质量选择父码字并启用预学习的子码字，通过融合预计算和增量校正改善速度与精度的折中；在 ADE20K 的受控比较中达到 43.33% mIoU，但仍低于精确注意力的 49.0%。

研究背景与动机¶

标准自注意力让每个查询与所有键交互，序列变长后，计算量随 token 数平方增长。 FlashAttention 通过分块和片上缓存避免保存完整注意力矩阵，主要解决内存访问问题，却没有取消这些两两点积。向量量化注意力走另一条路线：把相近的键映射到少量码字，先汇总对应的值，再让查询关注码字而非每个原始键。这样仍保留每个查询的独立输出，也不必像 token 丢弃那样直接删除某些键的信息，但键被代表向量替换后会引入近似误差。

固定码本的问题不是所有区域都同样不准确，而是精度分配与当前查询需要的信息不一致。一个查询块可能把大量注意力集中在少数键簇上，此时这些簇的粗量化误差尤其重要，其他簇的精细表示却很少影响输出。直接扩大码本可以降低误差，但会同时提高所有区域的分配与注意力开销。因此，本文要解决的不是给整个模型统一降位宽，而是在注意力计算过程中决定哪些键空间区域值得多花计算。

作者使用粗注意力本身提供重要性信号，再把预算投给最受关注的父码字。然而，FlashAttention 不会完整保存注意力权重，选中父码字后重新访问原始键也容易抵消收益；所以自适应策略必须与内核执行方式一起设计。 核心思路：提前学习并汇总父子码本，在每个查询块内部选择高注意力父码字，用子码字替换其中一部分粗表示，并在原有累加器上直接修正结果。

方法详解¶

整体框架¶

输入是某层某个注意力头的查询、键和值，输出仍是每个查询对应的注意力加权值，token 数不变。训练时先建立受约束的父子码本；推理时码字位置固定，只动态改变各查询块使用哪些子码字。整个前向过程由两个融合内核组成：第一个负责键的量化和聚合，第二个负责粗注意力、重要性选择及子码字校正。因此，“自适应”是按输入和查询块选择分辨率，不是在每次推理时重新训练聚类中心。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["训练键分布"] --> B["受约束的父子码本"]
    B --> C["全树融合预计算"]
    K["推理时的键和值"] --> C
    C --> D["查询块自适应校正"]
    Q["推理时的查询"] --> D
    D --> E["每个查询的注意力输出"]

关键设计¶

1. 受约束的父子码本：让精细表示能够低成本修正粗表示

每层的每个注意力头维护自己的码本，其中有 \(M_0\) 个父码字，每个父码字另有 \(C\) 个子码字。父子不是简单的互斥替代关系：某个键若已经离父码字足够近，细化后仍可以留在父码字上。训练时，先把键分配给最近的父码字，用指数移动平均更新父中心，再在该父码字及其子码字中寻找更好的表示，更新相应子中心。这使子码字描述的是一个父簇内部的细节，而不是重新与全局所有键竞争。

关键约束是父码字等于其全部子码字的算术平均。独立的移动平均更新通常破坏这个关系，因此作者在每次更新后进行质量加权投影；根据图 2，占有更多键的子码字更抵抗位移。这里“质量加权”描述投影修正的代价，不是把最终父子关系改成加权平均。线性点积因此具有一个有用性质：查询与父码字的点积，可以由它与全部子码字点积的平均值恢复。后续已经要计算子点积，就无需再次加载父码字来算旧权重，这正是约束服务于硬件执行的地方。具体投影闭式解在未随缓存提供的补充材料 A 中，本笔记不补写该公式。

2. 全树融合预计算：提前准备细化统计量，避免注意力中途停下来重做量化

第一个内核一次读取键和值，先在父码本中寻找最近邻，然后只检查所属父码字的子码字。一个键的比较范围由全体父子码字缩小为 \(M_0+C\) 个候选距离，而不是遍历 \(M_0(1+C)\) 个向量。键和值在父、子分配之间保留在寄存器中，并通过原子累加直接形成每个码字的值向量总和与键数量。这同时融合了最近邻分配和聚合，减少中间张量进入显存的次数。

必须区分两种统计量：父码字保留整个父簇的值总和与数量，子码字只记录确实比父码字更接近的那些键。因此，子统计量是后续需要从父表示中替换出去的子集，并不是一份应直接叠加到父输出上的新信息。系统会预计算整棵两层树的聚合，即使某些子码字最终不会被某个查询块使用。这看似浪费，却让所有查询块共享一次量化成本，并允许各自选择不同父码字，而不必中断注意力内核、回显存保存状态后再启动量化。

为让同一查询块更可能共享合理的细化决策，作者还按 Gilbert 空间填充曲线重排图像 token。其作用是让连续的一维内核块对应紧凑的二维图像区域，并不是删除 token 或增加一种注意力掩码。这个重排属于全树预计算和块级调度的配套组织方式；是否有独立收益需要补充材料 H，正文图 3 只展示其空间结构。

3. 查询块自适应校正：选择重要父簇，再把重复计算改成贡献替换

第二个内核先对所有父码字做粗注意力。由于一个码字代表多个键，分子使用该簇的值向量总和，分母必须乘上簇中键的数量；否则大小不同的簇会被错误地当成等量 token。内核沿用在线 softmax 的分子、分母和稳定最大值累加器，不保存完整的查询—码字注意力矩阵。父簇的重要性是该查询块给它的归一化注意力质量之和，而不是码字占有率或未归一化点积。按算法 2 第 8 行，其定义为：

\[ w_j(I)=\sum_{i\in I}\frac{A_{ij}n_j}{\bar Z_i}. \]

这里 \(I\) 是查询块，\(A_{ij}\) 是经过数值稳定处理的指数点积，\(n_j\) 是父簇键数，\(\bar Z_i\) 是对应查询的粗注意力分母。每个块独立选出得分最高的 \(P\) 个父码字，加载它们的子码字及提前准备的聚合量。当父码本能一次放入片上存储时，这个分母是精确的；较大父码本分块时，论文允许近似分母，其具体方案在未提供的补充材料 B 中。

细化不能直接把子注意力加上去，因为这些键已经在父注意力里贡献过一次。利用父子均值约束恢复父点积后，内核用“子权重减父权重”乘子簇聚合量，同时修正分子和分母。算法 2 第 15–18 行明确给出了减法；省略在线最大值平移后的机制可写成：

\[ \Delta A_{ic}=\exp(S_{ic})-\exp(S_{ip}). \]

\(S_{ic}\) 和 \(S_{ip}\) 分别是子、父点积；实际实现对二者使用同一个稳定最大值，并在最大值变化时重标定已有累加器。负校正项是合法的，因为它在移除旧表示，并不意味着最终注意力概率为负。没有迁移到子码字的键继续由父码字表示，没有被选中的父簇也完全保留粗表示，所以所有键仍通过某个码字贡献输出。

一个完整示例¶

考虑论文采用的一类配置 \(M_0=64,P=16,C=8\)，码本共保存 576 个父子码字。第一个内核为所有父簇及其子簇建立统计量，而某个查询块只关注 64 个父码字，再细化其中 16 个父码字的 128 个子码字。因此该块在注意力阶段使用 192 个父子码字交互，而不是遍历全部 576 个码字。这些数量由配置直接计算，只用于解释流程，不是新增实验结果。

设某个被选父簇中的一部分键更靠近某个子码字，它们起初已经被计入父簇总和。细化时只对这部分键对应的聚合值应用权重差，旧的父权重被扣掉，新的子权重被加上。另一个查询块可能选择另一组父簇，但可复用同一份预计算统计量。这说明方法节省的主要是每个查询的精细交互，而不是完全免去所有子簇的聚合工作。

损失函数 / 训练策略¶

分类和分割实验从预训练 Transformer 出发，替换注意力层后进行微调，码本通过在线 k-means 与指数移动平均更新。正文没有完整给出优化器、学习率或量化分配的梯度处理，不能把这套更新自行解释成某个特定直通估计器。表 3 的高效注意力比较使用相同的 30 epoch 训练方案，并报告三次随机种子的分割统计。图 5 是另一组速度—质量曲线，不应把它的点与表 3 混成同一训练设置下的精确数值。

扩散实验采用 LinFusion 的知识蒸馏目标、数据和超参数，但只训练 50k 步，对方方案为 100k 步。 AVQ 主要替换 SD1.5 UNet 的五个最外层自注意力块；另有一个配置同时替换五个下一层块，短序列的最内层块保留精确注意力。因此这不是把整套扩散模型全部改成 AVQ，也不能把局部内核加速率当成整个模型的加速率。

实验关键数据¶

主实验¶

下表精选原文表 3，任务为 ADE20K 语义分割，骨干为 DPT-Large；mIoU 越高越好，注意力内核耗时越低越好。正文说明使用相同的 30 epoch 比较方案，带误差的行是三次种子的均值与标准差；没有明确列出评估 split 名称，完整配方指向补充材料 F。

方法	mIoU（%，↑）	内核耗时（ms，↓）
Flash Attention-v2 基线	49.0	0.313
AVQ，32/8/8	43.33 ± 0.12	0.116
固定 VQ，128 码字	42.70 ± 0.08	0.103
固定 VQ，192 码字	43.04 ± 0.05	0.164
Swin，窗口 7	42.90 ± 0.09	0.108

AVQ 相对 128 码字的固定 VQ 提高 0.63 个百分点，但略慢；相对 192 码字的固定 VQ 提高 0.29 个百分点，同时更快。这些是相近预算比较，而非严格相同延迟比较。相对精确注意力，AVQ 仍少 5.67 个百分点，不能描述成无损替换。

消融实验¶

下表精选原文表 2 的内核融合分析，并非逐模块任务精度消融。输入为计时用张量，批量 \(B=4\)、头数 \(H=12\)、每头维度 \(D=64\)；无数据集 split，64k 对应 65,536 个 token。所有方法的注意力阶段都采用 FlashAttention 风格内核，未融合版本只在 VQ 预计算阶段使用经过 torch.compile 的 PyTorch，因此同码本的两行可以隔离融合收益。

配置	1k token 耗时（ms，↓）	64k token 耗时（ms，↓）
Flash Attention	0.21	847
固定 VQ，256 码字，未融合	0.35	20.7
固定 VQ，256 码字，融合	0.18	10.6
AVQ，64/8/8，融合	0.13	6.67
AVQ，64/16/8，融合	0.16	7.95

固定 VQ 在 64k 下由 20.7 ms 降到 10.6 ms，约为 1.95 倍加速，说明收益并非全部来自自适应选择。 AVQ 64/8/8 的约 127 倍是相对同表 Flash Attention 的长序列内核比值，不包含整模型其他层，也没有配套同长度任务精度证明。硬件型号与完整计时条件位于未提供的补充材料 E，不能把该数字推广到任意设备。

关键发现¶

正文 §5.2 定义 capture 为选中父簇所含键获得的真实注意力权重占比，而非 AVQ 自己的粗注意力占比。在 \(M_0=64,C=8\)、训练时全部展开的实验中，推理只选 \(P=16\) 时，两项任务的 capture 均超过 82%；作者称性能距最大值不足 0.3%，但正文未列出逐任务原始表值。
原文表 4 的 COCO 评估中，采用 ScaleCrafter 的 \(1024\times1024\) 生成下，AVQ 32/8/8 的 FID 为 35.39、CLIP 分数为 0.297、每步 UNet 前向为 133.48 ms；SD1.5 基线分别为 41.03、0.292、213.43 ms。 FID 越低越好，CLIP 越高越好；这里只能比较该蒸馏与高分辨率协议，不能据此断言所有生成任务都更优。
ImageNet-1k 的图 5 支持自适应码本在相近内核成本下优于固定码本的趋势。缓存中的曲线文字提取错位，故不从图中猜读精确准确率或自行补出坐标。

亮点与洞察¶

父子均值约束不是普通的聚类正则化，而是为恢复旧点积创造代数条件。表示学习的约束直接减少重访数据，使算法与内核设计相互支撑。
子码字不必整体取代父码字，保留父中心作为候选可避免把已表示良好的键强行移动。校正只作用于真正迁移的子集，因此“增加分辨率”与“增加重复贡献”被清晰分开。
把全树统计量预计算、把细化选择留给每个查询块，是本文重要的开销分工。它用可共享的工作换掉反复启动内核和读写中间状态，而非只追求纸面最少运算。

局限与展望¶

作者明确把围绕 AVQ 设计端到端高效架构留作未来工作，并承认各层的父码字数、展开数和子码字数可能应分别配置。更深码本和结合量化误差的选择规则是讨论中的方向，不是本文已经验证的结果。
读者观察：重要性由粗表示估计，未必能选中量化误差最大、却被粗注意力低估的区域。这与作者建议把量化误差纳入细化标准相呼应，但本文没有给出相应方案的定量增益。
读者观察：受控分割比较保留明显精度差距，而长序列的巨大加速来自内核计时。需要区分近似精度、算子延迟和端到端吞吐，不能用一项代替其余两项。
来源边界：缓存含完整正文与参考文献，但不含所引用的补充材料 A–H。部分正文公式提取损坏，本笔记只采用可由算法 2 核对的重要性和校正公式；投影闭式解、详细训练配方及曲线读数均未补造。

评分¶

新颖性：4/5。注意力质量驱动的层次细化与父贡献恢复形成了明确的联合设计。
实验充分度：3/5。覆盖分类、分割、生成与内核分析，但部分关键复现细节只在补充材料中。
写作质量：4/5。正文给出两段内核伪代码，较清楚地连接表示结构与实际计算。
价值：4/5。适合研究注意力近似和算子优化的结合，使用时必须接受任务精度与硬件依赖的约束。