트랜스포머의 어텐션은 토큰 수가 늘수록 계산량이 N의 제곱에 비례해 커진다. 이를 줄이기 위해 제안된 VQ-Attention은 키를 코드워드로 묶어 O(MN)으로 낮추지만, 코드북 용량을 모든 구간에 똑같이 적용한다는 한계가 있다. arXiv 프리프린트 2607.12789은 이런 고정 코드북 대신, 전방 계산에서 주의 질량이 큰 부모 코드워드만 골라 자식 코드워드로 더 세밀하게 나누는 Adaptive Vector-Quantized Attention(AVQ-Attention)을 제안한다.
이 방법은 어텐션 가중치를 중요도 신호로 쓴다. 먼저 부모 코드워드로 값을 집계하고, 쿼리 타일별로 중요도가 높은 상위 P개 부모를 선택한 뒤 미리 학습된 자식 코드워드를 추가한다. 부모는 자식들의 평균이어야 한다는 제약을 두어, 자식을 계산할 때 부모의 기여를 다시 불러와 교정할 수 있게 했다. 논문은 이 구조가 필요한 곳에만 세밀한 양자화를 적용하고, 덜 중요한 구간은 거친 근사로 남긴다고 설명한다.
구현은 custom Triton kernels와 Flash Attention의 tiled computation을 활용한다. 논문은 양자화, 값 집계, 중요도 산정, 자식 삽입, 부모 기여 교정을 하나의 타일 계산 안에서 처리해 중간 결과를 글로벌 메모리에 쓰고 다시 읽는 부담을 줄인다고 적었다. 전체 파이프라인은 VQ precompute와 attention의 두 커널로 이뤄지며, 표 1에서 제시한 대로 총 FLOPs는 O(N(M0+PC)D)로 유지된다.
실험은 ImageNet-1k, ADE20K, Stable Diffusion 1.5에서 이뤄졌다. ImageNet-1k에서는 ViT-Base를 N=785 토큰으로 4 epoch 미세조정했고, ADE20K에서는 DPT-Large를 N=901 토큰으로 10 epoch 미세조정했다. 동일한 조건에서 AVQ는 고정 코드북 VQ와 비교해 더 나은 정확도-효율 절충을 보였다고 보고됐다. ADE20K의 비교 표에서는 3개 시드 평균 mIoU가 AVQ 43.33±0.12%, flat VQ 42.70±0.08% 또는 43.04±0.05%로 제시됐고, 커널 시간은 AVQ 0.116ms, flat VQ 0.103ms와 0.164ms, Flash Attention-v2 baseline은 0.313ms였다.
긴 시퀀스에서는 속도 차이가 더 크게 나타났다. N=65,536에서 AVQ의 fused kernel 구성은 Flash Attention보다 81~127배 빠르다고 보고됐다. 또한 P를 바꾸는 추론 실험에서는 M0=64, C=8로 학습한 단일 모델을 P=4부터 32까지 평가했을 때 ADE20K mIoU가 41.44%에서 42.49% 수준, ImageNet 정확도가 79.23%에서 82.29% 수준으로 변했다. P=16에서는 실제 주의 질량의 82% 이상을 포착하면서 최대 성능과의 차이가 0.3% 이내였다고 논문은 밝혔다.
| 항목 | 값 | 근거 위치 |
|---|---|---|
| ADE20K 성능(3개 시드 평균) | AVQ 43.33 ± 0.12% mIoU, flat VQ 42.70 ± 0.08% 또는 43.04 ± 0.05% mIoU | Table 3 |
| ADE20K 커널 시간 | AVQ 0.116 ms, flat VQ 0.103 ms 또는 0.164 ms, Flash Attention-v2 baseline 0.313 ms | Table 3 |
| Elastic inference, P=16 | ADE20K 42.42% mIoU, capture 85.3%, kernel 0.20 ms; ImageNet 82.04% accuracy, capture 82.7%, kernel 0.14 ms | Table 6 |
| 대규모 시퀀스 시간(N=65,536) | AVQ fused kernel이 Flash Attention보다 81–127× 빠름 | Sec. 5.1 / Table 2 |
자료: STORIUM 정리
부가 분석도 제시됐다. 코드북 활용률은 flat VQ의 경우 M=512에서 90.4%였고, AVQ는 총 코드워드 수가 576일 때 97.0%, 1152일 때 95.4~95.5%를 유지했다. 쿼리 순서를 Gilbert space-filling curve로 재배열하는 실험에서는 ADE20K에서 raster order보다 약 0.2%포인트 나은 결과가 나왔고, random permutation은 ImageNet에서 약 1.8%포인트, ADE20K에서 약 0.8%포인트 성능 저하를 보였다. 논문은 코드북 학습에 EMA 기반 online k-means를 썼으며, DiVeQ나 dead-code 처리 같은 다른 양자화 기법은 더 나은 결과를 낼 수 있다고 덧붙였다.
저작권자 © STORIUM 무단전재 및 재배포 금지
