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

Conference: ECCV2026
Paper: ECCV Paper
Code: https://github.com/HVision-NKU/OneVAE
Area: Video Generation
Keywords: video VAE, finite scalar quantization, joint continuous-discrete optimization, multi-token quantization, first-frame enhancement

TL;DR

OneVAE alternates continuous and discrete reconstruction within a shared video VAE, combining progressive initialization and architectural improvements to raise PSNR from 27.22 to 28.15 in the unified-training ablation while supporting both latent representations at the same compression configuration.

Background & Motivation

Video generation first requires compressing high-dimensional pixels into a tractable representation, and tokenizer errors limit the reconstruction fidelity available to the subsequent generator. Diffusion models commonly use continuous latents, whereas purely discrete next-token prediction requires indices from a finite vocabulary. Videos contain motion as well as spatial detail, so temporal compression further amplifies information loss from quantization. VQ relies on nearest-neighbor codebook matching and can suffer codebook collapse; Cosmos uses finite scalar quantization (FSQ) to mitigate this problem, yet discrete video VAEs can still train slowly and reconstruct poorly.

The authors ask whether the latent space already learned by a continuous model can help optimize a discrete model, rather than simply enlarging its vocabulary. When only the intermediate quantizer is trained inside a pretrained continuous VAE, Table 1 reports PSNR of 26.46 for FSQ and 22.82 for VQ. This local experiment suggests that FSQ adapts more readily to the existing continuous representation under this fixed-backbone setting; it does not establish universal superiority over VQ. The resulting training idea is to retain a continuous reconstruction path so that the shared encoder and decoder need not rely entirely on approximate gradients through discrete quantization.

Stable optimization does not remove the capacity-allocation problems of high compression: a single token per spatial location limits expressiveness, and a causal model cannot use future frames to reconstruct its first frame. The paper therefore combines an optimization recipe with two architectural improvements rather than attributing every gain to a new quantization operator. Core Idea: preserve a learnable latent space through continuous reconstruction, then improve discrete video tokenization through progressive training, multi-token representations, and first-frame capacity allocation.

Method

Overall Architecture

The input is a video clip, which a shared 3D encoder compresses spatially and temporally into a channel vector at each latent location. The continuous path sends continuous representations to the decoder; the discrete path applies FSQ before using the same decoder to reconstruct the video. Training randomly selects one path per step rather than computing both reconstruction losses at every step. Model development begins with a low-compression continuous VAE and derives higher-compression and discrete configurations; first-frame enhancement and multi-token quantization can additionally address causal context and discrete capacity. Solid edges below denote reconstruction data flow, while dashed edges denote training initialization or supervision; the continuous path bypasses quantization rather than concatenating continuous and discrete outputs before decoding.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    Train["Unified Progressive Training"] -.-> Encoder["Shared 3D Encoder"]
    Video["Input Video"] --> Anchor["First-Frame Enhancement"]
    Anchor --> Encoder
    Encoder -->|Discrete Path| Quant["Multi-Token Quantization"]
    Encoder -->|Continuous Path| Decoder["Shared 3D Decoder"]
    Quant --> Decoder
    Train -.-> Decoder
    Decoder --> Output["Reconstructed Video"]
    Output -.->|Training Only: Reconstruction Supervision| Train

The diagram shows the reconstruction path with architectural enhancements; basic OneVAE does not require both enhancements to be enabled together. First-frame enhancement affects capacity allocation during encoding and decoding; its placement at the input indicates a change in compression allocation, not the generation of an extra reference image. At inference time, the application selects continuous or discrete mode without randomly sampling a training path. Both modes can share a model at the same compression ratio, but different ratios are progressively derived model configurations, not automatic switching to arbitrary compression rates.

Key Designs

1. Unified Progressive Training: use continuous reconstruction to support discrete optimization

