PaperlayerView exact revision on arXiv

Source-linked AI summary

Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation

Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, Il-Chul Moon

arXiv:2106.05527v5cs.LGcs.AIstat.ML

TL;DR

Diffusion training must reconcile likelihood estimation, which relies heavily on small diffusion times, with sample generation, which relies more on large times, while loss contributions remain imbalanced. The paper introduces Soft Truncation, which randomizes the truncation threshold across mini-batches and connects this procedure to general weighting and perturbed likelihood estimation. Across CIFAR-10, CelebA, CelebA-HQ 256x256, and STL-10, it reports strong FID results while preserving NLL performance.

  • Problem

    Diffusion models exhibit a trade-off between NLL and FID because density estimation and sample generation emphasize different diffusion-time regimes.

    Source: §1 · p. 1, §4.1 · p. 6, §6 · p. 9

  • Method

    Soft Truncation replaces the fixed smallest diffusion time ϵ with a randomly sampled threshold τ at each mini-batch update.

    Source: §1 · p. 1

  • Results

    Soft Truncation achieves FID 1.90 on CelebA with DDPM++, improving on the previous best FID of 2.92 by DDGM, while keeping NLL intact.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Randomizing truncation lets the score network learn across the diffusion-time range and reduces the imbalance that otherwise favors particular evaluation objectives.

    Source: §1 · p. 1, §4.2 · p. 6, §6 · p. 9

    The paper limits its detailed focus to linear SDEs, particularly VESDE and VPSDE, to maintain simplicity, although it states that Soft Truncation applies to any SDE of the proposed form.

    Source: §2 · p. 2

Abstract

from arXiv · show

Recent advances in diffusion models bring state-of-the-art performance on image generation tasks. However, empirical results from previous research in diffusion models imply an inverse correlation between density estimation and sample generation performances. This paper investigates with sufficient empirical evidence that such inverse correlation happens because density estimation is significantly contributed by small diffusion time, whereas sample generation mainly depends on large diffusion time. However, training a score network well across the entire diffusion time is demanding because the loss scale is significantly imbalanced at each diffusion time. For successful training, therefore, we introduce Soft Truncation, a universally applicable training technique for diffusion models, that softens the fixed and static truncation hyperparameter into a random variable. In experiments, Soft Truncation achieves state-of-the-art performance on CIFAR-10, CelebA, CelebA-HQ 256x256, and STL-10 datasets.

1. Introduction

Diffusion models face a trade-off between likelihood quality and sample quality because different training weightings favor NLL or FID. Soft Truncation randomizes the truncation threshold, resolving this trade-off while connecting to general weighting and Maximum Perturbed Likelihood Estimation.

  • Soft Truncation resolves the NLL–FID trade-off, matching FID-favorable models in FID while retaining the NLL level of NLL-favorable models.Source: §1 · p. 1
  • The paper studies Soft Truncation as a generally applicable technique for diffusion models that improves the balance between likelihood estimation and image generation.Source: §1 · p. 1
  • The method replaces the fixed truncation threshold ϵ with a newly sampled threshold τ for every mini-batch, training the score only on [τ, T].As τ varies, the score network estimates scores across the full diffusion-time range [ϵ, T], improving FID.Source: §1 · p. 1
  • Soft Truncation is equivalent in expectation to a diffusion model with a general weight, with the distribution of τ determining that weight function.This explains why the method can achieve FID performance comparable to FID-favorable training despite implementing only threshold randomization.Source: §1 · pp. 1–2
  • Soft Truncation optimizes a likelihood perturbed by τ and is therefore framed as Maximum Perturbed Likelihood Estimation for diffusion models.Source: §1 · p. 2

2. Preliminary

Continuous-time diffusion models represent data-to-noise corruption and generation with forward and reverse stochastic differential equations. Their score-learning objective uses a diffusion-time weighting function, with linear SDEs yielding Gaussian transitions and equivalent NCSN and DDPM losses.

  • Continuous-time diffusion models gradually perturb data x0 into noise xT through an SDE, while the reverse-time SDE defines the generative process.The reverse process has the same solution as the forward SDE when run backward in time.Source: §2 · p. 2
  • A parametrized generative SDE learns the stochastic process by starting from a noise distribution and integrating backward in time.Source: §2 · p. 2
  • Diffusion training minimizes a score loss whose weighting function λ(t) controls each diffusion time’s contribution.The score loss is equivalent, up to a θ-independent constant, to the continuous denoising NCSN loss.Source: §2 · p. 2
  • VESDE and VPSDE provide analytic Gaussian transition probabilities and are instances of linear diffusions with time-dependent drift and diffusion functions.The paper focuses on these linear SDEs for simplicity, although Soft Truncation applies to any SDE in the stated family.Source: §2 · p. 2
  • For linear SDEs, the denoising NCSN loss is equivalent to the continuous DDPM loss, making NCSN and DDPM interchangeable in this framework.Likelihood training is recovered when λ(t) = g2(t), the likelihood weighting.Source: §2 · pp. 2–3
  • The truncation parameter ϵ excludes unstable small diffusion times, where the score-matching error can diverge and Monte Carlo loss estimates become high variance.The variational-bound integrand remains extremely imbalanced even with positive truncation, motivating careful threshold selection.Source: Figure 1 · p. 3

