기술 요약: Attend Locally, Remember Linearly (ARL2)

1. 문제 정의

자기회귀 (AR) 비디오 확산은 인과적, 청크 단위 공식화 덕분에 스트리밍 및 상호작용형 비디오 생성의 선도적 패러다임으로 부상했습니다. 그러나 기존 AR 비디오 확산 모델은 Diffusion Transformer(DiT) 내에서 소프트맥스 셀프 어텐션에 의존합니다. 이러한 의존성은 근본적인 확장성 병목 현상을 초래합니다:

• 이차적 계산 및 선형 메모리: 표준 어텐션 메커니즘은 시퀀스 길이 $N$(비디오 프레임) 에 대해 $O(N^2)$의 계산 복잡도와 $O(N)$의 메모리 사용량을 초래합니다.
• KV 캐시 증가: 시간적 일관성을 유지하기 위해 이러한 모델은 생성된 블록마다 선형적으로 증가하는 키 - 값 (KV) 캐시를 누적합니다. 5 초 480p 비디오의 경우 KV 캐시는 모델 파라미터 자체를 초과하는 34GB 를 초과할 수 있습니다.
• 기존 해결책의 한계: 희소 어텐션, KV 캐시 양자화, 또는 제거와 같은 현재 솔루션들은 단계별 비용을 줄이거나 메모리를 제한하지만, 선형 메모리 증가를 동시에 제거하고 스트리밍 컨텍스트를 효과적으로 관리하지는 못합니다. 희소 어텐션은 여전히 증가하는 캐시에 의존하며, 제거는 과거 컨텍스트를 되돌릴 수 없게 폐기합니다.

핵심 과제는 사전 훈련된 AR 비디오 확산 모델을 **상수 메모리 ($O(1)$)**와 선형 시간 확장을 달성하면서도 프레임 내 상호작용의 공간적 충실도나 프레임 간 생성의 시간적 일관성을 희생하지 않는 하이브리드 아키텍처로 변환하는 것입니다.

2. 방법론: ARL2

저자들은 비디오 어텐션의 이질적 특성을 해결하기 위해 셀프 어텐션을 두 개의 서로 다른 브랜치로 분해하는 **ARL2(Attend Locally, Remember Linearly)**라는 하이브리드 어텐션 모듈을 제안합니다:

2.1 하이브리드 어텐션 분해

모델은 어텐션을 다음과 같이 분할합니다:

1. 프레임 내 브랜치 (Attend Locally): 각 프레임 내에서 원래의 양방향 소프트맥스 어텐션을 유지합니다. 이 브랜치는 시각적 충실도에 필수적인 세밀한 공간 상관관계와 지역적 의존성을 포착합니다. 프레임 내 모든 토큰이 결합적으로 디코딩되므로 추가 지연을 도입하지 않고 사전 훈련된 프로젝션을 재사용하여 공간적 세부 사항을 유지합니다.
2. 프레임 간 브랜치 (Remember Linearly): 증가하는 KV 캐시를 고정 크기 순환 상태($S$)로 대체합니다. 이 브랜치는 게이트드 델타 네트워크 (GDN) 메커니즘을 사용하여 장거리 시간적 컨텍스트를 유지합니다.
   • 적응: 사전 훈련된 $Q, K, V$ 프로젝션은 소프트맥스에 최적화되어 있으므로, 모델은 선형 순환 도메인에 표현을 적응시키기 위해 헤드별 그룹화된 $1\times1$ 특징 맵 ($\phi_q, \phi_k, \phi_v$) 을 도입합니다.
   • 상태 업데이트: 상태 $S_N$은 게이트드 규칙을 사용하여 업데이트됩니다: $S_N = \alpha_N S_{N-1} + \beta_N K'_N(V'_N - K'_N S_{N-1})$. 여기서 $\alpha$(포기 게이트) 와 $\beta$(학습률) 는 토큰별 학습된 프로젝션입니다.
   • 출력: 프레임 간 출력은 $O_{inter} = Q'_N S_{N-1}$로 계산됩니다.

2.2 메모리 상태 관리

순환 상태가 노이즈가 있는 중간 단계에 의해 손상되거나 프레임 내 비대칭성을 유발하는 것을 방지하기 위해 ARL2 는 두 가지 특정 메커니즘을 사용합니다:

• 프레임 단위 접근: 단일 프레임 내의 모든 토큰은 동일한 사전 업데이트 상태($S_{N-1}$) 를 쿼리합니다. 상태는 전체 프레임이 처리된 후에만 업데이트됩니다. 이는 현재 프레임의 모든 토큰에 대해 일관된 프레임 간 컨텍스트를 보장합니다.
• 클린 - 패스 업데이트: 다단계 디노이징 루프 동안 모델은 상태를 쿼리하지만 업데이트하지는 않습니다. 상태는 프레임의 최종 "클린" 패스가 생성된 후에만 업데이트됩니다. 이는 노이즈가 있는 중간 단계가 장거리 메모리를 오염시키는 것을 방지하고, 상태 쓰기를 디노이징 단계당 한 번에서 프레임당 한 번으로 줄입니다.

2.3 훈련 파이프라인: 2 단계 증류

처음부터 훈련하는 대신, 저자들은 사전 훈련된 Causal Forcing 모델을 변환하기 위한 효율적인 2 단계 증류 파이프라인을 제안합니다:

