VAE 없이 텍스트-이미지 모델 학습, Linum의 Pyramid-JiT가 픽셀 공간에서 새 기록

안녕하세요 HN, 저자입니다!

맥락을 설명하자면, 저희는 생성형 비디오 모델을 훈련하는 2인 연구소입니다. 목표는 제어 가능한 새로운 애니메이션 도구 세트를 만드는 것입니다 (관심 있으시면 여기에서 더 읽어보실 수 있습니다: https://www.linum.ai/about).

저희의 지난 텍스트-투-비디오 모델에서 훈련 및 추론 비용 측면에서 가장 큰 병목은 어텐션입니다. 비디오 모델은 토큰 밀도가 엄청납니다 (예: 몇 초짜리 클립에 110K 토큰). 만약 그 컨텍스트 윈도우를 더 공격적으로 압축할 수 있다면, 훨씬 적은 비용으로 더 큰 모델을 훈련하고, 오늘날 Seedance 같은 대형 모델과 달리 (실행 비용이 어마어마하죠) 프로슈머에게 합리적인 가격으로 제공할 수 있습니다.

전통적으로 이미지 및 비디오 모델은 두 개의 분리된 구성 요소를 가집니다: VAE (Variational Autoencoder)와 Diffusion Transformer (DiT). 이들은 별도로 훈련되며, 경험적으로 VAE는 16x16 토큰 감소를 넘어서기 어려워 보입니다.

여기서는 픽셀 공간으로 전환하고, VAE를 버리고, 32x32 토큰 감소 (컨텍스트 윈도우가 4배 더 작음)를 달성하면서도 훨씬 적은 훈련 샘플로 더 나은 전체 모델을 학습합니다.

핵심 주장은 "단순한 것이 더 낫다"입니다. 압축 문제를 더 강력한 Diffusion Transformer (DiT)에 넣을 수 있다면, 생성에 최적화된 "잠재 공간"을 학습하고 생성 품질을 해치지 않으면서 더 나은 압축을 얻을 수 있을 것입니다.

앞으로 몇 시간 동안 이 글을 수시로 확인할 예정이니, 아래에 질문을 남겨주세요. 그리고 최선을 다해 답변하겠습니다 […]

그 중 얼마나 많은 부분이 최적이 아닌 VAE 사용 때문인지 궁금하네요. 동일한 최적화를 VAE에도 적용할 수 있고, 제 직관으로는 더 나은 접근 방식으로 VAE를 직접 재훈련하면 (특히 이동 및 회전 불변성을 보장하고, 잠재 변수를 재조정/블러링하는 능력 포함) 총 연산 비용이 훨씬 더 최적이 될 것 같습니다.

저해상도 잠재 공간에서 드래프트를 사용하는 것과 동일한 이점을 누릴 수도 있습니다. 학습하기 더 쉽고 교란에 더 강건한 데이터 분포를 가지면서, 512x512 대신 동일한 연산으로 4096x4096 출력을 얻을 수 있습니다 (8배 VAE를 가정할 때).

또한 VAE가 단일 단계이고 훨씬 작을 수 있기 때문에 메인 LDM에 대해 많은 수의 확산/플로우 매칭 단계를 사용할 수 있다는 장점도 있습니다. VAE는 언어나 중요한 장면 이해를 처리할 필요 없이 지각적 압축만 하면 되니까요. 이는 프레임 간 압축에 의존하지 않고 생성 모델을 훈련하는 것을 매우 주저하게 되는 비디오 모델에 특히 중요해 보입니다.