AB-Attention은 영상 확산 트랜스포머의 어텐션 계산을 희소화한다

무엇인가

영상 확산 트랜스포머는 공간과 시간 의존성을 잡기 위해 3D 전체 어텐션을 사용하지만, 시퀀스 길이에 대해 비용이 제곱으로 늘어난다. 해상도와 프레임 수가 커질수록 이 비용이 커지고, 어텐션은 여러 transformer 레이어와 디노이징 스텝에서 반복되므로 추론의 주요 병목이 된다. 따라서 품질을 유지하면서 어텐션 계산을 줄이는 것이 더 길고 고해상도인 영상 생성의 핵심 문제다.

어떻게 동작하나

동적 블록 희소 어텐션은 현재 활성값에서 중요한 key/value 블록만 골라 정확한 어텐션을 계산해 비용을 줄인다. 그러나 논문은 세 가지 한계를 지적한다. 블록 중요도를 mean pooling 같은 거친 요약으로 매기면 토큰별 어텐션 차이가 섞여 순위가 틀어지고, 전역 k-means와 토큰 순열은 정확하지만 비싸다. top-k 선택은 버려진 블록을 어텐션 분자와 softmax 정규화 분모 모두에서 제거해 꼬리 질량을 잃는다. 선택 단위가 실행 타일이라 일부 key만 중요한 타일도 통째로 유지하거나 버린다. AB-Attention은 학습 없이 높은 희소성에서도 dense attention을 정확히 근사하려고 이 세 지점을 각각 IBCS, SIP, 타일 패킹 커널로 바꾼다.

무엇과 다른가

먼저 Intra-Block Clustering Selector(IBCS)는 쿼리를 실행 타일 높이에 맞춘 128토큰 블록으로 나누고, key/value는 4토큰의 훨씬 작은 블록으로 나눈다. key/value 블록을 작게 하면 oracle top-k가 유지하는 어텐션 질량이 128토큰일 때 59%에서 4토큰일 때 69%로 올라간다고 논문은 제시한다. 각 쿼리 블록 안에서만 한 번의 k-means 반복으로 8개 쿼리 클러스터를 만들고, 각 클러스터의 centroid와 쿼리 비중을 구한다. 이 centroid들을 사용해 key/value 블록의 pooled key에 대한 softmax 질량을 클러스터 비중으로 가중 평균해 블록 점수를 매긴다. 전역 클러스터링이나 토큰 순열 없이도 일부 클러스터만 주목하는 블록의 점수를 보존하려는 설계다.

어떻게 쓰나

SIP(Sampling with Inclusion-Probability Weighting)는 이 점수를 top-k 절단이 아니라 포함확률로 쓴다. 각 쿼리 블록에서 블록 j의 포함확률을 π_ij = min(1, λ_i S_ij)로 두고, Σπ_ij = s가 되도록 예산에 맞춘다. 점수가 높아 π=1인 블록은 항상 포함되고, 나머지 블록도 0이 아닌 포함확률을 가진다. systematic sampling으로 정확히 s개의 서로 다른 블록을 뽑고, 한 쿼리 블록의 모든 쿼리가 같은 샘플을 공유한다. 각 샘플 블록의 어텐션 분자와 분모 기여도를 1/π_ij로 가중하면, 이는 샘플 블록에 대해 logit을 −log π_ij만큼 이동한 일반 softmax와 같아진다. 고정 활성값에서 분자와 분모는 Horvitz–Thompson 비편향 추정이지만, 비율인 출력 자체는 일반적으로 편향될 수 있다고 논문은 명시한다. 같은 블록 크기에서 SIP의 어텐션 오차는 oracle top-k보다 41% 낮다고 보고한다.

전제와 한계

세 번째는 선택 세분도와 실행 효율을 분리하는 타일 실행이다. AB-Attention은 4토큰 key/value 블록처럼 실행 타일보다 작은 단위로 선택하되, 커널이 샘플된 블록을 인덱스로 모아 전체 tensor-core 타일로 패킹한다. 이때 역포함확률 logit 이동과 online softmax, value 누적을 한 번의 어텐션 패스에서 수행한다. 토큰을 재정렬하거나 재정렬된 시퀀스를 메모리에 만들 필요가 없고, 고정 예산에서 tensor-core 타일 계산 수를 늘리지 않으면서 선택만 더 세밀하게 한다. 선택기 자체도 여러 쿼리·key/value 블록을 배치 행렬 연산으로 점수화해 작은 후보마다 별도 연산을 피한다. 샘플링 계획은 4 디노이징 스텝마다 다시 만들고 그 사이 재사용하지만, 어텐션은 매 스텝 현재 활성값으로 다시 계산한다.

