기술 요약: Veda - 증류된 희소 어텐션을 통한 확장 가능한 비디오 확산

1. 문제 제기

고해상도, 장시간 비디오 생성을 위한 확산 트랜스포머 (DiT) 의 확장성은 현재 시공간 토큰 시퀀스 길이에 대한 자기 어텐션 메커니즘의 이차적 계산 및 메모리 복잡도 ($O(N^2)$) 에 의해 병목 현상을 겪고 있습니다. 희소 어텐션이 이를 완화할 이론적 경로를 제공하지만, 기존 방법들은 중요한 트레이드오프에 직면해 있습니다:

• 정적 패턴: 사전 정의된 희소 패턴 (예: 슬라이딩 윈도우) 을 사용하는 방법은 DiT 가 학습한 동적이고 헤드별인 어텐션 구조에 적응하지 못합니다.
• 동적 선택: 동적 마스크를 학습하는 방법 (예: VSA, VMOBA) 은 종종 부정확한 추정을 겪습니다. 이들은 확산 목적을 통한 암시적 감독과 평균 풀링과 같은 거친 통계에 의존하여, 완전 어텐션의 실제 기하학적 구조와 불일치를 초래합니다.
• 핵심 문제: 실증적 증거에 따르면 높은 희소성 수준 (≥90%) 에서 기존 접근 방식들은 생성 품질을 저하시켜 공간 왜곡, "물결무늬" 패턴, 그리고 시간적 깜빡임과 같은 구조적 아티팩트를 생성합니다. 본 논문은 이러한 열화가 희소성 자체에 기인한 것이 아니라, 희소 마스크가 완전 어텐션의 타일별 기하학에 정렬되지 못했기 때문이라고 주장합니다.

2. 방법론: Veda 프레임워크

Veda 는 극단적인 희소성 하에서도 완전 어텐션의 타일별 구조와 순위를 보존하도록 설계된 증류된 희소 어텐션 프레임워크입니다. 이는 타일 선택을 확산 목적의 암시적 부산물이 아닌 명시적 재구성 문제로 취급합니다.

2.1. 증류된 타일 점수화

희소성을 형성하기 위해 확산 손실에 의존하는 대신, Veda 는 완전 어텐션 백본으로부터 타일 수준의 어텐션 점수를 명시적으로 학습하는 경량 추정기를 사용합니다.

• 목표 구성: 참조 "오라클" 마스크는 쿼리 - 키 타일 영역에 대한 완전 어텐션 행렬에 최대 풀링 (max-pooling) 을 적용하여 구성됩니다. 평균 풀링 대신 최대 풀링을 선택한 이유는 배경 잡음에 의해 희석되기 쉬운 중요한 고주파 신호 피크를 보존하기 위함입니다.
• 통계 인식 추정기: 이러한 점수를 재구성하기 위해 Veda 는 각 타일에 대해 평균, 최대, 최소 통계를 연결하는 TripPool 설명자를 사용합니다. 이는 단순한 평균 풀링을 넘어 타일 표현을 풍부하게 합니다.
• 헤드별 투영: 추정기는 이질적인 의존성 패턴을 포착하기 위해 각 어텐션 헤드에 대해 별도의 MLP 투영을 사용합니다.
• 최적화: 모델은 예측된 타일 점수 분포를 완전 어텐션 참조와 정렬시키기 위해 행 단위 KL-발산 증류 목적 함수($L_{distill}$) 를 사용하여 훈련됩니다. 중요한 점은 추정기에 입력되는 백본 특징에 stop-gradient 연산이 적용된다는 것입니다. 이는 마스크 학습을 특징 학습과 분리하여 사전 훈련된 생성 매니폴드의 교란을 방지합니다.

2.2. 헤드 인식 타일링

어텐션 헤드가 시공간 의존성에서 상당한 이질성을 보인다는 점 (일부는 국소적 공간 상호작용을, 다른 일부는 장거리 시간적 의존성을 포착함) 을 인식하여, Veda 는 균일한 타일링 전략을 포기합니다.

• 구성 탐색: 각 레이어와 헤드에 대해 Veda 는 하드웨어 타일 크기를 인수분해하는 최적의 타일링 구성 $\pi_{l,h} = (p_t, p_h, p_w)$을 탐색합니다.
• 오프라인 선택: 최적의 구성은 보정 세트에서 완전 어텐션 출력과 희소 어텐션 출력 간의 근사 오차를 최소화함으로써 오프라인에서 선택되어, 타일링 전략이 각 헤드의 특정 구조적 요구 사항과 일치하도록 보장합니다.

2.3. 하드웨어 효율적 구현

이론적 FLOP 감소를 실제 월클럭 속도 향상으로 전환하기 위해, Veda 는 ThunderKittens DSL 을 사용하여 사용자 정의 타일 스킵 희소 어텐션 커널을 구현합니다.

