확산 변환기를 위한 효율적인 학습 가능 로그-선형 희소 어텐션

초록

확산 트랜스포머(DiT)는 시각적 생성 분야에서 최첨단 성능을 보여주지만, 2차적인 자기 주의력(self-attention) 비용으로 인해 긴 토큰 시퀀스로의 확장이 근본적으로 제한됩니다. 최근의 Top-K 희소 주의력 접근법은 토큰을 블록 단위 표현으로 압축하고 소수의 관련 키 블록을 선택하여 DiT의 계산량을 줄이지만, 여전히 (i) 압축된 토큰에 대한 2차 선택 비용과 (ii) 시퀀스 길이가 증가함에 따라 모델 품질 유지를 위해 필요한 K 값 증가 문제가 발생합니다. 우리는 이러한 비효율성이 단일 수준 설계, 즉 단일한 coarse 수준이 전역 구조를 표현하기에 부족하기 때문임을 확인했습니다. 본 논문에서는 계층적 구조를 활용하여 선택 비용과 주의력 비용을 2차에서 로그-선형 복잡도로 감소시키는, 극도로 긴 토큰 시퀀스를 위한 학습 가능한 희소 주의력 메커니즘인 로그-선형 희소 주의력(LLSA)을 소개합니다. LLSA는 계층적 Top-K 선택을 수행하여 이전 수준에서 발견된 인덱스로 희소 Top-K 선택을 점진적으로 적용하며, 주의력 계산 시 다양한 세분화의 더 적은 토큰을 사용하면서도 전역 컨텍스트를 보존하는 계층적 키-값 강화(Hierarchical KV Enrichment) 메커니즘을 도입합니다. 효율적인 학습을 지원하기 위해 순전파와 역전파 모두에서 희소 인덱스만을 사용하는 고성능 GPU 구현을 개발하여 밀집 주의력 마스크의 필요성을 제거했습니다. 패치화와 VAE 인코딩을 사용하지 않고 고해상도 픽셀 공간 이미지 생성에 대해 LLSA를 평가했습니다. LLSA는 256x256 픽셀 토큰 시퀀스에서 주의력 추론 속도를 28.27배, DiT 학습 속도를 6.09배 가속시키면서도 생성 품질을 유지했습니다. 결과는 LLSA가 긴 시퀀스 DiT를 효율적으로 학습하기 위한 유망한 방향을 제시함을 보여줍니다. 코드는 https://github.com/SingleZombie/LLSA에서 확인할 수 있습니다.

English

Diffusion Transformers (DiTs) set the state of the art in visual generation, yet their quadratic self-attention cost fundamentally limits scaling to long token sequences. Recent Top-K sparse attention approaches reduce the computation of DiTs by compressing tokens into block-wise representation and selecting a small set of relevant key blocks, but still suffer from (i) quadratic selection cost on compressed tokens and (ii) increasing K required to maintain model quality as sequences grow. We identify that their inefficiency is due to the single-level design, as a single coarse level is insufficient to represent the global structure. In this paper, we introduce Log-linear Sparse Attention (LLSA), a trainable sparse attention mechanism for extremely long token sequences that reduces both selection and attention costs from quadratic to log-linear complexity by utilizing a hierarchical structure. LLSA performs hierarchical Top-K selection, progressively adopting sparse Top-K selection with the indices found at the previous level, and introduces a Hierarchical KV Enrichment mechanism that preserves global context while using fewer tokens of different granularity during attention computation. To support efficient training, we develop a high-performance GPU implementation that uses only sparse indices for both the forward and backward passes, eliminating the need for dense attention masks. We evaluate LLSA on high-resolution pixel-space image generation without using patchification and VAE encoding. LLSA accelerates attention inference by 28.27x and DiT training by 6.09x on 256x256 pixel token sequences, while maintaining generation quality. The results demonstrate that LLSA offers a promising direction for training long-sequence DiTs efficiently. Code is available at: https://github.com/SingleZombie/LLSA

확산 변환기를 위한 효율적인 학습 가능 로그-선형 희소 어텐션

Trainable Log-linear Sparse Attention for Efficient Diffusion Transformers

초록

Support