WF-VAE : Amélioration de Video VAE par Flux d'Énergie Piloté par Ondelettes pour le Modèle de Diffusion Vidéo Latent

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

November 26, 2024

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

cs.AI

Résumé

Le Variational Autoencoder Vidéo (VAE) encode des vidéos dans un espace latent de faible dimension, devenant un composant clé de la plupart des Modèles de Diffusion Vidéo Latente (LVDM) pour réduire les coûts d'entraînement du modèle. Cependant, à mesure que la résolution et la durée des vidéos générées augmentent, le coût d'encodage des VAE Vidéo devient un goulot d'étranglement limitant dans l'entraînement des LVDM. De plus, la méthode d'inférence par bloc adoptée par la plupart des LVDM peut entraîner des discontinuités de l'espace latent lors du traitement de vidéos de longue durée. La clé pour résoudre le goulot d'étranglement computationnel réside dans la décomposition des vidéos en composants distincts et l'encodage efficace des informations critiques. La transformée en ondelettes peut décomposer les vidéos en plusieurs composants de domaine fréquentiel et améliorer significativement l'efficacité, nous proposons donc le Variational Autoencoder à Flux d'Ondelettes (WF-VAE), un autoencodeur qui exploite la transformée en ondelettes multi-niveaux pour faciliter le flux d'énergie à basse fréquence dans la représentation latente. De plus, nous introduisons une méthode appelée Causal Cache, qui maintient l'intégrité de l'espace latent lors de l'inférence par bloc. Comparé aux VAE Vidéo de pointe, le WF-VAE démontre des performances supérieures à la fois en termes de PSNR et de métriques LPIPS, atteignant un débit 2 fois plus élevé et une consommation de mémoire 4 fois inférieure tout en maintenant une qualité de reconstruction compétitive. Notre code et nos modèles sont disponibles sur 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.