기술 요약: 효율적인 비디오 생성을 위한 토큰 반경 어텐션 (Token Radius Attention)

1. 문제 정의

비디오 확산 트랜스포머(VDiTs)는 복잡한 공간적 및 시간적 의존성을 포착하기 위해 조밀한 3D 셀프 어텐션(dense 3D self-attention)에 의존하며 고충실도 비디오 생성의 주류 아키텍처로 자리 잡았습니다. 그러나 조밀한 어텐션은 비디오 토큰 수 $N$에 대해 이차적인 계산 비용($O(N^2)$)을 발생시키며, 이는 스케일링 및 배포를 위한 심각한 추론 병목 현상을 초래합니다.

이러한 비용을 완화하기 위한 기존 방식들은 주로 헤드 수준(head-level) 및 블록 수준(block-level) 방법으로 분류되는 희소 어텐션(sparse attention) 메커니즘을 채택합니다.

• 헤드 수준 방법은 개별 어텐션 헤드에 구조화된 공간적 또는 시간적 패턴을 할당합니다.
• 블록 수준 방법은 중요한 쿼리-키(query–key) 블록 쌍을 선택합니다.

핵심 한계점: 두 패러다임 모두 일반적으로 모든 쿼리 토큰에 대해 단일한 계산 예산(예: 고정된 유지 밀도 또는 반경)을 공유한다는 가정을 따릅니다. 이러한 "공유 할당(shared-allocation)" 가정은 셀프 어텐션의 본질적인 특성, 즉 각 쿼리에 대해 독립적으로 수행되는 Softmax 정규화와 충돌합니다. 결과적으로, 집중된 어텐션 분포(낮은 엔트로피)를 가진 쿼리는 너무 많은 키를 참조하도록 강제될 경우 계산을 낭비하게 되며, 분산된 분포(높은 엔트로피)를 가진 쿼리는 공유 예산이 너무 작을 경우 중요한 상호작용을 놓칠 수 있습니다. 이러한 불일치는 생성 품질을 저하시키거나 최적의 효율성을 달一个하지 못하게 만듭니다.

2. 방법론: 토큰 반경 어텐션 (Token Radius Attention, TRA)

저자들은 명시적인 쿼리별 키 랭킹(per-query key ranking) 없이도 토큰별 구조적 희소성을 구현하는 훈련 불필요(training-free) 프레임워크인 **토کن 반경 어텐션(TRA)**을 제안합니다. TRA는 VDiT 어텐션 패턴 분석에서 도출된 두 가지 핵심적인 경험적 통찰력을 기반으로 작동합니다.

통찰 I: 토큰별 어텐션 희소성

저자들은 동일한 레이어와 헤드 내에서도 쿼리 간에 "유지 밀도"(목표 어텐션 질량을 보존하기 위해 필요한 상위 순위 키의 비율)가 수 자릿수 차이로 변한다는 것을 관찰했습니다. 결정적으로, 이 유지 밀도와 어텐션 엔트로피 사이에는 강력한 **로그-선형 관계(log-linear relationship)**가 있음을 발견했습니다.

• 메커니즘: 높은 엔트로피는 더 큰 토큰 예산을 필요로 하는 분산된 어텐션 분포를 나타내고, 낮은 엔트로피는 더 적은 토큰을 필요로 하는 집중된 분포를 나타냅니다.
• 구현: TRA는 분석적 근사치($\hat{B} \approx \tau \exp(H_i) / N$)를 사용하여 쿼리 엔트로피를 토큰별 예산으로 직접 매핑함으로써, 비용이 많이 드는 쿼리별 정렬(sorting)이나 랭킹 과정을 제거합니다.

통찰 II: 토큰 반경 어텐션 패턴

저자들은 쿼리에 대한 지배적인 상호작용이 공간 영역에서 **쿼리 중심의 원형 이웃(query-centered circular neighborhoods)**을 형성한다는 점을 관찰했습니다. 정규화된 어텐션 확률은 2D 공간 거리와 함께 대략 지수적으로 감소합니다.

• 메커니즘: 임의의 상위-$k$ 키를 선택하는 대신, TRA는 스칼라 토큰 예산을 공간적 반경으로 변환합니다.
• 시간적 감쇠(Temporal Decay): 비디오 시퀀스를 처리하기 위해, 공간 반경은 시간적 거리 감쇠 규칙($\phi(\delta) = \exp(-\gamma \delta)$)에 의해 조정되어 규칙적인 시공간적 지원 구조를 생성합니다.

TRA 파이프라인

