Keep-or-Drop? Adaptive Tokenizer for Compact Video Representation

AI summary

Overview

Research area: Computer vision and generative modeling — specifically latent diffusion models (LDMs), video tokenizers/VAEs, and adaptive token sparsification for video representation.

Technical level: Advanced. The paper assumes familiarity with variational autoencoders, latent diffusion, flow matching, transformer encoder–decoder architectures, Gumbel–Softmax relaxation, and standard video-generation metrics (PSNR, SSIM, LPIPS, FVD).

Scope in one sentence: The paper introduces KATok, a transformer-based adaptive video VAE that learns per-token keep-or-drop decisions through a differentiable Gumbel–Softmax gate plus sparsity regularization, and pairs it with two diffusion-side strategies (joint content–position generation and cascaded mask-prior conditioning) to preserve spatial consistency when latents are aggressively sparsified.

What This Paper Is About

Conventional video VAEs used inside latent diffusion models compress video at a fixed ratio, so the number of latent tokens always scales with spatial and temporal resolution even though most video content is spatio-temporally redundant. Recent flexible-length tokenizers let a user choose the token budget, but that is flexibility, not adaptivity — the optimal count per clip still requires extra models or inference-time search. KATok instead learns to automatically decide which tokens to keep and which to drop from each individual clip, aiming for compact yet expressive latents that remain compatible with diffusion-based video generation.

Key Contributions

1. An adaptive VAE for token-wise compression. KATok learns token importance through soft attention gating (Gumbel–Softmax relaxation) plus an L1 sparsity penalty on the soft masks, so sample-adaptive compression emerges during training without any predefined token budget or inference-time search.

2. Two remedies for content–position misalignment under sparse tokenization. Because dropping tokens can disturb original spatio-temporal structure, the paper proposes (i) joint content–position prediction with timestep decoupling and (ii) a cascaded scheme in which a lightweight mask prior selects positions before content generation.

3. Efficiency gains in training and generation. The authors report 6.9x faster training and 3.2x faster inference than the transformer-based tokenizer baseline, while generating on the order of ten times fewer tokens.

4. An analysis of what drives token allocation. Quantitative correlation analysis shows that the number of tokens kept tracks temporal complexity more strongly than spatial complexity, and this relationship yields emergent controllability of motion and detail at generation time.

Main Findings

• Reconstruction with far fewer tokens (Panda-70M validation, ~5,000 videos). At 256²×16, KATok uses 366.24 tokens on average (64 channels) and reaches PSNR 31.24, LPIPS 0.04, SSIM 0.94, rFVD 5.12, versus OmniTokenizer-VAE at 5,120 tokens (PSNR 28.10, LPIPS 0.05, SSIM 0.88, rFVD 7.84) and ElasticTok-KL at 3,845.56 tokens (PSNR 30.52, LPIPS 0.06, SSIM 0.91, rFVD 12.37).

• Compression ratio grows with resolution. At 512²×32, KATok uses 1,554.24 tokens with compression ratio 253.00 (PSNR 33.23, rFVD 6.40), compared with OmniTokenizer-VAE's 36,864 tokens, compression ratio 96, PSNR 24.07 and rFVD 16.85. At 256²×16 the compression ratio is 134.21. ElasticTok is only available at 256² spatial resolution, so 512²×32 results are marked as not supported.

• Extreme sparsity in near-static content. A uniform white 256²×16 clip is reconstructed with as few as 28 tokens, and a 512² example uses 836 tokens — fewer than ElasticTok's 1,104 tokens at 256² — while achieving better reconstruction quality. On the ocean clip the token count drops from 1,104 to 836; indoors from 1,712 to 330.

• Token count tracks temporal complexity. Correlating expected token count with entropy measures gives Pearson r = 0.865 for temporal entropy, r = 0.618 for spatial entropy, and r = 0.877 for spatio-temporal entropy (the figure caption states the comparison as approximately 0.87 versus 0.62).

