PaperlayerView exact revision on arXiv

Source-linked AI summary

Temporal Alignment Guidance: On-Manifold Sampling in Diffusion Models

Youngrok Park, Hojung Jung, Sangmin Bae, Se-Young Yun

arXiv:2510.11057v1cs.LGcs.AI

TL;DR

Arbitrary guidance can push diffusion samples off the data manifold, where score errors accumulate and undermine fidelity. TAG uses a time predictor to estimate temporal mismatch and steer samples toward the desired manifold. The authors report reduced off-manifold behavior and improved sample quality across multiple scenarios.

  • Problem

    Arbitrary modifications to diffusion sampling can push samples into low-density regions, where unreliable score estimates accumulate errors and degrade fidelity.

    Source: §1 · p. 1, §2 · p. 3

  • Method

    TAG uses a time predictor to identify temporal mismatch and guides samples back toward the desired manifold during generation.

    Source: §1 · p. 1, §3.1 · p. 4, §5 · p. 9

  • Results

    TAG significantly reduces off-manifold behavior across multiple scenarios and improves generated-sample quality across domains and tasks.

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

  • Takeaways & Limitations

    The authors suggest TAG may be especially effective for real-world downstream tasks whose desired conditions vary in real time.

    Source: §5 · p. 10

    Once fidelity is high, further narrowing the time gap yields marginal or no quality improvement; more sophisticated time predictors may unlock additional gains.

    Source: A BROADER IMPACT AND LIMITATIONS · p. 20

Abstract

from arXiv · show

