Re-ttention: 주의 통계 재구성을 통한 초희소 시각적 생성

Re-ttention: Ultra Sparse Visual Generation via Attention Statistical Reshape

May 28, 2025

저자: Ruichen Chen, Keith G. Mills, Liyao Jiang, Chao Gao, Di Niu

cs.AI

초록

디퓨전 트랜스포머(DiT)는 비디오 및 이미지와 같은 고품질 시각적 콘텐츠를 생성하기 위한 사실상의 표준 모델로 자리 잡았습니다. 그러나 주요 병목 현상은 해상도와 비디오 길이에 따라 복잡도가 2차적으로 증가하는 어텐션 메커니즘입니다. 이러한 부담을 줄이기 위한 논리적인 방법 중 하나는 희소 어텐션(sparse attention)으로, 계산에 포함되는 토큰 또는 패치의 일부만을 사용하는 것입니다. 그러나 기존 기술은 극도로 높은 희소성 수준에서 시각적 품질을 유지하지 못하며, 심지어 무시할 수 없는 계산 오버헤드를 초래할 수 있습니다. 이러한 문제를 해결하기 위해, 우리는 Re-ttention을 제안합니다. Re-ttention은 디퓨전 모델의 시간적 중복성을 활용하여 어텐션 메커니즘 내의 확률적 정규화 편차를 극복함으로써 시각적 생성 모델을 위한 매우 높은 희소 어텐션을 구현합니다. 구체적으로, Re-ttention은 이전 소프트맥스 분포 기록을 기반으로 어텐션 점수를 재구성하여 극도로 높은 희소성 수준에서도 완전한 2차 어텐션의 시각적 품질을 유지합니다. CogVideoX 및 PixArt DiT와 같은 T2V/T2I 모델에 대한 실험 결과는 Re-ttention이 추론 중에 단 3.1%의 토큰만을 필요로 하며, FastDiTAttn, Sparse VideoGen, MInference와 같은 최신 방법들을 능가함을 보여줍니다. 또한, 우리는 지연 시간을 측정하여 H100 GPU에서 45% 이상의 종단 간 지연 시간 감소와 92% 이상의 셀프 어텐션 지연 시간 감소를 무시할 수 있는 오버헤드 비용으로 달성할 수 있음을 입증했습니다. 코드는 다음 링크에서 확인할 수 있습니다: https://github.com/cccrrrccc/Re-ttention

English

Diffusion Transformers (DiT) have become the de-facto model for generating high-quality visual content like videos and images. A huge bottleneck is the attention mechanism where complexity scales quadratically with resolution and video length. One logical way to lessen this burden is sparse attention, where only a subset of tokens or patches are included in the calculation. However, existing techniques fail to preserve visual quality at extremely high sparsity levels and might even incur non-negligible compute overheads. % To address this concern, we propose Re-ttention, which implements very high sparse attention for visual generation models by leveraging the temporal redundancy of Diffusion Models to overcome the probabilistic normalization shift within the attention mechanism. Specifically, Re-ttention reshapes attention scores based on the prior softmax distribution history in order to preserve the visual quality of the full quadratic attention at very high sparsity levels. % Experimental results on T2V/T2I models such as CogVideoX and the PixArt DiTs demonstrate that Re-ttention requires as few as 3.1\% of the tokens during inference, outperforming contemporary methods like FastDiTAttn, Sparse VideoGen and MInference. Further, we measure latency to show that our method can attain over 45\% end-to-end % and over 92\% self-attention latency reduction on an H100 GPU at negligible overhead cost. Code available online here: https://github.com/cccrrrccc/Re-ttention{https://github.com/cccrrrccc/Re-ttention}