BIGFix: Bidirectional Image Generation with Token Fixing

TL;DR

BIGFix introduces a self-correcting, bidirectional token-fixing framework for fast image and video generation by combining multi-token prediction with training-time random token injection. The method enables backtracking during sampling and uses Halton sequence-based token ordering with arccos scheduling to preserve 2D spatial structure while accelerating generation. Empirical results on ImageNet-256, CIFAR-10, UCF-101, and NuScenes show significant quality gains (e.g., improved FID/IS and FVD) and faster sampling (32 steps) compared with prior MaskGIT-like and autoregressive approaches. The approach offers a practical, scalable path to high-fidelity visual synthesis with strong correction capabilities and competitive speed, highlighting the value of robustness to perturbations in token contexts for multi-token generation.

Abstract

Recent advances in image and video generation have raised significant interest from both academia and industry. A key challenge in this field is improving inference efficiency, as model size and the number of inference steps directly impact the commercial viability of generative models while also posing fundamental scientific challenges. A promising direction involves combining auto-regressive sequential token modeling with multi-token prediction per step, reducing inference time by up to an order of magnitude. However, predicting multiple tokens in parallel can introduce structural inconsistencies due to token incompatibilities, as capturing complex joint dependencies during training remains challenging. Traditionally, once tokens are sampled, there is no mechanism to backtrack and refine erroneous predictions. We propose a method for self-correcting image generation by iteratively refining sampled tokens. We achieve this with a novel training scheme that injects random tokens in the context, improving robustness and enabling token fixing during sampling. Our method preserves the efficiency benefits of parallel token prediction while significantly enhancing generation quality. We evaluate our approach on image generation using the ImageNet-256 and CIFAR-10 datasets, as well as on video generation with UCF-101 and NuScenes, demonstrating substantial improvements across both modalities.

Figures (13)

• Figure 1: Self-correction: After the first 6 unmasking steps to create the overall structure, we proceed with two generations: without correction (top) and ours with correction (bottom). Correction solves structural errors such as supernumerary or missing elements.
• Figure 2: Schematics. During training (top, in blue), we perturb the input tokens by replacing a fraction $\alpha$ of the unmasked context tokens with randomly sampled ground-truth tokens from the same image. The random tokens originate from the correct image distribution but are arranged in an incorrect position. This augmentation not only enhances the model’s robustness to such errors during sampling but also enables the identification and correction of incompatible tokens during inference (bottom, in orange). While baseline methods must assume the input tokens are correct and unchangeable, BIGFix allows to refine its predictions iteratively.
• Figure 3: One-step reconstruction on ImageNet 256$\times$256 validation set. Evaluation of model accuracy and fidelity under corrupted token contexts.
• Figure 4: Random token correction: Top rows: random token injected at different values of $\alpha$, with no other tokens masked. Bottom rows: the model's prediction in a single step. The figure illustrates the model's ability to accurately correct tokens even when $\alpha$ exceeds the training range (i.e., $\textcolor{red}{\alpha > 0.2}$). Better visible when zoomed in.
• Figure 5: Patch tokens correction: Tokens from the green dotted square are copied and pasted into the solid green square, creating a corrupted image (second row). Only the tokens in the solid green square are altered, while all other tokens remain unchanged. BIGFix is able to correct these tokens in a single step (third row), demonstrating its ability to perform correction under out-of-distribution latent manipulations. Better visible when zoomed in.