WF-VAE: Melhorando o VAE de Vídeo por Fluxo de Energia Impulsionado por Wavelet para Modelo de Difusão de Vídeo 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

Resumo

O Autoencoder Variacional de Vídeo (VAE) codifica vídeos em um espaço latente de baixa dimensão, tornando-se um componente-chave na maioria dos Modelos de Difusão de Vídeo Latente (LVDMs) para reduzir os custos de treinamento do modelo. No entanto, à medida que a resolução e a duração dos vídeos gerados aumentam, o custo de codificação dos VAEs de Vídeo se torna um gargalo limitante no treinamento dos LVDMs. Além disso, o método de inferência por blocos adotado pela maioria dos LVDMs pode levar a descontinuidades no espaço latente ao processar vídeos de longa duração. A chave para lidar com o gargalo computacional está na decomposição dos vídeos em componentes distintos e na codificação eficiente das informações críticas. A transformada wavelet pode decompor vídeos em múltiplos componentes de domínio de frequência e melhorar significativamente a eficiência, portanto, propomos o Autoencoder de Fluxo Wavelet (WF-VAE), um autoencoder que aproveita a transformada wavelet em vários níveis para facilitar o fluxo de energia de baixa frequência na representação latente. Além disso, introduzimos um método chamado Causal Cache, que mantém a integridade do espaço latente durante a inferência por blocos. Em comparação com os VAEs de vídeo de última geração, o WF-VAE demonstra desempenho superior em ambas as métricas PSNR e LPIPS, alcançando uma taxa de transferência 2 vezes maior e um consumo de memória 4 vezes menor, mantendo uma qualidade de reconstrução competitiva. Nosso código e modelos estão disponíveis em 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.