# Improved Masked Image Generation with Token-Critic

arXiv: https://arxiv.org/abs/2209.04439 · HTML: https://mlsanity.com/paper/2209.04439/

## Key facts

- Cited by 64 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 41 references that are in the map.
- First posted on arXiv 9 Sep 2022.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Analyzing and Improving the Image Quality of StyleGAN](https://mlsanity.com/paper/1912.04958/) (2019)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://mlsanity.com/paper/1809.11096/) (2018)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Cascaded Diffusion Models for High Fidelity Image Generation](https://mlsanity.com/paper/2106.15282/) (2021)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Improved Precision and Recall Metric for Assessing Generative Models](https://mlsanity.com/paper/1904.06991/) (2019)
- [Vector Quantized Diffusion Model for Text-to-Image Synthesis](https://mlsanity.com/paper/2111.14822/) (2021)
- [Vector-quantized Image Modeling with Improved VQGAN](https://mlsanity.com/paper/2110.04627/) (2021)
- [NVAE: A Deep Hierarchical Variational Autoencoder](https://mlsanity.com/paper/2007.03898/) (2020)
- [Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions](https://mlsanity.com/paper/2102.05379/) (2021)
- [StyleGAN-XL: Scaling StyleGAN to Large Diverse Datasets](https://mlsanity.com/paper/2202.00273/) (2022)

## Built on by (most cited first, 30 of 64)

- [Muse: Text-To-Image Generation via Masked Generative Transformers](https://mlsanity.com/paper/2301.00704/) (2023)
- [Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation](https://mlsanity.com/paper/2310.05737/) (2023)
- [MAGVIT: Masked Generative Video Transformer](https://mlsanity.com/paper/2212.05199/) (2022)
- [An Image is Worth 32 Tokens for Reconstruction and Generation](https://mlsanity.com/paper/2406.07550/) (2024)
- [NaturalSpeech 3: Zero-Shot Speech Synthesis with Factorized Codec and Diffusion Models](https://mlsanity.com/paper/2403.03100/) (2024)
- [MaskGCT: Zero-Shot Text-to-Speech with Masked Generative Codec Transformer](https://mlsanity.com/paper/2409.00750/) (2024)
- [A Reparameterized Discrete Diffusion Model for Text Generation](https://mlsanity.com/paper/2302.05737/) (2023)
- [Copilot4D: Learning Unsupervised World Models for Autonomous Driving via Discrete Diffusion](https://mlsanity.com/paper/2311.01017/) (2023)
- [MaskBit: Embedding-free Image Generation via Bit Tokens](https://mlsanity.com/paper/2409.16211/) (2024)
- [Visual Prompt Tuning for Generative Transfer Learning](https://mlsanity.com/paper/2210.00990/) (2022)
- [Zero-shot spatial layout conditioning for text-to-image diffusion models](https://mlsanity.com/paper/2306.13754/) (2023)
- [Informed Correctors for Discrete Diffusion Models](https://mlsanity.com/paper/2407.21243/) (2024)
- [Vector Quantized Time Series Generation with a Bidirectional Prior Model](https://mlsanity.com/paper/2303.04743/) (2023)
- [Masked Audio Generation using a Single Non-Autoregressive Transformer](https://mlsanity.com/paper/2401.04577/) (2024)
- [Fine-Tuning Masked Diffusion for Provable Self-Correction](https://mlsanity.com/paper/2510.01384/) (2025)
- [Image Compression with Product Quantized Masked Image Modeling](https://mlsanity.com/paper/2212.07372/) (2022)
- [Halton Scheduler For Masked Generative Image Transformer](https://mlsanity.com/paper/2503.17076/) (2025)
- [Masked Image Modeling: A Survey](https://mlsanity.com/paper/2408.06687/) (2024)
- [STAR: Scale-wise Text-conditioned AutoRegressive image generation](https://mlsanity.com/paper/2406.10797/) (2024)
- [MaskSketch: Unpaired Structure-guided Masked Image Generation](https://mlsanity.com/paper/2302.05496/) (2023)
- [Next Block Prediction: Video Generation via Semi-Autoregressive Modeling](https://mlsanity.com/paper/2502.07737/) (2025)
- [A Pytorch Reproduction of Masked Generative Image Transformer](https://mlsanity.com/paper/2310.14400/) (2023)
- [AnyEnhance: A Unified Generative Model with Prompt-Guidance and Self-Critic for Voice Enhancement](https://mlsanity.com/paper/2501.15417/) (2025)
- [Revisiting Non-Autoregressive Transformers for Efficient Image Synthesis](https://mlsanity.com/paper/2406.05478/) (2024)
- [σ-GPTs: A New Approach to Autoregressive Models](https://mlsanity.com/paper/2404.09562/) (2024)
- [Iterative Predictor-Critic Code Decoding for Real-World Image Dehazing](https://mlsanity.com/paper/2503.13147/) (2025)
- [AdaNAT: Exploring Adaptive Policy for Token-Based Image Generation](https://mlsanity.com/paper/2409.00342/) (2024)
- [TaDiCodec: Text-aware Diffusion Speech Tokenizer for Speech Language Modeling](https://mlsanity.com/paper/2508.16790/) (2025)
- [StraIT: Non-autoregressive Generation with Stratified Image Transformer](https://mlsanity.com/paper/2303.00750/) (2023)
- [ECHOPulse: ECG controlled echocardio-grams video generation](https://mlsanity.com/paper/2410.03143/) (2024)

## Abstract (arXiv)

Non-autoregressive generative transformers recently demonstrated impressive image generation performance, and orders of magnitude faster sampling than their autoregressive counterparts. However, optimal parallel sampling from the true joint distribution of visual tokens remains an open challenge. In this paper we introduce Token-Critic, an auxiliary model to guide the sampling of a non-autoregressive generative transformer. Given a masked-and-reconstructed real image, the Token-Critic model is trained to distinguish which visual tokens belong to the original image and which were sampled by the generative transformer. During non-autoregressive iterative sampling, Token-Critic is used to select which tokens to accept and which to reject and resample. Coupled with Token-Critic, a state-of-the-art generative transformer significantly improves its performance, and outperforms recent diffusion models and GANs in terms of the trade-off between generated image quality and diversity, in the challenging class-conditional ImageNet generation.