PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yeonkyeong Lee, Hyunsung Go, Jongmin Kim, Sewoong Lim, Donghoon Lee

arXiv:2608.24293v1cs.CV

TL;DR

Video VAEs use fixed compression despite varying spatio-temporal complexity, motivating KATok’s adaptive continuous token selection. KATok combines sparse latent tokenization with joint or cascaded position prediction, achieving faster, high-quality video generation while exposing failure cases under extreme motion or sparsity.

  • Problem

    Fixed-compression video VAEs allocate capacity to redundant or uninformative regions, while flexible tokenizers require predefined budgets or inference-time search.

  • Method

    KATok is an adaptive VAE that learns token importance with differentiable sparsity and uses joint or cascaded content–position generation for sparse diffusion latents.

  • Results

    6.9× faster training and 3.2× faster inference are achieved than with the strong transformer-based tokenizer baseline, with high-quality generation using substantially fewer tokens.

  • Takeaways & Limitations

    Adaptive tokenization coupled with suitable diffusion design provides a framework for efficient and high-quality video generation.

  • Takeaways & Limitations

    High-frequency detail under large motion can become blurry, with PSNR approximately 19–20 on Panda-70M; fewer than 3 generated tokens can produce near-uniform clips.

Abstract

from arXiv · show

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.

1 Introduction

KATok addresses the inefficiency of fixed video tokenization by learning content-adaptive token selection, while adding position-prediction strategies to preserve spatial consistency in diffusion generation.

  • Fixed-length VAE tokenization allocates capacity to uninformative regions despite strong spatio-temporal redundancy in videos.
  • Flexible-length tokenizers allow user-controlled compression but still require predefined token counts or inference-time search, whereas adaptive tokenization selects tokens from content complexity.
  • KATok predicts per-token importance with a differentiable Gumbel–Softmax gating mechanism and sparsity regularization, allowing compact continuous latent representations to emerge during training.
  • Sparse tokenization can disrupt original token positions, causing content–position mismatch, unstable boundaries, and temporal inconsistencies during generation.
  • KATok proposes joint content–position prediction and cascaded position selection to mitigate spatial misalignment under sparse tokenization.
  • 6.9× faster training and 3.2× faster inference are achieved than with the strong transformer-based tokenizer baseline while using substantially fewer tokens.

2 Related Works

Prior work improves compact visual modeling through latent tokenizers, flexible token allocation, and token pruning, but KATok targets adaptive compression learned directly during training.

  • Autoencoder-based tokenizers compress high-dimensional visual inputs into latent spaces that make large-scale image and video generation computationally manageable.
  • Recent compression and temporal-modeling methods still scale poorly with spatial resolution and sequence length, leaving cross-frame redundancy underutilized.
  • Flexible-length tokenizers allocate different numbers of tokens according to content importance using dropout-based, causal, or iterative strategies.
  • KATok performs continuous differentiable selection within a transformer encoder instead of relying on predefined budgets or inference-time search.
  • Inference-only pruning and merging methods reduce transformer redundancy, whereas KATok learns sparsity during training for efficiency in both training and inference.

3 Adaptive Tokenizer for Video Representation

KATok combines differentiable keep-or-drop selection, shared decoder masking, sparsity regularization, and asymmetric decoding to produce compact yet expressive video latents.

  • The framework aims to adjust token count to content complexity while avoiding preset budgets or inference-time search.
  • KATok represents videos as patch embeddings and learns an effective token count Neff(X) through per-token masks while maintaining reconstruction fidelity.
  • The encoder predicts keep/drop logits, uses Gumbel–Softmax for differentiable soft masks during training, and applies hard masks to discard tokens at inference.
  • Sharing soft masks with decoder attention keeps token dropping consistent; removing this masking causes training instability and severe reconstruction degradation.
  • An ℓ1 penalty on soft masks minimizes the expected number of active tokens, with its weight annealed after an initial phase that permits nearly all tokens to remain.
  • The learned selection makes active-token count an outcome of training rather than a predefined budget, eliminating inference-time search and handcrafted dropout schedules.
  • The tokenizer combines reconstruction, KL, adversarial, sparsity, and representation losses, while Video-LPIPS captures temporal coherence and motion consistency.
  • Latent noise augmentation and vJEPA-2 representation alignment regularize decoding to stabilize training and improve generative fidelity.

4 Diffusion with Sparse Latent Tokens

