기술 요약: 효율적인 비디오 생성을 위한 이질적 단계 할당 (HSA)

1. 문제 제기

확산 트랜스포머 (DiT) 는 비디오 생성 분야에서 최첨단 성능을 입증했으나, 막대한 계산 비용이라는 단점을 안고 있습니다. 표준 추론 프로토콜은 시공간 토큰의 모든 토큰에 동일한 수의 탈노이즈 단계 ($T$) 를 적용합니다. 이러한 단일화 접근법은 시각 데이터의 고유한 비대칭성을 간과합니다. 배경 영역, 정적 패치, 거친 구조는 지각적으로 단순하여 세부 사항이 풍부한 전경 객체나 고동역학 영역보다 적은 탈노이즈 단계를 필요로 합니다. 또한 인간의 시선은 정적 구성 요소에 비해 매우 동적인 구성 요소의 품질 저하에 덜 민감합니다. 따라서 모든 토큰에 전체 탈노이즈 예산을 균일하게 적용하는 것은 계산상 낭비일 뿐만 아니라 지각적 우선순위와도 부합하지 않습니다.

단계별 특징 캐싱 (예: TeaCache) 과 같은 기존 가속화 방법들은 일반적으로 전역적 결정을 내립니다. 즉, 주어진 반복 단계에서 모든 토큰이 계산되거나 모든 토큰이 캐시된 출력을 재사용합니다. 이러한 방법들은 시공간적 또는 토큰 수준의 이질성을 활용하지 못하며, 다양한 동적 요구 사항을 가진 토큰의 집합이 아닌 단일 블록으로 시퀀스를 처리합니다.

2. 방법론: 이질적 단계 할당 (HSA)

저자들은 사전 훈련된 DiT 내에서 토큰별 탈노이즈 궤적을 할당하는 훈련 불필요 (training-free), 모델 무관 (model-agnostic) 추론 알고리즘인 **이질적 단계 할당 (HSA)**을 제안합니다. 핵심 아이디어는 토큰 시퀀스를 속도 동역학에 기반하여 그룹으로 분할하고, 더 느리게 진화하는 토큰에는 적은 단계를 할당하는 반면 기준 그룹에는 전체 일정을 유지하는 것입니다.

2.1 토큰 분할 및 스케줄링

$N$개의 토큰 집합은 $K$개의 서로소 그룹 ($G_1, \dots, G_K$) 으로 분할됩니다. 각 그룹 $G_k$는 총 단계 $T$의 약수여야 하는 단계 예산 $s_k$를 할당받습니다.

• 기준 그룹 ($G_1$): 전체 예산 $s_1 = T$를 할당받습니다. 이 토큰들은 모든 반복 단계에서 활성화됩니다.
• 감소된 그룹: 더 작은 예산 (예: $T/2, T/4$) 을 할당받습니다. 이 토큰들은 전역 반복 단계 $i$가 $i \pmod{T/s_k} = 0$을 만족할 때만 활성화됩니다.
• 동적 선택: 논문은 토큰을 그룹에 할당하는 네 가지 전략을 평가합니다. 토큰별 L1 상대 속도 예측 변화에 기반하여 토큰을 할당하는 동적 토큰 선택 전략 (배경과 같은 느리게 진화하는 영역을 식별) 이 일관되게 최상의 성능을 보입니다.

2.2 KV 캐시 동기화

이질적 스케줄링의 주요 과제는 서로 다른 단계 수를 가진 토큰이 주어진 벽시계 (wall-clock) 반복 단계에서 서로 다른 노이즈 수준 ($\sigma$) 에 존재한다는 점입니다. 자기 어텐션은 모든 토큰이 서로에게 어텐션해야 하므로 동기화 불일치가 발생합니다. HSA 는 KV 캐시 동기화 메커니즘을 통해 이를 해결합니다:

1. 활성 토큰: 활성 토큰의 부분집합 ($A_i$) 만 QKV 프로젝션을 통한 순전파 (forward pass) 를 수행합니다.
2. 캐시 업데이트: 활성 토큰에 대한 새로운 키 ($K$) 및 값 ($V$) 프로젝션이 레이어별 캐시에 기록되어 오래된 항목을 덮어씁니다.
3. 전체 컨텍스트 어텐션: 활성 토큰은 전체 캐시된 $K, V$ 행렬 (활성 토큰의 신선한 항목과 건너뜀 토큰의 오래되었지만 유효한 항목을 모두 포함) 에 대해 쿼리를 계산합니다.
4. 복잡도 감소: 이는 자기 어텐션 복잡도를 $O(N^2)$에서 $O(|A_i| \cdot N)$으로 줄입니다. 여기서 $|A_i|$는 $i$단계에서의 활성 토큰 수입니다. 건너뜀 토큰은 비활성 반복 단계 동안 쿼리 계산, 크로스 어텐션 또는 피드포워드 연산을 수행하지 않습니다.

2.3 캐시된 오일러 업데이트

