AVQ-Attention: Adaptive Vector-Quantized Attention

Background The O(N²) cost of attention over N tokens remains a fundamental bottleneck as sequences scale, and while Flash Attention's tiling solves the memory-bandwidth side of the problem, the underlying computation stays quadratic. Vector-Quantized (VQ) attention offers an alternative, clustering keys into M representative codewords to cut cost to O(MN).

배경 Attention의 O(N²) 연산 비용은 시퀀스 길이 N이 커질수록 transformer의 근본적인 병목이 되며, Flash Attention 같은 타일링 기법은 메모리 대역폭 문제는 해결해도 연산량 자체는 그대로 이차로 남긴다. 벡터-양자화(VQ) attention은 key를 M개의 대표 코드워드로 클러스터링해 비용을 O(MN)으로 낮추는 대안을 제시해 왔다.

Problem Existing VQ-attention applies the same codebook capacity everywhere in key space regardless of where attention mass actually concentrates. High-attention regions end up coarsely approximated and lose accuracy, while low-attention regions receive representational capacity they do not need.

문제 기존 VQ-attention은 attention 질량이 실제로 어디에 몰려 있는지와 무관하게 모든 key 공간 영역에 동일한 코드북 용량을 배정한다. 그 결과 attention이 집중되는 영역은 거칠게 근사되어 정확도를 잃고, attention이 거의 없는 영역은 불필요하게 세밀한 표현 용량을 낭비한다.

Prior limitations A fixed-codebook VQ-attention can only trade accuracy for cost by growing M uniformly, which wastes resolution on regions nobody attends to and grows O(MN) cost without regard to where it actually helps. Other efficient-attention approaches -- sparse patterns, token merging -- likewise tend to fix their structure statically rather than adapting to where attention mass concentrates.

기존 한계 고정-코드북 VQ-attention은 코드북 크기 M을 늘려 정확도를 높일 수 있지만, 그러면 관심 없는 영역까지 균일하게 세밀해져 O(MN) 비용이 목표 없이 커진다. 희소 attention 패턴이나 토큰 병합 같은 다른 효율화 기법들도 대체로 구조를 정적으로 고정해, 어디에 중요한 attention이 몰릴지는 반영하지 않는다.

Goal Design a VQ-attention variant that concentrates fine-grained quantization exactly where attention mass concentrates and leaves the rest coarse, while allowing the speed/accuracy tradeoff to be adjusted at inference time with no retraining.

목표 어떤 재학습도 없이 추론 시점에 속도-정확도를 조절할 수 있으면서, attention이 실제로 집중되는 key 공간 영역에만 세밀한 양자화를 배정하고 나머지는 거친 양자화로 남겨 두는 VQ-attention을 설계하는 것이다.

Method Starting from a small set of parent codewords, the method computes per-codeword importance from the attention mass each parent receives during the actual forward pass, then "spawns" pre-learned child codewords for only the top-𝒫 most important parents -- refining resolution exactly where it matters and leaving everything else at the coarser parent level. The entire process (importance scoring, child spawning, parent-to-child contribution replacement) is implemented in custom Triton kernels fused directly into a Flash-Attention-style tiled computation loop, so nothing needs to leave SRAM.

방법 작은 parent 코드워드 집합에서 시작해, 순전파 도중 각 코드워드가 받는 attention 질량으로 중요도를 계산하고, 가장 중요한 상위 𝒫개 parent에 대해서만 미리 학습된 𝒞개의 child 코드워드를 "spawn"해 그 영역의 해상도를 정밀화한다. 이 전 과정 — 중요도 계산, child spawn, parent-기여 대체 — 을 Flash Attention 스타일의 타일 계산 루프 안에 직접 넣는 커스텀 Triton 커널로 구현해, 추가 자료 구조를 SRAM 밖으로 유출시키지 않는다.

Key idea The central move is adapting codebook resolution itself, at inference time and with no retraining, to where attention importance actually lies -- much like a cartographer drawing a city map in fine street-level detail downtown while rendering sparsely visited outskirts as large blocks. Changing 𝒫 alone, with no retraining, turns one trained checkpoint into a speed/accuracy dial, which is the practical payoff.

핵심 아이디어 핵심은 코드북 해상도 자체를 추론 시점에, 별도 재학습 없이 attention 중요도에 따라 적응시킨다는 것이다 — 마치 지도를 그릴 때 사람이 몰리는 도심은 골목까지 세밀하게, 인적 드문 외곽은 큰 구획으로만 그리는 것과 비슷하다. 𝒫를 바꾸는 것만으로 재학습 없이 하나의 체크포인트에서 속도와 정확도를 오가는 다이얼을 얻는다는 점이 실질적 이점이다.

Validation Evaluated by replacing attention layers in pretrained transformers and fine-tuning briefly, on ImageNet-1k classification (ViT-Base) and ADE20K semantic segmentation (DPT-Large), and by applying the same substitution to Stable Diffusion 1.5's UNet self-attention under the LinFusion distillation protocol. Wall-clock kernel time is measured out to N=65,536 to confirm the predicted linear scaling, with the gain from fusing the VQ precompute step reported separately from the additional gain from adaptive resolution allocation.

검증 ImageNet-1k 분류(ViT-Base)와 ADE20K 시맨틱 분할(DPT-Large)에서 사전학습된 transformer의 attention 층을 교체해 짧게 미세조정한 뒤 평가하며, Stable Diffusion 1.5 UNet의 self-attention에 LinFusion 증류 프로토콜을 그대로 적용해 검증한다. N=65,536까지 실제 wall-clock 커널 시간을 측정해 이론적 선형 스케일링을 확인하고, VQ 사전계산 단계를 융합하는 것만으로 얻는 이득과 적응형 해상도 배분이 추가로 주는 이득을 분리해서 보고한다.