• 1 단계 (레이어별 증류 및 선택): 30 개의 어텐션 레이어 각각이 독립적으로 증류됩니다. 저자들은 레이어를 "하이브리드 - 복구 가능 (HR)"과 "하이브리드 - 민감 (HS)"으로 분류하기 위해 어텐션 복구율 (ARR) 분석을 도입합니다.
  • HR 차원: 주제 일관성, 배경 일관성, 모션 부드러움, 미적 품질 (복구율 >85%).
  • HS 차원: 이미지 품질 (민감하며, 복구율 약 25% 만 가능).
  • 레이어 선택: 이미지 품질에 중요한 레이어가 소프트맥스 어텐션을 유지하도록, HS 차원에 덜 민감한 레이어를 교체하는 것을 우선시하는 점수 매기기 메커니즘 ($p_\ell$) 이 사용됩니다.
• 2 단계 (공동 증류): 하이브리드 모델 (선택된 레이어가 교체된 상태) 은 고정된 교사 모델에 대해 엔드 - 투 - 엔드로 공동 미세 조정되어 예측된 속도 필드가 모든 레이어에 걸쳐 일관된 작동을 보장하도록 맞춥니다.

3. 주요 기여

1. 최초의 하이브리드 AR 비디오 확산: 이는 사전 훈련된 AR 비디오 확산 모델을 하이브리드 선형 어텐션 아키텍처로 변환하는 첫 번째 작업으로, 처음부터 다시 훈련하지 않고도 상수 메모리를 갖춘 스트리밍 생성을 가능하게 합니다.
2. 역할 인식 분해: 이 설계는 비디오 어텐션의 이질성을 명시적으로 존중하여 공간적 세부 사항을 위해 양방향 소프트맥스를 유지하면서 시간적 메모리를 위해 선형 순환을 사용합니다.
3. 안정적인 상태 관리: 프레임 단위 접근과 클린 - 패스 업데이트의 도입은 반복적 디노이징에 순환 선형 어텐션을 적용할 때 발생하는 특정 과제를 해결하여 컨텍스트 비대칭성과 메모리 손상을 방지합니다.
4. 효율적인 증류 전략: 성능 저하를 최소화하면서 효율성 향상을 극대화하는 민감도 기반 2 단계 증류 파이프라인입니다.

4. 실험 결과

저자들은 13 억 파라미터 모델을 사용하여 다양한 비디오 길이 (81 에서 921 프레임) 에 걸쳐 VBench 벤치마크에서 ARL2 를 평가했습니다.

• 효율성:
  • 메모리: 921 프레임의 경우 (예: 40.5GB 대 87.9GB) 최대 54% 의 메모리 감소를 달성했으며, 비디오 길이에 관계없이 상수 메모리 사용량을 유지했습니다. 반면, 기저 모델들은 1005 프레임에서 OOM(메모리 부족) 으로 인해 실패했습니다.
  • 속도: 921 프레임에서 최대 2.26 배의 월클록 속도 향상을 달성했습니다. KV 캐시 오버헤드가 지배적이 됨에 따라 비디오 길이가 증가함에 따라 속도 향상도 확대됩니다.
• 품질:
  • ARL2 (50% 교체): 전체 품질에서 소프트맥스 교사와 일치하거나 초과했으며, 시간적 깜빡임 (+0.99) 과 모션 부드러움 (+1.35) 에서 구체적인 개선을 보였습니다.
  • ARL2 (75% 교체): 이미지 품질에서 기저 모델 대비 약 1% 이내의 경쟁력 있는 품질을 유지하면서 최대 효율성 향상을 달성했습니다.
  • 비교: ARL2(50%) 는 다른 증류된 모델들을 능가했으며, 3.5 배 적은 파라미터로 훨씬 큰 MAGI-1(45 억) 과 일치했습니다.
• 절대 실험 (Ablations):
  • 하이브리드 브랜치를 제거하거나 토큰 단위 접근을 사용하면 성능이 저하되어 두 브랜치와 프레임 단위 상태 공유의 필요성이 확인되었습니다.
  • 헤드별 게이팅은 성능과 파라미터 효율성 사이의 가장 좋은 균형을 제공했습니다.

5. 중요성 및 주장

이 논문은 ARL2 가 자기회귀 비디오 확산의 확장성 병목 현상을 근본적으로 해결한다고 주장합니다. 공간 및 시간 모델링을 분리함으로써, 이전에는 메모리 제약으로 인해 불가능했던 장편 비디오 합성을 가능하게 합니다. 저자들은 그들의 접근 방식이 상수 메모리를 갖춘 스트리밍 비디오 생성을 허용하여 상호작용 및 실시간 애플리케이션에 적합하다고 강조합니다.

이 작업은 메모리 집약적인 시간 메커니즘을 선형 순환 상태로 대체하면서도 소프트맥스 어텐션의 고품질 공간 모델링을 유지할 수 있음을 보여줍니다. 제안된 2 단계 증류는 이러한 아키텍처 변경이 최소한의 학습 가능 파라미터 (<2%) 와 계산 비용 (약 156 H100 GPU 시간) 으로 사전 훈련된 모델에 효율적으로 전이될 수 있음을 입증하여, 확장 가능한 비디오 생성을 위한 실용적인 경로를 제시합니다.