1. 엔트로피-예산 매핑 (Entropy-to-Budget): 초기 밀집 "웜업(warm-up)" 단계 동안, TRA는 각 쿼리에 대한 어텐션 엔트로피를 계산합니다. 이 엔트로피는 특정 토큰 예산으로 매핑됩니다.
2. 예산-반경 변환 (Budget-to-Radius): 예산은 쿼리별 기본 반경으로 변환됩니다. 이 반경은 시간적 거리에 따라 각 프레임별로 조정되어 2D 디스크 마스크를 형성합니다.
3. 효율적 실행:
   • 타일 우선 레이아웃 (Tile-Major Layout): 비디오 토큰을 타일 우선 순서로 재배열하여, 반경 내에 있는 공간적 인접 토큰들이 시퀀스상에서 연속되도록 보장합니다.
   • 융합 CUDA 커널 (Fused CUDA Kernels): 커스텀 커널을 통해 거리 계산, 반경 비교, 블록 프루닝(pruning), 토큰 투표를 융합하여, 불규칙한 토큰 수준의 반경 마스크를 FlashInfer와 호환 가능한 규칙적인 블록 희소 마스크로 변환합니다.
   • 단계 간 재사용 (Cross-Step Reuse): 엔트로피는 웜업 단계에서만 계산되며, 노이즈 제거 단계 전반에 걸쳐 관찰된 엔트로피 패턴의 안정성을 활용하여 이후의 희소 단계에서 재사용됩니다.

3. 주요 기여

1. 토큰별 희소성의 발견: 본 논문은 유지 밀도가 쿼리에 따라 극적으로 변하지만 엔트로피와 예측 가능한 로그-선형 관계를 따른다는 점과, 지배적인 상호작용이 쿼리 중심의 반경 패턴을 따른다는 점을 밝혀냈습니다.
2. 토큰 반경 어텐션 (TRA): 명시적인 키 랭킹 없이도 엔트로피-예산-반경 파이프라인을 통해 쿼리별 계산 요구 사항을 규칙적인 시공간적 지원 구조로 변환하는 훈련 불필요 프레임워크를 제안합니다.
3. 시스템 공동 설계 (System Co-Design): 오버헤드를 최소러 기존 블록 희소 백엔드에서 효율적으로 실행할 수 있도록 반경 마스크 커널과 융합 엔트로피 커널을 공동 설계했습니다.
4. 포괄적 검증: TRA는 Wan2.1, Wan2.2, HunyuanVideo(1.3B에서 14B 파라미터 범위)의 7가지 구성에 대해 텍스트-비디오(T2V) 및 이미지-비디오(I2V) 작업 모두에서 검증되었습니다.

4. 실험 결과

TRA는 VBench 벤치마크 제품군에서 밀집 어텐션 및 세 가지 훈련 불필요 희소 베이스라인(SVG, SVG2, Radial)과 비교 평가되었습니다.

• 효율성: TRA는 어텐션 상호작용의 **9%~19%**만을 유지합니다. 테스트된 모든 구성에서 밀집 모델 대비 1.56배에서 2.05배의 속도 향상을 달성했습니다. 예를 들어, HunyuanVideo-13B의 경우 T2V에서 2.05배, I2V에서 2.02배의 속도 향상을 보였습니다.
• 품질: TRA는 경쟁력 있는 생성 품질을 유지하며, 종종 희소 방법들 중 가장 우수한 VBench 종합 점수를 기록했습니다. Wan2.1-14B T2V의 경우, 밀집 어텐션 점수에 거의 근접하면서도 1.75배의 속도 향상을 제공했습니다.
• 충실도: 정성적 결과에 따르면, TRA는 다른 희소 방법들에 비해 피사체의 정체성, 구조적 무결성 및 시간적 일관성을 더 잘 보존하며, 깜빡임(flickering)이나 왜곡과 같은 아티팩트가 적게 나타납니다.
• 절제 연구 (Ablation Studies): 엔트로피 기반 예산 매핑이나 반경 마스킹(1D 거리로 대체)을 제거하면 성능이 크게 저하되어, 토큰별 적응성과 구조적 공간 지원의 필요성을 확인시켜 줍니다.

5. 의의 및 주장

본 논문은 TRA가 현재의 희소 어텐션 패러다임의 근본적인 비효란성, 즉 공유 할당 예산과 토큰별 어텐션 요구 사항 사이의 불일치를 해결한다고 주장합니다. 엔트로피와 어텐션 희소성 사이의 내재적 관계를 활용함으로써, TRA는 추가 학습이나 복잡한 온라인 랭킹 메커니즘 없이도 유리한 품질-효율성 트레이드오프를 달성합니다.

저자들은 TRA를 캐싱이나 양자화와 같은 다른 가속화 방법과 결측적인 보완 기술로 위치시키며, 실행되는 어텐션 레이어 내의 이차적인 쿼리-키 상호작용 비용을 줄이는 데 집중합니다. 이 연구는 토큰 수준의 적응성이 구조화된 반경 기반 마스킹을 통해 효율적으로 실현될 수 있음을 보여주며, 고충실도 비디오 생성 모델을 배포하기 위한 확장 가능한 경로를 제시합니다.