1. 문제 정의 (Problem)

확산 모델 (Diffusion Models), 특히 비디오 생성 모델의 계산 효율성을 높이기 위해 **희소 - 선형 어텐션 (Sparse-Linear Attention, SLA)**이 제안되었습니다. 기존 SLA 는 어텐션 맵을 희소 (Sparse) 부분과 저랭크 (Low-rank) 부분으로 분해하여 희소 어텐션과 선형 어텐션을 병용하는 방식입니다. 그러나 SLA 에는 두 가지 주요 한계가 존재했습니다.

1. 휴리스틱한 라우팅 (Heuristic Routing): SLA 는 어텐션 가중치의 크기 (magnitude) 에 따라 계산 작업을 희소 브랜치와 선형 브랜치로 나누는 휴리스틱 방식을 사용합니다. 이는 최적의 분할이 아닐 수 있으며, 일부 가중치를 강제로 이동시켜도 저랭크 부분의 랭크를 줄이지 못하면서 희소성만 개선되는 비효율이 발생할 수 있습니다.
2. 분해 오차 불일치 (Mismatch in Decomposition): SLA 의 수식적 구현은 원래 의도했던 '희소 + 선형' 분해 ($P = P_1 + P_2$) 와 정확히 일치하지 않습니다. SLA 는 희소 어텐션 출력에 스케일링 오차가 발생하며, 이를 보정하기 위해 추가적인 선형 프로젝션 (Projection) 을 학습시켜야 하는데, 이는 학습을 어렵게 만들고 오차를 완전히 해결하지 못합니다.

2. 방법론 (Methodology: SLA2)

저자들은 위 한계를 해결하기 위해 SLA2를 제안했습니다. SLA2 는 학습 가능한 라우팅, 더 정확한 분해 공식, 그리고 양자화 인식 학습 (QAT) 을 결합합니다.

가. 학습 가능한 라우터 (Learnable Router)

• 동적 선택: 휴리스틱한 가중치 크기 기반 분할 대신, 어텐션 계산이 희소 브랜치인지 선형 브랜치인지 동적으로 결정하는 **학습 가능한 라우터 (R)**를 도입했습니다.
• 입력 및 구조: 라우터는 쿼리 (Q) 와 키 (K) 를 입력받아 마스크 (M) 를 생성합니다. $N \times N$ 크기의 전체 어텐션 행렬을 계산하는 비용 ($O(N^2)$) 을 줄이기 위해 인접 토큰에 평균 풀링 (Mean Pooling) 을 적용하고, 학습 가능한 선형 프로젝션 ($proj_q, proj_k$) 을 통해 최적의 분할을 학습합니다.
• SoftTop-k: 학습 중에는 미분 불가능한 Hard Top-k 대신 SoftTop-k를 사용하여 그래디언트 역전파를 가능하게 합니다. 추론 시에는 Hard Top-k 를 사용합니다.

나. 분해 일관성 있는 혼합 공식 (Faithful Formulation)

• 기존 SLA 의 스케일링 오차를 해결하기 위해, 희소 어텐션 출력 ($O_s$) 과 선형 어텐션 출력 ($O_l$) 을 **학습 가능한 비율 벡터 ($\alpha$)**로 직접 가중치하여 결합하는 방식을 채택했습니다.
• 수식: $O = \alpha \odot O_s + (1 - \alpha) \odot O_l$
• 이 방식은 원래 의도했던 어텐션 분해 ($P \approx \alpha \odot P_s + (1-\alpha) \odot P_l$) 와 정확히 일치하도록 하여, 추가적인 보정 프로젝션 없이도 정확한 출력을 보장합니다.

다. 양자화 인식 학습 (Quantization-Aware Training, QAT)

• 추가적인 속도 향상을 위해 희소 어텐션 브랜치에 저비트 (Low-bit) 양자화를 적용했습니다.
• Forward Pass: 저비트 (예: INT8, FP8) 정밀도로 어텐션을 계산하여 속도를 높입니다.
• Backward Pass: 오차 역전파는 원래 FP16 정밀도로 수행하여 양자화로 인한 정확도 저하를 최소화합니다. 이를 통해 모델이 양자화 오차에 적응하도록 훈련합니다.

