---
title: 'MaskGIT Sampler: Iterative Image Refinement'
url: https://www.emergentmind.com/topics/maskgit-sampler
type: topic
---

# MaskGIT Sampler: Iterative Image Refinement

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\) [2202.04200]. 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\) [2202.04200]. 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 [2510.04525].

## 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 [2202.04200]. The model is trained with a bidirectional transformer decoder that predicts randomly masked tokens by attending to tokens in all directions [2202.04200].

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 \([MASK]\) token. The paper samples a mask ratio via a schedule \(\gamma(r)\in(0,1]\), where \(r\in[0,1)\) is a random ratio, and roughly \(\lceil \gamma(r)N\rceil\) tokens are masked. The loss is the negative log-likelihood over masked positions:
\[
\mathcal{L}_{\text{mask}}
=
-\mathbb{E}_{Y \sim \mathcal{D}}
\Big[
\sum_{\forall i \in [1,N],\, m_i=1}
\log p(y_i \mid Y_{\text{masked}})
\Big].
\]
In words, the model is fed a partially masked token sequence and trained to predict the original tokens at masked positions using cross-entropy [2202.04200].

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 [2202.04200]. 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^{(0)}=[\text{MASK},\text{MASK},\ldots,\text{MASK}],
\]
and performs a non-autoregressive iterative refinement procedure [2202.04200]. Let \(Y^{(t)}\) denote the token sequence at iteration \(t\), with some positions masked and some already fixed.

At each iteration, the transformer is run on \(Y^{(t)}\) to obtain token distributions for every masked position,
\[
p^{(t)}\in\mathbb{R}^{N\times K}.
\]
For each masked position \(i\), the model gives a categorical distribution \(p_i^{(t)}\) over the \(K\) codebook entries, and a token is sampled:
\[
y_i^{(t)}\sim p_i^{(t)}.
\]
The sampled probability is used as a confidence score for masked positions, while confidence for unmasked positions is set to \(1.0\) [2202.04200]. The paper also notes that sampling uses temperature annealing to encourage diversity [2202.04200].

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 [2202.04200]. With \(T=8\), the model can go from all masked to fully generated in just 8 steps [2202.04200]. 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 \(\gamma(\cdot)\). The number to mask at iteration \(t+1\) is
\[
n=\left\lceil \gamma\!\left(\frac{t}{T}\right)N \right\rceil,
\]
where \(T\) is the total number of decoding iterations and \(N\) is the total number of latent positions [2202.04200]. The schedule is decreasing, with
\[
\gamma(0)\approx 1,\qquad \gamma(1)\approx 0.
\]

After sampling, MaskGIT keeps the most confident tokens and re-masks the rest. If \(c_i\) is the confidence score for token \(i\), then the new mask is defined by thresholding against the \(n\)-th sorted confidence:
\[
m_i^{(t+1)}=
\begin{cases}
1, & \text{if } c_i < \mathrm{sorted}_j(c_j)[n],\\
0, & \text{otherwise}.
\end{cases}
\]
Operationally, the algorithm sorts all confidences, keeps the top \(N-n\) most confident tokens, and masks the bottom \(n\) tokens for the next iteration [2202.04200]. The result is an iterative mechanism that locks in high-confidence predictions and revisits uncertain ones.

The schedule \(\gamma(r)\) is required to be continuous, in \([0,1]\), monotonically decreasing in \(r\), and to satisfy \(\gamma(0)\to 1\) and \(\gamma(1)\to 0\) [2202.04200]. 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 [2202.04200]. The reported ablation on ImageNet \(256\times256\) gives the following example results: cosine achieves FID \(6.06\), IS \(181.5\), \(T=10\); square gives FID \(6.35\), IS \(179.9\), \(T=10\); linear gives FID \(7.51\), IS \(113.2\), \(T=16\); and square root and logarithmic are worse [2202.04200]. The paper finds cosine works best overall and emphasizes that the schedule is crucial for quality [2202.04200].

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 [2202.04200]. 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_1\rightarrow y_2\rightarrow \cdots \rightarrow y_N.
\]
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 [2202.04200]. 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 [2202.04200].

The reported consequence is a substantial speedup. For a \(32\times 32\) token image, the paper describes “8 steps instead of 256,” and reports up to \(30\text{–}64\times\) faster decoding than VQGAN-style autoregressive decoding [2202.04200]. The abstract states that MaskGIT accelerates autoregressive decoding by up to \(64\times\) and significantly outperforms the state-of-the-art transformer model on the ImageNet dataset [2202.04200].

