1. 연구 배경 및 문제 정의 (Problem)

비디오 확산 모델 (Video Diffusion Models) 은 긴 시퀀스 길이와 $O(N^2)$ 복잡도를 가진 어텐션 연산으로 인해 추론 비용이 매우 높습니다. 이를 해결하기 위해 기존에는 학습이 필요 없는 (Training-free) 희소 어텐션 방법들이 사용되었으나, 생성 품질을 유지하면서 더 높은 희소성을 달성하기 위해서는 학습 가능한 (Trainable) 희소 어텐션이 필요합니다.

하지만 기존 학습 가능한 희소 어텐션 방법들은 다음과 같은 두 가지 주요 한계점을 가지고 있었습니다:

1. Top-k 및 Top-p 마스킹의 실패 (High Sparsity Failure):

   • Top-k: 어텐션 가중치 분포가 균일 (Uniform) 한 경우, 고정된 토큰 수만 선택하면 중요한 컨텍스트를 놓치게 됩니다.
   • Top-p: 어텐션 가중치가 심하게 편향 (Skewed) 되어 있는 경우 (예: Attention Sink 현상), 누적 확률 임계값을 매우 적은 토큰으로 만족시켜 정보 가치가 높은 다른 토큰들을 누락시킵니다.
   • 결과: 매우 높은 희소성 (>90%) 에서 중요한 어텐션 연산이 손실되어 생성 품질이 저하됩니다.

2. 파인튜닝 시 데이터 분포 불일치 및 손실 함수의 한계:

   • 기존 방법들은 실제 데이터로 수집된 프롬프트 - 비디오 쌍을 사용하여 표준 Diffusion Loss로 모델을 파인튜닝합니다.
   • 그러나 오픈소스 비디오 모델 (예: Wan2.1) 의 경우 원본 학습 데이터가 공개되지 않아, 파인튜닝 데이터가 원본 프리트레이닝 데이터 분포와 일치하지 않습니다.
   • 결과: 희소 어텐션 도입 여부와 상관없이, 데이터 분포 불일치로 인해 Diffusion Loss 만을 최적화하면 모델의 생성 품질 (미적 품질, 시각적 일관성 등) 이 원본 모델보다 현저히 떨어지는 현상이 발생합니다.

2. 제안 방법 (Methodology)

저자들은 SpargeAttention2를 제안하며, 이는 세 가지 핵심 구성 요소를 통해 위 문제들을 해결합니다.

2.1. 하이브리드 Top-k + Top-p 마스킹 (Hybrid Top-k+Top-p Masking)

단일 마스킹 규칙의 실패를 방지하기 위해 Top-k 와 Top-p 를 결합한 통합 마스킹 전략을 도입했습니다.

• 원리: 각 어텐션 행 (Row) 에 대해 Top-k(상위 k% 토큰) 와 Top-p(누적 확률 p% 도달 토큰) 를 각각 적용한 후, 두 결과의 **합집합 (Union)**을 최종 마스킹 영역으로 선택합니다.
  • 수식: $M_{ij} = 1 \iff j \in \text{Top-k}(\bar{P}_i, k\%) \cup \text{Top-p}(\bar{P}_i, p\%)$
• 효과: 균일한 분포에서는 Top-p 가 Top-k 의 한계를 보완하고, 편향된 분포에서는 Top-k 가 Top-p 의 한계 (Attention Sink 우세) 를 보완하여 고희소성에서도 안정적인 정보 보존을 가능하게 합니다.

2.2. 속도 기반 증류 손실 (Velocity Distillation Loss)

데이터 분포 불일치로 인한 성능 저하를 해결하기 위해, 데이터에 의존하는 Diffusion Loss 대신 증류 (Distillation) 방식을 적용했습니다.

• 구조: Frozen 된 전체 어텐션 (Full Attention) 모델을 Teacher, 희소 어텐션 모델을 Student로 설정합니다.
• 손실 함수: Student 모델이 Teacher 모델의 **속도 예측 (Velocity Prediction, $u(x_t, c_t, t)$)**을 직접 모방하도록 학습시킵니다.
  • $L_{VD} = \mathbb{E}[\|u_{\text{sparse}} - u_{\text{full}}\|^2]$
