PaperlayerView exact revision on arXiv

Source-linked AI summary

Improved Distributional Diffusion Models

Tommaso Martorella, Alexandre Galashov, Felix Krause, Stefan Andreas Baumann, Valentin De Bortoli, Arthur Gretton, Björn Ommer

arXiv:2609.37147v1cs.CVcs.LG

TL;DR

Few-step generation challenges DDMs because multi-particle training is costly and fixed scoring-rule settings impose one trade-off across noise levels. The paper defers particle expansion and schedules scoring-rule hyperparameters over time. The resulting stochastic generator's FID does not degrade from 4 to 50 steps, and the approach transfers to text-to-image generation.

  • Problem

    DDMs face O(m) denoiser-evaluation overhead from multi-particle training and fixed scoring-rule hyperparameters that impose one trade-off across noise levels.

    Source: §1 · p. 2

  • Method

    The method shares early transformer computation before expanding into particles and uses time-dependent scoring-rule schedules informed by posterior concentration.

    Source: §1 · p. 2

  • Results

    FID is non-increasing from 4 to 50 steps, and iDDM nearly halves 4-step MS-COCO FID versus a matched flow-matching baseline (41.05 vs. 78.20).

    Source: §1 · p. 2

  • Takeaways & Limitations

    Deferred expansion and adaptive objectives make stochastic few-step DDM generation practical at DiT scale, with reported transfer to text-to-image generation.

    Source: §1 · p. 2, §6 · p. 10

    The regime anchors are approximate at finite scale, and distilled and fastforward methods remain stronger at 1–2 steps.

    Source: §6 · p. 10

Abstract

from arXiv · show

Distributional Diffusion Models (DDMs) replace the standard mean-prediction denoiser with a \emph{distributional} denoiser trained via a scoring rule objective, learning a stochastic approximation to $p(x_1 \mid x_t)$ rather than its conditional mean. However, scaling DDMs to modern image-generation settings faces two obstacles: (i) multi-particle training incurs overhead that scales with the number of particles, (ii) DDMs use globally fixed scoring rule hyperparameters, forcing a single trade-off across sampling budgets. We mitigate these limitations by deferring particle expansion to late transformer layers, and the hyperparameter trade-off by introducing time-dependent scoring rule schedules informed by the dynamical regimes of~\citet{Biroli2024}. Combined with a DiT-based latent setup, these changes make DDM training practical on class-conditional ImageNet-$256^2$, achieving 4.48 FID at 4 steps and 2.38 at 50 steps with DiT-XL/2, from a single model trained from scratch in one stage, without a teacher, self-distillation or JVPs. The result is a stochastic few-step generator whose FID does not degrade as the sampling budget grows from 4 to 50 NFE, and the same recipe transfers to text-to-image generation. Code and pre-trained models available at https://github.com/CompVis/iDDM.

1 INTRODUCTION

Few-step generation is difficult because coarse discretizations leave substantial posterior variance that conditional-mean denoisers do not capture. The paper makes distributional denoising more practical by reducing multi-particle cost and adapting scoring-rule hyperparameters over time.

1 INTRODUCTION

  • At coarse discretizations, p(x1|xt) can have substantial variance, making conditional-mean denoisers an insufficient summary of the posterior distribution.The motivation is especially relevant to few-step generation, where high-fidelity dynamics are challenging to approximate.Source: §1 · p. 2
  • DDMs replace conditional-mean regression with a stochastic denoiser that approximates p(x1|xt), trained using a generalized energy score.The score balances fidelity and sample-population diversity through its exponent β and trade-off weight λ.Source: §1 · p. 2
  • DDM scaling is limited by O(m) separate denoiser evaluations and globally fixed scoring-rule hyperparameters that impose one fidelity–diversity trade-off across noise levels.The fixed trade-off is problematic because p(x1|xt) varies in concentration over the diffusion trajectory.Source: §1 · p. 2
  • Deferred particle expansion reduces the marginal cost of added particles, while time-dependent scoring-rule schedules adapt the objective to posterior concentration.The schedules emphasize diversity for broad posteriors and more concentrated predictions for narrow ones, following Biroli et al. (2024).Source: §1 · p. 2
  • FID is non-increasing from 4 to 50 steps, and iDDM nearly halves 4-step MS-COCO FID versus a matched flow-matching baseline (41.05 vs. 78.20).The ImageNet-2562 and text-to-image results use a DiT backbone and the same approach, respectively.Source: §1 · p. 2

