MaskGIT Sampler: Iterative Image Refinement

• MaskGIT Sampler is an iterative masked-token decoding approach that refines a compressed token grid using a bidirectional transformer.
• It employs a decreasing mask schedule and confidence-based re-masking to progressively improve token predictions over several iterations.
• The method achieves up to 64× faster decoding than autoregressive techniques and supports diverse tasks like inpainting and image extrapolation.

The MaskGIT sampler is an iterative masked-token decoding procedure for image synthesis in a discrete latent token space learned by a VQ-style tokenizer, where an image is compressed into a grid of tokens $Y=[y_i]_{i=1}^N$, each $y_i$ being a codebook index from a vocabulary of size $K$ (Chang et al., 2022). Rather than generating tokens left-to-right in raster-scan order, it uses a bidirectional transformer that predicts masked tokens in parallel and then refines them iteratively. In the original formulation, the model begins with all tokens masked, predicts all currently masked positions simultaneously, samples token values, estimates confidence, keeps the most confident predictions, re-masks the uncertain ones, and repeats for a fixed number of iterations $T$ (Chang et al., 2022). A later theoretical analysis characterizes this procedure as a sample-then-choose sampler whose Gumbel-top-$k$ selection induces an implicit temperature sampling effect, and derives related choose-then-sample alternatives (Hayakawa et al., 6 Oct 2025).

1. Discrete latent representation and masked-token modeling

MaskGIT samples in a discrete latent token space. An image is first compressed into a grid of tokens

$Y=[y_i]_{i=1}^N,$

where each $y_i$ is a codebook index from a vocabulary of size $K$, and $N$ is the number of latent positions (Chang et al., 2022). The model is trained with a bidirectional transformer decoder that predicts randomly masked tokens by attending to tokens in all directions (Chang et al., 2022).

For context, the training objective is Masked Visual Token Modeling. A binary mask

$m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$

indicates which tokens are replaced by a special $y_i$0 token. The paper samples a mask ratio via a schedule $y_i$1, where $y_i$2 is a random ratio, and roughly $y_i$3 tokens are masked. The loss is the negative log-likelihood over masked positions: $y_i$4 In words, the model is fed a partially masked token sequence and trained to predict the original tokens at masked positions using cross-entropy (Chang et al., 2022).

This training setup is directly relevant to sampling because the model is trained to fill in missing tokens from both left and right context, which enables parallel inference (Chang et al., 2022). A plausible implication is that the sampler is not merely a decoding heuristic layered on top of an autoregressive model; it is aligned with the model’s bidirectional masked-prediction objective.

2. Iterative decoding from an all-mask canvas

At inference time, MaskGIT starts from an all-mask canvas,

$y_i$5

and performs a non-autoregressive iterative refinement procedure (Chang et al., 2022). Let $y_i$6 denote the token sequence at iteration $y_i$7, with some positions masked and some already fixed.

At each iteration, the transformer is run on $y_i$8 to obtain token distributions for every masked position,

$y_i$9

For each masked position $K$0, the model gives a categorical distribution $K$1 over the $K$2 codebook entries, and a token is sampled: $K$3 The sampled probability is used as a confidence score for masked positions, while confidence for unmasked positions is set to $K$4 (Chang et al., 2022). The paper also notes that sampling uses temperature annealing to encourage diversity (Chang et al., 2022).

The overall transition is a gradual fill-and-refine loop: all positions are initially masked, the first prediction pass predicts all positions jointly from no visual content, only the most confident predictions are kept, and subsequent passes predict the remaining masked positions again conditioned on the fixed tokens (Chang et al., 2022). With $K$5, the model can go from all masked to fully generated in just 8 steps (Chang et al., 2022). The procedure is therefore described as a parallel refinement decoder rather than a left-to-right generator.

3. Confidence-based re-masking and schedule design

A defining feature of the MaskGIT sampler is that it does not simply fill masked positions once. After each prediction-and-sampling round, it decides how many tokens should remain masked in the next iteration through a masking schedule $K$6. The number to mask at iteration $K$7 is

$K$8

where $K$9 is the total number of decoding iterations and $T$0 is the total number of latent positions (Chang et al., 2022). The schedule is decreasing, with

$T$1

After sampling, MaskGIT keeps the most confident tokens and re-masks the rest. If $T$2 is the confidence score for token $T$3, then the new mask is defined by thresholding against the $T$4-th sorted confidence: $T$5 Operationally, the algorithm sorts all confidences, keeps the top $T$6 most confident tokens, and masks the bottom $T$7 tokens for the next iteration (Chang et al., 2022). The result is an iterative mechanism that locks in high-confidence predictions and revisits uncertain ones.

