PaperlayerView exact revision on arXiv

Source-linked AI summary

PONDER: Learning to Generate in the Right Order

Julius Erbach, Dominik Narnhofer, Bernt Schiele, Jan Eric Lenssen, Konrad Schindler

arXiv:2610.23531v1cs.CV

TL;DR

Generation order affects the context available to later predictions, yet common fixed, random, and locally greedy schedules do not learn how to make sample-dependent decisions or revise prior predictions. PONDER learns ordering and rethinking policies from completed-sample rewards while keeping generators frozen. Across generation and visual constraint solving, learned ordering improves reported metrics, and rethinking yields further gains.

  • Problem

    Generation order affects final sample quality, but fixed, random, or locally greedy schedules do not learn sample-dependent decisions that account for later consequences.

    Source: §1 · p. 2

  • Method

    PONDER trains a policy to choose positions to fill or previously filled positions to remove for regeneration, using completed-sample rewards with the generator frozen.

    Source: §1 · p. 2

  • Results

    2.78 versus 4.09 MAR-B ImageNet FID and 56.6% versus 3.9% hard-Sudoku accuracy with learned ordering improve to 2.70 FID and 96.7% accuracy with rethinking.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Across visual constraint solving, class-conditional generation, and text-to-image synthesis, PONDER improves observed metrics over the generator’s default sampling schedule without changing the generator.

    Source: §5 · p. 9

    Policy behavior depends strongly on the training reward, and rethinking uses a prescribed budget rather than allocating revisions on demand.

    Source: §5 · p. 10

Abstract

from arXiv · show

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\'echet 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.

1 INTRODUCTION

Generation order affects the context available to later predictions and therefore the final sample, but fixed, random, and locally greedy schedules leave important ordering decisions unresolved. PONDER learns position selection and revision from completed-sample rewards, improving generation and visual constraint solving with frozen generators.

1 INTRODUCTION

  • Because learned conditionals break the equivalence of different factorizations, generation order substantially affects sample quality and may need to adapt per sample.The central question is how to learn an order that yields the best complete sample.Source: §1 · p. 2
  • Fixed schedules ignore generated content, random schedules do not exploit context at inference, and greedy schedules rely on a one-step proxy.Greedy confidence or difficulty estimates can favor locally easy positions without accounting for how the decision affects later steps.Source: §1 · p. 2
  • Intermediate choices have delayed effects, and predictions made with limited context may become inconsistent as generation proceeds.This motivates treating generation as sequential decision-making and revisiting earlier predictions when more context is available.Source: §1 · p. 2
  • PONDER learns a policy to order generation and remove earlier predictions for subsequent regeneration.The policy selects where a frozen generator fills content next using final-sample rewards; rethinking decisions can be made with a controllable budget.Source: §1 · p. 2
  • FIDgrad converts a distribution-level objective into an image-level reward by estimating each sample’s first-order effect on Fréchet Inception Distance.The contribution supports training a generation policy to target distribution quality.Source: §1 · p. 2
  • 2.78 MAR-B ImageNet FID and 56.6% hard-Sudoku accuracy with learned ordering improve to 2.70 FID and 96.7% accuracy with rethinking.Learned ordering also raises MaskGen-KL-XL HPSv2 from 24.50 to 26.50; performance continues to improve as the inference-time rethinking budget increases while policy and generator remain fixed.Source: §1 · p. 3

2 RELATED WORK

Prior work studies masked generation, adaptive ordering, revision, reward optimization, and distributional objectives. PONDER applies learned ordering and revision to continuous-latent diffusion generators while keeping the generator frozen.

Masked generation with continuous representations

  • Masked generators range from discrete-token MaskGIT (Chang et al., 2022) to continuous-patch MAR (Li et al., 2024), and PONDER builds on continuous masked generators with frozen weights.FLUID (Fan et al., 2025) studies representation and ordering effects in text-to-image synthesis; MaskGen (Kim et al., 2025) supports both discrete and continuous representations.Source: §2 · p. 3

