AdaTok: choose a budget for an observed image
An interactive field guide
Variable-length reconstruction and deciding when to stop generating are different problems.
Before this article: TiTok · Evaluation
What was already possible
Section titled “What was already possible”A fixed-length tokenizer spends the same number of codes on a simple image and a detailed one. A prefix-decodable tokenizer can trade length for quality. AdaTok adds a policy that chooses among discrete budgets using features of the observed input image.
What changes, and why
Section titled “What changes, and why”Nested tail masking trains useful prefixes of a TiTok-S-style code sequence. The allowed lengths are 32, 64, 96, 128, 160, 192, 224, and 256. Each budget has a rank-16 LoRA decoder head. After tokenizer training, freeze it and train a small budget-selection MLP using relative advantages across the eight reconstruction choices.
Read from left to right. Select a module to connect its role, tensor shape, and calculation. Snippets are teaching pseudocode. Sources provide the paper method or audited implementation.
Full candidate sequence
Nested masks make multiple prefixes decodable.
all_codes,features = tokenizer.encode(image)
prefix = all_codes[:budget]Budget policy
Policy training uses frozen tokenizer reconstructions at all allowed budgets.
budget = allowed_budgets[policy(features).argmax()]Budget-specific decoding
Use the matching low-rank decoder adaptation.
image_hat = decoder(all_codes[:budget],lora_head=budget)Work through the mechanism
Section titled “Work through the mechanism”The demo lets you set signal complexity and code budget. Its analytic decoder keeps enough basis coefficients to represent a chosen frequency. It shows how a quality constraint could favor different budgets. It does not implement the learned selector, and the displayed 256/budget ratio is not measured acceleration.
The policy sees encoder features of the original image. During unconditional generation, that image does not yet exist. A fixed target length that matches the average selected length does not reproduce autonomous stopping. The paper leaves visual EOS for future work.
Change a parameter, inspect the computed values, then compare this small example with the model described above.
Analytical teaching model · controls change the example, not a trained research model.
Grey: target signal. Green: decoded teaching signal. No trained image model is used.
| Quantity | Value |
|---|---|
| Toy signal MSE | 0.000 |
| Signal samples | 48 |
| Fixed-length ratio 256 / budget | 2.000 |
Training is not sampling
Section titled “Training is not sampling”Tokenizer training precedes the roughly 330K-parameter policy. The policy compares multiple budgets per image, balancing MSE and token count through relative normalization and dynamic weights. Relative normalization does not guarantee an exact per-image rate-distortion optimum. Encoding still creates the full candidate sequence before selection.
Reported evidence, not a local reproduction
Section titled “Reported evidence, not a local reproduction”| Setting | Metric | Value | Interpretation |
|---|---|---|---|
| ImageNet-256, 50k validation, full 256 | rFID ↓ / PSNR ↑ | 1.31 / 18.42 | Fixed full budget |
| Adaptive, mean 118 tokens | rFID ↓ / PSNR ↑ | 1.50 / 17.78 | Observed-image reconstruction, not EOS generation |
| Fixed 128 tokens | rFID ↓ | 1.63 | Paper comparison for budget selection |
What this does not establish
Section titled “What this does not establish”The reported roughly 2.1× AR throughput uses a fixed 118-token target length, chosen to match the adaptive mean. It is not dynamic generation with a learned stop rule. The full 256-token MaskGIT gFID of 2.28 cannot be attached to the adaptive 118-token reconstruction row. No fixed official code revision was audited.
Check your understanding
Section titled “Check your understanding”Why cannot the same budget policy directly choose a length before unconditional generation?
It needs features of an observed source image. Generation has no such source, so a different stopping mechanism is required.
Sources and next steps
Section titled “Sources and next steps”Diagrams and snippets are teaching reconstructions. Paper reports, implementation observations, and analytical examples are distinguished. No research model runs in this page.