The schedule $T$8 is required to be continuous, in $T$9, monotonically decreasing in $k$0, and to satisfy $k$1 and $k$2 (Chang et al., 2022). The paper groups schedules into linear, concave, and convex families. Concave schedules include cosine, square, cubic, and exponential; convex schedules include square root and logarithmic (Chang et al., 2022). The reported ablation on ImageNet $k$3 gives the following example results: cosine achieves FID $k$4, IS $k$5, $k$6; square gives FID $k$7, IS $k$8, $k$9; linear gives FID $Y=[y_i]_{i=1}^N,$0, IS $Y=[y_i]_{i=1}^N,$1, $Y=[y_i]_{i=1}^N,$2; and square root and logarithmic are worse (Chang et al., 2022). The paper finds cosine works best overall and emphasizes that the schedule is crucial for quality (Chang et al., 2022).

The accompanying intuition is that concave schedules match image generation better: early iterations only need to get a few key tokens right, while later iterations refine many details (Chang et al., 2022). This suggests that schedule design is not only a convergence-control mechanism but also a structural prior on coarse-to-fine generation.

4. Relation to autoregressive decoding and practical operating regime

Autoregressive image transformers generate tokens in a fixed sequence, typically raster scan,

$Y=[y_i]_{i=1}^N,$3

The original MaskGIT paper identifies two drawbacks of this strategy: no parallelization across tokens and long sequence cost, since images have many more tokens than text (Chang et al., 2022). MaskGIT avoids this by predicting all tokens at once, refining only a small number of times, and using bidirectional attention so every token can condition on all currently known tokens (Chang et al., 2022).

The reported consequence is a substantial speedup. For a $Y=[y_i]_{i=1}^N,$4 token image, the paper describes “8 steps instead of 256,” and reports up to $Y=[y_i]_{i=1}^N,$5 faster decoding than VQGAN-style autoregressive decoding (Chang et al., 2022). The abstract states that MaskGIT accelerates autoregressive decoding by up to $Y=[y_i]_{i=1}^N,$6 and significantly outperforms the state-of-the-art transformer model on the ImageNet dataset (Chang et al., 2022).

The practical sampling-related choices highlighted in the paper are tightly constrained. Good values for the number of iterations $Y=[y_i]_{i=1}^N,$7 are around $Y=[y_i]_{i=1}^N,$8; the example uses $Y=[y_i]_{i=1}^N,$9; and an ablation shows that too many iterations can hurt diversity and quality (Chang et al., 2022). Cosine is the default and best-performing schedule in the reported experiments. The base sampler does not require top-$y_i$0, nucleus, or beam search by default, although classifier-based rejection sampling can improve score further and is not part of the base sampler (Chang et al., 2022).

Beyond unconditional generation, the original paper states that MaskGIT can be easily extended to various image editing tasks, such as inpainting, extrapolation, and image manipulation (Chang et al., 2022). A plausible implication is that the sampler’s masked-token formulation makes conditioning on partial observations a native operation rather than an afterthought.

5. Sample-then-choose analysis and implicit temperature sampling

A later analysis recasts MaskGIT as a post-hoc sampler for masked diffusion and masked image modeling (Hayakawa et al., 6 Oct 2025). In that treatment, one round of MaskGIT proceeds by first independently sampling a token at every masked position and then choosing the positions to keep using Gumbel-top-$y_i$1 on the sampled token confidences: $y_i$2 Here $y_i$3 is the “Gumbel temperature” used in the official implementation (Hayakawa et al., 6 Oct 2025). The analysis therefore describes MaskGIT as fundamentally a sample-then-choose algorithm.

The central theoretical claim is that MaskGIT implicitly performs temperature sampling (Hayakawa et al., 6 Oct 2025). Using the Gumbel-top-$y_i$4 trick, the position-selection step can be rewritten so that, conditional on sampled tokens, selection is equivalent to sampling without replacement with logits proportional to $y_i$5. The paper introduces

$y_i$6

and, under a large-$y_i$7 approximation, derives the approximate joint distribution

$y_i$8

The interpretation given in the paper is that MaskGIT implicitly samples tokens with exponent $y_i$9, producing a temperature-like sharpening effect (Hayakawa et al., 6 Oct 2025).

Because $K$0, finite $K$1 induces biased, sharper-than-original token sampling, and the paper argues that this helps explain why MaskGIT can degrade as the number of steps increases: with many rounds, this implicit temperature bias accumulates (Hayakawa et al., 6 Oct 2025). This reframes confidence-based unmasking as only part of the story; the induced token-distribution distortion is treated as a dominant factor in MaskGIT’s empirical behavior.

6. Choose-then-sample alternatives, partial caching, and adaptive unmasking

From the large-$K$2 approximation, the later paper derives the moment sampler, presented as asymptotically equivalent to MaskGIT but more tractable and interpretable (Hayakawa et al., 6 Oct 2025). Its one-round form is