3. Training and Evaluation of Diffusion Models in Practice

Diffusion training truncates small diffusion times to control divergent loss variance, but truncation creates a trade-off: small times support density estimation while large times determine global sample quality. The section derives the positive-truncation variational bound and shows how this imbalance affects training, likelihood, and fidelity evaluation.

3.1. The Need of Truncation

  • As t →0, the score-matching error diverges, so previous methods truncate training to [τ, T] with τ=ϵ>0 for stable Monte-Carlo estimation.The divergence follows from the transition-score denominator σ^2(t) converging to zero.Source: §3.1 · p. 3

3.2. Variational Bound With Positive Truncation

  • For positive truncation, the paper generalizes the variational log-likelihood bound, recovering the original NCSN objective when τ=0.The generalization is obtained using the data processing inequality and Girsanov theorem.Source: §3.2 · p. 3

3.3. A Universal Phenomenon in Diffusion Training: Extremely Imbalanced Loss

  • The loss remains extremely imbalanced across diffusion time, while insufficiently small ϵ loosens the likelihood bound and harms microscopic sample quality.This imbalance persists throughout training, making ϵ a sensitive hyperparameter.Source: Figure 2 · p. 3, §3.3 · p. 4

3.4. Effect of Truncation on Model Evaluation

  • NELBO and NLL decrease as ϵ decreases, but smaller truncation worsens FID because large diffusion times remain less optimized.The model with σmin=10−3 achieves the best FID in the controlled denoising comparison.Source: §3.4 · p. 4, Table 2 · p. 4
  • Large diffusion times dominate sampling and determine global context, whereas small diffusion times mainly refine microscopic image details.The supplied figures illustrate this distinction through reverse-drift behavior and regenerated samples.Source: §3.4 · pp. 4–5, Figure 3 · p. 4, Figure 5 · p. 4

4. Soft Truncation: A Training Technique for a Diffusion Model

Soft Truncation replaces a fixed diffusion-time cutoff with a sampled cutoff, addressing importance-sampling concentration near the smallest diffusion time while preserving likelihood-oriented training. The method also admits equivalent general-weight and perturbed-likelihood interpretations.

4. Soft Truncation: A Training Technique for a Diffusion Model

  • Soft Truncation varies the truncation time τ across mini-batch updates instead of fixing it at ϵ, reducing the need to search for an optimal cutoff.Each update samples τ from P(τ) and optimizes the corresponding truncated loss.Source: §4 · p. 5, §4.2 · p. 6

4.1. Monte-Carlo Estimation of Truncated Variational Bound with Importance Sampling

  • Importance sampling reduces estimation variance but concentrates Monte-Carlo diffusion time near t ≈ ϵ, leaving large diffusion times inaccurately estimated.The method targets this sampling imbalance rather than replacing importance sampling altogether.Source: §4.1 · pp. 5–6, Figure 7 · p. 6

4.2. Soft Truncation

  • With τ softened, training covers small and large diffusion times more fairly, alleviating per-step loss-scale imbalance.The limited range [τ, T] changes across mini-batches, providing broader score-network coverage.Source: §4.2 · p. 6

4.3. Soft Truncation Equals to A Diffusion Model With A General Weight

  • In expectation, Soft Truncation reduces to a diffusion model with a general weight determined by the distribution P(τ).The equivalence is expressed through the expected Soft Truncation loss and the induced weighting function.Source: §4.3 · p. 6

4.4. Soft Truncation is Maximum Perturbed Likelihood Estimation

  • Soft Truncation can be interpreted as Maximum Perturbed Likelihood Estimation, because sampled τ perturbs the likelihood objective through a truncated variational bound.It is not ordinary MLE unless τ is sufficiently small because the relevant inequality is otherwise not tight.Source: §4.4 · p. 7

