OneVAE: Joint Discrete and Continuous Optimization Helps Discrete Video VAE Train Better

TL;DR

OneVAE introduces a joint discrete-continuous video VAE framework that leverages continuous priors and structural improvements to enhance training stability, speed, and reconstruction quality, enabling unified multi-modal video representations.

Contribution

The paper proposes a novel joint optimization scheme for discrete and continuous VAEs, improving training efficiency and reconstruction quality in video encoding.

Findings

FSQ preserves pre-trained continuous VAE priors effectively.

Multi-token quantization improves PSNR by nearly 1 dB.

Joint discrete-continuous optimization achieves competitive performance.

Abstract

Encoding videos into discrete tokens could align with text tokens to facilitate concise and unified multi-modal LLMs, yet introducing significant spatiotemporal compression compared to continuous video representation. Previous discrete video VAEs experienced unstable training, long training time, and degraded reconstruction quality. Given the easier training and superior performance of continuous VAEs, an intuitive idea is to enhance discrete video VAEs by leveraging continuous VAEs. After rethinking the intrinsic link between discrete and continuous representations, we found that FSQ could effectively preserve pre-trained continuous VAE priors compared to other quantization methods. By leveraging continuous VAE priors, it converges several times faster than training from scratch and achieves superior performance at convergence. Meanwhile, two structural improvements are proposed. First, inspired by how continuous VAEs enhance reconstruction via enlarged latent dimensions, we introduce a multi-token quantization mechanism, which achieves nearly a 1 dB improvement in PSNR without compromising the token compression ratio. Second, to tackle reconstruction challenges in high-compression video VAEs, we strengthen first-frame reconstruction, enabling the causal VAE to leverage this information in subsequent frames and markedly improving the performance of 4 x 16 x 16 discrete VAEs. Furthermore, we propose a joint discrete-continuous optimization scheme that unifies the two paradigms and, for the first time, achieves competitive performance on both continuous and discrete representations within a single network. We name our method OneVAE to reflect this connection.