WF-VAE: 잠재 비디오 확산 모델을 위한 웨이블릿 주도 에너지 흐름에 의한 비디오 VAE 향상

WF-VAE: Enhancing Video VAE by Wavelet-Driven Energy Flow for Latent Video Diffusion Model

November 26, 2024

저자: Zongjian Li, Bin Lin, Yang Ye, Liuhan Chen, Xinhua Cheng, Shenghai Yuan, Li Yuan

cs.AI

초록

비디오 변이형 오토인코더(VAE)는 비디오를 저차원 잠재 공간으로 부호화하여, 대부분의 잠재 비디오 확산 모델(LVDMs)에서 핵심 구성 요소로 작용하여 모델 훈련 비용을 줄이는 역할을 합니다. 그러나 생성된 비디오의 해상도와 지속 시간이 증가함에 따라, 비디오 VAE의 부호화 비용은 LVDMs의 훈련을 제한하는 병목 현상이 됩니다. 또한, 대부분의 LVDMs에서 채택된 블록별 추론 방법은 긴 지속 시간 비디오를 처리할 때 잠재 공간의 불연속성을 유발할 수 있습니다. 계산 병목 현상을 해결하는 핵심은 비디오를 구별 가능한 구성 요소로 분해하고 중요 정보를 효율적으로 부호화하는 데 있습니다. 웨이블릿 변환은 비디오를 여러 주파수 영역 구성 요소로 분해하고 효율성을 크게 향상시킬 수 있습니다. 따라서 저주파 에너지 흐름을 잠재 표현으로 용이하게 하는 다단계 웨이블릿 변환을 활용하는 Wavelet Flow VAE (WF-VAE)를 제안합니다. 더불어, 블록별 추론 중에 잠재 공간의 무결성을 유지하는 방법인 인과 캐시(Causal Cache)를 소개합니다. 최신 비디오 VAE와 비교하여, WF-VAE는 PSNR 및 LPIPS 메트릭에서 우수한 성능을 보여주며, 경쟁력 있는 재구성 품질을 유지하면서 처리량이 2배 높고 메모리 소비량이 4배 낮습니다. 저희의 코드와 모델은 https://github.com/PKU-YuanGroup/WF-VAE에서 확인하실 수 있습니다.

English

Video Variational Autoencoder (VAE) encodes videos into a low-dimensional latent space, becoming a key component of most Latent Video Diffusion Models (LVDMs) to reduce model training costs. However, as the resolution and duration of generated videos increase, the encoding cost of Video VAEs becomes a limiting bottleneck in training LVDMs. Moreover, the block-wise inference method adopted by most LVDMs can lead to discontinuities of latent space when processing long-duration videos. The key to addressing the computational bottleneck lies in decomposing videos into distinct components and efficiently encoding the critical information. Wavelet transform can decompose videos into multiple frequency-domain components and improve the efficiency significantly, we thus propose Wavelet Flow VAE (WF-VAE), an autoencoder that leverages multi-level wavelet transform to facilitate low-frequency energy flow into latent representation. Furthermore, we introduce a method called Causal Cache, which maintains the integrity of latent space during block-wise inference. Compared to state-of-the-art video VAEs, WF-VAE demonstrates superior performance in both PSNR and LPIPS metrics, achieving 2x higher throughput and 4x lower memory consumption while maintaining competitive reconstruction quality. Our code and models are available at https://github.com/PKU-YuanGroup/WF-VAE.