기술 요약: Kaleido

문제 정의

비디오 디퓨전 트랜스포머(vDiT)는 고충실도 비디오 생성을 위한 최첨단 아키텍처로 부상했으나, 심각한 계산 병목 현상에 직면해 있습니다. 추론 과정은 많은 타임스텝에 걸친 반복적인 디노이징을 포함하며, 이 과정에서 셀프 어텐션(self-attention) 메커니즘이 전체 실행 시간의 약 68%를 차지하며 지배적인 비중을 가집니다. 비디오 해상도와 지속 시간이 증가함에 따라 토큰 시퀀스 길이가 길어지며, 이로 인해 셀프 어텐션의 복잡도는 제곱($O(N^2 \times d)$)으로 증가합니다.

기존의 가속 전략들은 주로 대규모 언어 모델(LLM)의 희소 어텐션(sparse attention) 기술을 응용하는데, 이는 의미적 희소성(semantic sparsity)을 기반으로 중요도가 낮은 어텐션 점수를 제거하는 방식에 의존합니다. 그러나 이러한 접근 방식은 비디오 데이터와 근본적으로 맞지 않습니다. 텍스트와 달리, 비디오 토큰은 모든 토큰이 공간적 일관성과 시간적 일관성에 기여하는 조밀하고 연속적인 시각적 의미를 담고 있습니다. 어텐션 상관관계를 무분별하게 제거하면 움직임의 불연속성이나 깜빡임(flickering) 현상이 발생합니다. 또한, 디스틸레이션(distillation)을 통해 디퓨전 타임스텝이 줄어들더라도 셀프 어텐션은 여전히 유일하고 지배적인 병목 구간으로 남기 때문에, 타임스텝 감소 기술만으로는 충분하지 않습니다.

방법론

Kaleido는 토큰 수준의 희소성에 의존하는 대신, vDiT의 잠재 공간(latent space) 내에 존재하는 **채널별 시공간 상관관계(channel-wise spatio-temporal correlations)**를 활용하는 알고리즘-하드웨어 공동 설계(algorithm-hardware co-design)를 통해 이러한 과제를 해결합니다.

1. 알고리즘: 채널별 재사용 (Channel-Wise Reuse)

핵심 통찰은 vDiT 모델이 로터리 위치 임베딩(RoPE)을 사용하여 토큰 채널을 시간($t$), 수평($x$), 수직($y$)의 세 가지 별도 그룹으로 분할한다는 점입니다. 이 그룹들은 서로 다른 주파수에서 위치 정보를 인코딩하며, 이는 특정 패턴의 어텐션 점수를 유도합니다.

• 상관관계 발견: 저자들은 특정 채널 그룹 내의 인접한 토큰들이 해당 인코딩 방향에 대해 직교(orthogonal)하는 차원을 따라 높은 유사성을 보인다는 것을 입증했습니다. 예를 들어, $t$-채널은 공간 차원($x$ 및 $y$)을 가로질러 높은 상관관계를 가지며, $x$-채널은 시간 차원을 가로질러 상관관계를 가집니다.
• 재사용 메커니즘: Kaleido는 연산을 건너뛰는 대신(sparse attention 방식처럼), 부분적인 결과값을 재사용합니다.
  • 유사도 체크: 연산 전, 시스템은 정의된 축을 따라 인접한 토큰 사이의 유사도($\Delta$)를 계산합니다.
  • 임계값 설정: 두 개의 임계값 $\theta_{th1}$과 $\theta_{th2}$가 재사용 전략을 결정합니다:
    • $\Delta < \theta_{th1}$인 경우: 전체 재사용(Full Reuse). 해당 토큰은 이전 토큰의 부분 어텐션 점수를 완전히 재사용합니다.
    • $\theta_{th1} \le \Delta < \theta_{th2}$인 경우: 부분 재사용(Partial Reuse). 해당 토큰은 이전 결과의 가장 중요한 비트(예: 상위 3비트)를 재사용하고 나머지 하위 비트만 계산합니다.
    • $\Delta \ge \theta_{th2}$인 경우: 전체 연산(Full Computation).
  • 적용 범위: 이 전략은 셀프 어텐션, 크로스 어텐션 및 MLP 레이어에 적용됩니다. MLP의 경우, 재사용은 가중치가 아닌 활성화 함수(activations)에 적용됩니다.

2. 하드웨어: 재구성 가능한 시스톨릭 가속기 (Reconfigurable Systolic Accelerator)

재사용 알고리즘으로 인해 발생하는 불규칙한 데이터 액세스 패턴은 표준 시스톨릭 어레이와 호환되지 않습니다. Kaleido는 이를 지원하기 위해 맞춤형 하드웨어 아키텍처를 도입합니다.

