Source-linked AI summary
Visual Autoregressive Generation with Next Token Imagination
Ruihang Jiang, Zongyao Zhao, Zhaolin Wang, Yuanwei Liu, Xiatian Zhu
TL;DR
Standard visual AR training supervises each position using only the ground-truth next token, although plausible alternative continuations may exist. NTI trains models to explore plausible non-ground-truth tokens using context-regulated imagination and snapshot-based context approximation. On ImageNet, it improves visual AR results, including reducing FID from 3.02 to 1.67 on the same AR-L backbone.
Problem
Standard visual AR training conditions on ground-truth prefixes and supervises only the ground-truth next token, despite plausible non-ground-truth continuations.
Source: §1 · p. 1Method
NTI uses Context-Regulated Imagination to assign probability to plausible non-ground-truth tokens and Snapshot-based Context Approximation to construct model-generated contexts without sequential rollout.
Source: §1 · p. 2Results
On the same AR-L backbone, NTI-L reduces FID from 3.02 to 1.67, and NTI improves results across model scales and tokenizer designs.
Source: §1 · p. 2, §4.2 · p. 7Takeaways & Limitations
NTI improves visual AR generation without modifying the tokenizer, generation order, or inference pipeline.
Source: §1 · p. 2
Abstract
from arXiv · showhide
The space of possible visual content is enormous, and arguably unbounded. Any training set, however large, therefore covers only a tiny fraction of it. Yet generative models are typically trained to fit only the images they are given. In visual autoregressive (AR) generation, for example, each token is predicted from the preceding ground-truth (GT) tokens and supervised solely by the GT next token. Confining learning to observed data in this way limits how much of the visual distribution a model can capture. In this work, we hypothesize that letting a model imagine beyond its training data, by exploring plausible visual realizations it has never observed, can ease this limitation. We use visual AR generation as a proof of concept, since its fixed, discrete token space makes imagination easier to steer and constrain. We propose Next Token Imagination (NTI), a method that allows the model to explore beyond the GT next token, subject to a carefully designed criterion that keeps this imagination bounded by the underlying visual distribution. Extensive experiments on ImageNet show that NTI substantially and consistently improves the fidelity of generated images, approaching state-of-the-art diffusion and AR models while using a smaller model and lower training cost.
1 Introduction
Standard visual AR training supervises each position from a ground-truth prefix using only the ground-truth next token, although plausible non-ground-truth continuations can lead to valid visual realizations. NTI introduces context-regulated imagination and snapshot-based context approximation; ImageNet experiments report improved fidelity across model scales and tokenizer designs.
1 Introduction
- Standard visual AR training predicts each next token from the ground-truth prefix and supervises only the corresponding ground-truth token.At inference, sampled non-ground-truth tokens can accumulate, moving the generated sequence away from the ground-truth sequence.Source: §1 · p. 1
- A non-ground-truth token can lead to another plausible visual realization, but off-ground-truth contexts can also produce structural distortion.The examples distinguish plausible and implausible continuations, so broadening the next-token space alone is insufficient.Source: §1 · pp. 1–2
- NTI uses Context-Regulated Imagination to assign probability to plausible non-ground-truth tokens and Snapshot-based Context Approximation to construct model-generated contexts without sequential rollout.Together, these mechanisms move training beyond the observed ground-truth sequence while regulating imagination.Source: §1 · p. 2
- On ImageNet, NTI reduces FID from 3.02 to 1.67 and improves IS from 256.20 to 312.72, with consistent gains across model scales.It also outperforms representative AR methods while remaining competitive with substantially larger or more specialized generative models.Source: §1 · p. 2
2 Preliminaries
Visual autoregressive generation encodes images as discrete token sequences, models their distribution with a causal generator, and decodes generated sequences back into images. Standard training predicts each ground-truth next token from its ground-truth prefix, whereas inference conditions on previously generated tokens.
2 Preliminaries
- Visual autoregressive generation encodes images into latent sequences, models their distribution with an autoregressive generator, and decodes generated sequences into images.The pipeline has three stages: tokenization, sequence generation, and decoding.Source: §2 · p. 3
- A visual tokenizer maps an image to a latent feature grid, assigns each feature to its nearest codebook entry, and decodes the quantized grid to reconstruct the image.The codebook entries are discrete indices used to represent images as tokens.Source: §2 · p. 3
- The causal generator models a raster-order sequence of codebook indices, and standard training predicts each ground-truth next token from its ground-truth context.At inference, it instead conditions on the previously generated tokens, forming a generated context.Source: §2 · p. 3
3 Next Token Imagination
NTI broadens visual autoregressive supervision beyond the ground-truth next token while regulating imagined targets by context. Controlled experiments motivate this approach, and SCA and CRI implement it without autoregressive rollout during training.
3.2 Training with Next Token Imagination
- NTI extends supervision to plausible non-GT tokens and uses SCA to construct model-generated contexts in a single forward pass instead of autoregressive rollout.SCA uses a periodically refreshed frozen snapshot of the current generator to approximate model-generated contexts while preserving a single training forward pass.Source: §3 · p. 4, §3.2 · p. 5
3.1 Understanding Restricted Next Token Space
- Non-GT contexts can remain visually plausible, although lower GTR and GTS generally correspond to larger perceptual deviations in generated images.Facial features blur progressively in the cat examples; dog examples retain coherent leg structures under moderate deviations before severe distortions emerge.Source: §3.1 · p. 4
- Under the same generated context, non-GT tokens achieve comparable or higher mean DINO scores than GT tokens, showing GT is not the only plausible target.The comparison spans different GTR values; snow leopard examples show both choices can produce visually plausible images.Source: §3.1 · p. 5
- Context-adaptive selection performs better than fixed selection as GTR decreases, while both choices degrade further and call for imagination to contract toward the GT token.At high GTR, fixed and context-adaptive selection perform similarly.Source: §3.1 · p. 5
3.2 Training with Next Token Imagination
- CRI combines codebook similarity with prediction consistency across GT and model-generated contexts to select and weight non-GT supervision targets.It relies more on model-generated contexts when their predictions remain consistent, and shifts toward the GT context as predictions diverge.Source: §3.2 · pp. 5–6
- CRI controls total imagination probability using GT-context confidence, probability assigned beyond the GT token, and prediction consistency, reducing imagination when uncertainty or divergence rises.NTI applies its target distribution at positions with model-generated contexts while retaining the standard AR objective for direct learning from GT sequences.Source: §3.2 · p. 6
4 Experiments
ImageNet experiments show that NTI improves raster-order visual AR generation, generalizes across tokenizers, and adds modest training cost. Ablations identify context-regulated supervision and a bounded replacement rate as important design choices.
4.1 Experimental Setup
- Class-conditional ImageNet-1K generation at 256 × 256 follows the ADM evaluation protocol, generating 50,000 images with classifier-free guidance and reporting FID and IS.Models use a pretrained MaskGIT-VQGAN tokenizer and causal Transformer variants trained for 400 epochs; NTI maximum replacement probability rises to 0.3, then falls to zero by epoch 320.Source: §4.1 · p. 6
4.2 Main Results
- NTI reduces FID from 3.02 to 1.67 and improves IS from 256.20 to 312.72 on the same AR-L backbone, with consistent gains across model scales.With 461M parameters and 400 training epochs, NTI-L also outperforms reAR-L, RAR-L, LlamaGen-L, and BitDance-B-1x in FID while remaining competitive with substantially larger or more expensive models.Source: Table 1 · p. 7, Figure 5 · p. 7, §4.2 · p. 7
- NTI improves FID from 3.94 to 3.66 on GigaTok and from 1.61 to 1.51 on AliTok without changing their tokenizers or generators.The reported gains extend NTI to non-standard visual AR configurations.Source: Table 2 · p. 7, §4.2 · p. 7
- NTI adds about 10% training overhead, compared with online two-pass construction that nearly doubles cost, and leaves the inference pipeline unchanged.On 8 NVIDIA A100 GPUs, vanilla AR, NTI, and online two-pass construction take 11.62, 12.80, and 22.90 minutes per epoch, respectively.Source: §4.2 · pp. 7–8
4.3 Ablation Studies
- Under the same SCA contexts, CRI reduces FID from 2.11 to 2.04; context-independent supervision and SCA with GT supervision yield 2.10 and 2.11 FID.Online scheduled sampling with CRI achieves 2.03 FID, only 0.01 lower than NTI, while SCA avoids its additional prediction pass.Source: Table 3 · p. 8, §4.3 · p. 8
- Loss weights λ = 0.5 and 1.0 yield comparable generation quality, so the experiments use λ = 0.5.The authors suggest AdamW's scale-free property may partly explain the limited sensitivity.Source: §4.3 · p. 8
- Raising peak replacement probability from 0.10 to 0.30 improves FID from 2.09 to 2.04, while 0.40 degrades it to 2.09.The authors use a peak replacement probability of 0.30.Source: §4.3 · p. 8
- With 80 training epochs, τ = 0.05 achieves 3.55 FID compared with 3.63 for τ = 0.10, and is used in all experiments.A smaller τ emphasizes codebook candidates more similar to the GT token.Source: §4.3 · p. 8
5 Related Work
Prior visual AR work has modified tokenizers, representations, prediction units, and generation order, while exposure-bias methods alter training contexts. NTI instead changes next-token supervision to permit context-dependent plausible continuations beyond the GT token.
Visual autoregressive generation
- Visual AR research has advanced tokenizers, representations, and prediction formulations, whereas NTI leaves these components and generation order unchanged.Prior approaches include compact tokenizers (Yu et al., 2024b; Xiong et al., 2025), flexible sequence dependencies (Bachmann et al., 2025; Wu et al., 2026), randomized order (Yu et al., 2025a; Pang et al., 2025), and alternative prediction units (Li et al., 2024; Tian et al., 2024; Ren et al., 2025).Source: §5 · p. 8
Exposure bias
- Exposure-bias methods introduce predicted or noisy contexts but retain GT-token supervision, while NTI changes supervision under off-GT contexts.Examples include scheduled sampling (Bengio et al., 2015), stochastic sampling (Lee et al., 2022), Gumbel-softmax mixing, and noisy contexts.Source: §5 · pp. 8–9
Supervision beyond the GT token
- Unlike token-relationship methods such as DIST2Loss, IAR, and DCPE, NTI regulates non-GT supervision according to the current AR context.DIST2Loss (Chung et al., 2026), IAR (Hu et al., 2025), and DCPE (Tang et al., 2025) derive supervision from token relationships without adapting it to that context.Source: §5 · p. 9
6 Conclusion
The conclusion frames GT-only next-token supervision as restricting valid continuations and presents NTI, with CRI and SCA, as a way to explore plausible alternatives. It reports improvements across tokenizer and generator designs without changing the tokenizer, generation order, or inference pipeline.
6 Conclusion
- The conclusion presents NTI with CRI and SCA as a way to move beyond GT-only targets and improve visual AR across tokenizer and generator designs.The stated approach does not modify the tokenizer, generation order, or inference pipeline.Source: §6 · p. 9
AI use statement
- The authors report using generative AI tools for research ideation, experiments, code, literature retrieval, manuscript preparation, and writing revision.They state that they reviewed and revised AI-assisted content, checked references, tested code, and take responsibility for the final work.Source: AI use statement · p. 9
A Additional Experimental Details
- Training and evaluation code and NTI-S/B/L inference checkpoints are publicly available in the project repository.The repository includes model configurations and instructions for sampling and computing FID and IS.Source: A Additional Experimental Details · p. 12
A.1 Dataset, Backbone, and Optimization
The appendix specifies the ImageNet-1K evaluation, tokenizer and Transformer configurations, and optimization settings. Its snapshot-based context construction uses cached SCA records and a scheduled replacement rate during training.
A.1 Dataset, Backbone, and Optimization
- ImageNet-1K evaluation uses class-conditional 256 × 256 generation, 50,000 samples, and FID and IS; the standard setup uses a frozen MaskGIT-VQGAN tokenizer.The tokenizer produces a 16 × 16 grid of 256 tokens from a 1,024-entry codebook.Source: §A.1 · p. 12
- Standard raster-order models train for 250,000 steps, approximately 400 epochs, with batch size 2,048 and AdamW.The learning rate warms up to 4 × 10^-4 and then follows a cosine decay to 10^-5.Source: §A.1 · p. 12
A.2 Snapshot-Based Context Construction
- The maximum replacement probability follows a cosine rise-and-decay schedule, peaking at 0.3 and reaching zero at 80% of training.For each sequence, the replacement ratio is sampled up to the scheduled maximum; the CRI affinity temperature is 0.05 by default.Source: §A.2 · p. 12
- A frozen snapshot refreshes SCA contexts every 10,000 steps; cached SCA records and GT sequences are mixed and processed in one trainable-generator forward pass.The cached contexts and CRI targets remain fixed between refreshes, and NTI loss is normalized across valid NTI positions.Source: §A.2 · pp. 12–13
A.3 Sampling and Efficiency Evaluation
This section specifies sampling settings, checkpoint evaluation, throughput measurement, and training-cost accounting for the efficiency comparisons.
A.3 Sampling and Efficiency Evaluation
- NTI sampling uses a power-cosine classifier-free guidance schedule with model-specific (w, γ) settings.The reported settings are (22, 2.75) for NTI-S, (14.5, 2.25) for NTI-B, and (14, 1.75) for NTI-L.Source: §A.3 · p. 13
- Figure 5 checkpoints span 50K to 250K optimization steps, or approximately 80 to 400 epochs.The final checkpoint is averaged over sampling seeds 42, 142, and 242; earlier checkpoints use seed 42.Source: §A.3 · p. 13
- Sampling throughput is timed on one A100-SXM4-80GB GPU at batch size 128, including token generation and tokenizer decoding.Image-file writing is excluded; measurements use KV caching and FP32 inference without autocast or compilation.Source: §A.3 · p. 13
- Training-cost estimates include optimization and periodic SCA construction, while the online comparator constructs contexts in two passes at every optimization step.NTI context-construction cost uses a frozen snapshot and is averaged across epochs between refreshes.Source: §A.3 · p. 13
B Analysis Details of Next Token Imagination
This section gives additional details for the two analysis experiments in Section 3.1, covering their controlled comparisons, evaluation metrics, and interpretation under the main-text GTR and GTS definitions.
B Analysis Details of Next Token Imagination
- The appendix details two Section 3.1 analysis experiments, including controlled comparisons, evaluation metrics, and interpretation using the main-text GTR and GTS definitions.Source: B Analysis Details of Next Token… · p. 13
B.1 Feasibility of Next Token Imagination
The analyses test how context deviations affect later generations and whether plausible non-GT targets can preserve class semantics. Greater perturbation generally increases image and suffix deviation, while suitable non-GT targets depend on context.
B.1 Feasibility of Next Token Imagination
- Lower GTR generally reduces PSNR and SSIM while increasing paired DINO distance and ∆NLL; higher GTS under stronger perturbations generally preserves greater image similarity and limits the suffix-NLL increase.Figure 6 reports PSNR, SSIM, paired DINO distance, and fixed-GT-suffix ∆NLL across requested GTS levels and measured GTR levels.Source: §B.1 · p. 14, Figure 6 · p. 15
- Greater context deviations increasingly affect subsequent generation, but perturbed contexts can still produce coherent generations beyond the GT sequence.The analysis compares continuations from perturbed and corresponding GT contexts using the same source, target position, and sampling randomness.Source: §B.1 · pp. 14–15
B.2 Regulation of Next Token Imagination
- Non-GT continuations achieve comparable or higher Class-DINO scores than GT-token continuations under generated contexts, indicating that the GT token is not the only choice preserving target-class semantics.Class-DINO measures feature similarity to real images from the target class.Source: §B.2 · p. 15
- Suitable non-GT targets depend on the current context; fixed and context-adaptive selection are compared using Class-DINO scores.Context-adaptive selection combines codebook cosine similarity with the model prediction.Source: §B.2 · pp. 15–16
C Qualitative Evidence of Expanded Visual Space
Qualitative samples indicate that NTI-L produces broader visual realizations than AR-L. The discussion distinguishes NTI’s context-regulated next-token supervision from approaches that vary generation order or prediction procedures.
C Qualitative Evidence of Expanded Visual Space
- NTI-L shows greater variation than AR-L in volcano and guitar-pick morphology, composition, appearance, arrangement, viewpoint, and surrounding-object interactions.Examples include crater lakes, ash plumes, lava eruptions, snow-covered volcanoes, and close-up volcanic structures.Source: C Qualitative Evidence of Expanded… · p. 16
- NTI-L hippopotamuses and wallabies span more poses and viewpoints than AR-L hippopotamuses, which are often partially submerged in similar configurations.NTI-L examples include frontal, side-view, standing, walking, crouching, and close-up configurations.Source: C Qualitative Evidence of Expanded… · p. 16
- Together with controlled analyses, the qualitative results show that NTI broadens the plausible visual realizations produced by a standard visual AR model at inference.Source: C Qualitative Evidence of Expanded… · p. 16
D Additional Discussion on Related Work
- The discussion compares NTI with methods that improve visual AR through generation order, tokenization, alternative prediction formulations, or additional training supervision.NTI extends next-token supervision beyond the GT token and regulates non-GT supervision according to the current context.Source: D Additional Discussion on Related Work · p. 16
D.1 Randomized Generation Order
- RandAR, RAR, and ARPG vary token order or support randomized parallel generation, whereas NTI retains raster order and changes supervision under a given context.RandAR uses position instructions, RAR decreases order randomization during training, and ARPG separates positional queries from content representations; NTI instead allows context-dependent plausible non-GT targets.Source: §D.1 · p. 16
D.2 1D and Causality-Aligned Tokenization
NTI modifies generator supervision rather than tokenization or the autoregressive prediction formulation, and its gains on GigaTok and AliTok complement tokenizer improvements. It retains discrete raster-order next-token prediction while adapting supervision to context.
D.2 1D and Causality-Aligned Tokenization
- NTI-L generates greater variation in appearance, composition, pose, and viewpoint than AR-L across the same ImageNet classes.The figure highlights broader morphology and composition for volcanoes and guitar picks, and richer pose and viewpoint diversity for hippopotamuses and wallabies.Source: Figure 7 · p. 17
- FID falls from 3.94 to 3.66 on GigaTok and from 1.61 to 1.51 on AliTok without changing their tokenizers or generator architectures.The results indicate that tokenizer improvements and context-regulated supervision are complementary.Source: §D.2 · p. 18
- Tokenization approaches improve compactness, representation quality, or compatibility with causal generation, whereas NTI changes supervision for a given representation.TiTok, GigaTok, FlexTok, AliTok, CaTok, SpectralAR, and Hita exemplify these tokenization directions.Source: §D.2 · pp. 16–18
D.3 Alternative Autoregressive Formulations
- Unlike methods that alter prediction space, prediction unit, or autoregressive factorization, NTI retains discrete raster-order next-token prediction and regulates supervision over the existing vocabulary.The contrasted alternatives include continuous-token, next-scale, flexible-entity, and masked or parallel prediction formulations.Source: §D.3 · p. 18
D.4 Training Regularization and Supervision
- NTI uses codebook affinity to identify plausible non-GT candidates, then GT- and model-generated-context predictions regulate their weights and overall supervision strength.This differs from token-aware supervision based on token relationships and generated-context methods operating under different sequence objectives and prediction spaces.Source: §D.4 · pp. 18–19
Related papers
This paper cites
- Language Models are Few-Shot Learners2020
- The Unreasonable Effectiveness of Deep Features as a Perceptual Metric2018
- Classifier-Free Diffusion Guidance2022
- Diffusion Models Beat GANs on Image Synthesis2021
- Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks2015
- Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think2024