기술 요약: NABLA - 효율적인 비디오 생성을 위한 이웃 적응형 블록 수준 어텐션 (Neighborhood Adaptive Block-Level Attention)

문제 정의

Diffusion Transformer(DiT)를 이용한 비디오 생성은 전체 셀프 어텐션(full self-attention) 메커니즘의 계산 복잡도로 인해 심각한 병목 현상에 직면해 있습니다. 잠재 확산 모델(latent diffusion models)에서 어텐션 메커니즘은 시공간적 토큰 수에 따라 이차적으로($O((T \cdot H \cdot W)^2)$) 증가합니다. 이러한 이차적 스케일링은 고해상도 및 장기 지속 비디오 처리를 매우 느리고 메모리 집약적으로 만듭니다. Sliding Tile Attention(STA)과 같은 정적 희소 어텐션(static sparse attention) 방식은 하드웨어 친화적인 로컬 윈도우를 활용하여 효율성을 높이지만, 경직성이라는 문제를 안고 있습니다. 정적 패턴은 동적이고 콘텐츠 특화적인 의존성을 포착하는 데 실패하여, 고해상도 생성 시 객체 중복이나 불충분한 전역 어텐션 커버리지와 같은 아티팩트를 유발할 수 있습니다. 반대로, 기존의 동적 희소 어텐션 방식은 복잡한 오프라인 프로파일링, 커스텀 CUDA 커널을 요구하거나 상당한 지연 시간 오버헤드를 발생시켜 실용성을 제한합니다.

방법론: NABLA

저자들은 커스텀 커널이나 보조 손실 함수 없이도 동적으로 콘텐츠 인식형 희소 어텐션 마스크를 구축하도록 설계된 메커니즘인 NABLA(Neighborhood-Adaptive Block-Level Attention)를 제안합니다. 이 방법은 PyTorch의 FlexAttention에 원활하게 통합되며 다음의 핵심 구성 요소를 통해 작동합니다:

1. 토큰 재정렬 (Fractal Flattening): 인접한 공간 토큰 간의 의미적 관계를 보존하기 위해 입력 시퀀스를 재정렬합니다. $P \times P$ 크기의 공간 패치 내 토큰들은 연속적인 시퀀스로 그룹화되는 반면, 시간 차원은 순서가 유지됩니다. 이를 통해 전역 비디오 차원과 관계없이 로컬 의미 구조가 유지되도록 합니다.

2. 적응형 마스크 구축 (블록 수준 풀링 및 CDF 임계값 설정):

   • 다운샘플링: 쿼리($Q$)와 키($K$)를 $N \times N$ 블록으로 재구성하고 평균 풀링하여 어텐션 맵 계산의 차원을 축소합니다.
   • 축소된 어텐션: 이러한 다운샘플링된 블록 표현 위에서 전체 어텐션 맵을 계산합니다.
   • CDF 임계값 설정: 축소된 어텐션 맵의 각 행에 대해 누적 분포 함수(CDF)를 계산합니다. 누적 확률이 임계값($1 - thr$)을 초과하는 블록만을 유지함으로써 마스크를 이진화합니다. 이 단계는 특정 입력 컨텍스트에 따라 각 어텐션 헤드에 대해 가장 중요한 블록을 동적으로 선택합니다.
   • 확장: 결과로 나온 이진 마스크를 원래의 토큰 해상도로 다시 확장하여, 어떤 $N \times N$ 블록의 어텐션을 계산할지 정의합니다.

3. STA와의 하이브리드 통합: 순수하게 적응형 희소성만 사용할 때 발생할 수 있는 경계 아티팩트를 완화하기 위해, NABLA는 선택적으로 Sliding Tile Attention(STA)과 결합됩니다. 최종 마스크는 NABLA 유래의 콘텐츠 인식 마스크와 정적 STA 윈도우 마스크의 합집합입니다 ($M = M_{\nabla} \lor M_{STA}$). 이 하이브리드 접근 방식은 정적 윈도우의 강력한 사전 지식(prior)을 활용하면서도 동적 선택의 유연성을 유지합니다.

주요 기여