Adaptive and learned generation orders; Revision during generation

  • PONDER learns ordering and revision for continuous-latent diffusion generators, whereas the described ordering methods for discrete tokens only reveal masked positions and keep revealed tokens fixed.PONDER also uses a prescribed removal budget, separating the decision to discard content from its later regeneration.Source: §2 · p. 3

Reward optimization for generative models

  • PONDER optimizes position-selection actions with rewards on completed images, rather than training the generator parameters as in DDPO (Black et al., 2024).Its policy selects positions to fill or remove, so training does not require likelihoods for generated continuous content.Source: §2 · p. 3

Distributional objectives

  • FIDgrad uses evolving population statistics to estimate a single sample’s first-order effect on FID and trains the policy without differentiating through the generator or feature extractor.This differs from FD-loss (Yang et al., 2026), which backpropagates through the current generated batch while maintaining a feature population.Source: §2 · p. 4

3 METHOD

PONDER controls the order in which a masked generator fills patches and can also select earlier patches for regeneration. Its policy is trained against rewards on completed samples while the generator remains frozen.

3 METHOD

  • PONDER learns a policy for ordering patch fills and optionally re-masking patches for regeneration, optimizing a task-specific target such as FID or human preference.The generator starts from an empty canvas and repeatedly fills masked regions.Source: §3 · p. 4

3.1 POLICY-GUIDED MASKED GENERATION

  • Each iteration evaluates the backbone on the current canvas, selects patches to unmask, and fills them; a prescribed schedule sets the number of patches while the policy chooses their locations.Visible patches provide context, and masked patches are represented by MASK tokens.Source: §3.1 · p. 4

3.2 ADAPTIVE PATCH SELECTION

  • A bidirectional-transformer policy maps generator features and start-noise vectors to logits, then samples patches sequentially without replacement.Within an iteration, selected patches are filled using the same context from the preceding iteration.Source: §3.2 · p. 5

3.3 PONDERING BY REVISITING PATCHES

  • The policy can select visible patches to mask; they remain masked until selected again and are regenerated with the later context.Unlike repeated resampling under the same context, rethinking can occur after other canvas regions have been filled.Source: §3.3 · pp. 5–6

3.4 GROUP-RELATIVE POLICY OPTIMIZATION

  • Training uses terminal rewards on completed samples and compares independently sampled policy trajectories that share the same initial canvas and generator noise.Only the ordering policy is updated; the generator remains frozen.Source: §3.4 · p. 6

3.5 DISTRIBUTION-LEVEL OPTIMIZATION WITH FIDGRAD

  • FIDgrad scores a generated image by its first-order effect on the current distribution’s FID, enabling distribution-quality rewards for policy training with a frozen generator.Its scores are terminal rewards, and the generated feature moments are tracked during training.Source: §3.5 · pp. 6–7

4 EXPERIMENTS

PONDER is evaluated on visual Sudoku, class-conditional ImageNet, and text-to-image tasks using frozen generators. Learned ordering improves reported task metrics, and rethinking further improves them as its inference-time budget increases.

ImageNet (MAR)

  • FIDgrad rewards a new image according to its first-order effect on FID, combining contributions from changes to feature mean and covariance.In the illustrated case, shifting the generated mean toward the reference contributes positively, while increasing variance that already exceeds the reference contributes negatively.Source: Figure 4 · p. 7

Shared experimental setup

  • Experiments cover MAR-B on ImageNet, MAR on Sudoku, and MaskGen-KL-XL on text-to-image synthesis, with all generators frozen.Policy training uses DDIM with 25 steps for MaskGen-KL-XL and MAR-B, and flow matching with 12 steps for Sudoku.Source: §4 · p. 7