2 BACKGROUND

The background introduces flow matching as transport from noise to data, then explains how DDMs model the conditional data distribution stochastically. It also reviews reverse-process regimes that motivate later schedule design.

2.1 FLOW MATCHING

  • Flow matching learns a time-dependent velocity field transporting samples from noise p0 to data p1, using an interpolant between noise and data.The model regresses on the conditional velocity; inference integrates the learned field with a numerical ODE solver.Source: §2.1 · pp. 2–3

2.2 DISTRIBUTIONAL DIFFUSION MODELS

  • DDMs add auxiliary noise ξ to predict samples from an approximation to p(x1|xt), rather than only its conditional mean.The generalized energy score is strictly proper for β∈(0, 2) and λ=1, while β=2 and λ=0 recover mean-targeting MSE.Source: §2.2 · p. 3

2.3 DYNAMICAL REGIMES OF THE REVERSE PROCESS

  • The reverse dynamics pass through undifferentiated, coarse-structure, and training-datum-attracted regimes, separated by speciation and collapse times.The regime analysis uses an Ornstein–Uhlenbeck forward process and motivates later schedule construction.Source: §2.3 · pp. 3–4

3 METHOD

The method reduces multi-particle DDM training cost by sharing early transformer computation, and adapts the scoring rule over time according to posterior evolution. Its schedules use dynamical regimes as design guidance.

3.1 EFFICIENT MULTI-PARTICLE TRAINING AND ξ CONDITIONING

  • Deferred expansion shares the early transformer layers across particles, reducing m=4 DDM training cost from roughly 4× flow matching to about 1.5×.The shared hidden state is replicated at layer ℓstart, where particle-specific noise is injected; all particles then pass through the same remaining blocks.Source: §3.1 · p. 4
  • Fixed-inner-width concatenation gives the lowest 4-step FID at roughly baseline parameter count, while adaptive normalization and register tokens break down at few steps.The compared ξ-injection families are concatenation, residual addition, adaptive normalization, and register tokens.Source: §3.1 · p. 4, Figure 3 · p. 5

3.2 TIME-DEPENDENT SCORING RULE SCHEDULES

  • Time-dependent (λ(t), β(t)) schedules link the scoring objective to the complexity of p(x1|xt), transitioning from distributional toward regression-like behavior.The motivation is that λ<1 underestimates posterior variance, which is beneficial when the posterior is concentrated but harmful when it is broad.Source: §3.2 · pp. 4–5

3.2.1 UNDERSTANDING THE POSTERIOR EVOLUTION THROUGH DYNAMICAL REGIMES

  • The schedule design maps Biroli et al. (2024)'s transition times to flow-matching time and favors distributional objectives before collapse.The regimes distinguish prior/manifold localization, an ambiguous mode-dominated window, and an unambiguous data-end regime.Source: §3.2.1 · p. 5
  • The default schedule uses linear shape functions sλ(t)=sβ(t)=1−t with βmin=0.1, matching the dynamical-regimes profile without requiring regime thresholds.The main ablation caption reports that CONFIG F achieves the best performance among the progressively augmented configurations.Source: §3.2.1 · p. 5, Table 1 · p. 6

4 EXPERIMENTS

Experiments evaluate architectural and objective ablations before comparing iDDM with few-step methods on ImageNet and text-to-image generation. The results show strong few-step performance and a single checkpoint whose FID does not worsen from 4 to 50 steps.

4.1.1 DEFERRED EXPANSION & ARCHITECTURAL IMPROVEMENTS

  • Fixed-inner-width concatenation is adopted; t-adaptive noise gating slightly reduces many-step performance but boosts few-step quality.Adaptive norms and token conditioning underperform, especially at few steps.Source: §4.1.1 · pp. 6–7

