We introduce PONDER, a method that learns in which order an image generator should fill in content and which of its earlier predictions to revise. Masked generative models construct a sample by repeatedly filling in subsets of missing variables. The order of these partial predictions determines the context available in subsequent steps and therefore affects the quality of the final sample. Common strategies are to choose the order randomly or based on local confidence estimates. Conditioned on the partially generated sample and the generator's features, our policy chooses which positions to fill next and/or which previously filled positions to mask for later regeneration, a mechanism we call rethinking. In this way, the generation process is able to adapt to the evolving context and also to revise earlier predictions. We optimize the policy using group-relative policy gradients, with rewards that measure constraint satisfaction, human preference, or distribution quality. For the latter, we introduce FIDgrad, an image-level reward derived from the (first-order) change in Fréchet Inception Distance caused by adding a generated sample to the current feature distribution. With the generators frozen, learned ordering reduces MAR-B ImageNet FID from 4.09 to 2.78 and raises the accuracy on hard visual Sudokus from 3.9% to 56.6%, surpassing uncertainty-based greedy ordering (38.7%). Rethinking further improves these results to 2.70 FID and 96.7% accuracy, while increasing the rethinking budget of a fixed policy yields continued improvements at inference time.
Figures & tables
Figure 1: Learning to generate and rethink. Left: a learned policy selects positions of patches to generate or mask for later regeneration, enabling revision of earlier predictions. Right: increasing the inference-time rethinking budget improves Sudoku-hard accuracy and MAR-B image quality, with each policy and generator held fixed.
Figure 2: Overview of Ponder . The policy processes frozen backbone features using transformer blocks and a linear head, then samples Kt positions without replacement. Selected masked positions are filled, while rethinking masks selected visible positions for later regeneration. A terminal reward trains only the policy through group-relative policy gradients, while the generator remains frozen.
Figure 3: An example of sequential generation with rethinking on a visual Sudoku instance. We visualize next patch log probabilities as background color and policy decisions as border colors.
Figure 4: Illustration of the mean and covariance component of the FIDgrad reward for a new image with features z . FID compares Gaussian approximations of generated and reference image distributions in Inception feature space ( Heusel et al., 2017b ) . Our reward measures the first order change in FID induced by adding the new image into the pool of generated images. a) The new sample z shifts the mean of the generated distribution towards the reference distribution. This contributes to a positive reward. b) The new sample increases the variance of the generated distribution, although the variance is already higher than the reference distribution. This contributes negatively to the reward. In general a sample affects both: mean and covariance.
Table 5Table 6
Figure 7: We visualize the successive generation of a MAR-B image. In the bottom row we visualize patches before they have been remasked and after after the same patch has been re-generated.
Table 8Table 9
MaskGen-KL-L
MAR-Sudoku ( R=9 )
HPSv2 ↑
Accuracy (%) ↑
Policy input
R=0
R=4
Overall
Hard
Medium
Easy
Generator features + noise
26.11
26.61
96.38
95.27
96.67
97.21
Generator features only
25.01
25.71
96.49
95.72
96.78
96.97
Latent state only
25.05
25.42
59.19
22.11
59.82
95.64
Table 5: Policy-input ablation with frozen MaskGen-KL-L and MAR-Sudoku. The benefit of observing generation noise depends on the generator and task: it helps MaskGen in this comparison, but is not required for strong Sudoku performance. Generator features remain useful in both settings.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Parameter
MAR-B
MaskGen-KL-XL
MAR-Sudoku
Architecture
Number of positions
256
32
81
Generator feature width
768
1280
256
Noise width
16
16
32
Input projection
784→768
1296→1024
288→256
Transformer depth
6
4
2
Appendix
Table 6: Dynamic-removal policy architectures and optimization settings.
Parameter
MAR-B
MaskGen-KL-XL
MAR-Sudoku
Architecture
Tokenizer
KL16 VAE
TA-TiTok-BL32 VAE
MNIST VAE
Token layout
16×16 grid
32-token sequence
9×9 grid
Number of tokens
256
32
81
Token dimension
16
16
32
Backbone
encoder–decoder
MM-DiT decoder
encoder–decoder
Appendix
Table 7: Generator architectures and dynamic-removal inference settings.
R
Overall (%) ↑
Hard (%) ↑
Medium (%) ↑
Easy (%) ↑
0
80.17
56.60
87.07
96.85
1
89.68
75.50
94.22
99.33
2
93.62
84.90
96.50
99.46
4
96.53
92.52
97.53
99.53
8
98.07
96.39
98.30
99.52
9
98.21
96.68
98.39
99.57
Appendix
Table 8: Test-time scaling of the Ponder policy trained with R=9 on MNIST Sudoku. Overall and hard-split exact-solution accuracy are measured on 30,000 held-out boards balanced across difficulty levels.
Position policy
R
Decisions
FID ↓
FDr-5 ↓
IS ↑
Random
0
256
4.09
13.69
181.22
Random
8
272
3.97
13.42
182.22
Random
32
320
3.80
12.82
185.90
Ponder
0
256
2.78
13.28
206.55
Ponder
1
258
2.78
13.22
206.13
Ponder
2
260
2.78
13.18
207.06
Appendix
Table 9: Class-conditional ImageNet results with the frozen MAR-B generator. The FIDgrad policy is trained with a rethinking budget of R=8 and evaluated at every listed test-time budget R . We generate 50,000 images (50 per class) with DDIM using 25 denoising steps and guidance scale 1.
Figure 10: MaskGen-KL-XL examples. Random ordering at R=0 (left) versus the HPSv2 policy at R=4 (right), using the same prompts and DDIM-25 sampler.
Figure 11: Additional MaskGen-KL-XL examples. Random ordering at R=0 (left) versus the HPSv2 policy at R=4 (right), using the same prompts and DDIM-25 sampler.
Figure 12: MAR-B ImageNet qualitative examples. Random-policy and Ponder samples with a rethinking budget of R=8 . Each class is shown as a random/ Ponder pair.
Figure 13: MAR-B rethinking trajectory. Each upper sequence shows successive partial canvases and the final image; the lower strip enlarges patches before removal and after regeneration. Colored boxes identify the revised locations.