• 비동기 실행: 커널은 NVIDIA Hopper 의 텐서 메모리 액세스 (TMA) 와 워프 특화를 활용합니다. 생산자 - 소비자 패러다임을 사용하여 데이터 이동을 계산과 분리합니다. 여기서 생산자 워프는 비연속적인 키/값 타일을 공유 메모리로 가져오고, 소비자 워프는 텐서 코어 연산을 실행합니다.
• 효율성: 이 설계는 밀집 행렬 수학 뒤에 메모리 지연을 숨겨 FlashAttention-3 의 메모리 FLOP 활용도 (MFU) 의 약 80% 를 달성합니다.
• 훈련 효율성: 특수 TileLang 커널은 두 번의 패스를 통해 정답 히트맵을 생성하여, 성능을 저해하지 않으면서 훈련 오버헤드를 줄이기 위해 희소 감독 (쿼리 타일의 무작위 부분집합) 을 가능하게 합니다.

3. 주요 기여

1. 실증적 통찰: 본 논문은 희소 비디오 확산에서의 생성 품질이 희소성 비율 자체가 아니라 완전 어텐션의 타일별 기하학에 대한 희소 마스크의 정렬에 의해 결정됨을 입증합니다.
2. 증류된 희소 어텐션: Veda 는 통계 인식 추정기와 헤드별 투영을 사용하여 추정 오차를 최소화하도록 타일 선택을 명시적 재구성 문제로 공식화하는 프레임워크를 도입합니다.
3. 헤드 인식 타일링: Video DiT 의 의존성 패턴의 이질성을 해결하기 위해 서로 다른 어텐션 헤드에 서로 다른 시공간 타일 인수분해를 할당하는 새로운 전략입니다.
4. 하드웨어 최적화 커널: 시퀀스 길이에 따라 거의 선형으로 확장되고 높은 하드웨어 활용도를 달성하는 사용자 정의 타일 스킵 커널로, 알고리즘적 희소성을 실질적인 지연 시간 향상으로 전환합니다.

4. 실험 결과

실험은 최대 720P 해상도와 최대 241 프레임의 시퀀스 길이에 걸쳐 Waver-T2V-12B 및 Wan2.1-T2V-14B를 포함한 대규모 비디오 확산 모델에서 수행되었습니다.

• 속도 향상: Waver-T2V-12B(720P, 241 프레임) 에서 Veda 는 5.1 배의 엔드투엔드 속도 향상(샘플 시간을 19.4 분에서 3.8 분으로 단축) 과 10.5 배의 자기 어텐션 속도 향상을 달성합니다. 어텐션 오버헤드는 92% 에서 50% 로 감소합니다.
• 확장성: 속도 향상은 시퀀스 길이가 증가함에 따라 커집니다. 245K 토큰에서 Veda 는 완전 어텐션의 이차적 성장 (1576.5ms) 에 비해 309.1ms 로 거의 선형적으로 확장되어 5.1 배의 속도 향상을 제공합니다.
• 품질 보존:
  • 인간 평가: Waver-bench 1.0 에서 90% 희소성으로 Veda 는 완전 어텐션과 지각적 동등성을 달성합니다. 95% 희소성에서도 기존 최첨단 VSA 방법 (VSA 가 더 낮은 희소성으로 작동할 때조차) 을 크게 능가합니다.
  • 아티팩트 감소: 높은 희소성에서 물결무늬 패턴과 시간적 깜빡임을 겪는 베이스라인과 달리, Veda 는 높은 시각적 충실도와 시간적 일관성을 유지합니다.
  • 정량적 지표: VBench 평가에 따르면 Veda 는 95% 희소성에서도 완전 어텐션 베이스라인과 비교 가능한 주제 일관성과 운동 매끄러움을 유지합니다.

5. 중요성 및 주장

본 논문은 Veda 가 구조적 무결성을 훼손하지 않고 고해상도 비디오 생성의 확장성 병목 현상을 성공적으로 해결했다고 주장합니다. 희소 어텐션을 휴리스틱 검색이나 암시적 학습 작업이 아닌 증류 문제로 재정의함으로써, Veda 는 심각한 품질 저하 없이 이전에 달성 불가능했던 공격적인 희소성 (최대 95%) 을 가능하게 합니다.

저자들은 이득이 단순히 이론적이지 않음을 강조합니다. 사용자 정의 커널은 희소성이 실제 월클럭 가속으로 전환되도록 보장합니다. 이 프레임워크는 시공간 해상도에 따라 유리하게 확장되는 모듈식 하드웨어 인식 솔루션으로 제시되어, 현재 하드웨어 제약 조건 하에서도 고충실도 장편 비디오 생성을 실현 가능하게 만듭니다. 이 연구는 향후 개선 사항으로 커널 융합 강화, 타임스텝 전반의 적응형 희소성 전략, 그리고 타일 점수의 시간적 캐싱에 초점을 맞출 수 있음을 시사합니다.