Visual Sudoku

  • The Sudoku evaluation classifies each cell and checks row, column, and box constraints on generated handwritten MNIST puzzles.The frozen MAR generator uses a flow-matching head over per-cell VAE latents, and one policy is trained across difficulty levels with a binary valid-solution reward.Source: §4 · p. 7
  • 96.7% hard-split Sudoku accuracy is reached with learned ordering and rethinking at R = 9, versus 56.6% with learned ordering alone.Learned ordering alone exceeds uncertainty-based greedy ordering at 38.7%; the policy trained at R = 9 reaches 98.6% at test-time R = 128.Source: §4 · pp. 7–8, Figure 5 · p. 8, Table 1 · p. 8

ImageNet (MAR)

  • 2.70 FID, 12.47 FDr-5, and 212.60 IS are reached on ImageNet with MAR-B at rethinking budget R = 32.Without rethinking, PONDER reduces FID from 4.09 to 2.78 and increases IS from 181.22 to 206.55 versus random ordering; FDr-5 improves from 13.69 to 13.28.Source: §4 · p. 8, Table 2 · p. 8, Figure 6 · p. 8

Text-to-image generation

  • At R = 32, MaskGen-KL-XL reaches 27.80 HPSv2 and 0.81 ImageReward; GenEval accuracy rises from 37.3% with random ordering to 41.1% with learned ordering at R = 0.At R = 4, HPSv2 is 26.95 and ImageReward is 0.73; GenEval accuracy is 42.6% for the learned policy versus 38.0% for random ordering.Source: §4 · pp. 8–9, Table 3 · p. 9, Table 4 · p. 9

5 CONCLUSION AND DISCUSSION

PONDER makes generation order a learned component for frozen masked generative models, with a policy that can also revise earlier predictions. Across three visual tasks, it improves observed metrics over default sampling, while reward dependence and fixed rethinking budgets remain limitations.

5 CONCLUSION AND DISCUSSION

  • PONDER learns the generation order of frozen masked generative models and can remove earlier predictions for regeneration as context changes.Across visual constraint solving, class-conditional generation, and text-to-image synthesis, it improves observed metrics over default sampling while leaving generators untouched.Source: §5 · p. 9
  • The policy adds a forward pass per generation iteration, a minor overhead relative to the generative backbone and scheduled denoising steps.Its behavior depends strongly on the training reward, and rethinking uses a prescribed budget rather than allocating revisions on demand.Source: §5 · p. 10

AI USE STATEMENT

The authors report using generative AI tools to assist with code, methodology feedback, literature identification, mathematical derivations, and manuscript revision; they state that they reviewed and tested the resulting work.

AI USE STATEMENT

  • The authors used generative AI tools for code implementation and debugging, methodology feedback, literature identification, FIDgrad derivations, and manuscript revision.They report reviewing the AI-assisted text, checking references and derivations, and testing AI-assisted code.Source: AI USE STATEMENT · p. 10

A.1.1 SHARED SETUP

The generators remain frozen during ordering-policy training. The shared setup uses matched generation noise within conditioning groups and Transformer policies with specified inputs and optimization settings.

A.1.1 SHARED SETUP

  • The generators remain frozen while their ordering policies are trained.Policy inputs generally concatenate frozen-generator features with Gaussian noise for the next fill; the policies are bidirectional pre-norm Transformer encoders with GELU activations, learned positional embeddings, and zero dropout.Source: §A.1.1 · p. 13
  • Trajectories within a conditioning group share generation noise indexed by position and previous fills, while policy actions use a separate random stream.Evaluations match noise across rethinking budgets and hold noise capacity fixed.Source: §A.1.1 · p. 13
  • Table 6 reports dynamic-removal policy architectures and optimization settings.The shared setup describes Transformer policy architecture and optimization details.Source: Table 6 · p. 13, §A.1.1 · p. 13
  • Policies use Adam with ϵ = 10^-8 and zero weight decay, with bfloat16 mixed precision during training.Each update contains 128 independent conditioning inputs, each expanded into a trajectory group.Source: §A.1.1 · p. 13

A.1.2 MAR-B

The MAR-B policy is trained with FIDgrad and a rethinking budget of R = 8. Moving feature-moment estimates support the reward, but the resulting updates need not decrease population FID at every step.