• Generation quality (gFVD, lower is better). Ours-Cascaded scores 23.19 on SkyTimelapse, 61.53 on UCF-101, and 160.84 on Kinetics-600; Ours-Joint scores 21.36, 73.16, and 193.78. OmniTokenizer-VAE scores 23.28, 100.00, and 206.58; ElasticTok-KL scores 95.53 on SkyTimelapse and 712.56 on UCF-101, with no Kinetics result reported.

• Convergence and throughput. Ours-Cascaded reaches gFVD 49.34 at 200k steps versus 82.31 for OmniTokenizer, using approximately 28.6% of the effective tokens on average under OmniTokenizer's 2×2 spatial patch setting. It reaches 73.81 at 80k steps, surpassing OmniTokenizer's final 200k-step result with roughly a 6.9x wall-clock speed-up. Against ElasticTok it is approximately 46.7x faster, reaching 203.51 at 30k steps versus 571.41 at 200k steps. Generation throughput is 15.71 videos/sec, which is 3.1x, 3.2x, and 3.7x faster than Ours-Joint (5.01), OmniTokenizer (4.91), and ElasticTok (4.24) using the DOPRI5 ODE solver.

• Sparse generation strategy progression (UCF-101, gFVD). Naive content-only flow matching reaches 95.69, joint content–position generation improves to 73.16, and cascaded mask-prior conditioning gives the best result at 61.53. Under the same tokenizer, removing latent regularization degrades gFVD from 101.23 to 161.23.

• Ablation of VAE components (stage-1 model, 100K iterations, max 514 tokens). Full model: PSNR 30.85, LPIPS 0.07, SSIM 0.93, 365.57 tokens. Removing latent regularization: 31.26 / 0.07 / 0.94 / 361.70. Removing asymmetric decoding: 29.61 / 0.11 / 0.92 / 377.80. Replacing Video-LPIPS: 29.28 / 0.11 / 0.91 / 404.00. Removing register tokens: 29.17 / 0.12 / 0.91 / 410.20. Removing the soft attention mask: 19.00 / 0.51 / 0.54 / 2.00 (training collapse attributed to register tokens alone remaining active). Removing Gumbel–Softmax: 18.91 / 0.52 / 0.53 / 2.00.

• Emergent controllability. Adjusting the number of tokens at generation time modulates motion and detail without retraining or extra conditioning: 200-token samples produce simpler, low-motion videos, while 400-token samples produce more dynamic and visually rich sequences (UCF-101).

Methodology in Plain English

A video is cut into spatio-temporal patches and encoded by a transformer into continuous latent tokens. A lightweight token-importance network scores each token with keep/drop logits, and a Gumbel–Softmax relaxation turns those logits into a differentiable soft mask, so the keep-or-drop decision can be trained end to end. The same soft mask is used to bias decoder attention logits (b_j = log(m̃_j + ε)), which the paper states is essential — without it, training collapses. An L1 penalty on the soft masks gives the sparsity objective, with the penalty weight annealed: a weak penalty for roughly the first 5k iterations, then progressively increased.

The tokenizer is trained with a weighted sum of reconstruction (L1 plus Video-LPIPS computed with spatio-temporal S3D features), KL, sparsity, adversarial, and representation-alignment losses. Latent regularization adds small Gaussian noise sampled uniformly from [0, 0.2] and an alignment loss against a frozen vJEPA-2 encoder; the authors note these slightly hurt reconstruction metrics but improve generation and convergence.

Architecturally, the encoder is a single-stream transformer with 3D rotary positional embeddings and two register tokens; the decoder is a double-stream design (adopted from FLUX) where learnable query tokens attend to masked latent tokens. Encoding uses coarse 16²×8 patches while decoding uses a finer 8²×4 grid — an asymmetric coarse-to-fine scheme that keeps the latent compact but the reconstruction detailed. Training proceeds in three stages: single-resolution 256²×16 (210K steps), multi-resolution (30K steps), and GAN-based fine-tuning (50K steps), with a total of 344M parameters and a batch size of 256.

