AVQ-Attention은 트랜스포머 모델의 어텐션 복잡도를 줄이는 새로운 방식입니다. 기존 벡터 양자화 어텐션(VQ) 방식의 단점을 개선하여 중요도가 높은 영역은 정밀하게, 낮은 영역은 대략적으로 양자화합니다. 트라이톤 커널을 사용해 플래시 어텐션과 유사한 효율성을 유지하며, 정확도와 효율성 간 균형을 개선했습니다.
AVQ-Attention은 어텐션 중요도에 따라 코드북 용량을 적응적으로 할당합니다. 작은 코드셋에서 시작하여 순방향 패스 과정에서 가장 중요한 코드를 식별하고, 미리 학습된 자식 코드워드로 개선하여 중요도가 높은 영역에서는 정밀한 양자화를, 그렇지 않은 영역에서는 대략적인 양자화를 유지합니다.
기존 벡터 양자화 어텐션과 동일한 O(MN) 복잡도를 유지하면서도 정확도와 효율성 측면에서 개선된 성능을 보입니다.