FSQ first maps continuous features into a bounded range, rounds each dimension, and can then convert the resulting integer tuple into a discrete index. Unlike VQ, it does not search for the nearest codebook entry for the entire vector, making it suitable for insertion into an existing continuous latent space. A compact restatement of the mechanism described in Section 3.1 is:

\[ q = \operatorname{Round}(\operatorname{Bound}(z)). \]

Here \(z\) denotes pre-quantization features and \(q\) denotes the per-dimension integer result, not an already flattened vocabulary index. The cached equations have extraction damage, so unverified bounds, sampling details, and piecewise conditions are not reconstructed; this expression only captures the operations explicitly described in the prose. The frozen-continuous-VAE experiment in Table 1, which trains only the quantizer, provides empirical motivation for this connection rather than a rigorous proof of latent-space similarity.

At each training step, \(R_{dis}\) controls the proportion of discrete-path updates, with the remaining updates using continuous reconstruction. Shared parameters receive two learning signals: discrete steps teach reconstruction after quantization, while continuous steps preserve the reconstruction connection between pre-quantization features and pixels. The authors suggest that the continuous path reduces optimization difficulties from relying solely on approximate discrete gradients; Section 4.2 explicitly presents the more accurate gradient-flow explanation as a hypothesis. This is neither distillation between two independently trained VAEs nor averaging two reconstructed outputs at deployment.

Progressive training first learns an \(8\times8\times4\) continuous VAE and then extends it to higher-compression configurations such as \(16\times16\times4\). At a fixed compression configuration, inserting FSQ and continuing joint training derives the discrete version without relearning all video reconstruction capabilities from scratch. These ratios follow the spatial, spatial, temporal order in Table 3; the prose also uses the temporal, spatial, spatial notation \(4\times16\times16\) for the same type of downsampling configuration. Unified training combines continuous initialization and dual-path optimization, so Table 2 cannot assign separate gains to the two mechanisms. For a primarily discrete reconstruction objective, the authors use a high \(R_{dis}\), but its exact value is not given in the readable main text. For a unified model intended to support both continuous and discrete representations, they explicitly set \(R_{dis}=0.5\).

2. First-Frame Enhancement: provide a more reliable starting point for causal decoding

A causal video VAE cannot use future frames to reconstruct the current frame, leaving the first frame particularly short of temporal context. Under high compression, losing text, boundaries, or textures in that frame also removes a reliable starting point for later reconstruction. By comparing causal and non-causal models, the authors associate poor early-frame reconstruction with this contextual asymmetry. First-frame enhancement therefore reduces compression only for the first frame, allocating more tokens per spatial unit so that subsequent causal reconstruction can use clearer structure and texture.

This does not repair the first frame with future video information or synthesize it using a separate teacher model. It changes capacity allocation across the sequence while retaining causal access to information. Section 3.3 claims an unchanged overall video compression ratio and token budget, but the readable main text does not fully specify the compensating allocation or exact token counts. The location and purpose of the enhancement are therefore clear, but a verified budget-preserving implementation cannot be reconstructed from these claims alone. Figure 3 illustrates improvements to both the first and later frames, while Table 3 reports overall reconstruction for the high-compression OneVAE-FE configuration.

3. Multi-Token Quantization: increase per-location capacity through compositional encoding

Conventional discrete encoding compresses the entire channel vector at a latent location into one token, allowing only one discrete selection regardless of local complexity. OneVAE divides that vector into multiple subvectors, using two equal parts in the experiments, and quantizes them into two tokens. The two parts use a shared codebook design to jointly describe the original location rather than requiring a larger independent codebook for the complete vector. Intuitively, two discrete choices can express more combinations, but they also produce more tokens per latent location.

Consequently, an equal token compression ratio must account for latent-grid size as well as tokens per location. Table 3 uses \(8\times8\times8\) for OneVAE-MT and \(8\times8\times4\) for ordinary OneVAE. Stronger temporal downsampling combined with two tokens per location explains the authors' capacity comparison under their stated equal token compression ratio. For finite causal clips, exact sequence length also depends on first-frame and boundary handling; this note does not infer an exact total token count from nominal ratios. This design differs from first-frame enhancement: one changes the discrete representation at each latent location, while the other reallocates capacity across temporal positions.