For downstream generation, the VAE supplies latents to a flow-matching model. To prevent content–position misalignment under heavy sparsification, the paper offers two options: jointly predicting content and position with decoupled noise schedules (sensitive to scheduler tuning), or a cascaded scheme in which a lightweight 8.3M mask-prior model first predicts which positions are active and feeds them as positional conditioning to the main generator. The cascaded approach is adopted as the default.

Why This Matters

Impact on research: The paper reframes tokenizer compression from a user-specified budget to a learned, content-dependent outcome, and it identifies a concrete failure mode — content–position misalignment — that arises specifically when adaptive tokenization is combined with diffusion. The finding that differentiability of the token gate is essential (removing it collapses training) and that token count correlates far more strongly with temporal than spatial entropy provides diagnostic signals for follow-up work on learned sparsity in generative models.

Real-world applications (derived from the paper's efficiency and controllability claims):

• Generative video content creation, where producing more videos per GPU-hour lowers serving cost.
• High-resolution and long-sequence video pipelines, where the paper shows the compression ratio rising from 134.21 at 256²×16 to 253.00 at 512²×32.
• Interactive generation tools where token count acts as a single knob for motion intensity and visual richness (200-token versus 400-token samples), without retraining or extra conditioning.
• Compression of highly redundant footage, such as static or near-uniform scenes reconstructed with as few as 28 tokens.

Industry relevance: The work comes from Kakao Corp., and its headline efficiency numbers — 6.9x faster training, 3.2x faster inference, 15.71 videos/sec, and roughly 10x fewer generated tokens per clip — are measured on 8 H200 GPUs with a global batch size of 256, targeting the cost profile that dominates industrial video generation.

Future Directions

• Reducing sensitivity to scheduler hyperparameters. The joint content–position approach improves spatial consistency but the paper states its performance is sensitive to the choice of scheduler hyperparameters and the optimal combination of the two noise schedules, typically requiring careful tuning.
• Extending the comparison across tokenizer paradigms. The main paper restricts baselines to transformer-based continuous-token (KL) VAEs; comparisons against CNN-based VAEs (Cosmos, LTX-Video) and VQ-based adaptive tokenizers (EVATok, AdapTok) are deferred to supplementary material, leaving cross-paradigm rankings open.
• Covering baselines at higher resolution. ElasticTok is only available at 256² spatial resolution, so the 512²×32 comparison cells are unsupported; broader coverage of adaptive baselines on high-resolution and long video would strengthen the claims.
• Better understanding of the mask prior's role. The cascaded strategy relies on a lightweight 8.3M mask prior; how its accuracy interacts with generation fidelity under more extreme sparsification is not reported.

Target Audience

Researchers and engineers working on video generation, latent diffusion, and visual tokenization who need a compact latent representation with adaptive, content-dependent token budgets. It is also relevant to practitioners optimizing inference and training cost for high-resolution or long-sequence video models, and to readers interested in how token-dropping ideas from efficiency literature (DynamicViT, EViT, ToMe) transfer from inference-time pruning to learned sparsity inside a trainable tokenizer.

Authors’ abstract

Latent diffusion models have emerged as a dominant framework for high-fidelity image and video synthesis, operating in compact latent spaces with variational autoencoders (VAEs) to enhance computational efficiency without compromising visual quality. However, conventional VAEs are suboptimal for video data as they employ fixed compression ratios that cannot adapt to the varying complexity of spatio-temporal content. We present KATok (Keep-or-Drop? Adaptive Tokenizer for Compact Video Representation), a transformer-based VAE that incorporates an adaptive token selector which is jointly learned with latent tokens. By evaluating each token's content-richness as keep-or-drop probability, the token selector effectively discards uninformative tokens, naturally allowing data-dependent compression. Applying adaptive tokenization to diffusion models may cause spatial misalignment, as token dropping can disturb the original spatio-temporal structure. To alleviate this issue, we propose two position-prediction strategies: cascaded and joint generation, to ensure spatial consistency. We empirically show that our model achieves strong reconstruction and generation quality at a state-of-the-art compression ratio. Further analysis on video data reveals that this improvement is primarily achieved by reducing spatio-temporal redundancy and removing uninformative tokens, as supported by both quantitative and qualitative results.