PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers

기술 요약: PISA - 효율적인 확산 트랜스포머를 위한 조각별 희소 어텐션 (Piecewise Sparse Attention)

1. 문제 정의

확산 트랜스포머(Diffusion Transformers, DiTs)는 고충실도 이미지 및 비디오 생성의 표준이 되었습니다. 그러나 이들의 확장성은 입력 시퀀스 길이($L$)가 증가함에 따라(더 높은 해상도나 긴 비디오의 경우) 셀프 어텐션 메커니즘의 이차 복잡도($O(L^2)$)에 의해 병목 현상이 발생합니다.

기존 솔루션들은 다음과 같은 상당한 한계에 직면해 있습니다:

• 희소 어텐션 ("유지 또는 폐기" - Keep-or-Drop): 비중요 키-값(KV) 블록을 버리는 방식은 계산을 가속화하지만, 높은 희소도 비율에서 컨텍스트 손실로 인해 성능 저하를 겪습니다. 또한, 선택에 따른 오버헤드가 이득보다 큰 짧은 시퀀스에서는 어려움을 겪습니다.
• 선형/하이브리드 어텐션: 어텐션 분포를 변경하는 방식(예: 피처 매핑이나 희소/선형 브랜치의 가산적 결합을 통한 방식)은 사전 학습된 모델의 내재적 정규화를 방해하는 경우가 많습니다. 이는 사전 학습된 가중치를 직접 재사용하는 것을 방지하며, 값비싼 재학습이나 미세 조정을 필요로 합니다.
• 하드 트렁케이션 (Hard Truncation): 블록을 직접적으로 드롭하는 것은 원래의 어텐션 분포로부터 급격한 편차를 유발하여 생성된 콘텐츠의 품질 저하를 초래합니다.

핵심 과제는 사전 학습된 가중치를 재학습 없이 그대로 사용하면서도, 전체 어텐션의 정확도를 유지하고 부-이차(sub-quadratic) 복잡도로 추론을 가속화할 수 있는 메커니즘을 설계하는 것입니다.

2. 방법론: 조각별 희소 어텐션 (PISA)

PISA는 전통적인 "유지 또는 폐기" 전략을 넘어, 새로운 "정확 또는 근사(exact-or-approximate)" 패러다임을 도입합니다. 이는 어텐션을 전체 어텐션 범위를 부-이차 복잡도로 커버하는 조각별 프로세스로 취급합니다.

핵심 통찰

저자들은 사전 학습된 DiT들을 분석하여 두 가지 결정적인 특성을 발견했습니다:

1. 분포적 안정성 (Distributional Stability): 비중요(uncritical) 블록들의 pre-softmax 어텐션 스코어는 0 또는 음수 값을 중심으로 하는 대칭적인 종 모양 분포를 보입니다. 이는 이들이 단순한 폐기가 아닌, 테일러 전개(Taylor expansion)를 통한 정확한 근사에 매우 적합함을 의미합니다.
2. 정규화 일관성 (Normalization Consistency): 학습 없이 호환성을 유지하기 위해서는, 근사치가 단순히 서로 다른 브랜치의 출력을 더하는 방식(이는 가중 합 규칙을 위반함)이 아니라, 소프트맥스 분자와 분모 내부에 내부적으로 통합되어야 합니다.

핵심 알고리즘

PISA는 KV 블록을 두 세트로 분할합니다:

1. 중요 블록 ($S_i$): 높은 살리언스(salience) 정보를 보존하기 위해 정확하게 계산됩니다.
2. 비중요 블록 ($U_i$): 블록 단위 테일러 전개를 사용하여 효율적으로 근사됩니다.

출력 $o_t$는 가중치 합산된 값의 집계(aggregation)를 정규화하여 도출됩니다:
$$o_t = \frac{N_t}{D_t}$$

여기서 분모($D_t$)와 분자($N_t$)는 다음과 같이 구성됩니다:

• 정확한 희소 항 (Exact Sparse Terms): 선택된 블록 $j \in S_i$에 대해 계산됩니다.
• 블록 단위 0차 근사 (Block-wise Zeroth-Order Approximation): 선택되지 않은 블록들에 대해, 블록 중심점 $\bar{k}_j$와 값의 합을 사용하여 근사됩니다.
• 전역 1차 보정 (Global First-Order Correction): 모든 미선택 블록에 대해 별도의 1차 항을 계산하는 대신(이는 메모리 제한 문제를 야기함), PISA는 **하이브리드 차수 근사(hybrid-order approximation)**를 채택합니다. 모든 블록에 걸친 블록 단위 1차 행렬의 평균인 전역 통계량 $\bar{H}$를 사용하여 전체 테일(tail) 세트에 대한 1차 그래디언트를 근사합니다.

분자 $N_t$에 대한 공식은 다음과 같습니다:
$$N_t = \sum_{j \in S_i} \sum_{n} e^{q_t k_{j,n}^\top} v_{j,n} + \sum_{j \in U_i} \alpha_{t,j} \left(\sum_{n} v_{j,n}\right) + q_t \bar{H} \sum_{j \in U_i} \alpha_{t,j}$$
여기서 $\alpha_{t,j} = \exp(q_t \bar{k}_j^\top)$ 입니다.

공분산 인식 라우팅 (Covariance-Aware Routing)

