MixTok: Adaptive Visual Tokenization with Global 1D Tokens and Sparse 2D Refinement

AI summary

Overview

Research area: Computer vision, specifically visual tokenization for image reconstruction and class-conditional image generation (discrete latent representations, vector quantization, masked generative transformers).

Technical level: Intermediate. The core idea is intuitive, but the paper assumes familiarity with VQ-VAE-style tokenizers, codebooks, MaskGIT-style masked generation, and FID-style metrics.

Scope in one sentence: The paper introduces MixTok, a hybrid visual tokenizer that pairs a compact 1D global token sequence with content-adaptively selected sparse 2D tokens for local detail refinement, evaluated on ImageNet 256×256 reconstruction and generation.

What This Paper Is About

Most visual tokenizers either lay tokens out on a uniform spatial grid (which preserves local detail but wastes capacity in visually simple regions) or compress an image into a short 1D latent sequence (which is compact but tends to lose fine-grained local detail). MixTok's goal is to get both properties at once: a compact global representation plus a small number of 2D tokens placed only where an image actually needs refinement. It also lets the number of these extra tokens vary per image, rather than fixing the same budget for every image at a given resolution.

Key Contributions

1. A hybrid tokenizer (MixTok): it combines a compact 1D global representation with sparse 2D tokens, where selected 2D features replace the corresponding 1D-derived base features and a shared decoder reconstructs the image from this mixed grid.
2. A content-adaptive router: a Reconstruction Grid Scorer ranks spatial locations by refinement priority, and a top-p Budget Controller decides how many 2D tokens each image receives, under a target average token budget.
3. A two-stage MaskGIT generation pipeline: a 1D MaskGIT generates global tokens, which guide selection of refinement locations and provide spatial context for a sparse 2D MaskGIT that generates only the selected tokens.
4. Empirical validation across backbones: results on ImageNet 256×256 for reconstruction (rFID 1.195 at roughly 128 content tokens) and generation (gFID 1.95 at roughly 131 tokens), plus a transferability test replacing the TiTok–LlamaGen pair with ALIT and its VQGAN backbone.

Main Findings

• Reconstruction quality at a compact budget: MixTok achieves an rFID of 1.1950 and a PSNR of 19.51 dB with approximately 128 content tokens per image on average. At a comparable token budget it improves rFID over TiTok S-128 (1.7222), TiTok BL-128 VQ (1.5063), and One-D-Piece L-256 (1.4406), and remains competitive with FlexTok d18-d28 (1.1422) and DetailFlow-16 (1.1370).
• Generation quality: MixTok reaches a gFID of 1.95 with approximately 131 tokens. At a similar token budget it improves gFID over TiTok-S-128 (2.00 to 1.95) while reducing sampling rounds from 64 to 16. Relative to DetailFlow-32, it reports a lower gFID (1.95 vs. 2.63) with roughly half as many tokens and fewer rounds (16 vs. 39).
• Latency: Under the paper's single-image latency protocol on one RTX 5090 at batch size 1, MixTok takes 270.1 ms and 1.452 T FLOPs per image, reported as 3.77× faster than the TiTok configuration and 2.48× faster than the DetailFlow configuration.
• Sparse refinement drives generation gains: starting from the same generated 1D tokens, sparse 2D refinement improves gFID from 4.33 to 1.95 with eight additional sampling rounds.
• Learned spatial ranking beats heuristics: holding each image's top-p budget from the Reconstruction Grid Scorer fixed, the Reconstruction Grid Scorer reaches the best rFID of 1.1950 (PSNR 19.5074), ahead of the Generation Grid Scorer (1.3479), gradient entropy (1.4012), random (1.5460 ± 0.0093), and uniform (1.6604) rankings.
• Per-image budgets help beyond mere variation: with the same Grid Scorer, top-p dynamic allocation at an average K of 96.1093 gives rFID 1.1950 and PSNR 19.5074, better than a fixed K = 96 (1.2052) and better than shuffled budgets that preserve the budget distribution but break its link to image content (1.2236 ± 0.0068).
• Transferability to other backbones: on the ALIT–VQGAN combination at a matched 128-token budget (64 1D + 64 2D), rFID improves from 2.9453 to 2.4720 and PSNR from 18.71 to 18.76 dB over ALIT-128.
• Distillation preserves refinement priorities: the Generation Grid Scorer, which sees only 1D base features during generation, reaches rFID 1.3479 and PSNR 18.4776, outperforming all tested heuristic rankings even without access to 2D features.
• Qualitative effect: selective 2D refinement recovers stripe patterns, curved boundaries, and sharp structures that are blurred or missing in the 1D-only base reconstruction, per the paper's qualitative figures.

Methodology in Plain English

MixTok is built on top of two existing, pretrained tokenizers. A 1D tokenizer (TiTok-L-32, 32 tokens) compresses the whole image into a short sequence that has no explicit spatial layout; a 2D tokenizer (LlamaGen VQ-16, a 16×16 grid) keeps tokens tied to image locations. The trick is combining them rather than choosing between them.