4.5. Choice of Truncation Probability Distribution

  • For VPSDE, k ≈ 1.0 provides a sample-quality sweet spot, whereas overly small k excludes τ values near ϵ and harms both generation and density estimation.The authors leave further optimization of the truncation distribution for future work.Source: §4.5 · p. 7

5. Experiments

Experiments evaluate Soft Truncation across datasets, architectures, diffusion SDEs, truncation levels, and probability distributions. The method improves the NLL–FID trade-off and achieves strong sample-generation results across several benchmarks.

FID by Iteration

  • Soft Truncation beats vanilla training after 150k iterations in CelebA FID by training steps.Figure 8 reports FID on the y-axis and training steps on the x-axis.Source: §5 · p. 8, Figure 8 · p. 8

Ablation Studies

  • Soft Truncation improves FID over vanilla training across weighting functions and outperforms the variance-weighted FID-favorable model on CIFAR-10 and ImageNet32 with P0.9.Paired experiments indicate better optimization even when average losses coincide.Source: §5 · pp. 8–9, Table 4 · p. 8
  • Soft Truncation boosts FID while maintaining density-estimation performance across score networks and diffusion strategies, including effectiveness on CelebA with variance weighting.The variance-weighting result does not hold on CIFAR-10 and ImageNet32.Source: Table 5 · p. 8, §5 · p. 9
  • Soft Truncation monotonically reduces both NLL and FID across ϵ values, unlike vanilla training's inverse NLL–FID trend.This also reduces the effort required to search over ϵ.Source: Table 6 · p. 8, §5 · p. 9
  • Soft Truncation significantly improves FID for 0.8 ≤ k ≤ 1.2 on VPSDE and also works with a nonlinear forward SDE.These results extend the observed effectiveness beyond one truncation distribution and beyond linear diffusion settings.Source: Table 7 · p. 8, Table 8 · p. 8, §5 · p. 9

Quantitative Comparison to SOTA

  • FID 1.90 with DDPM++ on CelebA exceeds the previous best FID of 2.92 by DDGM, while Soft Truncation keeps NLL intact across benchmark datasets.State-of-the-art sample-generation performance is reported on CIFAR-10, CelebA, CelebA-HQ, and STL-10.Source: §5 · p. 9, Table 9 · p. 9

6. Conclusion

The paper proposes Soft Truncation as a generally applicable diffusion-model training method that addresses imbalanced score learning across diffusion times.

  • Soft Truncation softens the truncation level at each mini-batch update by randomizing the smallest estimated diffusion time.The method is motivated by density estimation relying mostly on small diffusion times and sample generation on large diffusion times.Source: §6 · p. 9

A. Derivation

This section derives transition probabilities for linear SDEs, explains why denoising score loss diverges near zero diffusion time, and formulates general weighted diffusion losses.

A.1. Transition Probability for Linear SDEs

  • Linear SDE transition probabilities are Gaussian, with mean and covariance characterized by ODE systems and initial conditions.Source: §A.1 · p. 12

A.2. Diverging Denoising Loss

  • The denoising score loss diverges as t →0 because the conditional transition score becomes unbounded there.No locally Lipschitz continuous score function can converge to the resulting unbounded vector field.Source: §A.2 · p. 13

A.3. General Weighted Diffusion Loss

  • The denoising score loss can be rewritten using an appropriate weighting function by changing the order of integration.Source: §A.3 · p. 13

B. Theorems and Proofs

The theorems connect truncated diffusion objectives to weighted losses and restricted path measures, providing the formal basis for Soft Truncation.

  • The proof uses restricted forward and generative path measures together with data processing, KL chain rules, Girsanov’s theorem, and martingale properties.Source: B. Theorems and Proofs · p. 14
  • Theorem 1 establishes that a truncation-based objective corresponds to a weighted diffusion loss under stated regularity conditions.The weighting function is nondecreasing, nonnegative, absolutely continuous on [ϵ, T], and zero before ϵ.Source: B. Theorems and Proofs · p. 15
  • The theorem extends to continuous weighting functions and to weighting functions with finite jumps on [0, T].Source: B. Theorems and Proofs · p. 16
  • The truncation probability is normalized using g2(T), which is the normalizing constant of the unnormalized truncation probability.Source: B. Theorems and Proofs · p. 16

C.1. Additional Score Architectures: Unbounded Parametrization

The paper studies an unbounded score parametrization and a reciprocal variance-exploding SDE to improve behavior near zero diffusion time and restore a geometric variance progression from t = 0.

