Improved Video VAE for Latent Video Diffusion Model

Abstract: Variational Autoencoder (VAE) aims to compress pixel data into low-dimensional latent space, playing an important role in OpenAI's Sora and other latent video diffusion generation models. While most of existing video VAEs inflate a pretrained image VAE into the 3D causal structure for temporal-spatial compression, this paper presents two astonishing findings: (1) The initialization from a well-trained image VAE with the same latent dimensions suppresses the improvement of subsequent temporal compression capabilities. (2) The adoption of causal reasoning leads to unequal information interactions and unbalanced performance between frames. To alleviate these problems, we propose a keyframe-based temporal compression (KTC) architecture and a group causal convolution (GCConv) module to further improve video VAE (IV-VAE). Specifically, the KTC architecture divides the latent space into two branches, in which one half completely inherits the compression prior of keyframes from a lower-dimension image VAE while the other half involves temporal compression through 3D group causal convolution, reducing temporal-spatial conflicts and accelerating the convergence speed of video VAE. The GCConv in above 3D half uses standard convolution within each frame group to ensure inter-frame equivalence, and employs causal logical padding between groups to maintain flexibility in processing variable frame video. Extensive experiments on five benchmarks demonstrate the SOTA video reconstruction and generation capabilities of the proposed IV-VAE (https://wpy1999.github.io/IV-VAE/).

• The paper introduces an improved video VAE that integrates keyframe-based temporal compression and Group Causal Convolution to enhance video generation.
• The dual-branch design significantly improves model convergence and reduces spatial-temporal conflicts, achieving superior PSNR, SSIM, LPIPS, and FVD results.
• Extensive experiments on datasets like Kinetics and ActivityNet demonstrate that IV-VAE robustly handles varied resolutions, boosting realism and consistency.

Improved Video VAE for Latent Video Diffusion Model

Introduction

The paper presents innovations in video generation through the introduction of an Improved Variational Autoencoder (VAE) for Latent Video Diffusion Models (LVDMs). Variational Autoencoders, essential for compressing high-dimensional pixel video data into low-dimensional latent spaces, serve as foundational elements for models such as OpenAI's Sora. The traditional approach of increasing temporal compression via inflating 2D image VAEs into 3D structures reveals certain shortcomings. The paper identifies two critical issues: the detrimental effect of initializing with well-trained image VAEs on temporal compression enhancement and the imbalance caused by unidirectional causal interactions in frame processing.

Methodology

The authors introduce an innovative approach using a dual-branch architecture and Group Causal Convolution (GCConv) to address these limitations. This methodology results in an Improved Video VAE (IV-VAE) design:

1. Keyframe-Based Temporal Compression (KTC): This architecture divides the latent space into distinct branches. One branch adopts a lower-dimensional image VAE for keyframe compression, while the other employs 3D GCConv for temporal compression, enhancing convergence speed and reducing spatial-temporal conflicts.
2. Group Causal Convolution (GCConv): GCConv allows for efficient inter-frame information exchange within groups while maintaining causal temporal relationships across groups, thereby ensuring temporal flexibility and reducing frame performance imbalance. The method depends on intra-group standard convolution and inter-group causal padding to maintain causal dependencies.
3. Enhancements for High Resolutions: Introduction of dilated convolutions and additional attention mechanisms expand receptive fields, improving high-resolution video processing capabilities.

Experimental Results

Extensive tests across various benchmarks reveal the superiority of IV-VAE, achieving state-of-the-art (SOTA) results. The experimental outcomes highlight:

• Performance Metrics: IV-VAE demonstrates significant improvements over traditional and contemporary methods across PSNR, SSIM, LPIPS, and FVD metrics on several datasets such as Kinetics-600 and ActivityNet.
• Resolution Handling: IV-VAE exhibits robustness and efficiency over varying resolutions (e.g., 480P to 1080P), outperforming existing methodologies by a substantial margin, especially in higher resolutions due to its enhanced temporal motion perception capabilities.
• Video Generation Capabilities: When applied in a generation context using diffusion models like Latte, IV-VAE offers unprecedented results in terms of realism and consistency, particularly evident on datasets like Kinetics-400 and SkyTimelapse.

Implementation Insights

The proposed system leverages a UNet-based architecture with nuanced modifications such as RMSNorm usage and strategic channel dimension adjustments. Key implementation details include:

• Training regimes with particular attention to KL divergence, Mean Absolute Error (MAE), and perceptual loss (LPIPS).
• Use of a cache mechanism over traditional overlap mechanisms to enhance memory usage and maintain semantic integrity during long video reconstructions.

Conclusion

The authors contribute significant innovations to video encoding via IV-VAE, which combines keyframe-oriented initialization with advanced causal convolution strategies. By meticulously enhancing temporal compression prerequisites and reducing inter-frame inconsistencies, this paper not only provides empirical evidence of improved video VAE performance but also opens pathways for future research in higher-resolution video generation and enhanced model architectures. While the paper effectively addresses current challenges, it also notes the potential for future exploration into alternative architectures like Dit or Mamba to further optimize the video VAE performance landscape.