AVQ-Attention: Adaptive Vector-Quantized Attention

Conference: ECCV2026
Paper: ECCV official page · Paper PDF
Area: Model Compression
Keywords: vector-quantized attention, adaptive codebook, hierarchical codebook, kernel fusion, attention acceleration

TL;DR

AVQ-Attention starts with coarse codeword attention and selectively activates learned children for the parents receiving the most attention within each query tile; fused precomputation and incremental correction improve the speed–quality trade-off, reaching 43.33% mIoU in a controlled ADE20K comparison, still below exact attention's 49.0%.

Background & Motivation

Standard self-attention compares every query with every key, making computation quadratic in sequence length. FlashAttention avoids storing the full attention matrix and reduces memory traffic through tiling, but does not eliminate those pairwise dot products. Vector-quantized attention instead assigns similar keys to representative codewords, aggregates their values, and lets queries attend over that smaller set. Each query keeps its own output, and keys are represented rather than simply discarded, although replacing keys introduces approximation error.

A fixed codebook does not necessarily allocate precision where the current queries need it. If a query tile puts most of its attention on a few key clusters, errors inside those clusters matter more than fine distinctions elsewhere. Increasing the entire codebook improves resolution but also increases assignment and attention costs across the board. The objective here is therefore not uniform low-bit model quantization: it is deciding which regions of key space deserve extra computation during attention.

The authors use coarse attention itself as the importance signal and refine the most attended parent codewords. However, FlashAttention does not materialize all attention weights, and revisiting raw keys after selection could erase the savings, so the adaptive rule must fit the execution strategy. Core idea: learn and preaggregate a parent–child codebook, select important parents independently inside each query tile, and replace the relevant coarse contributions directly in the existing attention accumulators.

Method

Overall Architecture

Inputs are the queries, keys, and values of an attention head; outputs remain one weighted value vector per query, with no reduction in token count. Training learns a constrained parent–child codebook, while inference freezes codeword positions and changes only which children each query tile uses. The forward pass has two fused kernels: one assigns keys and aggregates values, and the other computes parent attention, selects refinements, and corrects the output with child attention. Adaptation therefore means input- and tile-dependent resolution, not retraining cluster centers during inference.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Training key distribution"] --> B["Constrained parent–child<br/>codebook"]
    B --> C["Fused full-tree<br/>precomputation"]
    K["Inference keys and values"] --> C
    C --> D["Query-tile adaptive<br/>correction"]
    Q["Inference queries"] --> D
    D --> E["Attention output per query"]

Key Designs

1. Constrained parent–child codebook: make fine representations cheap to substitute for coarse ones

Each attention head in each layer has \(M_0\) parent codewords and \(C\) additional children per parent. Children supplement their parent rather than necessarily replacing it: a key already well represented by the parent can remain there after refinement. Training first assigns keys to their nearest parent and updates parent centroids with exponential moving averages, then searches among that parent and its children for a better representation and updates the child centroids. Children thus model distinctions inside one parent cell instead of competing over the entire key space.

The crucial constraint makes a parent equal the arithmetic mean of all its children. Independent moving-average updates generally violate this relation, so a mass-weighted projection restores it; Figure 2 explains that children assigned more keys resist displacement more strongly. Mass weighting governs the projection cost, not a replacement of the final arithmetic-mean constraint by a weighted mean. Linearity then lets the kernel recover a query's parent dot product by averaging its dot products with all that parent's children. Since child dot products are already needed for refinement, the kernel avoids loading the parent again to recompute its old weight. The projection's closed-form expression is deferred to Supplement A, which is absent from the cache, so it is not reconstructed here.

2. Fused full-tree precomputation: prepare refinement statistics without interrupting attention

The first kernel reads keys and values, finds each key's nearest parent, and checks only that parent's children. Assignment requires distances to \(M_0+C\) candidates per key rather than an exhaustive search through \(M_0(1+C)\) codewords. Keys and values remain in registers between the two assignment stages, while atomic additions accumulate a value-vector sum and key count for each codeword. Fusing assignment and aggregation reduces intermediate memory traffic, rather than relying on asymptotic savings alone.

There are two distinct statistics to keep straight: a parent contains the values and counts of its entire original cluster, while children contain only keys closer to them than to the parent. Child aggregates identify subsets whose coarse contributions can subsequently be replaced; they are not extra information that can simply be added to the parent output. The kernel precomputes aggregates for the entire two-level tree, including children a particular query tile will never select. That extra shared work avoids suspending attention, writing its state to global memory, and launching another quantization step after parent selection. It also allows different query tiles to choose different refinements while reusing the same aggregates.