First, the 1D tokens are expanded into a 16×16 spatial feature grid and projected into the decoder's feature space, giving a "base" representation of the image. The 2D tokenizer's features already live in that same space, so a 2D feature can be substituted into the grid at any chosen location. A small convolutional network called the Reconstruction Grid Scorer looks at the base grid, the full 2D feature grid, and the base reconstruction, and assigns each of the 16×16 locations a refinement score. A top-p rule then converts those scores into an actual selection: the locations are ranked, and the chosen set is the shortest prefix of that ranking whose cumulative softmax probability reaches a threshold p. When probability mass concentrates on a few locations, fewer tokens are used; when it is spread out, more are used. The threshold p is calibrated once on a subset of training images to hit a target average budget and then fixed, so individual images still get different token counts. The selected 2D features replace the base features at those locations, and one shared decoder reconstructs the image from the mixed grid.

Training proceeds in stages. A short alignment stage (2.5 epochs) teaches the spatial expansion and feature adapter to match 1D features to the decoder space while everything else stays frozen. Then joint training with a fixed budget of K = 96 refinement tokens (20 epochs) trains the scorer, the encoding and decoding paths, and the discriminator (adversarial training starts at epoch 1.5 after a 0.01-epoch warm-up). Finally, top-p fine-tuning (6 epochs, with T = 2 and p = 0.91) adapts the model to content-dependent budgets. Losses combine L1 reconstruction, LPIPS perceptual loss, GAN loss, DINOv2 and CLIP feature losses, and routing penalties, with the 1D encoder and quantizer frozen throughout.

For generation, a two-stage pipeline is used. A pretrained TiTok MaskGIT generates the 1D tokens, which are expanded into the base grid. Because the real 2D features are unavailable at generation time, a separate Generation Grid Scorer is trained by distillation to predict refinement priorities from the 1D base grid alone. The same top-p rule picks the locations. A fine-tuned Halton-MaskGIT then generates 2D tokens only at those locations, with an Embedding Adapter supplying continuous context at the unselected positions.

Why This Matters

Impact on research. The paper argues that token allocation should adapt to content in two ways at once — where extra capacity goes and how much each image gets — rather than fixing a uniform grid or a single per-image token count. It provides a concrete mechanism (learned grid scoring plus top-p budget control) and ablation evidence that both components contribute independently: learned ranking beats random, uniform, and gradient-entropy heuristics, and content-dependent budgets beat both fixed and shuffled budgets at the same average. It also shows the recipe transfers beyond its own backbone pair, since the ALIT–VQGAN combination improves at a matched 128-token budget.

Real-world applications

• Generative image and content systems: shorter token sequences and fewer sampling rounds translate directly into cheaper inference, which matters for interactive or high-volume image synthesis.
• Bandwidth-constrained transmission: a compact discrete representation of an image with variable per-image token counts is relevant to progressive or adaptive-rate image delivery.
• Media storage and retrieval: compact global latents plus sparse local detail tokens provide a representation usable for compression-style pipelines and latent-space search.
• Downstream vision modeling: tokenized latents of this kind serve as the modeling space for subsequent generative or predictive models that operate on discrete tokens rather than raw pixels.

Industry relevance. The reported latency and FLOPs figures (270.1 ms and 1.452 T/image on one RTX 5090 at batch size 1, versus 1018.4 ms for the TiTok configuration) frame the work as an efficiency argument, not only a quality argument. Because it builds on released pretrained checkpoints and publishes code and configuration files, the approach is relatively easy to graft onto existing tokenizer stacks.

Future Directions

• Sparse computation in the decoder and generator. The authors note that although MixTok retains and generates fewer discrete content tokens, reconstruction still requires dense 2D encoding, and the second-stage generator still operates on the full spatial grid. Removing that dense computation is stated as future work.
• Broader token budgets. Evaluation covers only a narrow range of average token budgets, so how the top-p allocation behaves at very low or very high budgets is an open question.
• Higher resolutions and more diverse data. The paper evaluates only ImageNet at 256×256; the authors list higher resolutions and more diverse visual data as targets.
• Extension of adaptive global–local allocation beyond this setup. Whether the same router design benefits other tokenizer families, and how it interacts with richer conditioning or different visual domains, remains untested here.

Target Audience

Researchers and engineers working on visual tokenization, discrete latent representations, image compression, and masked generative models; practitioners who need to trade off token budget against reconstruction or generation fidelity in vision systems; and readers with a background in deep generative modeling who want a concrete example of content-adaptive token allocation with ablation evidence for each design choice.

Authors’ abstract

Efficient visual tokenization requires balancing representation compactness with the preservation of fine-grained image content. Conventional grid-based tokenizers allocate tokens uniformly across spatial locations despite non-uniform visual complexity, whereas compact 1D tokenizers aggregate global information but can struggle to preserve local details. We introduce MixTok, a hybrid visual tokenizer that combines a compact 1D global representation with sparse 2D tokens for local refinement. A content-adaptive router ranks candidate locations by refinement priority and determines the number of 2D tokens for each image under a target average token budget. At the highest-scoring locations, 2D token features replace the spatial features recovered from the 1D representation, and a shared decoder reconstructs the image from the resulting mixed features. On ImageNet 256 x 256, MixTok achieves highly competitive reconstruction and generation performance, with an rFID of 1.195 and a gFID of 1.95 at respective average budgets of approximately 128 and 131 content tokens per image.