WF-VAE: Potenziare il Video VAE tramite il Flusso di Energia Driven da Wavelet per il Modello di Diffusione Video Latente

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

November 26, 2024

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

cs.AI

Abstract

Il Video Variational Autoencoder (VAE) codifica i video in uno spazio latente a bassa dimensione, diventando un componente chiave della maggior parte dei Modelli di Diffusione Video Latenti (LVDM) per ridurre i costi di addestramento del modello. Tuttavia, all'aumentare della risoluzione e della durata dei video generati, il costo di codifica dei Video VAE diventa un collo di bottiglia limitante nell'addestramento dei LVDM. Inoltre, il metodo di inferenza a blocchi adottato dalla maggior parte dei LVDM può portare a discontinuità dello spazio latente durante il processamento di video di lunga durata. La chiave per affrontare il collo di bottiglia computazionale risiede nella decomposizione dei video in componenti distinte e nella codifica efficiente delle informazioni critiche. La trasformata wavelet può decomporre i video in multiple componenti nel dominio delle frequenze e migliorare significativamente l'efficienza, pertanto proponiamo il Wavelet Flow VAE (WF-VAE), un autoencoder che sfrutta la trasformata wavelet a più livelli per facilitare il flusso di energia a bassa frequenza nella rappresentazione latente. Inoltre, introduciamo un metodo chiamato Causal Cache, che mantiene l'integrità dello spazio latente durante l'inferenza a blocchi. Rispetto ai VAE video all'avanguardia, il WF-VAE dimostra prestazioni superiori sia in termini di PSNR che di metriche LPIPS, raggiungendo un throughput 2 volte superiore e un consumo di memoria 4 volte inferiore pur mantenendo una qualità di ricostruzione competitiva. Il nostro codice e i modelli sono disponibili su 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.