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) кодирует видео в пространство низкой размерности, становясь ключевым компонентом большинства моделей латентной диффузии видео (LVDM), направленных на снижение затрат на обучение модели. Однако с увеличением разрешения и продолжительности создаваемых видео стоимость кодирования видео VAE становится ограничивающим фактором при обучении LVDM. Более того, метод блочного вывода, принятый в большинстве LVDM, может привести к разрывам в латентном пространстве при обработке видео большой продолжительности. Ключ к преодолению вычислительного узкого места заключается в разложении видео на отдельные компоненты и эффективном кодировании важной информации. Вейвлет-преобразование может разложить видео на несколько компонентов в частотной области и значительно повысить эффективность. Мы предлагаем Вейвлет-потоковый вариационный автоэнкодер (WF-VAE), автоэнкодер, использующий многоуровневое вейвлет-преобразование для облегчения потока энергии низкой частоты в латентное представление. Кроме того, мы представляем метод, названный "Причинный кэш", который поддерживает целостность латентного пространства во время блочного вывода. По сравнению с передовыми видео VAE, WF-VAE демонстрирует превосходную производительность как по метрикам PSNR, так и LPIPS, достигая удвоенного уровня пропускной способности и вчетверо меньшего потребления памяти при сохранении конкурентоспособного качества восстановления. Наш код и модели доступны по адресу 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.