라. 2 단계 학습 전략

1. Stage 1: 라우터 (R) 와 비율 벡터 ($\alpha$) 를 초기화하기 위해 풀 어텐션 (Full Attention) 과의 MSE 손실을 최소화하며 학습합니다.
2. Stage 2: 전체 확산 모델을 엔드 - 투 - 엔드 (End-to-End) 로 미세 조정 (Fine-tuning) 합니다. 이 단계에서는 라우팅 로직을 추론과 일치시키며 (Hard Top-k 사용), 모델이 고희소성 환경에서도 성능을 유지하도록 적응시킵니다.

3. 주요 기여 (Key Contributions)

1. SLA2 아키텍처 제안: 기존 SLA 의 휴리스틱 라우팅과 분해 불일치 문제를 해결한, 학습 가능한 라우터와 분해 일관성 있는 혼합 공식을 갖춘 새로운 희소 - 선형 어텐션 방법론을 제안했습니다.
2. QAT 통합: 저비트 양자화를 양자화 인식 학습 (QAT) 방식으로 통합하여, 비디오 생성 품질을 저하시키지 않으면서 어텐션 연산 속도를 추가로 가속화했습니다.
3. 성능 입증: 비디오 확산 모델 (Wan2.1) 에서 기존 방법론들을 능가하는 효율성과 품질을 입증했습니다.

4. 실험 결과 (Results)

실험은 Wan2.1-1.3B와 Wan2.1-14B 비디오 생성 모델에서 수행되었습니다.

• 효율성 (Efficiency):
  • 97% 어텐션 희소성을 달성했습니다. (선형 브랜치를 고려한 실제 계산 절감률은 약 96.7% 입니다.)
  • 18.6 배의 어텐션 실행 시간 속도 향상 (Speedup) 을 기록했습니다.
  • 97% 희소성에서도 90% 희소성을 가진 기존 베이스라인 (VMoBA, VSA, SLA) 보다 더 빠른 속도를 보였습니다.
• 품질 (Quality):
  • 97% 희소성에서도 **전체 어텐션 (Full Attention, 0% 희소성)**과 비교했을 때 비디오 생성 품질이 저하되지 않았으며, 오히려 일부 지표 (Imaging Quality, Overall Consistency 등) 에서 베이스라인보다 우수한 성능을 보였습니다.
  • 특히 14B 모델에서도 97% 희소성에서 4.35 배의 엔드 - 투 - 엔드 지연 시간 단축을 달성했습니다.
• Ablation Study:
  • QAT 를 사용하지 않을 경우 생성 품질이 저하됨을 확인하여 QAT 의 효과성을 입증했습니다.
  • 학습 가능한 라우터가 휴리스틱 Top-k 라우터보다 성능이 뛰어났습니다.

5. 의의 (Significance)

SLA2 는 확산 모델, 특히 고해상도 비디오 생성 모델의 계산 비용을 획기적으로 줄이면서도 생성 품질을 유지하는 실용적이고 효율적인 솔루션을 제공합니다.

• 이론적 개선: 어텐션 분해의 수학적 불일치를 해결하고, 휴리스틱이 아닌 학습 기반의 최적 라우팅을 도입하여 모델의 표현력을 높였습니다.
• 실용적 가치: 18 배 이상의 속도 향상과 97% 의 희소성은 고비용인 비디오 생성 모델을 실시간 또는 저비용 환경에서 배포할 수 있는 가능성을 열었습니다.
• 확장성: QAT 와 결합된 설계는 하드웨어 제약을 받는 환경에서도 고품질 생성을 가능하게 하는 새로운 표준을 제시합니다.

결론적으로, SLA2 는 희소 어텐션과 선형 어텐션의 장점을 최적화하고 양자화 기술을 효과적으로 접목하여, 차세대 효율적인 생성형 AI 모델의 핵심 기술로 자리 잡을 것으로 기대됩니다.