PONDER: Learning to Generate in the Right Order
Abstract Overview
PONDER is a policy-learning framework for masked image generation that learns which patches to generate next and which earlier predictions to remove and regenerate later. The method treats generation order as a sequential decision problem and keeps the underlying generator frozen, training only a position-selection policy from terminal rewards on completed samples. The paper also introduces FIDgrad, an image-level reward derived from the first-order effect of a sample on Fréchet Inception Distance, so that a policy can be optimized toward distribution-level quality. Experiments cover visual Sudoku, class-conditional ImageNet generation, and text-to-image generation, showing that learned ordering and optional rethinking improve performance across these settings.
Novelty
The paper's main novelty is learning generation order and revision decisions for frozen continuous-token masked generators, rather than relying on fixed, random, or confidence-based schedules. A second distinctive contribution is FIDgrad, which converts a distribution-level FID objective into a per-sample reinforcement-learning reward for training the ordering policy.
Results
With frozen generators, learned ordering reduces MAR-B ImageNet FID from 4.09 to 2.78, and rethinking further improves it to 2.70 while also improving FDr-5 and Inception Score. On hard visual Sudoku, accuracy rises from 3.9% with random order and 38.7% with SRM greedy ordering to 56.6% with learned ordering and 96.7% with rethinking. For MaskGen-KL-XL text-to-image generation, PONDER improves HPSv2 from 24.50 to 26.50 at R=0 and to 27.80 at R=32, with gains also reported on ImageReward and GenEval.
Key Points
- PONDER learns sample-dependent patch ordering and can revisit earlier predictions through a controllable rethinking budget while leaving the generator weights fixed.
- FIDgrad provides an image-level reward based on the first-order change in FID, enabling reinforcement learning toward distribution quality without differentiating through the generator.
- Empirical gains are shown across three tasks—ImageNet generation, visual Sudoku, and text-to-image synthesis—and test-time increases in rethinking budget continue to improve performance with the same trained policy.