• 효율적인 콘텐츠 인식 마스킹: NABLA는 블록 단위 풀링과 CDF 임계값 설정을 통해 동적으로 희소 마스크를 구축하며, 장거리 의존성을 보존하고 복잡한 시공간 관계를 처리하는 데 있어 STA와 같은 고정된 희소 패턴보다 뛰어난 성능을 보입니다.
• 하드웨어 불가지론적 구현: 본 방법은 표준 PyTorch 연산(FlexAttention)을 사용하여 구현되었으며, 커스텀 CUDA 커널이나 외부 프로파일링 단계를 필요로 하지 않습니다. 이는 기존의 훈련 및 추론 파이프라인에 원활한 통합을 보장합니다.
• 이단계 가속화: 학습과 추론 모두에 초점을 맞춘 다른 연구들과 달리, NABLA는 DiT 모델의 사전 훈련(pre-training/fine-tuning)과 추론 단계를 모두 가속화합니다.
• 하이브리드 견고성: NABLA와 STA의 결합은 계산 효율성과 고품질 생성 사이의 균형을 효과적으로 맞추며, 순수 정적 또는 순수 동적 방식의 한계를 해결합니다.

실험 결과

저자들은 두 가지 주요 작업에 대해 NABLA를 평가했습니다: 대규모 사전 훈련된 모델의 미세 조정(fine-tuning)과 소규모 모델의 처음부터 시작하는 사전 훈련(pre-training).

1. 미세 조정 (Wan 2.1 14B 720p):

• 속도: NABLA는 전체 어텐션 대비 추론 시간에서 최대 2.7배의 속도 향상을 달계했습니다.
• 품질: 이 방법은 CLIP, VBench, FVD 전반에 걸쳐 베이스라인 품질 지표와 일치하는 성능을 보였습니다. 특히 높은 희소도 수준(예: 90%)에서도 NABLA는 42.08의 CLIP 점수(베이스라인 42.06 대비)와 83.17의 VBench 총점(베이스라인 83.16 대비)을 유지했습니다.
• 인간 평가: 50명의 참가자를 대상으로 한 사이드 바이 사이드 비교에서, 80% 희소도 조건에서 베이스라인과 NABLA 간의 지각적 품질(의미적 정렬, 시각적 품질, 움직임의 자연스러움)에 통계적으로 유의미한 차이가 없음을 확인했습니다.

2. 사전 훈련 (Custom 2B DiT 512²):

• 속도: 사전 훈련 중 NABLA는 반복(iteration) 시간을 10.9초에서 7.5초로 단축하여 1.46배의 가속을 달성했습니다.
• 수렴: NABLA 모델은 전체 어텐션 모델에 비해 더 낮은 검증 손실(validation loss)과 더 빠른 수렴 속도를 보여주었습니다.

3. 정적 방법과의 비교:

• 순수 STA 구성은 특히 다중 객체 및 공간적 관계를 처리할 때 의미적 점수(VBench)의 저하를 보였습니다.
• NABLA와 NABLA+STA 하이브리드는 모든 객관적 지표에서 베이스라인 수준의 성능을 유지하며, 우수한 적응성을 입증했습니다.

의의 및 주장

본 논문은 NABLA가 계산 효율성과 생성 품질 사이의 트레이드오프를 해결함으로써 효율적인 비디오 생성을 위한 새로운 최첨단(SOTA)을 확립한다고 주장합니다. 그 의의는 다음과 같습니다:

• 확장성: 전역적 일관성을 희생하지 않으면서 계산량(FLOPs)을 줄여 고해상도 비디오 합성을 가능하게 합니다.
• 실용성: 특수 하드웨어나 복잡한 엔지니어링(커스텀 커널)에 의존하지 않는 "드롭인(drop-in)" 솔루션을 제공하여, 기존 파이프라인에 즉시 도입할 수 있습니다.
• 적응성: 정적 패턴이 실패하는 비디오 콘텐츠의 동적인 특성을 성공적으로 해결하는 동시에, 다른 동적 방식들이 가진 지연 시간 오버헤드를 피합니다.

저자들은 NABLA를 메모리 중심 프레임워크를 보완하는 기술로 위치시키며, 특히 대규모 비디오 모델을 자원이 제한된 장치에 배포하기 위해 어텐션 메커니즘의 계산 복잡성을 타겟팅하고 있습니다.