$K$3

The paper proves

$K$4

The stated interpretation is that, in the regime $K$5, the moment sampler is asymptotically close to MaskGIT, makes the hidden temperature effect explicit, and is more interpretable because position choice depends on a moment quantity $K$6 rather than on sampled token realizations (Hayakawa et al., 6 Oct 2025).

This analysis motivates a broader choose-then-sample formulation: $K$7 The order is first to decide the next positions $K$8, and then to sample their tokens (Hayakawa et al., 6 Oct 2025). The paper proves that if $K$9 always and $N$0, then the generated $N$1, which it presents as an unbiasedness result for one-by-one CTS when the model marginals are exact (Hayakawa et al., 6 Oct 2025).

Two implementation consequences follow in the later paper. First, CTS enables partial caching for bidirectional transformers. For a selected set $N$2 of positions to unmask, split it into disjoint parts $N$3, with $N$4. For $N$5, tokens are sampled directly from $N$6; for $N$7, the updated conditional is approximated by running the transformer only on positions in $N$8, feeding sampled $N$9 for $m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$0, keeping $m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$1 on $m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$2, and reusing cached key-value vectors for positions not in $m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$3 (Hayakawa et al., 6 Oct 2025). The paper reports that the total attention cost is about $m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$4 times the original full computation, that this approximates an “intermediate” sampling step and effectively gives longer trajectories with only moderate extra cost, and that the empirical gains depend on hardware: it helps on A6000 but can be less beneficial on faster H100 GPUs where overheads dominate (Hayakawa et al., 6 Oct 2025).

Second, the paper formalizes adaptive selection of unmasking positions as a trade-off between exploitation and exploration. In a two-round CTS setting it decomposes the approximation error into

$m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$5

and further bounds it via three terms interpreted as exploitation, spatial dispersion, and exploration (Hayakawa et al., 6 Oct 2025). The paper then proposes a hybrid method that merges an exploration-oriented ordering $m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$6 and an exploitation-oriented ordering $m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$7 by taking the first $m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$8 indices from $m=[m_i]_{i=1}^N,\qquad m_i\in\{0,1\}$9 and the rest from $y_i$00, with Halton used for exploration, moment-based ordering for exploitation, and a merge ratio scheduled over time so the sampler shifts from exploration early to exploitation later (Hayakawa et al., 6 Oct 2025).

The empirical findings reported in that work are that Moment tracks MaskGIT closely, Temp almost matches MaskGIT performance, Random and Halton provide useful baselines, Moment+Cache yields some performance boost at a given latency, U-Moment performs well as an unbiased index-selection method in language generation, Hybrid combines complementary strengths of Halton and U-Moment, Hybrid+Cache improves both trade-off and runtime, and overall speedup is roughly $y_i$01 in the language setting (Hayakawa et al., 6 Oct 2025). The paper also reports that lower-temperature methods may improve perplexity but collapse entropy, which it uses to motivate the exploration-exploitation tension and the unbiased or hybrid approach (Hayakawa et al., 6 Oct 2025).

7. Conceptual interpretation and common points of confusion

In the original formulation, the MaskGIT sampler is best understood as a coarse-to-fine global refinement process: early passes decide global structure, later passes fill in details, uncertainty is revisited rather than committed immediately, and the model can use full bidirectional context at every iteration (Chang et al., 2022). This differs from the common assumption that non-autoregressive generation simply means “predict everything once.” In MaskGIT, parallel prediction is coupled to iterative confidence-based correction.

A second common point of confusion concerns what is being sampled. The sampler does not operate directly in pixel space. It outputs a fully filled token grid, which is then decoded back to pixels via the VQ decoder (Chang et al., 2022). Thus, both the speed and the refinement behavior arise from inference in a compressed discrete latent space.

A third issue is whether MaskGIT’s effectiveness should be attributed only to confidence-based ordering. The later theoretical analysis argues that MaskGIT’s performance is driven largely by its implicit temperature sampling, not only by confidence-based ordering (Hayakawa et al., 6 Oct 2025). This does not negate the operational role of confidence-based re-masking in the original algorithm; rather, it suggests that the sampler’s behavior is jointly shaped by ordering, token sharpening, and schedule design.

Taken together, the two papers establish the MaskGIT sampler as an iterative masked-token decoder defined by all-mask initialization, parallel prediction, confidence-based token retention, and a decreasing mask schedule (Chang et al., 2022), while also showing that this procedure admits a more explicit probabilistic interpretation through sample-then-choose and choose-then-sample formulations, implicit temperature sampling, and asymptotically equivalent alternatives such as the moment sampler (Hayakawa et al., 6 Oct 2025).