SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention
Created by
Haebom
Category
Empty
저자
Jintao Zhang, Haoxu Wang, Kai Jiang, Shuo Yang, Kaiwen Zheng, Haocheng Xi, Ziteng Wang, Hongzhou Zhu, Min Zhao, Ion Stoica, Joseph E. Gonzalez, Jun Zhu, Jianfei Chen
개요
Diffusion Transformer (DiT) 모델, 특히 비디오 생성에서 긴 시퀀스 길이와 이차 복잡도로 인해 어텐션 지연이 주요 병목 현상으로 작용합니다. 본 논문에서는 어텐션 가중치를 중요, 중요하지 않음, 무시할 수 있는 세 가지 범주로 분류하여 희소 어텐션(sparse attention)과 선형 어텐션(linear attention)을 결합하는 SLA (Sparse-Linear Attention)를 제안합니다. SLA는 중요 가중치에 O(N^2) 어텐션을, 중요하지 않은 가중치에 O(N) 어텐션을 적용하고, 무시할 수 있는 가중치를 건너뜁니다. SLA는 단일 GPU 커널로 이러한 계산을 통합하여 DiT 모델의 어텐션 계산량을 최대 20배 줄여 생성 품질 저하 없이 상당한 가속을 달성합니다. 실험 결과, SLA는 어텐션 계산량을 95% 감소시키면서 최종 생성 품질을 유지하며, 기존 방법보다 뛰어난 성능을 보입니다. 또한 효율적인 GPU 커널을 구현하여 어텐션 계산에서 13.7배, 비디오 생성에서 2.2배의 속도 향상을 보였습니다.