The paper trains flow-matching models on sparse latent tokens while addressing the spatial misalignment caused by adaptive token removal. It compares joint content–position prediction with cascaded mask-prior conditioning, which separates position selection from content generation.

  • Flow matching transports Gaussian noise to clean latent tokens by learning a continuous-time velocity field along linearly interpolated states.The objective matches the predicted velocity to the displacement between clean latents and noise.
  • Adaptive token removal can misalign surviving tokens with their original locations, causing unstable boundaries and temporal inconsistencies.
  • 4.1 Joint content–position generation with timestep decoupling: Joint generation predicts latent contents and original positions together, optimizing content and position flow-matching objectives with a position-loss weight.Separate noise schedulers decouple content and position denoising dynamics.
  • 4.1 Joint content–position generation with timestep decoupling: The joint approach improves spatial consistency but is sensitive to scheduler hyperparameters and typically requires careful tuning.
  • 4.2 Cascaded mask-prior conditioning: The cascaded approach first predicts a binary token-selection mask, then conditions the main flow model on positional embeddings from the selected positions.Ground-truth positions provide training conditioning, while mask-prior predictions provide inference conditioning.
  • 4.2 Cascaded mask-prior conditioning: Mask-prior conditioning is adopted by default because it is more robust and achieves higher generative fidelity under heavy sparsification.Joint position prediction remains a strong ablation.

5 Experiments

Experiments show that KATok reconstructs and generates videos with substantially fewer tokens while preserving or improving quality. Its adaptive selection concentrates tokens on informative dynamics, scales efficiently with resolution, and supports controllable generation, though key architectural components are necessary for stable training and performance.

  • Reconstruction Quality: 31.24 PSNR and 5.12 rFVD are achieved at 256² × 16 using 366 tokens, versus 5,120 for OmniTokenizer-VAE and 3,846 for ElasticTok-KL.At 512² × 32, KATok reaches 33.23 PSNR and 6.40 rFVD with 1,554 tokens, compared with OmniTokenizer-VAE’s 24.07 PSNR and 16.85 rFVD using 32,768 tokens.
  • Adaptive Compression: Adaptive token selection retains motion-rich regions while suppressing static or homogeneous areas, using only 28 tokens for an entire uniform-white 256² × 16 clip.At higher spatial resolution, token dropping becomes stronger within the same scene, while reconstruction quality remains perceptually favorable.
  • Adaptive Compression: Effective token usage correlates most strongly with spatio-temporal entropy (r=0.877) and temporal entropy (r=0.865), more than spatial entropy (r=0.618).This indicates that token allocation tracks video dynamics particularly closely, while spatial complexity alone is a weaker predictor.
  • Video Generation: KATok generates 366 tokens per clip on average versus Omni-VAE’s fixed 5,120-token budget, while consistently outperforming Omni-VAE across datasets.Under the same tokenizer, naive flow matching, joint content–position generation, and cascaded mask-prior conditioning yield progressively better generation quality.
  • Video Generation: Ours-Cascaded reaches 49.34 gFVD at 200k steps versus OmniTokenizer’s 82.31, and reaches 73.81 at 80k steps, giving a ∼6.9× wall-clock speed-up.It uses only ∼28.6% of OmniTokenizer’s effective tokens under the compared spatial patch setting.
  • Video Generation: Increasing generation-time token counts produces more dynamic and visually detailed videos without retraining or additional conditioning.Using fewer tokens yields simpler, low-motion sequences, while higher allocations modulate motion intensity and visual richness.
  • Ablation: Asymmetric decoding improves PSNR from 29.61 to 31.26 while keeping average token counts similar, from 377.80 to 361.70.The change increases only decoder reconstruction granularity, using finer decoder patches than the encoder.
  • Ablation: Removing Video-LPIPS increases token usage from 377.80 to 404.00 and reduces reconstruction quality, while removing soft attention masking or Gumbel–Softmax causes severe degradation or training collapse.Without soft attention masking, only two register tokens remain active, producing PSNR 19.00 and SSIM 0.54.

6 Conclusion

KATok adaptively keeps or drops tokens according to content complexity, while joint and cascaded position strategies address sparse-token content–position misalignment.

  • KATok selectively keeps or drops tokens according to content complexity for compact, expressive video representations.Differentiable attention gating and sparsity loss balance fidelity and efficiency.
  • Joint content–position generation with timestep decoupling and cascaded mask-prior conditioning mitigate misalignment from sparse latent tokens.
  • The proposed adaptive tokenization framework supports efficient, high-quality video generation when paired with suitable diffusion design.

Supplementary Material for “Keep-or-Drop?

The supplementary material identifies the paper and its authorship and lists adaptive tokenization and video diffusion among its keywords.

  • The paper is titled “Adaptive Tokenizer for Compact Video.”
  • Yeonkyeong Lee, Hyunsung Go, Jongmin Kim, Sewoong Lim, and Donghoon are listed as authors.
  • The authors are affiliated with Kakao Corp., Republic of Korea.

