Soft Tail-dropping for Adaptive Visual Tokenization

TL;DR

Soft Tail-dropping Adaptive Tokenization (STAT) introduces a content-adaptive 1D visual tokenizer that outputs a variable-length latent token sequence with per-token keep probabilities. By enforcing a monotonic decreasing importance profile and aligning token counts with a perceptual complexity proxy, STAT achieves efficient, adaptive tokenization that pairs well with vanilla autoregressive image generation. The two-stage training (flexible prefix reconstruction and content-adaptive token allocation) yields high reconstruction quality with fewer tokens and strong generation performance, competitive with diffusion models while using simpler architectures. The approach demonstrates strong performance on ImageNet-1k for reconstruction and generation, extends to video, and provides insights into token-budget-aware modeling for scalable multimodal autoregressive systems.

Abstract

We present Soft Tail-dropping Adaptive Tokenizer (STAT), a 1D discrete visual tokenizer that adaptively chooses the number of output tokens per image according to its structural complexity and level of detail. STAT encodes an image into a sequence of discrete codes together with per-token keep probabilities. Beyond standard autoencoder objectives, we regularize these keep probabilities to be monotonically decreasing along the sequence and explicitly align their distribution with an image-level complexity measure. As a result, STAT produces length-adaptive 1D visual tokens that are naturally compatible with causal 1D autoregressive (AR) visual generative models. On ImageNet-1k, equipping vanilla causal AR models with STAT yields competitive or superior visual generation quality compared to other probabilistic model families, while also exhibiting favorable scaling behavior that has been elusive in prior vanilla AR visual generation attempts.

Figures (11)

• Figure 1: (a) STAT adaptively allocates token counts based on image complexity (the outlined figures), using fewer tokens for simple images and more for complex ones. When paired with a vanilla autoregressive model, STAT enables (b) high-quality text-conditional image generation and (c) the best FID in class-conditional generation on ImageNet among visual tokenizers.
• Figure 2: (a) Overview of STAT. A ViT encoder followed by a vector-quantizer (VQ) layer produces discrete latent tokens along with a keep probability for each token position. Probabilistic token dropping is applied to obtain a masked latent sequence for reconstruction, while two regularization losses enforce content-adaptive allocation and a decreasing importance prior. (b) Image generation with STAT. A special End-of-Sequence (EoS) token enables adaptive-length autoregressive generation, where the EoS position is determined by a threshold with the keep-probability profile predicted by STAT.
• Figure 3: Token dropping strategies in two training stages.
• Figure 4: Keep-probability curves and corresponding reconstructions across samples of varying complexity. More complex images receive higher token counts, and the learned allocation correlates strongly with the JPEG file sizes of the inputs.
• Figure 5: Ablation Study on the two loss terms for decreasing importance prior and the content-adaptive prior. The decreasing importance prior is necessary for learning the monotonic structure for AR generation, while the content-adaptive prior is crucial for successfully learning token allocation strategy.