Diffusion models have achieved remarkable success as generative models. However, even a well-trained model can accumulate errors throughout the generation process. These errors become particularly problematic when arbitrary guidance is applied to steer samples toward desired properties, which often breaks sample fidelity. In this paper, we propose a general solution to address the off-manifold phenomenon observed in diffusion models. Our approach leverages a time predictor to estimate deviations from the desired data manifold at each timestep, identifying that a larger time gap is associated with reduced generation quality. We then design a novel guidance mechanism, `Temporal Alignment Guidance' (TAG), attracting the samples back to the desired manifold at every timestep during generation. Through extensive experiments, we demonstrate that TAG consistently produces samples closely aligned with the desired manifold at each timestep, leading to significant improvements in generation quality across various downstream tasks.

1 INTRODUCTION

Guidance helps diffusion models target desired properties, but modifying the reverse process can push samples off the data manifold and degrade fidelity. The paper introduces Temporal Alignment Guidance (TAG) to steer samples toward higher-density regions, reporting improved quality across domains and tasks.

1 INTRODUCTION

  • Arbitrary modifications to the reverse process can push samples into low-density regions, where score errors accumulate and outputs deviate from the true data manifold.The paper identifies this as the off-manifold phenomenon.Source: §1 · p. 1
  • TAG represents time as a probability distribution over possible values and steers samples toward higher-density regions where learned scores are reliable.This differs from approaches relying on fixed timesteps in the reverse process.Source: §1 · pp. 1–2
  • TAG significantly improves generated-sample quality across multiple domains and tasks, with experiments reporting state-of-the-art results.The overview illustrates conditional-generation examples for ImageNet, molecules, and CelebA.Source: §1 · p. 2, Figure 1 · p. 2
  • The paper identifies off-manifold behavior across multiple scenarios and presents TAG with theoretical guarantees as a general framework for guiding samples toward the desired manifold.Its contribution summary also reports extensive experiments across domains and tasks.Source: §1 · pp. 1–2

2 OFF-MANIFOLD PHENOMENON IN DIFFUSION MODELS

Off-manifold behavior arises when sampling departs from the distributions on which the learned score is reliable. The section describes several settings where external guidance or discretization can produce such departures and degrade generation.

Controlling by external guidance

  • Training-free guidance can approximate conditional sampling with an unconditional model, but its per-timestep guidance can move samples away from the learned data manifold.The guidance uses a target estimate approximated with Tweedie’s formula.Source: §2 · p. 3

Multi-conditional guidance

  • Naively combining independent conditional score functions does not yield the multi-conditional score function, complicating guidance for multiple target properties.The section gives the relation ∇x log p(x|c1, c2) ≠ ∇x log p(x|c1) + ∇x log p(x|c2).Source: §2 · p. 3

Few-step generation

  • Reducing function evaluations with the probability-flow ODE can introduce discretization errors that accumulate during reverse sampling and cause off-manifold behavior.The paper points to Appendix B.5 for further details.Source: §2 · p. 3

Degradation of sample quality in low-density regions

  • In low-density regions, the learned score becomes unreliable, and its approximation errors can compound across generation steps to degrade sample quality.The score is trained to guide samples toward high-density regions of the noisy data distribution at each timestep.Source: §2 · p. 3

Controlling by external guidance

  • In a toy Gaussian-mixture example, adding an external drift at every reverse-process step produces samples far from the original distribution.The example is described as illustrating how off-manifold behavior can harm diffusion sampling.Source: §2 · p. 3

3 METHOD

TAG uses timestep information to estimate temporal misalignment and correct samples at each reverse step, including when task guidance or perturbations alter the sampling path. The section combines this mechanism with conditional guidance, gives theoretical analysis, and reports improved quality under increasing corruption.

3.1 TEMPORAL ALIGNMENT GUIDANCE (TAG)

  • A lightweight time predictor classifies timestep labels, and its gradient estimates TLS, whose field directs samples toward high-probability regions when conventional scores struggle off-manifold.Source: §3.1 · p. 4

3.2 IMPROVING GUIDANCE WITH TAG

  • In training-free guidance, TAG adds a time-misalignment correction alongside property guidance, treating time as another condition in conditional sampling.Source: §3.2 · pp. 4–5

3.3 THEORETICAL ANALYSIS OF TAG

  • TLS combines score functions across timesteps, attracting samples toward the target-time manifold while repelling competing manifolds, especially when target-time density is low.Source: §3.3 · p. 5

3.4 UNDERSTANDING TAG UNDER CORRUPTED REVERSE PROCESS

  • As TAG strength ω increases, FID and IS improve while Time-Gap decreases; TAG also alleviates FID degradation as corruption σ increases.Table 1 varies TAG strength across corruption levels; Figure 3 compares FID across corruption levels with and without TAG.Source: Table 1 · p. 6, Figure 3 · p. 6, §3.4 · p. 6

3.3 THEORETICAL ANALYSIS OF TAG

  • Under mild assumptions, TAG reduces the upper bound on total-variation distance between the generated-sample distribution and the target distribution.Source: §3.3 · p. 6

3.4 UNDERSTANDING TAG UNDER CORRUPTED REVERSE PROCESS

  • The Time-Gap metric quantifies temporal deviation during corrupted reverse sampling; lower values indicate closer temporal-manifold alignment and correlate with improved generation quality.Source: §3.4 · p. 6

4 EXPERIMENTS

The experiments evaluate TAG on standard guidance benchmarks, multi-conditional generation, few-step sampling, and large-scale text-to-image tasks. Across these settings, the authors report improved fidelity or alignment, including under stronger guidance and fewer sampling steps.

4.1 TFG BENCHMARK

  • TAG improves fidelity while maintaining conditioning across most TFG benchmark tasks, and remains robust where several comparison methods degrade under external-guidance drift.The benchmark uses six pretrained models across image restoration, conditional generation, molecular generation, and audio synthesis.Source: §4.1 · p. 7, Table 2 · p. 7

4.2 MULTI-CONDITIONAL GUIDANCE

  • TAG significantly outperforms independent guidance combinations across multi-conditional tasks; single and unconditional time predictors match or exceed multi-conditional performance.The evaluations include molecular-property combinations and CelebA attribute combinations.Source: §4.2 · p. 8, Table 5 · p. 8

4.3 FEW-STEP GENERATION

  • 41.1% lower FID in single-step CIFAR10 sampling; TAG also boosts sample quality across accelerated sampling step counts, particularly with fewer steps.The authors treat discretization error from skipped timesteps as external noise that TAG can mitigate in practice.Source: Table 4 · p. 8, §4.3 · p. 9

4.4 LARGE-SCALE TEXT-TO-IMAGE GENERATION

  • In large-scale text-to-image tests, adding TAG to DAS increases final reward and reduces Time-Gap, while TFG+TAG improves style adherence and reduces Time-Gap.The tests cover reward alignment with aesthetic, CLIP, and combined objectives, plus style transfer.Source: §4.4 · p. 9, Table 6 · p. 9

4.5 ABLATION STUDY

  • The ablations examine predictor training steps, guidance strengths under added noise, scaling to 50k samples, Time-Gap correlations with image-quality scores, and predictor robustness.The authors also report hyperparameter-sensitivity and additional-baseline analyses in the appendices.Source: §4.5 · p. 9

5 CONCLUSION AND FUTURE WORKS

The conclusion presents TAG as a time-predictor-based way to measure and reduce off-manifold deviation during diffusion generation. Experiments report reduced deviation across multiple scenarios, while further domains and predictor architectures are left for future work.

5 CONCLUSION AND FUTURE WORKS

  • A time-prediction mechanism measures Time-Gap to identify off-manifold phenomena, and TAG steers samples toward the desired manifold at each timestep.The authors describe TAG as a guidance mechanism for reducing the measured gap.Source: §5 · p. 9
  • Experiments show that TAG significantly reduces off-manifold phenomena across multiple scenarios.The authors identify real-time-varying desired conditions as a setting where TAG could be especially effective.Source: §5 · p. 10
  • The authors identify reinforcement-learning tasks and discrete diffusion models as promising domains for future investigation.Source: §5 · p. 10

A BROADER IMPACT AND LIMITATIONS

The paper notes possible misuse of improved generative models and a limit to further quality gains from narrowing Time-Gap at high fidelity. Its background sections review diffusion models and score-based sampling.

A BROADER IMPACT AND LIMITATIONS

  • Improved generative models may be misused to generate harmful images, unsafe compounds, or audio, and the authors call for stronger safeguards.The paper also reports that narrowing Time-Gap yields only marginal or no quality improvement once sample fidelity is high.Source: A BROADER IMPACT AND LIMITATIONS · p. 20

B.1 DIFFUSION MODELS

  • Diffusion models sample from data distributions through a forward noising process and learn a score function to denoise noisy data.The background describes conditional diffusion as sampling from a posterior given a condition, using a conditional score function.Source: §B.1 · pp. 20–21

B.2 SCORE BASED DIFFUSION MODEL

  • The background covers score-based diffusion formulations, including VP-SDE and VE-SDE processes and probability-flow ODE sampling.It notes that a discretized probability-flow ODE sampler can be interpreted as DDIM and used to skip network evaluations.Source: §B.2 · p. 21

B.3 TRAINING-FREE GUIDANCE

This section details training-free TAG guidance, including multi-condition time-predictor variants, timestep-specific manifolds, and analyses of guidance and sampling. It also completes a Gaussian-posterior proof and presents TAG representations for discretized and continuous-time diffusion models.

B.3 TRAINING-FREE GUIDANCE

  • Training-free guidance uses Tweedie-denoised x0 estimates for property conditioning, while TAG adds time prediction and reparameterizes samples to reduce off-manifold drift.The time predictor estimates the plausible timestep given a noisy sample and condition; the guided score aligns each sampling step with the property and time.Source: §B.3 · pp. 22–23
  • For multiple conditions, TAG avoids a predictor per combination by sequentially reparameterizing samples with a single-condition or unconditional time predictor, extending to k > 2 conditions.A separate multi-condition predictor fully models condition effects but becomes infeasible for complex or high-dimensional combinations.Source: §B.3 · pp. 23–24

B.4 MANIFOLD ASSUMPTION

  • The correct manifold at timestep t is defined as points whose density pt(x) is at least threshold ϵt; a sample is off-manifold when it falls outside it.Further theoretical understanding of this definition is left for future work.Source: §B.4 · p. 24

B.5 FEW STEP GENERATION

  • Skipping diffusion-model evaluations at some timesteps causes discretization error, which higher-order terms can reduce; combining TAG with a higher-order solver is left for future work.The appendix discusses this in the context of few-step generation.Source: §B.5 · p. 26

C.1 UPPER BOUND BY EXTERNAL DRIFT

  • Even with a perfectly known score, external guidance can cause worst-case distributional divergence between guided and original reverse-process outputs.Proposition C.1 bounds the total variation distance between the final distributions.Source: §C.1 · p. 27

C.2 PROOF OF PROPOSITION B.1

  • The Proposition B.1 proof models temporal likelihood through a property loss, assumes a prior on xt, and derives the posterior using a first-order likelihood expansion.For a Gaussian posterior, the proof notes that the MAP estimate equals the expectation.Source: §C.2 · p. 28

C.3 PROOF OF PROPOSITION B.2

  • The Proposition B.2 proof frames the two-condition derivation as hierarchical Bayesian inference, with intermediate samples reflecting c1 before c2 and flat priors in its posterior setup.Source: §C.3 · pp. 28–29
  • For Gaussian posteriors, the MAP estimate equals the expectation, completing the proof.Source: §C.3 · p. 29

C.4 PROOF OF THEOREM 3.3

  • For discretized diffusion timesteps, the i-th-step TAG is represented by rearranging terms using p_tot := sum_j p_j(x).Source: §C.4 · p. 29

C.5 CONTINUOUS TIME LIMIT OF TAG

  • For continuous-time diffusion models, Theorem C.2 states that the TLS score can be decomposed.Source: §C.5 · p. 30

C.6 PROOF OF PROPOSITION 3.4

Proposition 3.4 states that TAG changes the timestep energy barrier map, with the proof deriving the result by integrating score terms. The following section sets up a Langevin-dynamics analysis using the JKO gradient-flow perspective.

C.6 PROOF OF PROPOSITION 3.4

  • TAG changes the energy barrier map U_k(x) = −log p_k(x) to Φ_k(x) at each timestep.The proof identifies the gradient of the potential with the score and integrates the resulting relation, up to additive constants.Source: §C.6 · p. 30

C.7 FORMAL VERSION OF THEOREM 3.5 WITH ITS PROOF

  • The JKO scheme frames the Langevin Fokker–Planck equation as gradient flow of KL divergence under the Wasserstein-2 metric.This provides the basis for analyzing convergence of TAG-modified correction sampling.Source: §C.7 · p. 30
  • The proof introduces original and TAG-modified Langevin dynamics at timestep t_k for the convergence analysis.The displayed dynamics are not fully specified in the supplied passages.Source: §C.7 · p. 30

Fokker-Plank equation and gradient flow

The appendix analyzes how TAG modifies Langevin gradient flow and derives conditional convergence guarantees. It also distinguishes TAG from prior timestep-correction and regularization approaches.

Fokker-Plank equation and gradient flow

  • When the expectation of A(t) is strictly positive, the analysis says TAG helps escape a low-density region faster than original Langevin dynamics.Lower escape time is associated with faster convergence toward a high-density region and accelerated initial convergence.Source: Fokker-Plank equation and gradient flow · pp. 31–32
  • TAG reduces the upper bound on expected low-density-region escape time by a factor of 1 + η^2.The corollary also states that continuous modified Langevin flow reduces KL divergence to the target measure.Source: Fokker-Plank equation and gradient flow · p. 32
  • The formal guarantee assumes score-approximation error decreases monotonically with density, so error grows as samples move farther from the true distribution.The stated rationale is that the score network is trained on samples from p_k and rarely encounters points where p_k(x) is approximately zero.Source: Fokker-Plank equation and gradient flow · p. 33
  • Under Assumption C.11 and Theorem C.5’s assumptions, TAG’s modified Langevin correction improves the convergence guarantee for small t_0.The theorem compares total-variation error for the modified and original reverse processes after simulating the correction until time s.Source: Fokker-Plank equation and gradient flow · p. 33

D.1 RELATED WORKS

  • Unlike prior timestep approaches, TAG learns a dedicated predictor for p(t | x_t) and uses its score to pull samples toward their temporal manifold.The paper presents this as yielding convergence guarantees and empirical gains; it contrasts with timestep embedding or schedule adjustments and time-input perturbations (San-Roman et al., 2021).Source: §D.1 · p. 34
  • Unlike fine-tuning or trajectory-based control approaches, TAG requires neither offline history nor multiple iterations, and can handle conditions that change per generation.The comparison notes that fine-tuning for practical downstream tasks is costly when target conditions vary in real time, while control approaches usually require multiple trajectories.Source: §D.1 · p. 34

D.2 COMPARISON WITH BASELINE METHODS

This section contrasts TAG with methods that alter schedules, timesteps, or score inputs, emphasizing TAG’s adaptive temporal correction. It also presents comparisons with baseline approaches and discusses TAG as a low-overhead alternative to fine-tuning.

D.2 COMPARISON WITH BASELINE METHODS

  • TAG’s sample-specific temporal correction differs from fixed schedule strategies, timestep reassignment in TCS, time perturbations in TSG and SG, and heuristic exposure-bias adjustments.TAG keeps the solver’s current timestep and uses the posterior over possible timesteps rather than a single reassignment.Source: §D.2 · p. 35
  • TAG differs from CDM by using a time-classifier gradient to attract samples toward the desired timestep, rather than estimating output log density and approximating the score.For the multi-condition setting, Table 7 reports no improvement over the original model from input perturbation under off-manifold phenomena.Source: §D.2 · p. 35, Table 7 · p. 37
  • TAG adds a separate temporal-likelihood-score correction that strengthens with distance from the manifold, unlike PC sampling and NCSN’s reliance on the score ∇x log pt(x).The section connects this adaptiveness to robustness when off-manifold score estimates are error-prone.Source: §D.2 · p. 36
  • Unlike MultiDiffusion’s fusion of paths for spatial control, TAG applies temporal alignment to mitigate drift from combined standard guidance terms.Its temporal gradient can be approximated with predictors conditioned on single conditions or an unconditional predictor.Source: §D.2 · p. 36
  • TAG trains only a lightweight time predictor in minutes and applies an inference-time correction without changing the base model’s weights.The section contrasts this with fine-tuning approaches that require task-specific labeled data, architecture modifications, and hours of optimization.Source: §D.2 · p. 36

Baseline experiments

  • Table 8 comparisons report that TAG improves DPS performance while other methods struggle.The section also states that TAG outperforms score-based correction baselines and consistently improves over exposure-bias baselines.Source: §D.2 · pp. 35–36, Table 8 · p. 37

E.1 TOY EXPERIMENT

The toy experiment uses a two-Gaussian mixture with an external drift to examine TAG’s predictor size. The reported results indicate that TAG performs well with a predictor smaller than the score network.

E.1 TOY EXPERIMENT

  • The toy data comprise 40,000 samples from a mixture of two Gaussians, diffused with a 100-timestep DDPM VP-SDE and external drift v(x, t) = −0.01x.The mixture components are centered at (10, 10) and (−10, −10), each with identity covariance and weight one half.Source: §E.1 · p. 37
  • The toy diffusion model and time predictor are trained with 3-layer and 5-layer MLPs, respectively, for 5,000 full-batch epochs on one RTX 3090 GPU.Source: §E.1 · p. 37
  • TAG performed well with a toy-experiment predictor smaller than the score network, and the main experiments use a SimpleCNN predictor much smaller than the UNet backbone.The authors cite Table 9 for the toy ablation and Table 2 for the main-experiment comparison.Source: §E.1 · p. 37

E.2 CORRUPTED REVERSE PROCESS

This section reports corrupted reverse-process results and training-free guidance setups across several tasks, as well as time-predictor performance and large-scale image-alignment experiments. It also reports TAG’s FID and IS improvements in unconditional few-step generation when fewer model evaluations are used.

E.2 CORRUPTED REVERSE PROCESS

  • With σ=0.2 fixed, stronger time guidance consistently increased generation quality until performance saturated.The appendix presents a grid search over guidance strength ω.Source: §E.2 · p. 38, Table 10 · p. 38
  • At σ=0.05, 0.1, 0.2, and 0.3, lowest-FID TAG strengths were w=0.2, 1.25, 4.5, and 200.0, respectively; suitable TAG significantly improved generation quality.The comparison is between the original diffusion process and TAG at different noise levels.Source: §E.2 · p. 39, Table 12 · p. 39

E.3.1 LABEL GUIDANCE

  • For label guidance, 50,000-sample CIFAR-10 evaluation yielded improved, benchmark-consistent FID scores; target-label accuracy measures generation validity.The benchmark also evaluates ImageNet, while its rapid experiments use fewer samples.Source: §E.3.1 · p. 39, Table 13 · p. 40

E.3.2 GUASSIAN DEBLURRING / E.3.3 SUPER-RESOLUTION

  • For Gaussian deblurring, training-free guidance uses an L2 loss between blurred estimates and target; super-resolution uses training-free guidance to restore high-resolution images, with FID and LPIPS evaluating fidelity and conditioning effects.Source: §E.3.2 · p. 40, §E.3.3 · p. 40

E.3.4 MULTI-CONDITIONAL GUIDANCE

  • Multi-conditional guidance targets combinations of gender, hair color, or age; naive score summation can create off-manifold artifacts, and evaluation uses KID and conjunction accuracy.The experiments use CelebA-HQ images and assess whether generated samples satisfy all target attributes.Source: §E.3.4 · pp. 40–41

E.3.5 MOLECULAR GENERATION / E.3.6 AUDIO GENERATION

  • Molecular guidance minimizes property gaps and reports atom stability and MAE; audio declipping and inpainting use L2 guidance and are evaluated with FAD and DTW.Molecular targets include six QM-9 quantum properties; audio tasks repair clipped signals or reconstruct missing signal regions.Source: §E.3.5 · pp. 41–42, §E.3.6 · p. 42

E.4 TIME PREDICTOR

  • Lightweight time predictors use CNNs for image/audio and EGNNs for molecules; predictions are accurate for t<600 but errors rise near T on CIFAR-10 and molecular data, while better checkpoints improve TAG metrics.The checkpoint comparison contrasts 10K and 30K training steps on CIFAR-10.Source: §E.4 · pp. 43–44, Table 16 · p. 44

E.5 LARGE-SCALE TEXT-TO-IMAGE GENERATION

  • Large-scale image experiments apply reward alignment to aesthetic, CLIP, and combined objectives, and style transfer by matching intermediate CLIP-feature Gram matrices.Reward experiments report mean reward and Time-Gap; style transfer is evaluated with Style Score and CLIP Score.Source: §E.5 · pp. 44–45

F.1 FEW STEP UNCONDITIONAL GENERATION

  • TAG significantly improves FID and IS when fewer evaluation steps are used in unconditional generation.The comparison reports 50,000 samples; the supplied passages do not give the metric values.Source: §F.1 · p. 45, Table 17 · p. 45
  • Fewer function evaluations are associated with more severe off-manifold effects during the reverse diffusion process.The experiment focuses on few-step generation, where discretization error occurs.Source: §F.1 · p. 45
  • The experiment uses CIFAR-10 DDPM and DDPM sampling, with NFE denoting the number of diffusion-model evaluations during reverse generation.The setup follows the paper’s Section 4.3 experiment.Source: §F.1 · p. 45

F.2 TIME GAP

This section defines Time-Gap as a measure of temporal deviation and examines its behavior across timesteps and in relation to image-quality metrics. It also reports predictor-size comparisons and a limitation in Time-Gap’s correlation with FID.

F.2 TIME GAP

  • Time-Gap is the average absolute difference between the predicted timestep index and the true timestep index for a sample.The section says a lower Time-Gap indicates samples are closer to their expected temporal manifold.Source: §F.2 · p. 46
  • A 1.48 M-parameter SimpleCNN matches the time-gap performance of a 17.38 M-parameter UNet encoder.The comparison is reported across unconditional and conditional models and training checkpoints in Table 18.Source: §F.2 · p. 48, Table 18 · p. 48
  • As NFE increases, Time-Gap decreases while FID decreases and IS increases, supporting Time-Gap as a measure of off-manifold phenomena.The comparison varies NFE from 1 to 50 using 50,000 CIFAR10-DDPM samples with DDIM sampling.Source: §F.2 · p. 48, Figure 10 · p. 48
  • When the reverse process is already good and Time-Gap is small, it often loses correlation with FID.The authors suggest that improving the time predictor may reduce this limitation.Source: §F.2 · p. 48
  • The authors use Time-Gap to quantify off-manifold phenomena, not to propose a new metric.Source: §F.2 · p. 48

G VISUALIZATIONS OF GENERATED SAMPLES

The section presents qualitative examples for four experimental settings—DPS, DPS+TAG, TFG, and TFG+TAG—across image, molecule, and face-generation tasks. It specifies the dataset configurations and target conditions used for those examples.

CIFAR-10

ImageNet

QM9

CelebA-HQ

Loading 2510.11057v1…