The authors reorder image tokens along a Gilbert space-filling curve so contiguous query tiles correspond to compact spatial regions. Queries sharing a refinement decision are then more likely to attend to related regions; this changes data organization, not token count or the attention mask. The reordering is a supporting part of the precomputation and tile organization rather than a separate selection mechanism. Its isolated task benefit requires Supplement H; the main paper's Figure 3 establishes the spatial arrangement, not an ablation result.

3. Query-tile adaptive correction: select important clusters and replace, rather than duplicate, their contributions

The second kernel first computes attention over all parents. Because a codeword represents multiple keys, the numerator uses the sum of their value vectors and the denominator includes their count. Treating every cluster as one equally weighted token would otherwise misrepresent clusters of different sizes. The computation maintains the online-softmax numerator, denominator, and stable maximum instead of storing the entire query–codeword attention matrix. Importance is the normalized attention mass that a query tile assigns to a parent, not cluster occupancy or an unnormalized dot product. Algorithm 2, line 8, specifies:

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

Here \(I\) is the query tile, \(A_{ij}\) is the stabilized exponential dot product, \(n_j\) is the parent-cluster count, and \(\bar Z_i\) is the coarse-attention denominator for that query. Each tile chooses its own top-\(P\) parents and loads their children together with the precomputed child aggregates. The denominator is exact when all parents fit in on-chip memory; larger tiled parent codebooks may use an approximation whose details are deferred to unavailable Supplement B.

Adding child attention without correction would count the reassigned keys twice, since they already contributed through their parent. After recovering parent logits using the mean constraint, the kernel multiplies child aggregates by the difference between child and parent weights and updates both accumulators. The subtraction is explicit in Algorithm 2, lines 15–18; suppressing the common online-maximum shift, the mechanism is:

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

The two logits belong to the child and parent respectively; implementation uses the same stabilizing maximum for both and rescales existing accumulators whenever that maximum changes. A negative correction is valid because it removes an old contribution, not because the final attention probability becomes negative. Keys that do not move remain represented by the parent, and unselected parents retain their complete coarse contribution. All keys therefore still contribute through a codeword, although their original exact attention weights are not preserved.

A Worked Example

For the paper's configuration \(M_0=64,P=16,C=8\), the stored tree contains 576 parent and child codewords. The first kernel builds statistics for the full tree, but an individual query tile attends to 64 parents and then refines 16 of them with 128 children. Its attention stage therefore processes 192 parent/child codeword interactions per query rather than traversing all 576 codewords. These counts are arithmetic illustrations of the configuration, not additional measured results.

Suppose some keys in a selected parent cluster are better represented by one child. Their values are already present in the parent aggregate, so refinement applies a weight difference only to that child's subset: subtract its old parent-weighted contribution and insert its child-weighted contribution. Another query tile can select different parents without rebuilding the aggregates. The main saving is selective fine-grained interaction per query, not avoiding every unused child's aggregation.

Loss & Training

Classification and segmentation start from pretrained transformers, replace their attention layers, and fine-tune while learning codebooks through online k-means and exponential moving averages. The main paper does not fully specify optimizers, learning rates, or gradient handling through assignments; it would be unjustified to infer a particular straight-through estimator. Table 3 uses the same 30-epoch recipe across efficient-attention alternatives and reports segmentation statistics from three seeds. Figure 5 is a separate speed–quality sweep and should not be merged with Table 3 as if every point had the same training setting.

Diffusion experiments use LinFusion's knowledge-distillation objective, data, and hyperparameters, but train for 50k steps rather than its 100k steps. AVQ mainly replaces the five outermost self-attention blocks of the SD1.5 UNet; one configuration also replaces five blocks at the next level, while the shortest-sequence innermost blocks retain exact attention. This is not a replacement of the entire diffusion model, and an isolated attention-kernel speedup is not the model's end-to-end speedup.

Key Experimental Results

Main Results

Selected rows from the paper's Table 3 evaluate DPT-Large semantic segmentation on ADE20K. Higher mIoU and lower attention-kernel latency are better; the comparison uses an identical 30-epoch recipe, with marked rows reporting mean and standard deviation over three seeds. The main text does not name the evaluation split explicitly and defers the complete recipe to Supplement F.

Method	mIoU (%, ↑)	Kernel time (ms, ↓)
Flash Attention-v2 baseline	49.0	0.313
AVQ, 32/8/8	43.33 ± 0.12	0.116
Flat VQ, 128 codewords	42.70 ± 0.08	0.103
Flat VQ, 192 codewords	43.04 ± 0.05	0.164
Swin, window 7	42.90 ± 0.09	0.108

AVQ gains 0.63 percentage points over 128-codeword VQ but is slightly slower; it gains 0.29 points over 192-codeword VQ while running faster. These are comparable-budget comparisons, not precisely matched latencies. Its 5.67-point deficit to exact attention is substantial, so the result should not be described as lossless replacement.

Ablation Study