근사 오차를 최소화하기 위해, PISA는 공분산 인식 Top-k 전략을 사용합니다. 블록 선택 점수는 어텐션 스코어뿐만 아니라 블록의 공분산 행렬 노름($M_j = \|H_j - \bar{H}\|_2$)을 포함합니다. 높은 의미적 관련성이나 높은 이질성(heterogeneity)을 가진 블록은 정확한 계산을 위해 선택되며, 더 매끄러운(smoother) 블록은 테일러 전사로 위임됩니다.

하드웨어 구현

PISA는 온라인 소프트맥스 과정 내에서 정확한 계산과 근사 스캐닝을 인터리빙(interleaving)하는 **융합 커널(fused kernel)**을 통해 구현됩니다(FlashAttention 스타일). 파이프라인은 다음과 같습니다:

1. 사전 계산 (Pre-computation): 블록 중심점, 값의 합, 그리고 전역 통계량 $\bar{H}$를 계산합니다.
2. 단계 1 (정확 - Exact): 선택된 블록을 정확한 어텐션을 위해 SRAM으로 로드합니다.
3. 단계 2 (근사 - Approximation): 0차 전개를 사용하여 미선택 블록들에 대한 집계된 블록 벡터를 스캐닝합니다.
4. 단계 3 (보정 - Correction): 전역 1차 보정 항을 주입합니다.

이 설계는 메모리 효율성을 보장하며, 모든 블록에 대해 별도의 행렬을 로드하는 데서 오는 병목 현상을 방지합니다.

3. 주요 기여

1. 새로운 조각별 희소 어텐션: 정확한 희소 계산과 블록 단위 근사를 통합하여 부-이차 복잡도로 전체 어텐션 범위를 커버할 수 있는 학습 불필요(training-free) 메커니즘을 제공합니다.
2. 하이브리드 차수 근사 체계: 블록 단위 0차 전개와 전역 1차 보정을 결합한 전략입니다. 이는 표준 희소 어텐션에 비해 무시할 만한 계산 비용만으로 정확도를 크게 높입니다.
3. 공분산 인식 라우팅: 어텐션 점수와 블록 공분산 노름을 모두 기준으로 블록을 선택함으로써 근사 발산을 최소화하는 유도된 방식입니다.
4. 최첨단 성능: 재학습 없이 다양한 작업(비디오 및 이미지 생성)에서 우수한 효율성과 품질을 입증했습니다.

4. 실험 결과

저자들은 비디오 생성 모델(Wan2.1-1.3B/14B, Hunyuan-Video)과 이미지 생성 모델(SD 3.5, FLUX.1)에 대해 PISA를 평가했습니다.

• 비디오 생성:

  • Wan2.1-14B: 87.5% 희소도에서 밀집 어텐션(dense attention) 대비 1.91배의 속도 향상(1468초 $\to$ 818초)을 달성하면서도, 밀집 어텐션과 비교하여 품질 지표(VBench, PSNR, SSIM)를 유지하거나 약간 개선했습니다.
  • Hunyuan-Video: 일관된 품질 보존과 함께 2.57배의 속도 향상(1651초 $\to$ 641초)을 달성했습니다.
  • 비교: 기존의 SpargeAttn 및 SVG2와 같은 희소 방식보다 속도와 품질(특히 미세한 디테일 및 시간적 일관성 보존) 측면에서 모두 뛰어난 성능을 보였습니다.

• 이미지 생성:

  • FLUX.1-dev: 85% 희소도에서 1.2배의 속도 향상(8.32초 $\to$ 6.87초)을 달성했습니다. 특히, PISA는 밀집 어텐션 및 SpargeAtt는 모두보다 인간 선호도 지표(ImageReward, HPSv2)와 유사도 지표(PSNR, SSIM)에서 더 높은 점수를 기록했습니다.
  • SD 3.5: 밀집 어텐션과 비슷하거나 더 나은 FID 점수를 유지하면서 상당한 속도 향상(예: SD 3.5-M에서 1.4배)을 보여주었습니다.

• 커널 효율성:

  • PISA는 다양한 시퀀스 길이(4K ~ 32K)와 밀도에 걸쳐 FlashAttention-2/3 및 SpargeAttn보다 레이턴시 측면에서 일관되게 우수한 성능을 보였습니다.
  • 짧은 시퀀스에서 성능이 저하되는 SpargeAttn과 달리, PISA는 4K 토큰에서도 속도 우위를 유지했습니다.

5. 의의 및 주장

본 논문은 PISA가 "유지 또는 폐기"에서 "정확 또는 근사"로의 패러다임 전환을 나타낸다고 주장합니다. 주요 의의는 다음과 같습니다:

• 속도-품질 간극 해소: PISA는 무시할 만한 오버헤드로 거의 손실 없는 근사를 달 Achieves 함으로써 추론 속도와 생성 품질 사이의 간극을 효과적으로 메웁니다.
• 학습 불필요 호환성: 근사를 소프트맥스 정규화 프레임워크에 직접 통합함으로써, PISA는 선형 또는 하이브리드 어텐션 방식에 필요한 값비싼 재학습이나 미세 조정을 생략하고 사전 학습된 가중치를 직접 재사용할 수 있게 합니다.
• 확장성: 고해상도 이미지 및 장기 비디오 생성 모두를 위한 통합된 솔루션을 제공하며, 생성된 콘텐츠의 충실도를 희생하지 않고 이차 복잡도 병목 현상을 해결합니다.

저자들은 자신들의 방법이 블록 단위 테일러 전개와 희소 어텐션을 사전 학습된 모델의 내재적 분포를 보존하는 통합된 프레임워크로 최초로 결합한 방식임을 강조합니다.