• 재구성 가능한 프로세싱 엘리먼트 (PEs): PE 어레이는 네 가지 실행 모드를 지원하도록 추가 레지스터와 멀티플렉서가 보강되었습니다. 이 모드들을 통해 PE는 단일 타일 내에서 전체 연산, 전체 재사용, 그리고 부분 재사용(하프 비트 연산)을 동적으로 처리할 수 있습니다. 또한 PE는 온칩 메모리에서 다시 로드하지 않고도 이전 프레임으로부터 부분 결과를 누적할 수 있습니다.
• 데이터 디스패처 (Data Dispatcher): 낮은 PE 활용도를 유발하는 불규칙한 희소성을 완화하기 위해 전용 데이터 디스패처가 도입되었습니다. 이는 다음으로 구성됩니다:
  • 스케줄러 (Scheduler): 유사한 재사용 패턴(이진 문자열로 표현됨)을 가진 토큰들을 클러스터로 그룹화합니다.
  • 매처 (Matcher): 클러스터 내의 호환 가능한 입력 채널들을 쌍으로 묶어 데이터 스트림을 병합함으로써, PE 어레이에 공급되는 인접 행들이 동시에 처리될 수 있도록 하여 활용도를 극대화합니다.

주요 기여

1. vDiT의 체계적 특성 규명: 본 논문은 vDiT 잠재 공간의 시공간 상관관계를 체계적으로 설명하고, RoPE 기반의 채널 파티셔닝이 독특한 어텐션 패턴의 근원임을 밝힌 최초의 연구입니다.
2. 채널별 재사용 알고리즘: 채널 수준에서 부분 결과를 재사용함으로써 중복 연산을 건너뛰는 경량 알고리즘입니다. 토큰 단위의 병합(merging)과 달리, 이는 픽셀 수준의 정밀도를 보존하여 이전의 공동 설계 모델들보다 훨씬 높은 생성 품질(PSer > 17 dB 개선)을 달면합니다.
3. 공동 설계된 가속기 아키텍처: 불규칙한 희소성을 처리할 수 있는 맞춤형 PE와 데이터 디스패처를 갖춘 재구성 가능한 시스톨릭 어레이를 통해, PE 활용도를 약 69%에서 85%로 향상시켰습니다.
4. 포괄적 평가: 네 가지 주요 vDiT 모델(HunyuanVideo, Wan, CogVideoX, TurboDiffusion)을 대상으로 최첨단 GPU 및 가속기와 비교하여 광범위한 평가를 수행했습니다.

결과

16nm 공정 기술 노드와 128x128 PE 어레이 구성을 기준으로 평가한 결과는 다음과 같습니다:

• 성능: Kaleido는 최첨단 가속기(예: AdapTiV, Exion) 대비 최대 5.9배, Nvidia A100 GPU 대비 6.6배 빠른 속도를 달합니다.
• 에너지 효율: A100 베이스라인 대비 최대 16.0배의 에너지 절감을 실현합니다.
• 생성 품질: Kaleido는 HunyuanVideo에서 29.9 dB, Wan에서 24.5 dB의 PSNR 점수를 기록하며 우수한 비디오 품질을 유지하며, 이는 Cambricon-D와 같은 근사 기반 베이스라인보다 17 dB 이상 높은 수치입니다. 또한 이미지 품질 지표(PSNR, SSIM, LPIPS)에서 소프트웨어 최적화된 희소 어텐션 방식(예: MInference, SVG)보다 뛰어난 성능을 보였습니다.
• 확장성: 비디오 해상도가 480p에서 720p로 높아질수록 성능 격차가 더 커지는데, 이는 고해상도일수록 알고리즘이 활용하는 시공간 상관관계가 더욱 강력해지기 때문입니다.

의의

본 논문은 Kaleido를 비디오 생성 가속화의 필수적인 진화 단계로 자리매김합니다. 단순히 LLM 기반의 희소 어텐션 기술을 비디오에 전이하는 것은 시각 데이터의 연속성 때문에 효과적이지 않다고 주장합니다. "제거(pruning)"에서 vDiT의 고유한 채널별 상관관계에 기반한 "재사용(reusing)"으로 초점을 전환함으로써, Kaleido는 출력 품질을 저하시지 않으면서도 계산 병목 현상을 해결할 수 있는 원칙적인 접근 방식을 제시합니다. 이 연구는 차세대 비디오 디퓨전 모델의 효율성을 확보하기 위해 알고리즘-하드웨어 공동 설계가 필수적임을 입증합니다.