AB-Attention, 비디오 확산 트랜스포머의 희소 어텐션을 더 정확하게

비디오 확산 트랜스포머는 높은 품질의 영상을 생성하지만, 프레임 수와 해상도가 증가하면 추론 비용도 빠르게 커집니다. 특히 3D 전체 어텐션은 많은 비디오 토큰 사이의 상호작용을 모두 계산하기 때문에 이차적인 비용이 발생합니다. 동적 블록 희소 어텐션은 계산량을 줄일 수 있지만, 희소성을 지나치게 높이면 밀집 어텐션과의 품질 차이가 커지는 문제가 있습니다.

AB-Attention은 단순히 더 많은 블록을 제거하는 대신, 적은 계산으로 밀집 어텐션을 더 충실하게 추정하는 데 초점을 둡니다. 기존 방식은 거친 요약값으로 키·밸류 블록의 중요도를 평가하거나, 비용이 큰 전역 클러스터링을 사용합니다. 또한 하나의 타일을 통째로 유지하거나 버리는 경우가 많아, 실제로는 타일 내부의 일부 키만 중요한 상황을 제대로 활용하기 어렵습니다. 버려진 블록은 softmax에 포함되지 않기 때문에 최종 어텐션 분포가 달라질 수도 있습니다.

첫 번째 구성 요소인 Intra-Block Clustering Selector(IBCS)는 각 쿼리 블록 안에서만 쿼리를 클러스터링합니다. 토큰 순서를 바꾸지 않으므로 전역 클러스터링 구조를 만드는 비용을 피할 수 있습니다. 동시에 하나의 블록에 속한 쿼리들이 서로 다른 키·밸류 블록을 필요로 할 수 있다는 점을 반영해, 단일 요약값보다 세밀한 중요도 추정을 제공합니다.

두 번째 구성 요소는 Sampling with Inclusion-Probability Weighting(SIP)입니다. 중요도가 낮은 블록을 임계값 아래에서 일괄적으로 잘라내는 대신, 추정된 중요도에 비례해 블록을 샘플링합니다. 이후 각 블록의 포함 확률 역수를 이용해 결과를 재가중합니다. 따라서 선택되지 않은 블록을 영구적으로 무관하다고 가정하는 대신, 한 번의 희소 계산으로 밀집 어텐션을 추정하는 효과를 노립니다. 희소성이 높을수록 이런 보정이 잘못된 제거의 누적 영향을 줄이는 데 중요할 수 있습니다.

논문은 GPU 실행 경로도 함께 설계했습니다. 일반 타일보다 작은 키·밸류 블록을 모아 완전한 Tensor Core 타일로 구성함으로써, 더 세밀한 블록 선택을 Tensor Core의 실행 형태와 결합합니다. 논문에 따르면 이 방식은 Tensor Core 계산 비용을 추가하지 않으면서 선택 단위를 세분화할 수 있습니다. 이 점은 이론적인 FLOPs 감소뿐 아니라 실제 커널 효율까지 고려했다는 의미가 있습니다.

Wan2.1과 HunyuanVideo를 대상으로 한 실험에서 AB-Attention은 어텐션 계산의 8992%를 건너뛰었습니다. 같은 비교에서 기준 방법들은 1.83.6배 더 많은 어텐션을 계산했습니다. 또한 가장 강력한 기준선보다 PSNR을 최대 4.4 dB 높였고, 엔드투엔드 추론에서는 최대 2.23배의 속도 향상을 기록했다고 보고합니다.

이 연구의 의미는 비디오 확산 어텐션 가속을 거친 블록 절단에서 확률적으로 보정된 근사 계산으로 확장했다는 데 있습니다. 중요도 추정, 샘플링 보정, GPU 친화적 블록 수집을 하나의 추론 방법으로 묶어 희소성과 정확도의 균형을 개선하려는 접근입니다. 다만 현재 결과는 Wan2.1과 HunyuanVideo, 그리고 특정 실험 조건에 기반합니다. 실제 속도와 품질 향상은 GPU, 시퀀스 길이, 목표 희소율, 생성 작업에 따라 달라질 수 있으므로 다양한 모델과 환경에서의 추가 검증이 필요합니다.

출처: Hugging Face Daily Papers