OneVAE: Joint Discrete and Continuous Optimization Helps Discrete Video VAE Train Better
Abstract
Encoding videos into discrete tokens can enable concise andunified multi-modal LLMs. Previous discrete video VAEs suffer from un-stable training, long training time, and degraded reconstruction quality.We revisit the relationship between continuous and discrete VAEs andfind that bridging discrete and continuous representations improves dis-crete token learning. Based on this insight, we propose a unified progres-sive training framework that (i) jointly optimizes continuous and discretereconstructions within a single network, and (ii) progressively derives afamily of VAEs at different compression ratios, leading to faster conver-gence and better final performance. Furthermore, leveraging this unifiedtraining, a single VAE can achieve competitive performance for bothcontinuous and discrete representations. Meanwhile, we propose two ar-chitectural improvements to further boost the performance of discreteVAEs. First, inspired by the use of enlarging continuous latent dimensionto boost reconstruction, we propose multi-token quantization, improvingPSNR by nearly 1 dB at the same token compression ratio. Second, weintroduce first-frame enhancement, which uses a lower-compression firstframe as an anchor in causal VAEs, alleviates the limited-context is-sue of the first frame and significantly improves reconstruction underhigh-compression settings (e.g., 4 × 16 × 16). The code is available athttps://github.com/HVision-NKU/OneVAE.