1 Implementation Details

KATok uses a transformer-based adaptive VAE with sparse token selection, multi-stage training, and diffusion components designed to preserve spatio-temporal coherence.

  • Tokenizer Model Architecture: KATok is a transformer-based adaptive VAE comprising linear patchification, a single-stream encoder, a double-stream decoder, and an adaptive token selector.The model follows a ViT-B configuration with 344M parameters and uses 3D rotary positional embeddings.
  • Decoder Design: The decoder uses learnable query tokens, 3D positional embeddings, and soft attention masking to reconstruct spatial layouts from sparse latents.The masking implementation shifts attention scores using the logarithm of soft mask values.
  • Adaptive Token Selection: The selector predicts Gaussian latent parameters and Gumbel keep/drop logits from encoder outputs, while register tokens remain active.Soft masks, sparsity, KLD, reconstruction, perceptual, and adversarial losses are combined with stage-specific weights.
  • Asymmetric Coarse-to-Fine Decoding: The encoder uses coarse 16^2 × 8 patches, while the decoder reconstructs on a finer 8^2 × 4 grid for asymmetric coarse-to-fine decoding.The finer decoder grid compensates for details discarded by coarse encoding.
  • Multi-Stage Training: Training proceeds through 210K iterations of single-resolution pretraining, 30K multi-resolution iterations, and 50K adversarial fine-tuning iterations.Stage 3 uses reconstruction-based GAN training to improve perceptual fidelity and sharpness.
  • Diffusion Model Architecture: The cascaded diffusion model predicts occupancy over a fixed 3D grid before content generation, separating position selection from content generation.Its mask prior uses a SiT transformer and is trained jointly with the content diffusion model.

2 Additional Quantitative Results

Extended comparisons show that KATok remains competitive across tokenizer paradigms while using substantially fewer tokens and achieving strong controlled generation quality.

  • The controlled generation comparison pairs all tokenizers with an identical SiT-XL backbone and training protocol, isolating tokenizer effects.
  • KATok remains Pareto-leading on Kinetics-600 and Pareto-competitive on the harder 512² Epic-Kitchens benchmark.
  • 7× fewer tokens deliver 2.6–4.0 dB higher PSNR than adaptive VQ tokenizers across Kinetics-600 and Epic-Kitchens.
  • 61.53 gFVD with 366 tokens outperforms OmniTokenizer-VAE at 100.00 gFVD with 5,120 tokens in the controlled UCF-101 comparison.

3 Additional Qualitative Results

Qualitative reconstructions and generations show adaptive token allocation across diverse scenes, with token budgets controlling motion complexity and visual detail.

  • VAE Reconstruction: Adaptive reconstruction preserves spatio-temporal details while discarding redundant regions across diverse scenes.
  • Video Generation with Mask Visualization: The model allocates more tokens to complex motion, fine details, cloud textures, horizon details, and large camera or object motion.
  • Video Generation with Mask Visualization: Cascaded generation produces temporally coherent and visually consistent sequences across UCF-101, SkyTimelapse, and Kinetics-600.

4 Extended Analyses and Visualizations

Extended analyses characterize efficiency, mask quality, token-budget sensitivity, domain behavior, failure cases, and the relationship between redundancy and adaptive sparsification.

  • Mask Prior Analysis: The mask prior adds only 8.3M parameters, or 1.2% of the 686M content model, and 4.72% to total sampling cost.
  • Mask Prior Analysis: Predicted masks achieve 91.80 FVD on Kinetics-600 mask sequences, evaluated against a random-mask upper bound using two 2,048-clip validation halves.
  • Behavior under Domain Shift and Extreme Conditions: KATok behaves consistently under domain shift, continuing to track motion and texture complexity on Kinetics-600 and Epic-Kitchens.
  • Behavior under Domain Shift and Extreme Conditions: High-frequency detail combined with large motion produces blurry reconstructions in moving fine-texture regions, with PSNR approximately 19–20.
  • Behavior under Domain Shift and Extreme Conditions: Fewer than 3 generation tokens, outside training, leave insufficient capacity and produce near-uniform, single-color clips.
  • Image Reconstruction on ImageNet: Reducing ImageNet patch size from 32² to 16² increases redundancy, causing the selector to retain 229 of 256 tokens and drop homogeneous regions.
  • Token Count and Compression: On Panda-70M at 256² × 16, the model averages 366 of 512 tokens and reaches a 134× compression ratio.
  • Sparsity Sensitivity: Mild sparsity with λ_sparse ≤0.01 retains 370–400 tokens while maintaining PSNR ≥28.67 and rFVD ≤50.38.
Loading 2608.24293v1…