Source-linked AI summary
Are We Overthinking One-Step Generation?
Woobin Park, Jaeyeon Kim, Yilun Du, Sitan Chen, Bohyung Han
TL;DR
Diffusion and flow models typically require multiple network evaluations, and existing one-step approaches often rely on complex training machinery. Priming Models first train on reconstruction, then perform feature-space distributional post-training. Across image generation, text-to-image, and world modeling, they improve compute-quality trade-offs, including a state-of-the-art one-step ImageNet FID of 1.19 with approximately 10× less training compute than Drifting.
Problem
Diffusion and flow models rely on multiple network evaluations, while many one-step approaches require complex training machinery.
Source: §1 · p. 1Method
Priming Models pretrain a generator on reconstruction, then refine its one-step noise-to-sample mapping by matching generated and real distributions in feature space.
Source: §2.1 · pp. 3–4Results
Across ImageNet, text-to-image, and world modeling, Priming improves compute-quality trade-offs; ImageNet FID is 1.19 among one-step generators with approximately 10× less compute than Drifting.
Source: §1 · p. 2Takeaways & Limitations
Reconstruction pretraining is a minimal design principle for preparing generators for distributional post-training.
Source: §1 · p. 2Priming Models rely on pretrained feature encoders, and whether that dependence can be eliminated remains an open question.
Source: §5 · p. 10
Abstract
from arXiv · showhide
Diffusion and flow matching are now dominant frameworks for generative modeling in continuous domains, but they typically require multiple network evaluations at inference. This has motivated extensive research on one-step generative modeling primarily built on diffusion and flow, though existing approaches often rely on complex training machinery. In this work, we introduce Priming Models, a simple two-stage framework for training one-step generators. Motivated by the difficulty of distributional training from scratch, Priming Models are first trained on a reconstruction objective, e.g., pixel-space denoising at a fixed noise level, before undergoing feature-space distributional post-training. Our central scientific finding is that this simple reconstruction stage is already enough to entirely alleviate the difficulty of distributional training. Empirically, across class-conditional image generation, text-to-image generation, and world modeling, Priming Models significantly advance the Pareto frontier of training compute vs. generation quality. On ImageNet at $256 \times 256$ resolution, we achieve a state-of-the-art FID of $1.19$ among one-step generators using approximately one-tenth the training compute of prior methods with comparable FID. We open-source our codebase at https://github.com/JaeyeonKim01/priming.
1 INTRODUCTION
Diffusion and flow models offer strong generation but rely on iterative sampling, while one-step approaches often require complex training machinery. Priming Models use reconstruction pretraining before feature-space distributional post-training and report strong compute-quality results across three visual domains.
1 INTRODUCTION
- Diffusion and flow models use iterative sampling, while many one-step approaches require substantial extra training machinery.Source: §1 · pp. 1–2
- Priming Models first use a reconstruction objective, including fixed-level denoising, then post-train as one-step generators with feature-space distribution matching.Source: §1 · p. 2
- FID 1.19 is state-of-the-art among one-step ImageNet generators, with approximately 10× less training compute than Drifting.The comparison excludes FD-loss, which reports FID 0.72 by directly optimizing Fréchet distance.Source: §1 · p. 2
- Priming substantially outperforms one-step MeanFlow on GenEval and DPG while reducing training compute by a factor of 1.9.Source: §1 · p. 2
- In world modeling, Priming Models outperform Drifting with approximately half the training compute and achieve performance comparable to multi-step flow matching.Source: §1 · p. 2
- Reconstruction pretraining accelerates distributional post-training by aligning the network Jacobian and making output refinement more effective.Source: §1 · p. 2
2 PRIMING MODELS
The section motivates Priming Models as a response to the difficulty of distributional training from scratch and presents a two-stage reconstruction-then-distributional recipe. Its analysis associates higher pre-training Jacobian response with lower post-training online FID.
2 PRIMING MODELS
- The section identifies distributional optimization from scratch as challenging and reports that fixing style embeddings yields only marginal image differences across Gaussian noise draws.Varying the style embeddings still produces diverse samples.Source: §2 · p. 3
2.1 PRIMING MODELS: TRAINING RECIPE
- Priming Models train the same network first to reconstruct clean samples from corrupted inputs, then match generated and real distributions in frozen feature space.After pretraining, Gaussian noise is fed directly to the network as a one-step generator.Source: §2.1 · pp. 3–4
- A single fixed corruption level suffices for pretraining; the simplest setting uses s=0.8, bringing inputs near Gaussian noise while retaining information about clean targets.The pretraining stage also accommodates patch masking and reconstruction in a pretrained RAE representation space.Source: §2.1 · pp. 3–4
- The feature-space objective is flexible: optimal transport is the main choice, while Sinkhorn divergence and MMD surrogates are also supported, and mean pairwise distance is effective.The loss is estimated on minibatches and backpropagated through the frozen feature encoder composed with the generator.Source: §2.1 · p. 4
- Priming starts from poor initial one-step samples, unlike recent work that uses distributional losses to improve already strong generators, and requires neither prescribed trajectories nor a pretrained diffusion or flow model.The authors contrast their approach with prior work (Feng et al., 2026; Yang et al., 2026).Source: §2.1 · p. 4
2.2 PRIMING MODELS: UNDERSTANDING
- Across four models, higher normalized Jacobian response before post-training is associated with lower online FID afterward.At matched post-training updates, longer reconstruction pretraining is also associated with lower FID; the 160-epoch perceptual-loss model achieves a lower FID.Source: Figure 4 · p. 5, §2.2 · p. 5
3 EXPERIMENTS
Priming Models are evaluated on ImageNet class-conditional generation, text-to-image generation, and world modeling. Across these domains, the results show improved quality–compute trade-offs, including strong one-step performance and a two-step text-to-image extension.
3.1 IMAGENET: CLASS-CONDITIONAL GENERATION
- ImageNet: Priming-L reaches FID 1.19 and FDr5 7.10 at 238.95 EFLOPs, while all four configurations improve the quality–compute Pareto frontier over evaluated baselines.Evaluation uses one network pass per image; FDr5 averages normalized Fréchet distances across five feature spaces.Source: Table 1 · p. 6, §3.1 · pp. 6–7
- ImageNet post-training rapidly improves samples: although pretrained outputs are initially noisy, clean images appear within 1,250 updates and FID continues improving through 25,000 updates.The reported samples show no obvious signs of memorization, and precision and recall are comparable to baselines.Source: §3.1 · p. 7
3.2 TEXT-TO-IMAGE GENERATION
- Text-to-image: logit-normal Priming scores 0.818 on GenEval and 82.16 on DPG-Bench, outperforming one-step MeanFlow and approaching four-step MeanFlow with one network evaluation.The two-step extension scores 0.844 on GenEval and 82.81 on DPG-Bench, competitive with four-step MeanFlow.Source: §3.2 · pp. 7–8, Table 2 · p. 8
3.3 WORLD MODELING
- World modeling: Priming beats DriftWorld and 32-step flow matching on every reported quality metric except RT-1 FVD, using 59.7% less compute on RT-1 and 74.4% less on Bridge-V2 than DriftWorld.The models predict each frame in a single network evaluation and do not use DriftWorld’s motion weighting or self-forcing techniques.Source: Table 3 · p. 8, §3.3 · pp. 8–9
4 ADDITIONAL EXPERIMENTS
These experiments test alternatives in Priming Models’ pretraining and post-training designs. They report outcomes for reconstruction variants, losses, and feature spaces.
4.1 EXAMINING PRETRAINING DESIGN CHOICES
- Masking-based pretraining, which replaces 80% of patches with independent Gaussian noise, yields FID 1.87 after 160 pretraining epochs and 12,500 post-training steps.Source: §4.1 · p. 9
- Distilling clean-image representations from a frozen RAEv2 teacher yields FID 2.34 after 12,500 distributional post-training steps.A 300M-parameter student encoder learns representations from corrupted inputs and is composed with the frozen teacher decoder.Source: §4.1 · p. 9
- Without perceptual loss, terminal EMA FID worsens from 1.38 to 2.07 and FDr5 from 11.23 to 16.13 after 25,000 matched OT updates.The authors hypothesize that external-feature supervision helps capture useful data-manifold structure.Source: §4.1 · p. 9
- At fixed pretraining noise levels, s = 0.8 gives the lowest terminal EMA FID (1.65), followed by s = 0.7 (1.76), s = 0.9 (1.97), s = 0.5 (2.46), and s = 1.0 (65.08).The authors interpret this as supporting a balance between retaining clean-image information and keeping noise close to pure Gaussian noise.Source: §4.1 · p. 9
4.2 EXAMINING POST-TRAINING DESIGN CHOICES
- At 25,000 updates, OT gives terminal EMA FID 1.38, compared with 1.95 for pairwise distance and 1.85 for negative Laplacian-kernel similarity.The authors report that Priming works with multiple distributional losses, with OT yielding the best FID among those tested.Source: §4.2 · p. 9
- Lightweight feature spaces yield FID 1.65 and FDr5 6.13, versus 1.38 and 11.23 for the full-feature OT reference; the FDr5 result is not independent held-out-feature evidence.Inception, SigLIP, and ViT-MAE are components of FDr5.Source: §4.2 · pp. 9–10
- FD-loss is another distributional post-training objective tested with a pretrained Priming-B checkpoint, with setup and results reported in Appendix E.It accumulates feature statistics across steps while backpropagating through the current batch.Source: §4.2 · p. 10
5 CONCLUSION
Priming Models use reconstruction pretraining followed by distributional post-training, and the paper reports strong performance across visual domains. The conclusion identifies reliance on pretrained feature encoders and remaining performance gaps as open issues.
- Priming Models are a simple two-stage framework for end-to-end training of one-step generators.Source: §5 · p. 10
- The authors identify reliance on pretrained feature encoders as an open question, despite demonstrating the method with multiple encoder sets.Source: §5 · p. 10
- The relationship between reconstruction pretraining and Jacobian alignment warrants deeper investigation, and closing the gap with multi-step diffusion and flow models remains a challenge.Source: §5 · p. 10
- The authors hope the results open the door to simpler and more scalable approaches to one-step generative modeling.Source: §5 · p. 10
A RELATED WORK
Prior one- and few-step methods learn trajectories or flow maps, distill teachers, or directly match distributions. Priming Models instead use reconstruction pretraining to initialize direct one-step distribution matching from scratch.
Category 1: Trajectory straightening and flow-map learning
- Trajectory-straightening methods modify couplings or learn finite-time flow transitions, while Priming Models use reconstruction initialization rather than recovering a prescribed trajectory or flow map.Flow-map methods supervise a time-dependent transport or consistency relation.Source: A RELATED WORK · p. 16
Category 2: Teacher-guided distillation
- Teacher-guided distillation transfers a pretrained diffusion model’s knowledge, whereas Priming Models train from scratch and require no pretrained teacher.Distillation approaches inherit the cost and machinery needed to obtain and query that teacher.Source: A RELATED WORK · p. 16
Category 3: Distributional loss (from scratch)
- Direct distribution matching has precedents in GANs, MMD generators, and Sinkhorn-trained implicit models, but Priming Models apply it after reconstruction pretraining.Earlier work is described as limited in empirical scope.Source: A RELATED WORK · p. 16
- Priming Models first teach the generator to correct corrupted inputs toward clean images, then refine it with a distributional objective into a one-step generator.They do not optimize direct distribution matching from random initialization.Source: A RELATED WORK · p. 16
Category 3-1: Drifting Models
- Priming Models use an OT-based surrogate without the additional mechanisms described for Drifting and W-Flow, and all four reported models share the same post-training configuration.The comparison notes learned random style embeddings and classifier-free guidance in both other frameworks, plus drift-field normalization in Drifting; Priming does not use these mechanisms.Source: A RELATED WORK · p. 17
Priming
- The authors hypothesize that reconstruction pretraining reduces the need for extra training mechanisms, but their experiments do not isolate those mechanisms’ individual contributions.They also investigate whether reconstruction can be as simple as denoising at one fixed noise level without noise-level conditioning.Source: A RELATED WORK · p. 17
B DISTRIBUTIONAL POST-TRAINING DETAILS
The post-training details specify a stop-gradient distributional objective and its finite-batch implementation. Conditions and feature branches are matched separately, with branch scaling and a prescribed update procedure.
B DISTRIBUTIONAL POST-TRAINING DETAILS
- Post-training matches generated and real distributions in feature space using distributional objectives such as Sinkhorn divergence or Wasserstein distance.The section covers the stop-gradient loss, finite-batch form, conditional feature construction, normalization, and optimization algorithm.Source: B DISTRIBUTIONAL POST-TRAINING DETAILS · p. 17
B.1 WHY THE STOP-GRADIENT OBJECTIVE MINIMIZES A DIVERGENCE
- The stop-gradient objective has the same parameter gradient as a symmetric distributional discrepancy, while removing a constant term and duplicate generated self-interaction derivatives.The real self-interaction is constant with respect to parameters, and symmetry makes the two generated self-interaction derivatives equal.Source: §B.1 · pp. 17–18
B.2 FINITE-BATCH GRADIENT SURROGATES
- Finite-batch OT uses a Sinkhorn-approximated transport plan, while the MMD surrogate uses a symmetric-kernel loss.The stated Laplacian and energy-distance variants use k(a, b) = exp(−∥a−b∥2) and k(a, b) = −∥a−b∥2, respectively.Source: §B.2 · p. 18
B.3 CONDITIONAL AND MULTI-FEATURE CONSTRUCTION
- For each selected condition, the construction draws real targets and two independent latent batches; gradients pass only through active generated samples.Conditions may be ImageNet classes, text captions, or observation–action histories, and each condition-feature-position pair is matched separately.Source: §B.3 · p. 19
B.4 FEATURE SCALING
- Feature branches are divided by a common branch scale, using fixed calibration for ImageNet or a detached minibatch scale for text-to-image and world modeling.ImageNet scales are calibrated from median pairwise distances among 2,000 training images; the other settings use pair-RMS scaling.Source: §B.4 · p. 19
B.5 ALGORITHM AND DISTRIBUTED EXECUTION
- Each update samples conditions, real targets, and independent active and reference noises, generates both batches, then evaluates and combines branch losses.The active outputs carry gradients; the reference outputs are stop-gradient, and each branch’s detached scale is applied before its finite-batch objective.Source: §B.5 · p. 20
C EXPERIMENTAL DETAILS
The ImageNet configurations share training settings across four Priming-B/L variants. The appendix specifies their initialization, batches, feature transport, and evaluation procedure.
C.1 CLASS-CONDITIONAL GENERATION ON IMAGENET
- All four Priming-B/L configurations share the ImageNet training hyperparameters listed in Table 5.The 80- and 160-epoch variants start from two checkpoints of the same reconstruction run and are post-trained separately.Source: Table 5 · p. 21, §C.1 · p. 21
- Each update samples 16 distinct ImageNet classes and uses 1,024 differentiable generated images, 1,024 generated references, and 2,048 real images.Classwise distributional losses are computed locally within each GPU.Source: §C.1 · p. 21
- ImageNet distributional matching uses 84 frozen ResNet-MAE feature outputs and 13 ConvNeXt outputs, with each encoder family receiving half the total weight.Feature scales are calibrated once from median pairwise distances among 2,000 training images; separate transport problems are solved for each class and feature output.Source: §C.1 · p. 21
- Evaluation uses the final post-training EMA checkpoint, generates 50 images per class for 50,000 images total, and reports the mean across seeds 1, 2, and 3.These are repeated sampling evaluations of one trained checkpoint per configuration.Source: §C.1 · p. 21
C.2 TEXT-TO-IMAGE GENERATION
The text-to-image appendix details reconstruction and distributional training, evaluation, and a multi-step variant. The world-modeling appendix describes action-conditioned next-frame prediction and its autoregressive evaluation.
C.2 TEXT-TO-IMAGE GENERATION
- Text-to-image training uses a 258M-parameter MiniT2I MM-JiT-B/16 generator at 512 × 512, with frozen FLAN-T5-Large caption encoding.Reconstruction pretraining uses CC12M followed by reconstruction fine-tuning on a curated 120k-image mixture.Source: §C.2 · p. 21
- The fixed-noise text-to-image configuration uses s = 0.8 in both reconstruction stages, while the alternative uses a reweighted logit-normal corruption schedule.The alternative schedule is used through CC12M pretraining and curated-mixture fine-tuning.Source: §C.2 · p. 22
- One-step text-to-image post-training uses feature-space OT over pooled SigLIP, ViT-MAE, and Inception-V3 features, with equal branch weights and pairwise-RMS normalization.Each update uses 64 captions; each caption has 32 generated samples, half gradient-carrying and half stop-gradient, and one real image.Source: §C.2 · p. 22
- The multi-step text-to-image variant trains across noise levels (1.0, 0.75, 0.5, 0.25) in two 4,000-update phases, totaling 8,000 updates.It uses the same feature-space OT objective as the one-step variant.Source: §C.2 · p. 22
- Text-to-image evaluation uses all 553 GenEval prompts and all 1,065 DPG-Bench prompts; Table 6 reports training hyperparameters and durations in updates.Training-compute totals include reconstruction training and the subsequent method-specific stage.Source: §C.2 · p. 22, Table 6 · p. 22
C.3 WORLD MODELING
- World-modeling experiments predict the next frame from four observations and aligned actions, operating on 256 × 256 frames encoded into SD3-VAE latents.The same task and action-conditioned U-Net architecture are used for RT-1 and Bridge-V2.Source: §C.3 · p. 22
- World-model post-training starts from reconstruction-pretrained EMA weights and matches generated samples to real next frames using OT and normalized frozen-feature branches.Each update uses 32 conditioning sequences, 16 gradient-carrying samples, 16 stop-gradient samples, and one real target per condition.Source: §C.3 · p. 23
- Autoregressive evaluation inserts predicted frames into subsequent visual histories and re-anchors context on real trajectories between 8-frame chunks.Priming Models and DriftWorld use one network evaluation per predicted frame, versus 32 for the flow-matching baseline.Source: §C.3 · p. 23
C.4 EXPERIMENTAL DETAILS FOR SECTION 2.2
The experiments measure how reconstruction pretraining shapes the model’s Jacobian response to distributional correction directions. Across additional ResNet-MAE outputs, stronger alignment accompanies lower post-training FID, but the association is not uniform across representations.
C.4 EXPERIMENTAL DETAILS FOR SECTION 2.2
- The analysis compares four Priming-B initializations—20, 80, or 160 reconstruction-pretraining epochs, with or without perceptual loss—using matched Gaussian inputs across models.For each of eight fixed ImageNet classes, 512 images were generated.Source: §C.4 · p. 23
- The study estimates OT correction subspaces from gradients across 13 frozen feature outputs, retaining 64 PCA directions for each checkpoint, class, and feature-output pair.The outputs comprise 11 ResNet-MAE features and two ConvNeXt features.Source: §C.4 · p. 23
- A measures normalized Jacobian response along the estimated OT correction subspace relative to overall Jacobian scale, and is compared with online-weight FID after 12,500 post-training updates.The response estimate uses PCA-subspace and full-feature-space directions, with classwise scores averaged equally across eight classes.Source: §C.4 · p. 23
- Across ten additional ResNet-MAE outputs, alignment increases with longer reconstruction pretraining alongside lower post-training FID, while ConvNeXt global-mean alignment varies little and does not follow the FID ordering.The authors conclude that the association extends across multiple feature outputs but is not uniform across representations.Source: §C.4 · pp. 23–24
D ADDITIONAL RESULTS ON IMAGENET
Additional ImageNet results report comparable precision and recall to baselines and no observed signs of memorization in the tested Priming models. The FD-loss appendix defines a direct feature-moment matching objective and notes that it uses marginal rather than conditional distribution matching.
D ADDITIONAL RESULTS ON IMAGENET
- Priming Model instances exhibit precision and recall comparable to baselines on ImageNet.Table 8 reports the comparison; its caption says all methods use one-step generation except JiT-H.Source: D ADDITIONAL RESULTS ON IMAGENET · p. 25, Table 8 · p. 25
- No signs of memorization were observed in Priming-B or Priming-L-160 among 20 visually selected queries per model using same-class CLIP ViT-L/14 retrieval.Each query was compared against 1,300 training images per class, with four nearest images shown.Source: D ADDITIONAL RESULTS ON IMAGENET · p. 25
E.1 FD-LOSS OBJECTIVE
- FD-loss matches generated and real feature distributions through their first two moments, using the squared Fréchet distance between corresponding moment-matched Gaussians.The objective is equivalently the squared 2-Wasserstein distance between those Gaussians.Source: §E.1 · p. 28
- FD-loss propagates gradients through generated feature moments while holding real moments fixed, and performs marginal rather than conditional distribution matching.Accumulated moment estimates allow optimization without putting the full feature population in one minibatch.Source: §E.1 · p. 28
E.2 FD-LOSS IMAGENET EXPERIMENTS
The ImageNet experiments test FD-loss after reconstruction pretraining, finding improved distributional metrics with longer pretraining and better results for perceptual reconstruction at fixed corruption. The appendix also provides qualitative ImageNet examples for Priming-B and Priming-L.
E.2 FD-LOSS IMAGENET EXPERIMENTS
- The FD-loss experiments use Priming-B for pretraining ablations and Priming-L for pretraining-duration tests, generating 256 × 256 ImageNet-1K images in one network evaluation.The models have 115.1M and 405.1M parameters, respectively; the experiments are separate from the main-text OT-post-trained results.Source: §E.2 · p. 28
- FDr6 falls from 7.49 to 4.95 with perceptual supervision at fixed s = 0.8, and from 8.30 to 5.57 under the tested logit-normal schedule.All configurations use 80 reconstruction-pretraining epochs followed by 20 FD-loss epochs; fixed s = 0.8 outperforms the tested logit-normal schedule under either reconstruction objective.Source: §E.2 · p. 28, Table 9 · p. 29
- With perceptual reconstruction fixed, s = 0.8 gives the lowest FDr6 among tested values from 0.6 to 1.0; at s = 1, performance deteriorates sharply.At s = 1, the corrupted input is independent of its paired clean image, so matching the eventual Gaussian input alone is insufficient for the better initialization reported here.Source: §E.2 · p. 28
- FID decreases from 1.07 to 0.94, FDr6 decreases from 5.46 to 3.29, and recall increases from 0.53 to 0.60 as Priming-L reconstruction pretraining grows from 20 to 160 epochs.FD post-training remains fixed at 20 epochs, and precision is approximately constant.Source: §E.2 · p. 29, Table 10 · p. 29
F.1 IMAGENET
- The appendix presents qualitative ImageNet examples for Priming-B and Priming-L, each pretrained for 160 epochs, using the same ten Gaussian noise vectors for each class.The listed examples include animal and non-animal classes.Source: §F.1 · pp. 29–30
F.2 TEXT-TO-IMAGE GENERATION
The appendix presents qualitative one-step text-to-image comparisons with a MeanFlow baseline, autoregressive world-model rollouts, and a style-embedding diversity comparison. The captions specify the conditioning and visual layouts for these examples.
F.2 TEXT-TO-IMAGE GENERATION
- Text-to-image examples compare one-step Priming samples with logit-normal pretraining against one-step MiniT2I MeanFlow samples using classifier-free guidance scale 12.0.The comparison uses selected captions and corresponding sample indices within four-sample groups.Source: §F.2 · p. 38, Figure 7 · p. 38, Figure 8 · p. 39
F.3 WORLD MODELS
- World-model rollouts show ground-truth frames above predictions, with the observed context at left and later predicted frames uniformly spaced across the rollout.Predictions are generated autoregressively while conditioning on the provided action sequence, with examples on RT-1 and Bridge-V2.Source: Figure 13 · p. 44
F.4 STYLE EMBEDDING
- With ImageNet class and learned style embedding fixed, varying Gaussian inputs produces highly similar sample appearance and composition, indicating that the embedding accounts for much of the variation.The comparison covers Drifting Pixel-L and W-Flow latent XL/2 across ten samples per row.Source: Figure 14 · p. 45
Related papers
This paper cites
- Learning Transferable Visual Models From Natural Language Supervision2021
- The Unreasonable Effectiveness of Deep Features as a Perceptual Metric2018
- Score-Based Generative Modeling through Stochastic Differential Equations2020
- Rethinking the Inception Architecture for Computer Vision2015
- Diffusion Models Beat GANs on Image Synthesis2021
- Diffusion Policy: Visuomotor Policy Learning via Action Diffusion2023