Prism: 2K 공동 영상·오디오 생성 학습을 위한 동적 희소 어텐션

AI 요약

Prism은 고해상도 영상·오디오 생성 모델을 학습할 때 전체 어텐션이 만드는 계산량과 학습 신호 분산 문제를 줄이려는 프레임워크입니다. 전체 어텐션은 토큰 수가 늘수록 비용이 제곱에 비례해 커집니다. 해상도가 높아지면 중복된 토큰에도 어텐션이 퍼져 유용한 정보에 집중하기 어렵고, 사전학습 모델의 특성도 흐트러질 수 있다고 저자들은 설명합니다. 기존 희소 어텐션 기법은 학습 없이 추론을 빠르게 하는 데 초점을 두거나, 영상과 오디오의 상호작용 구조를 반영하지 못한다고 지적합니다. 영상·오디오 상호작용은 소리가 발생하는 영역 주변에 집중되는 경향이 있다는 점이 Prism의 설계 출발점입니다.

지역별 정보에 맞춘 블록 구성

Prism은 토큰을 시공간 매크로 존(macro-zone)으로 나눕니다. 각 존에서 영상 특징의 채널 방향 분산과 오디오-영상 교차 어텐션(Audio-to-Video Cross-Attention)의 특징 노름을 계산합니다. 앞의 값은 영상 내용이 어느 방향으로 얼마나 변하는지 나타내고, 뒤의 값은 오디오가 각 영상 영역에 얼마나 강하게 영향을 주는지 나타냅니다. 두 신호를 함께 사용해 존마다 블록 모양을 조정합니다. 영상 변화가 빠르거나 오디오와 영상의 결합이 강한 축은 더 세밀하게 나눕니다. 그 결과 블록 안의 토큰이 의미적으로 더 가까워지고, 블록 단위 특징이 영상 정보와 공동 영상·오디오 상호작용을 함께 담도록 설계했습니다.

블록 선택에는 하이브리드 전략을 적용해 질의별 희소성도 동적으로 정합니다. 학습 설정은 Prism Block Sparse Attention(BSA)과 동적 블록 모양 할당을 켜는 방식입니다. 설정 예시에는 BSA 희소성 값 0.75, 0.80, 0.85, 0.93과 하이브리드 블록 선택을 위한 CDF 임계값 0.20이 제시되어 있습니다. 실험 결과로는 전체 어텐션 대비 학습 속도 2.5배 향상과 더 높은 생성 품질을 보고합니다.

데이터 준비와 학습 절차

학습 데이터는 클립 하나당 한 행을 담은 CSV로 지정합니다. 행에는 WAN VAE 잠재 표현 경로, 원본 영상과 오디오 경로, 캡션, 프레임 수, 영상 크기, 프레임률을 넣습니다. 잠재 표현은 영상마다 한 번 오프라인으로 추출합니다. 이어 데이터셋 도구로 720p, 1080p, 2K 해상도 버킷을 만들고, 잠재 표현 파일을 검사해 프레임·높이·너비 정보를 채울 수 있습니다. 버킷별 샘플 수 히스토그램도 출력합니다.

학습기는 이 CSV를 읽어 각 샘플의 공간·시간 버킷을 정합니다. 공간 버킷은 패치 수를 기준으로 가능한 높이·너비 조합을 만들고, 샘플의 종횡비와 가장 가까운 조합을 선택합니다. 시간 버킷은 최소 49프레임부터 시작해 설정된 간격으로 늘리며, 각 샘플에 들어맞는 가장 긴 버킷을 배정합니다. 시간 길이는 4k+1 형태를 따릅니다. 분산 버킷 배치 샘플러는 같은 버킷에 속한 샘플끼리 배치를 구성해 데이터 병렬 랭크 사이의 텐서 모양 불일치를 피합니다.

학습 설정은 최적화·병렬화·체크포인트·희소 어텐션을 담은 상위 YAML과 사전학습 모델, 손실 가중치, 노이즈 일정, 데이터 로더 설정을 담은 해상도별 YAML로 나뉩니다. 기본 학습률은 1e-5이며 워밍업은 200스텝입니다. 옵티마이저는 AdamW, 가중치 감쇠는 0.01, 그래디언트 클리핑 한도는 1.0으로 설정합니다. FSDP와 시퀀스 병렬화를 사용하며, 전체 모델 학습이 기본값입니다.

해상도와 하드웨어 요구 사항

Prism은 720p, 1080p, 2K에서 학습과 추론을 지원합니다. 720p 학습에는 NVIDIA 80GB GPU 최소 32개가 필요하고, 1080p와 2K 학습에는 각각 최소 64개가 필요합니다. 추론은 720p에서 80GB GPU 한 개로 실행할 수 있으며, 1080p와 2K는 80GB GPU 4개 이상을 권장합니다. 720p 추론에는 DeepSpeed 실행 스크립트를, 1080p와 2K에는 모델 파라미터를 GPU 여러 대에 나누는 FSDP 실행 스크립트를 안내합니다. 저장소에는 학습·추론 코드와 미리보기 체크포인트가 포함되며, 설치에는 Python 3.10 이상과 CUDA 12.4 이상이 필요합니다.

원문: Hugging Face / 번역·요약: Trawling