The practical sampling-related choices highlighted in the paper are tightly constrained. Good values for the number of iterations \(T\) are around \(8\text{–}12\); the example uses \(T=8\); and an ablation shows that too many iterations can hurt diversity and quality [2202.04200]. Cosine is the default and best-performing schedule in the reported experiments. The base sampler does not require top-\(k\), nucleus, or beam search by default, although classifier-based rejection sampling can improve score further and is not part of the base sampler [2202.04200].

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 [2202.04200]. 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 [2510.04525]. 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-\(k\) on the sampled token confidences:
\[
\begin{enumerate}[label=(MG\arabic*)]
\item Independently sample \(x_i\sim p_i\) and a standard Gumbel noise \(\xi_i\) for each \(i\in[N]\).
\item Choose \((i_1,\ldots,i_k)=\mathrm{argtop}_k{}_{i\in[N]}\{\log p_i(x_i)+\alpha\xi_i\}.
\item Return the indices \(i_1,\ldots,i_k\) and samples \(x_{i_1},\ldots,x_{i_k}\).
\end{enumerate}
\]
Here \(\alpha>0\) is the “Gumbel temperature” used in the official implementation [2510.04525]. The analysis therefore describes MaskGIT as fundamentally a sample-then-choose algorithm.

The central theoretical claim is that MaskGIT implicitly performs temperature sampling [2510.04525]. Using the Gumbel-top-\(k\) 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 \(\alpha^{-1}\log p_i(x_i)\). The paper introduces
\[
\beta:=1+\frac{1}{\alpha},
\]
and, under a large-\(N\) approximation, derives the approximate joint distribution
\[
P\{x_{i_\ell}, i_\ell^*=i_\ell \mid i_1^*=i_1,\ldots,i_{\ell-1}^*=i_{\ell-1}\}
\approx
\frac{p_{i_\ell}(x_{i_\ell})^\beta}
{\sum_{i\in[N]\setminus I_{\ell-1}}\lVert p_i\rVert_\beta^\beta}.
\]
The interpretation given in the paper is that MaskGIT implicitly samples tokens with exponent \(\beta\), producing a temperature-like sharpening effect [2510.04525].

Because \(\beta=1+1/\alpha\), finite \(\alpha\) 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 [2510.04525]. 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-\(N\) approximation, the later paper derives the moment sampler, presented as asymptotically equivalent to MaskGIT but more tractable and interpretable [2510.04525]. Its one-round form is
\[
\begin{enumerate}[label=(MM\arabic*)]
\item Let \((i_1,\ldots,i_k)=\mathrm{argtop}_k{}_{i\in[N]}\{\log\lVert p_i\rVert_\beta^\beta+\xi_i\}\) with i.i.d. standard Gumbel noise \((\xi_i)_{i=1}^N\).
\item Independently sample \(x_i\sim p_i^\beta/\lVert p_i\rVert_\beta^\beta\) for each \(i\in\{i_1,\ldots,i_k\}\).
\item Return the indices \(i_1,\ldots,i_k\) and samples \(x_{i_1},\ldots,x_{i_k}\).
\end{enumerate}
\]
The paper proves
\[
d_\mathrm{TV}(p_\mathrm{moment},p_\mathrm{MaskGIT})
\le
5\sqrt{\frac{k^2\lvert S\rvert^{1/\alpha}}{N}\left(1+\sqrt{\log^+\!\left(\frac{N}{k^2\lvert S\rvert^{1/\alpha}}\right)}\right)}.
\]
The stated interpretation is that, in the regime \(N\gg k^2\), 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 \(\lVert p_i\rVert_\beta^\beta\) rather than on sampled token realizations [2510.04525].

This analysis motivates a broader choose-then-sample formulation:
\[
\begin{enumerate}[label=(CTS\arabic*)]
\item Sample \(J\subset[D]\setminus I\) for where to unmask. Its distribution is denoted as \(\pi(\cdot|I,\bm{x}_I)\).
\item \(x_j\sim p_j\) with \(p_j\propto p_{j|I}(\cdot|\bm{x}_I)^\gamma\) for each \(j\in J\).
\end{enumerate}
\]
The order is first to decide the next positions \(J\), and then to sample their tokens [2510.04525]. The paper proves that if \(\lvert J\rvert=1\) always and \(p_{i|I}=q_{i|I}\), then the generated \(\bm{x}\sim q_\text{data}\), which it presents as an unbiasedness result for one-by-one CTS when the model marginals are exact [2510.04525].

Two implementation consequences follow in the later paper. First, CTS enables partial caching for bidirectional transformers. For a selected set \(I\) of positions to unmask, split it into disjoint parts \(I=A\cup B\), with \(A\cap B=\emptyset\). For \(i\in A\), tokens are sampled directly from \(p_{i|U}(\cdot|\bm{x}_U)\); for \(i\in B\), the updated conditional is approximated by running the transformer only on positions in \(I\), feeding sampled \(x_i\) for \(i\in A\), keeping \(M\) on \(i\in B\), and reusing cached key-value vectors for positions not in \(I\) [2510.04525]. The paper reports that the total attention cost is about \(1+\lvert I\rvert/D\) 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 [2510.04525].

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
\[
D_\mathrm{KL}(q\,\Vert\,p)
=
D_\mathrm{KL}(q_I\,\Vert\,\prod_{i\in I}q_i)
+
\mathbb{E}_{\bm{x}_I\sim q_I}
\left[
D_\mathrm{KL}(q_{I^c|I}(\cdot|\bm{x}_I)\,\Vert\, \prod_{i\notin I}q_{i|I}(\cdot|\bm{x}_I))
\right],
\]
and further bounds it via three terms interpreted as exploitation, spatial dispersion, and exploration [2510.04525]. The paper then proposes a hybrid method that merges an exploration-oriented ordering \(\bm{i}\) and an exploitation-oriented ordering \(\bm{j}\) by taking the first \(m\) indices from \(\bm{i}\) and the rest from \(\bm{j}\), 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 [2510.04525].

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 \(1.5\text{–}2\times\) in the language setting [2510.04525]. 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 [2510.04525].

## 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 [2202.04200]. 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 [2202.04200]. 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 [2510.04525]. 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 [2202.04200], 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 [2510.04525].

Source: https://www.emergentmind.com/topics/maskgit-sampler