Selected rows from Table 2 analyze kernel fusion rather than task-accuracy ablations of each module. The timing workload uses batch size \(B=4\), \(H=12\) heads, and head dimension \(D=64\); there is no dataset split, and 64k denotes 65,536 tokens. Every attention stage uses a FlashAttention-style kernel, while unfused VQ uses torch.compile'd PyTorch only for precomputation, making the matched-codebook rows informative about fusion.

Configuration	1k tokens (ms, ↓)	64k tokens (ms, ↓)
Flash Attention	0.21	847
Flat VQ, 256 codewords, unfused	0.35	20.7
Flat VQ, 256 codewords, fused	0.18	10.6
AVQ, 64/8/8, fused	0.13	6.67
AVQ, 64/16/8, fused	0.16	7.95

For flat VQ at 64k, fusion reduces latency from 20.7 to 10.6 ms, approximately a 1.95× speedup; adaptivity does not account for all the gains. AVQ 64/8/8 is approximately 127× faster than the same table's Flash Attention at that length, but this is an isolated kernel ratio, not whole-model latency or evidence of matched task accuracy at 64k. Hardware and complete timing conditions are deferred to unavailable Supplement E, limiting portability of the absolute numbers.

Key Findings

• Section 5.2 defines capture as the fraction of true attention weight assigned to keys in the selected parent clusters, not the fraction under AVQ's own coarse approximation. With \(M_0=64,C=8\) and full spawning during training, selecting \(P=16\) at inference captures over 82% on both tasks; the authors report performance within 0.3% of its maximum, without supplying the per-task table in the main text.
• Under Table 4's COCO protocol with ScaleCrafter at \(1024\times1024\), AVQ 32/8/8 reports FID 35.39, CLIP 0.297, and 133.48 ms per UNet forward step, versus 41.03, 0.292, and 213.43 ms for SD1.5. Lower FID and higher CLIP are better; these findings are specific to this distillation and high-resolution evaluation protocol, not a universal generation-quality claim.
• Figure 5 supports a better speed–quality trade-off than flat VQ on ImageNet-1k. The cached plot text is misaligned, so exact point coordinates and accuracies are not guessed from the extraction.

Highlights & Insights

• The parent–child mean constraint is an algebraic enabler for recovering an old logit, not merely a clustering regularizer. It directly links representation learning to reduced data movement in the kernel.
• Keeping the parent as a candidate avoids forcing well-represented keys onto a child. Correcting only the reassigned subset cleanly separates increased resolution from duplicate contributions.
• Full-tree preaggregation and tile-local refinement put shared work and query-specific work in different places. This spends some reusable computation to avoid repeated launches and intermediate memory traffic rather than optimizing operation counts in isolation.

Limitations & Future Work

• The authors leave end-to-end architectures built around AVQ to future work and note that codebook sizes and refinement budgets may need layer-specific settings. Deeper hierarchies and quantization-error-aware refinement are proposed directions, not demonstrated improvements.
• Reader interpretation: coarse importance can miss a region whose approximation error is large but whose coarse attention underestimates its relevance. This motivates the authors' suggestion to combine attention importance with quantization error, but no quantitative benefit for that extension is established here.
• Reader interpretation: the segmentation comparison retains a notable quality gap, whereas the largest speedup comes from long-sequence kernel timing. Approximation quality, operator latency, and end-to-end throughput must be evaluated separately.
• Evidence boundary: the cache contains the complete main paper and references but not the cited Supplements A–H. Several extracted equations are damaged; only importance and correction expressions checkable against Algorithm 2 are used, without inventing projection formulas, missing training details, or plot coordinates.

Related Work & Insights

• vs Transformer-VQ: both compress key interactions into codewords; AVQ additionally allocates fine codeword capacity per query tile and improves the assignment/aggregation kernel. Its contribution combines adaptive resolution with memory-traffic optimization rather than merely renaming a fixed codebook.
• vs FlashAttention: FlashAttention optimizes exact attention's memory access; AVQ reduces computation by approximating keys and reuses the online-softmax execution structure. Their techniques are compatible, but a FlashAttention-style kernel does not remove AVQ's approximation error.
• vs ToMe and sparse attention: those approaches merge tokens or restrict interactions, while AVQ keeps the query count and represents all keys at different resolutions. Keeping a global contribution path is not equivalent to preserving every exact attention weight.

Rating

• Novelty: 4/5. Importance-driven hierarchical refinement and parent-contribution recovery form a coherent joint design.
• Experimental Thoroughness: 3/5. Classification, segmentation, generation, and kernel analysis are covered, but important reproduction details are supplement-only.
• Writing Quality: 4/5. Two kernel algorithms clearly connect representation structure to execution.
• Value: 4/5. Useful for combining attention approximation with kernel optimization, subject to task-quality and hardware constraints.