A.1.2 MAR-B

  • The MAR-B policy is trained with FIDgrad and a rethinking budget of R = 8, using a constant learning rate for 2,000 updates followed by a linear decay over 2,000 updates.The learning rate decays to 2 × 10^-6.Source: §A.1.2 · p. 13
  • For a generated feature z, the FIDgrad perturbation changes the generated mean by d and covariance by dd^T − Σu to first order, yielding the Eq. 7 reward under Proposition 1's assumptions.Here d = z − µu, with generated moments (µu, Σu) and fixed real moments (µr, Σr).Source: §A.1.2 · p. 13
  • FIDgrad rewards are computed from current moments before they are updated, and the moment estimates use an exponential moving average with α = 0.001.The estimates are initialized from 50,048 separately generated images; one trajectory is selected uniformly from each group of eight to estimate batch moments.Source: §A.1.2 · p. 14
  • The moving moment estimates and within-group reward normalization produce a surrogate policy-training objective, so updates need not decrease population FID at every step.Source: §A.1.2 · p. 14

A.1.3 MASKGEN-KL-XL

The appendix reports evaluation and inference settings across MaskGen-KL-XL, MAR-Sudoku, and MAR-B, alongside test-time scaling results and qualitative examples. Sudoku accuracy rises with larger rethinking budgets, while qualitative figures compare policy-guided and random-order samples and illustrate patch revision.

A.1.3 MASKGEN-KL-XL

  • MaskGen-KL-XL policies are evaluated on 553 official GenEval prompts disjoint from Flow-GRPO training prompts, using four images per prompt and the original binary GenEval score.Separate policies are trained against HPSv2.1 and dense GenEval reward.Source: §A.1.3 · p. 14

A.1.4 MAR-SUDOKU

  • MAR-Sudoku training uses R = 9 and terminal reward 1 for a correct completed grid and 0 otherwise, with partially filled grids sampled across clue counts.The clue-count distribution has maximum density at zero and zero density at 81; clue locations are sampled randomly.Source: §A.1.4 · p. 14

A.2 GENERATOR ARCHITECTURES AND INFERENCE SETTINGS

  • Table 7 summarizes generator architectures and dynamic-removal inference settings, while the listed parameter counts exclude tokenizers or VAEs, the MaskGen text encoder, and the Sudoku reward classifier.Both MaskGen-KL-XL policies use the architecture shown in the table.Source: §A.2 · p. 14, Table 7 · p. 15

A.3 SUDOKU TEST-TIME SCALING

  • At R = 9, Sudoku accuracy is 98.21%, 98.02%, and 97.94% across three evaluation seeds, versus 80.17%, 79.96%, and 80.29% at R = 0.For the first seed, the improvement is 18.04 percentage points.Source: §A.3 · p. 15, Table 8 · p. 15

A.4 MAR-B TEST-TIME SCALING

  • MAR-B's Table 9 evaluates the FIDgrad policy trained at R = 8 across listed test-time budgets using 50,000 ImageNet images and DDIM with 25 denoising steps.FDr-5 averages normalized Fréchet distances across ConvNeXt, DINOv2, MAE, SigLIP, and CLIP representations, excluding the Inception reward space.Source: Table 9 · p. 16, §A.4 · p. 16

A.5 MASKGEN PREFERENCE TEST-TIME SCALING

  • MaskGen preference evaluation uses all 3,200 HPSv2 benchmark prompts, with learned-policy budgets sharing the same start noise and scores reported for HPSv2, PickScore, and ImageReward.Source: §A.5 · p. 16

A.6 QUALITATIVE EXAMPLES

  • Figures 10 and 11 compare random-order and policy-guided MaskGen samples, Figure 12 pairs random and PONDER MAR-B samples, and Figure 13 shows patch removal and regeneration.The MaskGen examples use the same prompts and DDIM-25 sampler; the MAR-B pairs use R = 8.Source: §A.6 · p. 16, Figure 10 · p. 17, Figure 11 · p. 18
Loading 2610.23531v1…