실험은 Wan2.1-T2V-1.3B, Wan2.1-T2V-14B, HunyuanVideo-T2V-13B에서 720p 영상 81프레임(Wan2.1) 또는 129프레임(HunyuanVideo)을 생성하며 Penguin Benchmark 600개 프롬프트를 쓴다. 비교 대상은 RadialAttention, SVG2, DFSAttn이고, dense와 같은 프롬프트·시드 영상 대비 PSNR, SSIM, LPIPS와 VBench 다섯 차원을 측정한다. 모든 방법은 첫 transformer 블록과 50 디노이징 스텝 중 첫 10스텝을 dense로 유지한다. AB-Attention은 key/value 블록 4, 쿼리 블록 128, 쿼리 클러스터 8, 4스텝마다 샘플링 계획 재구성을 사용해 Wan2.1 두 모델에서 89.16%, HunyuanVideo에서 92.35%의 어텐션 희소성을 달성한다. 주요 결과로 AB-Attention은 모든 모델에서 PSNR, SSIM, LPIPS가 가장 좋고, 가장 강한 베이스라인보다 PSNR을 1.6~4.4dB 높인다. 어텐션의 89~92%를 건너뛰는데도 비교 베이스라인은 AB-Attention보다 1.8~3.6배 많은 어텐션을 계산한다. 단일 H20에서 end-to-end 속도는 1.84~2.23배 빨라 모든 모델에서 베이스라인을 앞서며, HunyuanVideo 생성 시간을 4399초에서 1975초로 줄인다. HunyuanVideo에서 가장 가까운 경쟁자인 DFSAttn은 2.07배 속도를 내지만 dense 어텐션의 18.71%를 계산해 AB-Attention의 7.65%보다 2.4배 많은 어텐션을 쓴다. VBench의 미학 품질, 배경 일관성, 화질, 모션 부드러움, 피사체 일관성에서도 dense에 가깝게 유지된다고 보고한다.

절제 실험은 Wan2.1-T2V-14B와 Penguin Benchmark에서 무작위 100개 프롬프트, 동일한 89.16% 희소성에서 수행된다. 쿼리 클러스터링을 위치 기반 16쿼리 구간 평균으로 바꾸면 PSNR이 0.91dB 떨어지고, SIP를 빼는 변형은 3.99~7.25dB 떨어진다. 포함확률 가중을 제거하면 4.00dB 떨어져 top-k 수준이 되고, 점수까지 버린 균일 샘플링은 추가로 3.25dB 떨어지며 LPIPS가 두 배 이상 나빠진다. 실행 타일 단위 선택은 1.02dB 낮아지고, 선택 블록을 전체 타일에서 4토큰까지 줄일수록 PSNR이 꾸준히 오르며 패킹 덕분에 tensor-core 계산은 늘지 않는다. 모든 변형이 세 지표에서 AB-Attention보다 나쁘다.

개발자에게 이 논문은 비디오 DiT 추론 서빙에서 학습 없이 어텐션 병목을 줄이는 설계와 수치를 제공한다. 특히 높은 희소성에서 dense 품질을 추정하려는 SIP와 작은 블록을 타일로 패킹하는 커널 분리가, 기존 block-sparse 백엔드에 어떤 식으로 통합할 수 있는지 참고할 만하다. 다만 저자들이 밝힌 한계는 커널이 실험에 쓴 H20 GPU용이고 현재 forward inference만 지원한다는 점, 평가가 세 개 text-to-video diffusion transformer에 한정된다는 점이다. 향후 과제로 더 많은 하드웨어 커널 지원, 더 넓은 모델 평가, backward를 포함해 AB-Attention 아래에서 fine-tuning할 수 있는 trainable 구현을 제시한다. 실무 적용 전에는 대상 GPU와 커널 지원, BF16/FlashAttention-3 백엔드, 샘플링 계획 재구성 주기, dense 유지 구간을 확인해야 한다.