A Worked Example

Consider a training clip of \(17\times256\times256\): a continuous VAE first learns to compress and reconstruct it. After deriving the discrete model, the same clip passes through the shared encoder; a step assigned to the discrete path applies FSQ before decoding. With multi-token quantization, each latent location is split into two channel groups that produce separate tokens, not two separate output videos. On another step assigned to the continuous path, encoded features bypass discrete quantization and train reconstruction through the shared decoder. If first-frame enhancement is enabled, the first frame receives a less compressed representation, but later reconstruction still cannot access future frames. After training, reconstruction evaluation uses \(33\times512\times512\) inputs; for AR generation, a generator instead predicts discrete tokens that the decoder converts into video. This example connects the input specifications and mechanisms stated in the paper; it is not an additional experiment and assumes no unreported latent-grid boundary rules.

Loss & Training

Training uses WebVid-10M and AdamW with \(\beta_1=0.5\) and \(\beta_2=0.9\). FP16 reduces memory overhead, and the exponential moving average of model parameters uses a decay rate of 0.999. The reconstruction objective includes pixel and perceptual losses, with LPIPS weighted by 0.1. Training initially uses 17-frame clips at \(256\times256\) per frame, followed by decoder fine-tuning on 33-frame videos after convergence. Temporal PatchGAN adversarial training starts after a 20K-step warm-up to avoid adding unstable adversarial signals at the beginning. These are implementation choices rather than evidence that one loss term explains all gains; the main text does not provide the complete learning-rate, loss-weight, and compute configuration required for independent reproduction. Figure 4 reports approximately 5-fold faster convergence in training iterations, not a measured end-to-end wall-clock speedup.

Key Experimental Results

Main Results

Table 3, page 12: Panda70M test-set reconstruction at \(33\times512\times512\), resizing the shorter side to 512 before center cropping. Higher PSNR and SSIM are better; lower LPIPS and FVD are better. FVD here measures reconstructed-video distribution quality, not the generation gFVD below. Only discrete configurations are listed here; compression follows the original table, and MT uses a different latent-grid downsampling ratio from the basic model.

Discrete Model	Compression	PSNR	SSIM	LPIPS	FVD
Cosmos	\(8\times8\times4\)	30.34	0.9159	0.0524	82.86
OneVAE	\(8\times8\times4\)	30.67	0.9228	0.0528	108.93
OneVAE-MT	\(8\times8\times8\)	31.80	0.9399	0.0450	78.35
Cosmos	\(16\times16\times4\)	26.80	0.8531	0.1260	350.87
OneVAE	\(16\times16\times4\)	27.40	0.8690	0.1063	336.48
OneVAE-FE	\(16\times16\times4\)	28.11	0.8772	0.0814	213.04

Basic OneVAE does not beat Cosmos on every metric at the lower-compression setting: PSNR is higher, but LPIPS and FVD are worse. Adding FE to high-compression OneVAE raises PSNR from 27.40 to 28.11 and lowers FVD from 336.48 to 213.04, demonstrating the practical benefit of first-frame allocation. The paper states that the Cosmos \(16\times16\times4\) comparison is adapted from its \(16\times16\times8\) configuration, rather than being a fully native, identically trained setting. Continuous mode also improves: at \(8\times8\times4\), the same table reports PSNR of 32.48 for OneVAE and 31.12 for CV-VAE, but their latent channel counts are 6 and 4, respectively, so capacity differences matter.

Ablation Study

Table 2, page 11: unified-training and architecture ablations; Section 4.1 gives the general evaluation protocol, but Table 2 does not separately identify its compression configuration. The table's Feature Enhancement refers to first-frame enhancement in the prose and is abbreviated FE here; UT includes both continuous initialization and dual-path training.