C.1. Additional Score Architectures: Unbounded Parametrization

  • Standard sθ(xt, t) and sθ(xt, σ(t)) parametrizations cannot estimate the data score as t →0, whereas an unbounded parametrization can do so asymptotically.This explains why sθ(xt, log σ(t)) is successful for score estimation near zero diffusion time.Source: §C.1 · p. 17
  • UNCSN++ introduces η(t) in sθ(xt, η(t)) to focus score estimation near t ≈0.Source: §C.1 · p. 17
  • Under importance-weighted time sampling, the η-distribution becomes uniform, supporting the proposed parametrization’s focus near small diffusion times.Source: §C.1 · p. 17

C.2. Additional SDE: Reciprocal VESDE

  • VESDE’s geometric-variance approximation fails to converge to x0 in distribution because its minimum variance remains nonzero as t →0.Source: §C.2 · p. 18
  • Proposition 2 states that no SDE has the stochastic process {xt}t∈[0,T], defined by the cited transition probability p0t(xt|x0), as its solution.Source: §C.2 · p. 18
  • RVESDE introduces a diffusion process with geometric variance progression starting from t = 0, at the expense of reciprocated time 1/t.Source: §C.2 · p. 18

D.1. Experimental Details

The experiments specify training, evaluation, likelihood-estimation, sampling, and FID protocols across resolutions and model configurations.

D.1. Experimental Details

  • Training uses a learning rate of 0.0002, 5,000 warmup iterations, gradient clipping at 1, and dataset-dependent truncation settings.UNCSN++ uses σmin = 10^-3, NCSN++ uses σmin = 10^-2, and other truncation choices vary by ImageNet32 and non-ImageNet32 settings.Source: §D.1 · p. 18
  • Batch sizes range from 16 to 192 across 32 × 32 through 1024 × 1024 datasets, using multiple NVIDIA RTX-3090 machines for larger workloads.The 256 × 256 and 1024 × 1024 experiments each consume nearly 120 GB of GPU memory.Source: §D.1 · p. 18
  • NLL is evaluated with instantaneous change of variables using after-correction reporting in the main paper, with uniform dequantization applied to the data.The default integration range is ϵ = 10^-5 to T = 1, and the after-correction procedure follows Kim et al. (2022).Source: §D.1 · pp. 18–19
  • Sampling uses the Predictor-Corrector algorithm with resolution-specific signal-to-noise ratios and 1,000 or 2,000 iterative steps.VESDE uses reverse diffusion and annealed Langevin dynamics, whereas VPSDE uses Euler-Maruyama without a corrector.Source: §D.1 · p. 19
  • FID uses modified Inception V1 for CIFAR-10 and clean-FID with Inception V3 for the remaining datasets.The clean-FID implementation reports higher scores than the original FID calculation.Source: §D.1 · p. 19

E.1. Ablation Study on Reconstruction Term

The reconstruction-term ablation separates likelihood estimation from sample generation, while the appendix provides full Soft Truncation ablations across settings.

E.1. Ablation Study on Reconstruction Term

  • Including the reconstruction term improves NLL/NELBO but sacrifices sample generation performance.This comparison is reported for CIFAR-10 trained with DDPM++ (VP).Source: §E.1 · p. 19
  • With random τ, the reconstruction-based variational bound is not tight to negative log-likelihood, whereas omitting reconstruction yields a bound tight up to a constant.This tightness holds regardless of τ and motivates the name Maximum Perturbed Likelihood Estimation (MPLE).Source: §E.1 · p. 19

E.2. Full Tables

  • Full Soft Truncation results include ablations over weightings, architectures, diffusion SDEs, σmin, Pk, and combinations with normalizing flows.These experiments span CIFAR-10, ImageNet32, and CelebA.Source: §E.2 · p. 19, Table 11 · p. 20, Table 12 · p. 20

E.3. Generated Samples

The paper illustrates the sampling process and reports non-cherry-picked generated samples across several image datasets and resolutions.

E.3. Generated Samples

  • Figures 11–16 present non-cherry-picked generated samples from the trained model, with Figures 11–15 showing random samples of CIFAR-10, CelebA, LSUN Bedroom, LSUN Church, and FFHQ 256.Source: §E.3 · p. 20, Figure 11 · p. 23, Figure 12 · p. 24
  • Additional generated-sample figures document Soft Truncation ablations involving σmin, Pk, and diffusion combined with a normalizing flow.These ablations are listed for CIFAR-10 trained with DDPM++.Source: Table 13 · p. 21, Table 14 · p. 21, Table 15 · p. 21
  • Figure 10 depicts image generation by denoising with the predictor-corrector sampler.The figure shows the image-creation process from the trained model.Source: §E.3 · p. 20, Figure 10 · p. 22
Loading 2106.05527v5…