Results At N=65,536, fused AVQ-attention configurations are 81-127x faster than Flash Attention, and fusing the VQ precompute step alone (even for flat, non-adaptive VQ) already buys roughly 2x. AVQ beats flat VQ-attention at matched kernel cost on both ImageNet-1k and ADE20K (ADE20K mIoU 43.33% vs. flat-VQ 42.70%, Swin 42.90%, NATTEN 40.99%), and on Stable Diffusion 1.5, extended to an unseen 1024² resolution, it matches or beats LinFusion and ToMe-SD on FID/CLIP at a lower UNet forward time.

결과 N=65,536에서 융합된 AVQ-attention 설정은 Flash Attention 대비 81-127배 빠르며, flat VQ의 경우도 VQ 사전계산을 커널에 융합하는 것만으로 약 2배의 속도 향상을 얻는다. ImageNet-1k·ADE20K 모두에서 동일 attention 커널 비용 대비 flat VQ-attention보다 높은 정확도를 얻고(ADE20K mIoU 43.33% vs flat-VQ 42.70%, Swin 42.90%, NATTEN 40.99%), Stable Diffusion 1.5에서는 1024² 미학습 해상도로 확장해도 LinFusion·ToMe-SD보다 낮은 UNet forward 시간에 대등하거나 나은 FID·CLIP 점수를 유지한다.

Comparison At matched kernel cost, it outperforms flat VQ-attention, Swin, NATTEN, Linformer, and Performer (the last of which is additionally not Flash-Attention-compatible, a further practical disadvantage). Against Flash Attention itself, the reported speedup is shared between the VQ idea and the engineering effort of fusing it fully into the tile loop -- and the paper is honest about this, reporting the fusion-only gain and the adaptive-allocation gain separately rather than conflating them into one headline number.

비교 동일 커널 비용에서 flat VQ-attention·Swin·NATTEN·Linformer·Performer를 모두 앞서며, 특히 Performer는 Flash-Attention과 호환되지 않아 실질적 비교에서 더 불리하다. Flash Attention 자체와 비교하면 VQ라는 아이디어와 그것을 타일 루프에 완전히 융합하는 엔지니어링이 기여를 나눠 갖는데, 저자들은 이를 분리해서 "융합 자체의 이득"과 "적응형 배분의 이득"을 각각 보고한다는 점에서 비교가 정직하다.

Significance From an efficient-AI standpoint, this shows that "allocate representational capacity by importance" can be delivered with no extra kernel-level overhead once flat VQ is already fused into a tiled loop -- a general design principle applicable to attention approximation broadly, and a useful reference point for other adaptive-allocation methods, including KV-cache compression. It is worth being explicit, though, that the paper itself says nothing about causal language modeling or a cached, cross-step KV-cache setting.

의의 효율적 AI 관점에서 이 논문은 "중요도에 따라 표현 용량을 적응적으로 배분한다"는 원리가 커널 융합과 결합될 때 추가 오버헤드 없이 공짜로 얻어질 수 있음을 보여준다 — 이는 attention 근사 전반에 적용 가능한 일반적 설계 원리로서, KV 캐시 압축을 포함한 다른 적응형 할당 방법들의 참고 사례가 된다. 다만 이 논문 자체는 인과적 언어 모델링이나 KV 캐시 세팅을 전혀 다루지 않는다는 점은 분명히 해 둘 필요가 있다.

Limitations The authors themselves note the optimal M₀/𝒫/𝒞 likely varies by layer ("some layers may need far fewer codewords"), but automatic per-layer configuration is not attempted and is left as future work. Codebooks require a short fine-tuning pass on a pretrained transformer rather than being usable training-free, and deeper (more than one level) hierarchies are proposed but never tested.

한계 𝑀₀·𝒫·𝒞의 최적값이 층마다 다를 가능성이 높다고 저자들 스스로 언급하지만("일부 층은 훨씬 적은 코드워드만 필요할 수 있다"), 층별 자동 설정은 시도되지 않고 향후 과제로 남아 있다. 코드북은 사전학습된 transformer를 몇 epoch 미세조정해 얻어야 하므로 완전한 training-free 방법은 아니며, 2단계 이상의 더 깊은 계층 구조는 제안만 되고 실험되지 않았다.

Future work The authors suggest that captured attention mass and quantization error -- two signals AVQ-attention already exposes -- could enable layer-wise architecture search without full retraining. They also note that adding one more level of hierarchy costs only a linear (𝒫𝒞) increment while multiplying codebook resolution, and leave this as the natural next extension.

향후 과제 저자들은 캡처된 attention 질량과 양자화 오차라는, AVQ-attention이 이미 노출하는 두 신호를 이용해 재학습 없이 층별 아키텍처 탐색을 할 수 있으리라 제안한다. 또한 계층을 한 단계 더 깊게 하면 추가 비용은 선형(𝒫𝒞)인 반면 코드북 해상도는 곱셈적으로 늘어난다는 점을 지적하며, 이를 다음 자연스러운 확장으로 남겨 둔다.

Released resources No separate code repository link was found in the paper -- no public implementation beyond the arXiv listing was located.

자원 공개 논문 자체에 별도의 코드 저장소 링크는 확인되지 않았다 — arXiv 게재 외 공개 구현체는 발견되지 않았다.