Config	PSNR	SSIM	LPIPS
Baseline trained from scratch	27.22	0.8673	0.1029
UT	28.15	0.9039	0.0756
UT + MT	28.87	0.9207	0.0683
UT + FE	28.88	0.8973	0.0669
UT + MT + FE	29.35	0.9183	0.0611

UT improves PSNR by 0.93 over the baseline; adding MT or FE to UT reaches 28.87 or 28.88, respectively. Enabling both gives the best PSNR and LPIPS, but the highest SSIM remains 0.9207 with MT alone rather than 0.9183 for the full configuration. These findings support complementary benefits, not a claim that adding modules monotonically improves every metric.

Key Findings

The paper also evaluates downstream generation, since reconstruction quality alone cannot establish that tokens are easy to model. Table 4, page 13: UCF-101 at \(16\times128\times128\); the AR architecture matches LARP, but generator parameter counts differ. Starred gFVD uses reconstructed ground-truth videos as the reference, whereas ordinary gFVD uses original ground-truth videos; these protocols must not be mixed in rankings.

Method	Tokenizer Parameters M	Generator Parameters M	gFVD, Reconstructed Reference	gFVD, Original Reference
LARP-L	173	632	43	57
OneVAE	135	775	37	62

OneVAE scores 37 against reconstructed references, better than LARP-L's 43, but scores 62 against original references, worse than 57. Its larger vocabulary embedding adds 143M generator parameters, making this experiment stronger evidence for AR compatibility than for a strictly superior generation model. The original OneVAE tokenizer was not trained on UCF-101 or at \(128\times128\) resolution; this does not mean that the downstream AR generator was not trained on that dataset.

Highlights & Insights

• The continuous path can provide auxiliary supervision for discrete training rather than merely an alternative deployment format. This makes existing continuous-tokenizer training reusable.
• Multi-token quantization emphasizes compositional representations over vocabulary expansion alone, but latent-grid size and sequence length must be reported together to assess capacity gains fairly.
• First-frame enhancement targets the asymmetry of causal context instead of uniformly increasing capacity for every frame; its exact budget still requires implementation-level verification.

Limitations & Future Work

• Evidence boundary: the available source is the main paper, which directs MCL-JCV results to supplementary material; unread results are not supplied here.
• Reproduction boundary: equation extraction is damaged, and the readable main text does not fully specify first-frame budget compensation, exact vocabulary settings, or all training hyperparameters.
• Causal-explanation boundary: Table 2 does not isolate continuous initialization from dual-path training, and latent-feature visualization cannot independently establish the gradient mechanism.
• Evaluation boundary: better reconstruction need not lower original-reference gFVD; future comparisons should control vocabulary, total token count, and generator parameters.

Related Work & Insights

• Compared with VQ / FSQ: OneVAE primarily reorganizes training and representational capacity around FSQ rather than introducing an entirely new alternative to nearest-neighbor quantization.
• Compared with Cosmos: Cosmos provides multiple continuous or discrete tokenizers, while OneVAE shares a model across both representations at one compression configuration; its basic version still loses on some metrics.
• Compared with CV-VAE: a continuous video latent space becomes an initialization and training resource for a discrete tokenizer, not just an end product; latent-channel capacity remains relevant to comparisons.
• Compared with LARP: OneVAE demonstrates next-token video generation with a general-purpose tokenizer, but its larger vocabulary and differing gFVD references limit direct ranking claims.

Rating

• Novelty: 4/5. Joint training, progressive derivation, and capacity allocation form a video-tokenizer recipe whose contribution is primarily system design.
• Experimental Thoroughness: 4/5. Reconstruction, module ablations, and AR generation are covered, but key training factors and capacity variables are not fully disentangled.
• Writing Quality: 4/5. The central argument is clear, while compression-dimension ordering and first-frame budget accounting could be more explicit.
• Value: 4/5. Useful for reusing continuous VAEs to train discrete video representations, subject to additional budget and downstream-generation verification.