추가적인 모델 평가 없이 건너뜀 토큰의 잠재 상태를 진행시키기 위해 HSA 는 캐시된 오일러 업데이트를 사용합니다.

• 각 토큰은 가장 최근 활성 단계에서 예측된 속도 ($\hat{v}_n$) 를 저장합니다.
• 모든 전역 반복 단계에서 모든 토큰 (활성 및 건너뜀) 은 동일한 증분 오일러 단계 $(\sigma_{i+1} - \sigma_i)$만큼 진행됩니다.
• 활성 토큰은 freshly 계산된 속도를 사용하고, 건너뜀 토큰은 캐시된 속도를 재사용합니다.
• 이 업데이트는 분기 (branching) 를 피하고 GPU 효율성을 유지하면서 모든 $N$개 토큰에 대한 단일 텐서 연산으로 공식화됩니다.

2.4 위상 인식 캐싱 윈도우

초기 및 후기 탈노이즈 단계가 각각 전역 구조와 미세한 세부 사항에 중요하다는 점을 인식하여, HSA 는 캐싱을 궤적의 중간 부분으로 제한합니다. 첫 번째와 마지막 $m$단계는 모든 토큰에 대해 전체 캐시 없는 순전파를 실행하여, 오래된 근사치로 인해 구조적 무결성과 세부 사항 정교화가 훼손되지 않도록 보장합니다.

3. 주요 기여

1. HSA 알고리즘: 속도 동역학에 기반하여 시공간 토큰 그룹에 이질적 탈노이즈 궤적을 할당하는 완전한 훈련 불필요, 플러그앤플레이 추론 패러다임.
2. KV 캐시 동기화: 활성 토큰에 대해 전체 $N$개 토큰 어텐션 컨텍스트를 유지하면서 비활성 토큰을 완전히 우회하여, 반복당 어텐션 비용을 $O(|A_i| \cdot N)$으로 줄이는 메커니즘.
3. 캐시된 오일러 업데이트: 캐시된 속도를 사용하여 건너뜀 토큰을 진행시키는 간단한 텐서 연산으로, 추가 모델 평가 없이 생성 품질을 유지합니다.
4. 종합 평가: 다양한 추론 예산에 걸쳐 견고성을 입증하는 네 가지 토큰 그룹화 전략 및 네 가지 예산 사전 설정에 대한 분석.

4. 실험 결과

저자들은 Wan-2.1-1.3B 및 Wan-2.2-A14B(텍스트-비디오 및 이미지-비디오)와 LTX-2(오디오-비디오) 에서 HSA 를 평가했습니다.

• 품질-런타임 트레이드오프: HSA 는 균일한 흐름 매칭 (Flow Matching, FM) 과 TeaCache 와 같은 최첨단 캐싱 방법에 비해 우월한 파레토 프론티어를 달성합니다.
• 공격적 가속: 이점은 공격적인 가속 영역 (50% 및 25% 런타임) 에서 가장 두드러집니다. 25% 런타임에서 HSA 는 VBench 벤치마크에서 기저 모델들을 크게 능가합니다 (총 점수: FM 대비 79.87% 대 75.68%, TeaCache 대비 78.33%).
• 차원 견고성: 기저 방법들은 엄격한 예산 하에서 "치명적인 차원 붕괴"(운동 부드러움 또는 시간적 일관성과 같은 특정 VBench 차원의 심각한 저하) 를 겪는 반면, HSA 는 모든 16 개 평가 차원에서 전체 예산 기준을 견고하게 추적합니다.
• 이미지-비디오 (I2V): HSA 는 I2V 작업에서 이미지 조건부 정렬 및 시각적 충실도를 성공적으로 유지하는 반면, 감소된 단계 ($T=10$) 에서의 균일한 FM 기저 모델들은 심각한 구조적 붕괴를 보입니다.
• 오프라인 프로파일링 불필요: 모델 출력 변동을 추정하기 위해 값비싼 오프라인 프로파일링이 필요한 TeaCache 와 달리, HSA 는 오프라인 튜닝이나 파인튜닝 없이 작동합니다.

5. 중요성 및 주장

이 논문은 HSA 가 시각 데이터의 지각적 및 동적 복잡성과 계산 자원을 정렬함으로써 DiT 추론의 근본적인 비효율성을 해결한다고 주장합니다. 토큰을 전역적으로가 아닌 개별적으로 처리함으로써 HSA 는 공격적인 가속과 일반적으로 연관되는 품질 저하 없이 상당한 속도 향상을 달성합니다.

저자들은 HSA 가 다양한 토큰 그룹화 전략으로 구체화될 수 있는 유연한 프레임워크임을 강조합니다. 현재 동적 선택이 가장 강력한 결과를 산출하지만, 이 프레임워크는 더 정교한 그룹화 전략을 통해 추가적인 개선을 예상합니다. 이 연구는 시공간 데이터의 고유한 이질성을 활용하여 고품질 비디오 생성을 민주화하는 새로운 클래스의 추론 알고리즘에 대한 토대를 마련합니다.