WF-VAE: Mejorando VAE de Video mediante Flujo de Energía Impulsado por Wavelets para el Modelo de Difusión de Video Latente

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

November 26, 2024

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

cs.AI

Resumen

El Autoencoder Variacional de Video (VAE) codifica videos en un espacio latente de baja dimensión, convirtiéndose en un componente clave de la mayoría de los Modelos de Difusión de Video Latente (LVDM) para reducir los costos de entrenamiento del modelo. Sin embargo, a medida que aumenta la resolución y la duración de los videos generados, el costo de codificación de los VAE de Video se convierte en un cuello de botella limitante en el entrenamiento de los LVDM. Además, el método de inferencia por bloques adoptado por la mayoría de los LVDM puede llevar a discontinuidades en el espacio latente al procesar videos de larga duración. La clave para abordar el cuello de botella computacional radica en descomponer los videos en componentes distintos y codificar eficientemente la información crítica. La transformada wavelet puede descomponer videos en múltiples componentes en el dominio de frecuencia y mejorar significativamente la eficiencia, por lo tanto, proponemos el Autoencoder de Flujo Wavelet VAE (WF-VAE), un autoencoder que aprovecha la transformada wavelet de múltiples niveles para facilitar el flujo de energía de baja frecuencia hacia la representación latente. Además, presentamos un método llamado Causal Cache, que mantiene la integridad del espacio latente durante la inferencia por bloques. En comparación con los VAEs de video de última generación, WF-VAE demuestra un rendimiento superior en métricas tanto de PSNR como de LPIPS, logrando una mayor capacidad de procesamiento y un menor consumo de memoria, manteniendo al mismo tiempo una calidad de reconstrucción competitiva. Nuestro código y modelos están disponibles en 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.