4.1.2 TRAINING OBJECTIVE IMPROVEMENTS

  • Linear scoring-rule schedules give the best overall 4- and 50-step trade-off; β-only scheduling with λ = 1 and an early switch gives strong FID but leaves sample spread unconstrained afterward.The paper does not adopt the β-only setting as its default.Source: §4.1.2 · p. 7
  • An initial β value of 0.1 gives the best trade-off across sampling budgets, while jit time sampling leads to the best training-distribution results.Jit sampling biases training time toward initial noise, where distributional losses matter more.Source: §4.1.2 · pp. 7–8

4.1.1 DEFERRED EXPANSION & ARCHITECTURAL IMPROVEMENTS

  • Under CFG, fewer distributional layers improve quality while reducing training cost to 1.5× FM cost versus 4× for baseline DDM.The figure reports 2–4 distributional layers as best under optimal guidance and around 2–5 as a good efficiency-quality trade-off.Source: §4.1.1 · p. 6, Figure 4 · p. 8

4.2 COMPARISON WITH PREVIOUS METHODS

  • At matched training compute, iDDM improves 4-step FID from 26.70 to 14.33, while FM remains slightly better at 50 steps, 4.97 versus 5.62.The comparison accounts for iDDM’s higher per-step training cost.Source: §4.1.2 · p. 8
  • 4-step FID falls from 43.68 to 13.13 and 50-step FID from 10.04 to 4.57 versus naïve DDM in the controlled ImageNet comparison.Models were retrained from scratch at B-scale with comparable settings and relevant hyperparameters re-swept.Source: §4.2 · p. 8
  • 4.48 FID at 4 steps and 2.38 at 50 steps; one from-scratch, single-stage iDDM checkpoint has non-increasing FID across that range.At XL/2 scale after 200 epochs.Source: §4.2 · p. 8, Table 4 · p. 9
  • MS-COCO FID improves from 78.20 to 41.05 at 4 steps and from 26.33 to 19.40 at 8 steps versus matched flow matching.The transfer uses a 1.6B-parameter text-to-image model.Source: §4.2 · p. 9

5 RELATED WORK

The paper situates iDDM among approaches that speed diffusion sampling, learn the diffusion posterior, and train few-step generators. It distinguishes iDDM by directly scoring populations and targeting multi-step generation.

5 RELATED WORK

  • iDDM modifies the training objective rather than the sampler, complementing distillation, improved samplers, and parallel or speculative methods [Luhman & Luhman, 2021; Sauer et al., 2025; Yin et al., 2024; Salimans & Ho, 2022; Liu et al., 2024; Song et al., 2023; Xu et al., 2024; Salimans et al., 2024; Lu et al., 2022; 2025; Zheng et al., 2023; Karras et al., 2022; Shih et al., 2023; Chen et al., 2024; De Bortoli et al., 2025a].It can also integrate with Universal Inverse Distillation by replacing the diffusion loss with a distributional one [Kornilov et al., 2025].Source: §5 · p. 9
  • Unlike posterior-enriching and adversarial or MMD-based approaches, iDDM builds on generalized energy scoring rules to address DDM training cost and static hyperparameters [De Bortoli et al., 2025b; Shen et al., 2025].Related alternatives include variance-aware denoisers, conditional adversarial generators, and representation-space diversity methods [Ho et al., 2020; Nichol & Dhariwal, 2021; Bao et al., 2022a;b; Ou et al., 2025; Xiao et al., 2022; Xu et al., 2024; Aiello et al., 2024; Galashov et al., 2025; Wang & He, 2025; Yang et al., 2026].Source: §5 · pp. 9–10
  • Few-step alternatives include consistency and moment-matching methods, while iDDM shares a multi-particle philosophy with IMM and focuses on 4–50 inference steps.MFM also targets p1|t but uses consistency along auxiliary ODEs; its emphasis is reward alignment, whereas iDDM uses direct scoring-rule training for multi-step generation.Source: §5 · p. 10

6 CONCLUSION

The conclusion presents deferred particle expansion and time-adaptive scoring objectives as complementary changes that make stochastic few-step DDM training feasible at DiT scale. It also states the method’s scope limits and points to the materials for reproducing it.