• 효과: 파인튜닝 데이터의 품질이나 분포에 구애받지 않고, 원본 모델의 생성 행동을 유지하면서 희소 어텐션 구조에 적응할 수 있게 하여 생성 품질 저하를 방지합니다.

2.3. 효율적인 커널 구현

• FlashAttention 기반의 CUDA 커널을 사용하여 블록 단위 (Block-wise) 희소 어텐션을 구현했습니다.
• 마스킹된 영역의 행렬 곱셈과 Softmax 계산을 효율적으로 건너뛰어 실제 추론 속도 향상을 달성합니다.

3. 주요 기여 (Key Contributions)

1. 심층 분석: 확산 모델의 희소 어텐션에서 Top-k/Top-p 가 실패하는 시점, 학습 가능 방법이 더 높은 희소성을 달성할 수 있는 이유, 그리고 Diffusion Loss 기반 파인튜닝의 한계를 체계적으로 분석했습니다.
2. SpargeAttention2 제안:
   • 고희소성 환경에서 강건한 하이브리드 Top-k+Top-p 마스킹 전략.
   • 생성 품질을 보존하기 위한 속도 기반 증류 (Velocity Distillation) 파인튜닝 목표 함수.
   • 효율적인 GPU 커널 구현.
3. 성능 달성: 비디오 생성 품질을 저하시키지 않으면서 매우 높은 수준의 희소성과 속도 향상을 달성했습니다.

4. 실험 결과 (Results)

Wan2.1 (1.3B 및 14B 모델) 을 기반으로 한 실험에서 SpargeAttention2 는 기존 방법들 (VSA, VMoBA, SLA 등) 을 압도적으로 능가했습니다.

• 희소성 및 속도:
  • 95% 어텐션 희소성 달성.
  • 어텐션 연산 속도 16.2 배 향상.
  • 엔드 - 투 - 엔드 비디오 생성 속도 4.7 배 향상 (14B 모델 기준, 3043 초 → 650 초).
• 생성 품질:
  • 전체 어텐션 (Full Attention) 모델과 비교했을 때, 이미지 품질 (IQ), 일관성 (OC), 미적 품질 (AQ), 비주얼 리워드 (VR), VQA 점수 등 모든 지표에서 유사하거나 더 우수한 성능을 보였습니다.
  • 반면, 기존 희소 어텐션 방법들은 높은 희소성에서 품질이 크게 저하되거나 속도가 느렸습니다.
• Ablation Study:
  • 하이브리드 마스킹 (Top-k+Top-p) 이 단일 마스킹보다 품질이 우수함을 확인.
  • 학습 (Trainable) 이 필수적임을 확인 (학습 없는 경우 품질 급감).
  • 증류 손실 (Distillation Loss) 이 표준 Diffusion Loss 보다 분포 불일치 상황에서 훨씬 효과적임을 확인.

5. 의의 및 결론 (Significance)

SpargeAttention2 는 비디오 확산 모델의 계산 비용 문제를 해결하는 데 있어 중요한 이정표가 됩니다.

1. 실용성: 오픈소스 모델의 원본 학습 데이터가 없는 상황에서도, 증류 기법을 통해 원본 모델의 품질을 유지하면서 고도로 최적화된 희소 어텐션을 적용할 수 있음을 증명했습니다.
2. 효율성과 품질의 트레이드오프 해소: 기존에는 희소성을 높이면 품질이 떨어지는 경향이 있었으나, 본 방법은 95% 의 희소성을 달성하면서도 품질 저하 없이 4 배 이상의 생성 속도 향상을 이루어냈습니다.
3. 미래 지향성: 긴 시퀀스 처리가 필요한 비디오 생성 및 기타 대규모 시퀀스 모델 분야에서 학습 가능한 희소 어텐션의 새로운 표준을 제시합니다.

요약하자면, SpargeAttention2 는 하이브리드 마스킹과 증류 기반 파인튜닝을 결합하여, 기존 방법들의 한계를 극복하고 고희소성 환경에서도 고품질 비디오 생성이 가능한 획기적인 솔루션을 제시한 논문입니다.