6 CONCLUSION

  • Deferred population expansion shares early transformer layers and injects particle-specific noise late, reducing m = 4 DDM training cost from roughly 4× flow matching to about 1.5×.Time-dependent scoring schedules and improved training-time sampling provide time-adaptive objectives for few-step generation.Source: §6 · p. 10
  • Regime anchors are approximate at finite scale, and distilled and fastforward methods remain stronger at 1–2 steps.The reported gains hold on ImageNet-2562 and transfer to text-to-image; smooth schedules are more robust than hard steps.Source: §6 · p. 10
  • The paper describes iDDM as using publicly available datasets and standard transformers with minimal modifications, with training, evaluation, and hyperparameter details in the appendices.Code and model checkpoints are released to support reproduction.Source: REPRODUCIBILITY STATEMENT · p. 11

A MODEL CONFIGURATION

  • The appendix provides training-configuration tables for latent-space ImageNet models and for FM and iDDM.The ImageNet table covers DiT-B/2 and DiT-XL/2; the FM/iDDM table specifies zero-based block indices.Source: Table 5 · p. 18, Table 6 · p. 19

B.1 IMAGENET 2562

The ImageNet appendix reports multi-metric comparisons, component ablations, posterior-diversity diagnostics, and training-compute comparisons. Its evaluation protocol specifies sample counts and shared settings for FID.

B.1 IMAGENET 2562

  • FID uses 50,000 generated samples against full ImageNet training-set statistics, with shared Inception-V3 features and reference statistics; CFG is swept from 1.0 to 6.0.The paper characterizes results across heterogeneous model families as context rather than a controlled comparison.Source: §B.1 · p. 20
  • Adding each component independently to CONFIG B improves 4-step FID: ξ-conditioning 17.51→17.21, t-sampling 17.51→17.34, and (λ, β) schedule 17.51→15.77, with gains not dependent on the other components.Improved ξ-conditioning improves FID, FD-DINOv2, KID, KDD, and precision at both step counts, while 4-step recall declines from 0.154 to 0.139.Source: §B.1 · pp. 20–21
  • At 4 steps, iDDM improves over DDM on FID (43.6→13.2), FD-DINOv2 (801.8→318.8), KID (0.0337→0.0030), KDD (1.854→0.335), Inception Score (40.3→162.5), precision (0.704→0.778), and recall (0.164→0.226).At matched training steps, iDDM matches or exceeds the FM reference on every reported metric at both 4 and 50 steps.Source: §B.1 · p. 21, Table 9 · p. 21
  • For fixed xt, both distributional models show FID increasing with t; iDDM is more diverse than DDM at low t and less diverse at high t.The authors describe this as consistent with a schedule that shifts toward regression as the posterior concentrates, and caution that the diagnostic does not measure posterior accuracy.Source: §B.1 · p. 22, Table 11 · p. 22
  • iDDM’s per-step training cost is 1.5× flow matching’s, compared with 4× for DDM; its total training budget is about 17.6× lower than iMF’s and 1.67× lower than MeanFlow’s.The throughput benchmark used 4×A100-80GB and DiT-XL/2; the compute-matched setup aligns FM and iDDM budgets up to checkpoint granularity.Source: §B.1 · pp. 21–22, Table 10 · p. 22

B.2 TEXT-TO-IMAGE

The text-to-image appendix specifies the model, training, and evaluation setup, while the dynamical-regimes appendix explains how OU-based regime anchors relate to the flow-matching denoising problem. It presents those anchors as schedule-design aids, not universal ordering guarantees.

B.2 TEXT-TO-IMAGE

  • The text-to-image setup uses a frozen FLUX.2 autoencoder and Qwen3-VL text tower; FM and iDDM share a backbone, with iDDM adding gated noise channels in six widened blocks.Both models train on a recaptioned COYO subset; iDDM uses a generalized energy-kernel objective with four predictions per example. Evaluation includes COCO-FID, CLIP similarity, GenEval, T2I-CompBench, and HPS v2.1 at 4, 8, and 16 Euler steps.Source: §B.2 · pp. 22–23, Table 12 · p. 23

C DYNAMICAL REGIMES OF DIFFUSION MODELS

  • The appendix frames schedule construction around the conditional law induced by isotropic Gaussian corruption after normalization, rather than equivalence between the OU reverse SDE and the flow-matching probability path.The OU forward process described is the constant-rate VP-SDE case.Source: C DYNAMICAL REGIMES OF DIFFUSION MODELS · p. 23, §C.1 · p. 23

C.1 SNR COORDINATE AND OU/FM MAPPING

  • At matched SNR, OU and flow-matching corruptions define the same normalized denoising problem, not the same stochastic process, reverse dynamics, path measure, or raw-coordinate marginal entropy.The conditional law depends on the normalized observation z and scalar ρ; the appendix maps between corruption times by equating SNRs.Source: §C.1 · pp. 23–24

C.2 REGIME ANCHORS

  • The schedule anchors include speciation, collapse, and separation criteria; their ordering is a measured dataset property, not a universal theorem.Collapse is computed numerically, while the separation criterion uses a high-probability noise-ball containment condition around training points.Source: §C.2 · pp. 24–25

C.3 IMAGENET LATENT ESTIMATES AND SCHEDULE CONSTRUCTION

ImageNet latent diagnostics set the transition ranges for time-varying scoring schedules, while deferred particle conditioning and local scoring choices address practical design decisions for distributional denoising.

C.3 IMAGENET LATENT ESTIMATES AND SCHEDULE CONSTRUCTION

  • ImageNet schedule anchors use centered first-channel REPA-E latents from 50,000 training examples, with covariance, entropy-balance, and pairwise-distance diagnostics.The piecewise schedule is 1 below ρs, 0 above ρsep, and interpolated linearly between them.Source: §C.3 · p. 25, Table 14 · p. 26

D ON SCHEDULING λ(t) AND β(t)

  • The covariance-bias criterion favors assigning the greatest shrinkage to the highest SNRs, but does not determine a unique smooth profile or terminal-FID optimum.The regime anchors motivate transition locations; the precise profile and separate λ/β schedules are empirical and ablated.Source: D ON SCHEDULING λ(t) AND β(t) · pp. 25–26

E ξ CONDITIONING

  • After population expansion, each particle’s noise is injected through a conditioning mechanism, with a shared adaptive gate initialized to identity scaling.The gate maps timestep and, for class-conditional generation, class conditioning to a per-element multiplicative modulation.Source: E ξ CONDITIONING · p. 27

E.1 CONDITIONING MECHANISMS

  • Fixed-inner concatenation outperforms Ada-Norm and token conditioning in the few-step regime, while its proposed signal-preservation explanation remains a hypothesis.It keeps the residual stream widened but the attention and MLP inner computations at their original width.Source: §E.1 · pp. 27–28

F KERNEL SPATIAL MODES

  • The local kernel averages per-position scores and is strictly proper for per-position marginals, but it decouples cross-spatial dependencies.The authors use it by default for practical performance improvements.Source: F KERNEL SPATIAL MODES · pp. 28–29

G TRAINING-TIME t SAMPLING

Training-time sampling allocates emphasis across noise levels, and the section relates DDMs to other kernel-score approaches. The authors report that jit sampling improves few- and many-step FID.

G TRAINING-TIME t SAMPLING

  • 72.5% of jit training times fall in ρ∈[ρc, 1], versus 35.2% for uniform and 46.0% for logit-normal(0, 1), and this allocation improves both few- and many-step FID.The jit mode is at t≈0.25 (ρ≈0.1), in Regime IIIa, where the posterior is mode-dominated but still broad.Source: G TRAINING-TIME t SAMPLING · p. 29

H MMD LINEAGE

  • DDMs and IMM belong to one kernel-score family with different right-hand arguments; Consistency Models are a single-particle, first-moment limit in the IMM analysis.The generalized kernel score reduces to a Dirac-target MMD objective when λ=1.Source: H MMD LINEAGE · p. 29

I COMPUTE REQUIREMENTS

Deferred population expansion runs the shared transformer trunk once before replicating particle states, then processes the expanded batch in the final layers.

I COMPUTE REQUIREMENTS

  • The first ℓstart layers run at batch size B; after expansion and noise conditioning, the final H layers run at effective batch size B×m.The models were trained on 4×H200 for B-sized models and 8×H200 for XL-sized models.Source: I COMPUTE REQUIREMENTS · pp. 29–30

Related papers

Loading 2609.37147v1…