ArXiv: 2510.11057

🎯 Pitch

Diffusion models fall apart when external guidance pushes them off the data manifold, but a simple time predictor can detect—and correct—these deviations during sampling. By adding a gradient that pulls samples back to the right "time" on the manifold, Temporal Alignment Guidance (TAG) slashes FID by nearly half under heavy noise while boosting downstream task accuracy, all without retraining.


1. Executive Summary

This paper introduces Temporal Alignment Guidance (TAG), a general corrective mechanism that steers samples back onto the correct data manifold at every timestep of the diffusion reverse process by leveraging a learned time predictor to estimate temporal deviations—operationalized as a gradient term ∇x log p(t|x) that pulls samples toward high-density regions of the marginal distribution at the current timestep. Evaluated across the TFG benchmark with models including CIFAR10-DDPM, ImageNet-DDPM, and Molecule-EDM, TAG consistently improves both fidelity and conditioning accuracy when combined with standard training-free guidance methods like DPS and TFG, yielding relative improvements up to 47.9% in FAD on audio declipping and 40.3% in MAE on molecular polarizability control. The correction is particularly effective under strong external guidance or severe perturbations—reducing FID from 410.1 to 223.2 at noise level σ = 0.3—while providing diminishing returns on problems where samples already reside near the desired manifold, establishing that TAG's benefit is largest precisely when off-manifold drift is most severe.

2. Context and Motivation

The Core Problem: Diffusion Models Drift Off-Manifold Under Arbitrary Guidance

The fundamental problem this paper addresses is deceptively simple: when you modify the reverse diffusion process to steer samples toward desired properties, the generated samples often fall off the data manifold, producing unrealistic outputs. This "off-manifold phenomenon" (Section 2) is not a niche edge case—it arises systematically across virtually every practical application where diffusion models are used for controlled generation beyond their original training distribution.

To understand why this matters, we need to examine how diffusion models typically operate. During standard unconditional generation, the reverse process follows a carefully learned trajectory through noise space. At each timestep tt, the model's score function xlogpt(x)\nabla_x \log p_t(x) points toward regions of high probability density for that specific noise level. The forward diffusion process guarantees that these marginal distributions pt(x)p_t(x) have full support—that is, pt(x)>0p_t(x) > 0 everywhere in the data domain (Appendix B.4)—because the added Gaussian noise smooths out any low-dimensional structure. However, the learned score function is only reliable in regions where it encountered sufficient training data: areas of high pt(x)p_t(x). When an external guidance term pushes a sample into low-density regions where pt(x)0p_t(x) \approx 0, the score estimate degrades catastrophically.

The paper identifies this as a score approximation error problem (Section 2, paragraph on "Degradation of sample quality in low-density regions"):

"if an external force vv drives a sample to the low density region pt(xt)0p_t(x_t) \approx 0, the score function logpt(xt)\nabla \log p_t(x_t) estimated by the diffusion model becomes unreliable, as it is trained on noisy data that assumes the forward process is intact at the given timestep."

Critically, this error compounds over timesteps. An error at timestep tt means the sample arrives at an incorrect position for timestep t1t-1, making the score estimate at t1t-1 even less reliable, which produces an even larger error, and so on. This compounding effect explains why naively adding guidance terms—even principled ones—can produce samples that diverge dramatically from realistic outputs.

Why This Problem Is Important: The Ubiquity of Guided Generation

The off-manifold problem is important not because it affects a narrow technical corner case, but because guided generation is the dominant paradigm for making diffusion models useful in practice. The paper catalogs several concrete scenarios where the problem manifests (Section 2):

1. Training-free guidance (Section 2, "Controlling by external guidance"). This is perhaps the most practically important setting. Given an unconditional diffusion model, practitioners often want to condition it on external signals without retraining. The standard approach is to modify the reverse SDE from its original form:

dx=[f(x)g2(t)xlogqt(x)]dt+g(t)dwˉtdx = \left[f(x) - g^2(t)\nabla_x \log q_t(x)\right] dt + g(t)d\bar{w}_t

to include an additional guidance term v(x,c,t)v(x, c, t):

dx=[f(x)g(t)2(xlogqt(x)+v(x,c,t))]dt+g(t)dwˉtdx = \left[f(x) - g(t)^2 \left(\nabla_x \log q_t(x) + v(x, c, t)\right)\right] dt + g(t)d\bar{w}_t

For training-free guidance specifically, this term takes the form:

v(xt,c,t)=xtlogp(cx^0)v(x_t, c, t) = \nabla_{x_t} \log p(c|\hat{x}_0)

where x^0\hat{x}_0 is a denoised estimate obtained via Tweedie's formula (Equation 4). While mathematically motivated—this approximates sampling from the conditional distribution p(xc)p(x|c) using only an unconditional model—the approach has a fundamental flaw: each guidance step can push the sample away from the manifold of noisy data that the score function was trained on. The paper cites Shen et al. (2024) as directly observing that "this extra guidance in each timestep make samples far from the original learned data manifold."

The practical implications are severe. Training-free guidance enables applications like image inpainting, super-resolution, molecular property control, and audio restoration without task-specific retraining—but if the guidance degrades sample quality, the outputs become useless regardless of how well they satisfy the conditioning objective.

2. Multi-conditional guidance (Section 2, "Multi-conditional guidance"). When users want to control multiple attributes simultaneously—for example, generating a face that is both "female" and "young"—the natural approach is to combine guidance terms for each condition. However, as the paper states:

xlogp(xc1,c2)xlogp(xc1)+xlogp(xc2)\nabla_x \log p(x|c_1, c_2) \neq \nabla_x \log p(x|c_1) + \nabla_x \log p(x|c_2)

The left side is the true multi-conditional score. The right side is what practitioners actually do: sum independent conditional scores. This approximation is not merely inexact—it can be categorically wrong, producing samples that satisfy neither condition well or that achieve the conditions at the cost of severe realism degradation. The paper cites Du et al. (2023) as identifying this issue directly.

3. Few-step generation (Section 2, "Few-step generation"). Diffusion models are notoriously slow, requiring hundreds or thousands of sequential denoising steps. Accelerated samplers like DDIM reduce the number of function evaluations by skipping timesteps, but this introduces discretization errors in the probability flow ODE integration. Each skipped step means the sample trajectory deviates from the true reverse path, producing off-manifold drift even without any external guidance. This is a pure inference-time efficiency problem: users want fast generation, but speed comes at the cost of quality.

4. A unifying theme: all practical modifications break the careful calibration of the learned score. Whether the perturbation comes from external guidance, multi-condition combination, or discretization error, the mechanism of failure is the same. The diffusion model's score function was trained under the assumption that at timestep tt, the input xtx_t would be a noisy version of a real data point—that is, xtx_t would lie in a high-probability region of pt(x)p_t(x). Any modification to the reverse process violates this assumption, and the model has no mechanism to recover.

Where Prior Approaches Fall Short

The paper identifies several families of prior work that attempt to address related problems, each with specific limitations:

Score-based correction methods (Langevin dynamics, predictor-corrector sampling). Song & Ermon (2019) and Song et al. (2021b) proposed using Langevin dynamics steps—short MCMC chains that use the learned score function to refine samples—as correctors between predictor steps. The intuition is appealing: if the score function is accurate, repeatedly following it should push samples toward high-density regions. However, the paper makes a critical empirical observation (Appendix D.2):

"even an accurate score estimate can struggle to guide samples out of inherently flat probability landscapes. Indeed, our empirical findings in Appendix D.2 show that corrector sampling becomes ineffective, sometimes degrade the sample quality under external guidance."

The fundamental issue is that in low-density regions, the score function itself is unreliable (it was trained on data that assumes the forward process is intact), and even when it is accurate, the gradient may be too weak to escape flat regions quickly. The paper quantifies this: when applying Langevin dynamics on top of DPS on CIFAR-10, FID is 226.8 versus 217.1 for DPS alone (Table 8)—the correction actually hurts in the presence of off-manifold drift.

Exposure bias methods. Several works (Ning et al., 2023; 2024; Li et al., 2024a) address the train-test mismatch in diffusion models—the fact that during training, the model always sees clean forward-process samples, but during inference, it sees its own (potentially erroneous) outputs. Ning et al. (2023) add random perturbations during training to make the model robust to its own errors. Ning et al. (2024) scale the vector norm of model outputs as a heuristic correction. Li et al. (2024a) identify variance across sample batches to adjust time information. The paper's experiments (Tables 7 and 8) show that while these methods help in standard settings, they degrade or fail to help under strong external guidance. Input perturbation with η=0.10\eta = 0.10 on CIFAR-10 DPS produces FID 376.6 and accuracy 25.4% versus the baseline of 332.0 and 28.5% (Table 7). Epsilon Scaling achieves FID 186.0 (better than DPS at 217.1) but at the cost of reduced accuracy (53.0% vs. 57.5%) (Table 8). These methods are essentially heuristic patches that do not address the root cause: the sample has left the manifold where the score function is trustworthy.

Time perturbation methods (TSG, Self-Guidance). Sadat et al. (2024) and Li et al. (2024b) exploit the score model's sensitivity to its time input by perturbing tt (e.g., evaluating at t±δt \pm \delta) to derive contrastive guidance signals. The idea is that if the model's output changes significantly when the time input is slightly altered, the current sample may be temporally misaligned. However, as shown in Table 2, Timestep Guidance produces FID 393.2 and accuracy 9.4% on CIFAR-10 DPS—far worse than the baseline—and Self-Guidance produces FID 205.4 with accuracy 51.6%, improving fidelity slightly but at the cost of conditioning accuracy. These methods rely on local perturbation signals that may be too weak or too noisy when the sample is far off-manifold.

Time Correction Sampler (TCS). Jung et al. (2024) train a time predictor to classify which timestep a sample belongs to, then use the predicted timestep t~=argmaxϕ(xt)\tilde{t} = \arg\max \phi(x_t) to directly modify the diffusion step: instead of using the score at the current timestep tt, TCS evaluates the score at the predicted timestep t~\tilde{t} and adjusts the noise schedule accordingly. This is a "hard" reassignment: the sample is treated as if it were at a different point in the diffusion process. The paper's experiments (Table 2) show TCS degrades severely under external guidance—FID 213.4 and accuracy 29.4% on CIFAR-10 DPS, versus 217.1 and 57.5% for DPS alone—suggesting that hard temporal reassignment is too brittle when the guidance term systematically pushes samples off-manifold.

Fine-tuning approaches (RL-based, ControlNet, IP-Adapter). Methods like DDPO (Fan et al., 2023), AlignProp (Clark et al., 2024), ControlNet (Zhang et al., 2023), and IP-Adapter (Ye et al., 2023) adapt diffusion models for downstream tasks through additional training. However, as the paper notes (Appendix D.2):

"fine-tuning diffusion models for practical downstream tasks is highly costly where target condition vary in real-time"

These approaches require task-specific data collection, model architecture modifications, and hours of gradient-based optimization. They cannot handle scenarios where the target condition changes for every generation—a common requirement in interactive applications.

Classification Diffusion Models (CDM). Yadin et al. (2024) also use a timestep classifier, but for a fundamentally different purpose: to estimate density ratios and approximate score functions. The paper demonstrates (Appendix D.2) that CDM's theoretical framework can be derived as a special case of TAG's more general decomposition (Theorem 3.3), but CDM does not use the temporal gradient as a corrective force during sampling.

How This Paper Positions Itself

TAG is positioned as a general, lightweight, inference-time corrective mechanism that addresses the root cause of off-manifold drift rather than patching its symptoms. The key conceptual distinction is that TAG introduces a new gradient term—the Time-Linked Score (TLS), xlogp(tx)\nabla_x \log p(t|x)—that actively pulls samples back toward the correct temporal manifold at each step.

This is fundamentally different from all prior approaches:

  • Unlike Langevin correction, TAG does not rely solely on the (potentially degraded) score function xlogpt(x)\nabla_x \log p_t(x). The TLS provides an independent signal derived from a separate time predictor, which is trained to recognize temporal identity even when the score function is unreliable. Theorem 3.3 shows that the TLS is a linear combination of score functions from multiple timesteps, providing both attractive forces toward the correct manifold and repulsive forces away from incorrect ones.

  • Unlike TCS's hard reassignment, TAG maintains the current timestep for the primary score evaluation while adding the TLS as a soft correction. The decomposition in Equation 12 shows this is a weighted combination of all timestep-specific score functions, providing a more robust and adaptive correction than a single reassignment.

  • Unlike exposure bias heuristics, TAG is principled: it is derived from a probabilistic reinterpretation of the timestep variable as a conditioning signal rather than a fixed input. The correction term xlogp(tx)\nabla_x \log p(t|x) emerges naturally from treating time as an additional condition in a Bayesian framework (Section 3.2).

  • Unlike fine-tuning methods, TAG requires only training a lightweight time predictor (a SimpleCNN with ~1.48M parameters—8.5% the size of a UNet encoder—or a modified EGNN for molecular data) on the straightforward task of classifying timesteps from noisy samples. As the paper emphasizes in Appendix E.4, this training completes in minutes on minimal compute, and the resulting predictor transfers across guidance objectives without retraining.

The paper's central claim—and the motivation for the entire framework—is that temporal misalignment is the fundamental mechanism behind off-manifold degradation. By explicitly measuring and correcting temporal alignment at each step, TAG provides a unified solution that works across the diverse scenarios cataloged in Section 2: external guidance, multi-conditional generation, few-step sampling, and even extreme adversarial perturbations (the corrupted reverse process experiments in Section 3.4). This unification is what distinguishes TAG from the patchwork of task-specific fixes that characterize prior work.

The paper does not claim to eliminate the need for careful guidance design or to make any arbitrary guidance term safe. Rather, it provides a general robustness mechanism that makes diffusion models more tolerant of the inevitable score approximation errors that arise whenever the standard reverse process is modified. The Time-Gap metric (Definition F.1) operationalizes this: by measuring the average absolute deviation between predicted and actual timesteps, it provides a quantitative diagnostic for how severely off-manifold a generation process has become, and TAG's effect can be directly measured as a reduction in this gap.

3. Technical Approach

3.1 Reader Orientation

This paper introduces Temporal Alignment Guidance (TAG), a lightweight plug-in corrective mechanism that can be added to any diffusion model's reverse process to actively pull samples back toward the correct data manifold at every denoising step. TAG solves the off-manifold problem—where external guidance terms push generated samples into low-probability regions where the learned score function becomes unreliable—by training a small auxiliary time predictor that estimates which timestep a noisy sample belongs to, then using the gradient of this prediction as an additional force that steers the sample back to regions of high probability density for its current noise level.

3.2 Big-Picture Architecture (Diagram in Words)

The TAG framework consists of four interconnected components that augment the standard diffusion reverse process:

  1. Base Diffusion Model — A pretrained unconditional or conditional diffusion model (e.g., CIFAR10-DDPM, ImageNet-DDPM, Molecule-EDM, Stable Diffusion v1.5) that provides the standard score function xlogpt(x)\nabla_x \log p_t(x) at each timestep. This is the existing model that the user already has and wants to guide; TAG does not modify or retrain it.

  2. Time Predictor — A lightweight auxiliary neural network (SimpleCNN for images/audio, modified EGNN for molecules) trained to classify which timestep tt a noisy sample xtx_t was drawn from. Given an input sample at any point in the reverse process, the time predictor outputs a probability distribution over all possible timesteps, and TAG computes the gradient of the log-probability of the correct current timestep: xlogpϕ(tx)\nabla_x \log p_\phi(t | x). This gradient—the Time-Linked Score (TLS)—becomes the corrective signal.

  3. External Guidance Module (optional, task-specific) — The user's chosen guidance method that modifies the reverse process to steer samples toward desired properties. This could be DPS (Diffusion Posterior Sampling), TFG (Training-Free Guidance), multi-conditional guidance, reward alignment via DAS, or style transfer objectives. TAG is designed to work on top of any such guidance, not to replace it.

  4. TAG Correction Step — At each reverse diffusion timestep, after the external guidance (if any) has been applied, TAG adds the TLS as an additional gradient term: the sample is updated by xtxt+ωtxlogpϕ(txt)x_t \leftarrow x_t + \omega_t \cdot \nabla_x \log p_\phi(t | x_t), where ωt\omega_t is a timestep-dependent guidance strength. This corrected sample is then fed into the standard diffusion denoising step.

Information flow: Starting from random noise xTN(0,I)x_T \sim \mathcal{N}(0, I) → at each timestep tt: (1) apply external guidance if present, producing an intermediate xtx_t that may have been pushed off-manifold → (2) compute the TLS from the time predictor → (3) add the TAG correction xtxt+ωtxlogpϕ(txt)x_t \leftarrow x_t + \omega_t \cdot \nabla_x \log p_\phi(t | x_t) → (4) apply the standard reverse diffusion step using the base model's score function xlogpt(xt)\nabla_x \log p_t(x_t) → (5) proceed to t1t-1 → repeat until t=0t=0.

3.3 Roadmap for the Deep Dive

  • First, the off-manifold problem formalization — how external drift creates distributional divergence (Proposition C.1), establishing the quantitative gap that TAG must close.
  • Second, the Time Predictor — its architecture, training objective, and why it can recognize temporal identity even when the score function is unreliable.
  • Third, the Time-Linked Score (TLS) — the core gradient term, its decomposition via Theorem 3.3 revealing how it combines attractive and repulsive forces from all timesteps, and why this decomposition explains TAG's robustness.
  • Fourth, the TAG algorithm itself — the precise sampling procedure (Algorithm 1), the guidance schedule ωt\omega_t, and how TAG integrates with external guidance in both single-condition and multi-condition settings.
  • Fifth, the theoretical analysis — Proposition 3.4 on energy barrier reshaping, Theorem 3.5 on improved convergence guarantees, and the JKO scheme connection explaining why TAG accelerates escape from low-density regions.
  • Sixth, the corrupted reverse process experiments (Section 3.4) — a controlled study that isolates TAG's corrective mechanism by adding artificial noise, demonstrating the direct relationship between Time-Gap reduction and generation quality improvement.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a method paper whose core idea is that temporal misalignment—samples drifting to timesteps inconsistent with their noise level—is the root cause of off-manifold degradation, and that a learned time predictor can provide gradient signals that actively correct this misalignment at every step.


The Off-Manifold Problem: Formalizing How External Guidance Causes Distributional Divergence

The paper begins by formalizing the problem mathematically through the lens of stochastic differential equations. The standard reverse diffusion process (Anderson, 1982) is defined by the SDE:

dx=[f(x)g2(t)xlogqt(x)]dt+g(t)dwˉtdx = \left[f(x) - g^2(t)\nabla_x \log q_t(x)\right] dt + g(t)d\bar{w}_t

where f(x)f(x) is the drift coefficient from the forward process, g(t)g(t) is the diffusion coefficient controlling noise scale, xlogqt(x)\nabla_x \log q_t(x) is the true score function of the marginal distribution at time tt, and dwˉtd\bar{w}_t is a standard Wiener process in reverse time.

What it computes: This SDE describes how a sample xx evolves backward in time from pure noise at t=Tt = T to a clean data point at t=0t = 0. At each infinitesimal step, the deterministic drift term f(x)g2(t)xlogqt(x)f(x) - g^2(t)\nabla_x \log q_t(x) pushes the sample toward higher-density regions of qt(x)q_t(x) while the stochastic term g(t)dwˉtg(t)d\bar{w}_t adds noise for exploration. If the score function is perfectly known and the process is simulated exactly, the final distribution at t=0t = 0 exactly matches the data distribution.

Why this form: The forward process (which adds noise to data) has a known time-reversal formula. Anderson (1982) proved that any forward SDE can be reversed if the score function of the marginal distribution is available. This SDE is that time-reversed process—it is the unique stochastic process whose marginals match the forward process in reverse order.

When an external guidance term v(x,c,t)v(x, c, t) is added (for example, to steer samples toward a target property cc), the reverse SDE becomes:

dx=[f(x)g(t)2(xlogqt(x)+v(x,c,t))]dt+g(t)dwˉtdx = \left[f(x) - g(t)^2 \left(\nabla_x \log q_t(x) + v(x, c, t)\right)\right] dt + g(t)d\bar{w}_t

What it computes: The same reverse process, but with an additional drift term g2(t)v(x,c,t)-g^2(t) v(x, c, t) that pushes samples toward satisfying condition cc. For training-free guidance, v(x,c,t)=xlogp(cx^0)v(x, c, t) = \nabla_x \log p(c|\hat{x}_0) where x^0\hat{x}_0 is a denoised estimate of the clean data obtained via Tweedie's formula:

x^0=xt+(1αˉt)xtlogp(xt)αˉt\hat{x}_0 = \frac{x_t + (1 - \bar{\alpha}_t)\nabla_{x_t} \log p(x_t)}{\sqrt{\bar{\alpha}_t}}

where αˉt\bar{\alpha}_t is the cumulative noise schedule parameter defined by the forward process (for VP-SDE, αˉt=exp(120tβ(s)ds)\bar{\alpha}_t = \exp(-\frac{1}{2}\int_0^t \beta(s) ds)). The Tweedie estimate x^0\hat{x}_0 is the conditional expectation E[x0xt]\mathbb{E}[x_0 | x_t]—given a noisy observation xtx_t, this formula estimates what the clean data likely was. The guidance term then computes the gradient of some loss function c(A(x^0),c)\ell_c(A(\hat{x}_0), c) where AA is a property predictor (a classifier or analytic function) and c\ell_c measures discrepancy between the predicted and target properties.

Why this form: Using Tweedie's formula avoids having to run the full reverse process to evaluate the guidance—it provides an instantaneous estimate of the clean sample at every step, enabling gradient-based optimization of properties that are only defined on clean data.

The key theoretical insight is Proposition C.1, which bounds the damage this external guidance can cause:

dTV2(p0,p~0)KL(p0,p~0)120Txg(t)2pt(x)v(x,c,t)22dxdtd^2_{TV}(p_0, \tilde{p}_0) \leq \text{KL}(p_0, \tilde{p}_0) \leq \frac{1}{2} \int_0^T \int_x g(t)^{-2} p_t(x) \|v(x, c, t)\|^2_2 \, dx \, dt

where p0p_0 is the distribution of samples from the unmodified reverse process, p~0\tilde{p}_0 is the distribution from the guided reverse process, dTVd_{TV} is the total variation distance (a measure of distributional difference bounded between 0 and 1), and KL is the Kullback-Leibler divergence.

What it computes: An upper bound on how much the guided generation distribution can differ from the original data distribution. The bound depends on the integral over all timesteps of the expected squared magnitude of the guidance term, weighted by the inverse diffusion coefficient. Intuitively: the stronger the guidance signal vv, and the earlier in the process it is applied (when g(t)g(t) is smaller), the more the final distribution can diverge.

Why this form: The proof (Appendix C.1) uses Girsanov's theorem, which provides the Radon-Nikodym derivative between two SDE path measures that differ by a drift term. The KL divergence between path measures decomposes into the time-integrated expected squared drift difference (Equation 55). Data processing inequality and Pinsker's inequality then relate the final distributional divergence to the path-level divergence. Crucially, this bound does not assume the score function is perfect—it only bounds the additional error from the guidance term. If the base score function is already imperfect (as it always is in practice), the actual divergence can be larger.

This proposition establishes the fundamental tension: guidance is necessary for controlled generation, but any guidance term mathematically guarantees some distributional shift. TAG's role is to reduce this shift by providing a corrective force that counteracts off-manifold drift without eliminating the useful guidance signal.


The Time Predictor: Architecture, Training, and Why It Works

The time predictor is the workhorse of TAG. It is a neural network ϕ\phi trained to perform a seemingly simple task: given a noisy sample xtx_t, predict which timestep tt it came from. Formally, the time predictor outputs a probability distribution over TT discrete timesteps (in practice, T=1000T = 1000 for most models):

p^ϕ(xt)=softmax(fϕ(xt))[0,1]T\hat{p}_\phi(x_t) = \text{softmax}(f_\phi(x_t)) \in [0, 1]^T

where fϕf_\phi is the neural network logits and p^ϕ(xt)t\hat{p}_\phi(x_t)_t is the predicted probability that the input corresponds to timestep tt.

Training objective. The time predictor is trained with standard cross-entropy loss:

Ltp(ϕ)=Et,x0[logp^ϕ(xt)t]\mathcal{L}_{tp}(\phi) = -\mathbb{E}_{t, x_0}\left[\log \hat{p}_\phi(x_t)_t\right]

where tt is uniformly sampled from {1,...,T}\{1, ..., T\}, x0x_0 is drawn from the training data distribution, and xtx_t is generated by the forward diffusion process: xt=αˉtx0+1αˉtϵx_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon with ϵN(0,I)\epsilon \sim \mathcal{N}(0, I).

What it computes: The negative log-likelihood of the true timestep tt under the predictor's output distribution. The expectation averages over random data points and random timesteps. Minimizing this objective encourages the predictor to output high probability for the correct timestep and low probability for all others.

Why this form: Cross-entropy is the proper scoring rule for classification—it is minimized exactly when the predicted distribution matches the true conditional distribution p(txt)p(t | x_t). The paper explicitly chooses classification over regression because "overlapping supports of pt(x)p_t(x) and ps(x)p_s(x)" means that a given noisy sample may be consistent with multiple timesteps (Appendix E.4). A regression objective would force the predictor to commit to a single timestep even when ambiguity exists, producing unreliable gradients. The softmax classification output naturally handles this ambiguity by assigning partial probability mass across nearby timesteps.

Training details (Appendix E.4, Table 15). For image and audio data, the time predictor uses a SimpleCNN architecture: four convolutional layers with channel sizes (32, 64, 128, 256), each followed by ReLU activation and 2×2 average pooling, finishing with a linear layer that outputs logits for all TT timesteps. This architecture has approximately 1.48M parameters—roughly 8.5% the size of a typical UNet encoder (17.38M parameters) and substantially lighter than the full diffusion backbone. For molecular data, a modified Equivariant Graph Neural Network (EGNN, Satorras et al., 2021) processes node features, edge features, and spatial coordinates through message-passing layers, with the concatenated features passed through a feed-forward network to produce timestep logits. In conditional settings (where the time predictor must account for target properties), learned embedding vectors for conditions are concatenated before the final linear layer. Training uses the Adam optimizer with learning rate 1×1041 \times 10^{-4} for 300K iterations on most datasets (600K for ImageNet). Batch sizes range from 128 (molecules, audio, Cat) to 256 (CIFAR10, CelebA-HQ) to 1024 (ImageNet). Training completes in "minutes on a minimal computational resources" (Appendix D.2).

Why this architecture choice: The time predictor is intentionally kept lightweight for two reasons. First, the timestep classification task is easier than full denoising—the predictor only needs to recognize noise-level statistics, not reconstruct fine details. Second, TAG must compute the gradient xlogpϕ(tx)\nabla_x \log p_\phi(t | x) at every reverse step, so a lightweight predictor minimizes the per-step overhead. The paper empirically validates this choice: Table 18 shows that the 1.48M-parameter SimpleCNN matches or exceeds the time-gap performance of a 17.38M-parameter UNet encoder across multiple training checkpoints, confirming that the task does not require a large model.

Why the time predictor can succeed where the score function fails. This is the critical insight that makes TAG work. The score function xlogpt(xt)\nabla_x \log p_t(x_t) is trained by denoising score matching: given a noisy sample at a known timestep, the model must predict the noise that was added. This training procedure assumes the input xtx_t was produced by the forward process from some clean data point. When external guidance pushes a sample off-manifold, this assumption breaks—the sample no longer looks like a forward-process sample at any timestep. The score function output becomes unreliable because the model is extrapolating from its training distribution.

The time predictor, in contrast, is trained on the forward process distribution directly. It learns to recognize the statistical signature of each noise level—the variance, correlation structure, and relationship to clean data manifolds—independently of any denoising task. Critically, even when a sample is pushed off-manifold by external guidance, it still retains some noise-level characteristics (the total noise variance, the degree of blurring or corruption). The time predictor can leverage these residual statistical signatures to estimate which timestep the sample "most resembles," even if the score function for that sample would be unreliable. This independence between the time predictor's training signal and the score function's reliability is what makes the TLS a robust corrective gradient.


The Time-Linked Score (TLS): Definition and Decomposition

The TLS is the mathematical centerpiece of TAG. It is defined as:

TLS(x,t):=xlogp(tx)\text{TLS}(x, t) := \nabla_x \log p(t | x)

where p(tx)p(t | x) is the posterior probability that a given noisy sample xx originated from timestep tt, under the forward diffusion process distribution.

What it computes: The gradient with respect to the sample xx of the log-posterior probability of timestep tt. Intuitively, if the time predictor assigns low probability to the correct timestep tt (indicating the sample is temporally misaligned), then moving the sample in the direction of xlogp(tx)\nabla_x \log p(t | x) should increase the probability that it belongs to timestep tt—that is, it should move the sample back toward regions where samples at timestep tt typically reside. The TLS points toward the high-density region of pt(x)p_t(x) indirectly, by optimizing the temporal identity rather than the density directly.

Why this form: The TLS is designed to provide a corrective gradient that is reliable even when the score function is not. The posterior p(tx)p(t | x) depends on all timestep-specific densities through Bayes' rule: p(tx)=p(xt)p(t)kp(xtk)p(tk)p(t | x) = \frac{p(x | t)p(t)}{\sum_k p(x | t_k)p(t_k)}. This means the TLS is influenced by the density at all timesteps, not just the current one. If the sample has been pushed into a low-density region of pt(x)p_t(x) but a high-density region of some other ps(x)p_s(x), the TLS will recognize this temporal mismatch and provide a gradient that pushes the sample away from ss's manifold and toward tt's manifold.

The paper proves a decomposition theorem that reveals the full structure of this gradient:

Theorem 3.3 (TLS Decomposition). For discrete diffusion timesteps [t1,t2,...,tn][t_1, t_2, ..., t_n], the Time-Linked Score for a target timestep tit_i can be expressed as:

xlogp(tix)=kiptk(x)ptot(x)(xlogpti(x)xlogptk(x))\nabla_x \log p(t_i | x) = \sum_{k \neq i} \frac{p_{t_k}(x)}{p_{tot}(x)} \left(\nabla_x \log p_{t_i}(x) - \nabla_x \log p_{t_k}(x)\right)

where ptk(x)p_{t_k}(x) is the marginal density at timestep tkt_k, ptot(x)=jptj(x)p_{tot}(x) = \sum_j p_{t_j}(x) is the total density across all timesteps, and xlogptk(x)\nabla_x \log p_{t_k}(x) are the score functions at each timestep.

What this equation means operationally: The TLS for timestep tit_i is a weighted sum of pairwise differences between the score function at tit_i and the score functions at every other timestep. For each competing timestep tkt_k:

  • The weight ptk(x)ptot(x)\frac{p_{t_k}(x)}{p_{tot}(x)} is the relative probability that the sample came from timestep tkt_k. When the sample is far from manifold tit_i but close to manifold tkt_k, this weight becomes large—the correction becomes stronger precisely when it is most needed.

  • The difference xlogpti(x)xlogptk(x)\nabla_x \log p_{t_i}(x) - \nabla_x \log p_{t_k}(x) points from the tkt_k-manifold toward the tit_i-manifold. The first term xlogpti(x)\nabla_x \log p_{t_i}(x) attracts the sample toward high-density regions at the correct timestep. The second term xlogptk(x)-\nabla_x \log p_{t_k}(x) repels the sample away from high-density regions at incorrect timesteps.

Why this decomposition matters: It reveals that the TLS is not a simple "point toward the correct timestep" gradient—it is a competitive mechanism that simultaneously pulls toward the correct manifold and pushes away from all incorrect ones. This explains why TAG is more robust than hard temporal reassignment (like TCS): instead of committing to a single alternative timestep, TAG considers the full posterior distribution and derives a gradient that accounts for all possible temporal misalignments.

The paper extends this to continuous time (Theorem C.2):

xlogp(tx)=xlogpt(x)γsxlogps(x)ds\nabla_x \log p(t|x) = \nabla_x \log p_t(x) - \int \gamma_s \nabla_x \log p_s(x) ds

where γs=ps(x)pk(x)dk\gamma_s = \frac{p_s(x)}{\int p_k(x) dk}. In the continuous limit, the TLS is the score function at the target timestep minus the expected score function under the posterior distribution over all timesteps. This has an elegant interpretation: the TLS removes the "average" score direction (which may be pointing toward incorrect manifolds) and retains only the direction that is specific to the target timestep.

How the TLS is approximated in practice. The true densities pt(x)p_t(x) and their score functions are unknown. The time predictor ϕ\phi provides an approximation: xlogp(tx)xlogp^ϕ(x)t\nabla_x \log p(t | x) \approx \nabla_x \log \hat{p}_\phi(x)_t. Since the time predictor is trained with cross-entropy, its output is a consistent estimator of the true posterior p(tx)p(t | x). The gradient is computed via standard automatic differentiation through the time predictor network.


The TAG Algorithm: Integration into the Reverse Process

The full TAG sampling procedure is described in Algorithm 1. At its core, TAG modifies each reverse diffusion step by adding a single correction:

xtxt+ωtxlogpϕ(txt)x_t \leftarrow x_t + \omega_t \cdot \nabla_x \log p_\phi(t | x_t)

before the standard denoising step. The corrected sample is then processed by the base diffusion model as usual:

xt1ReverseStep(xt,xlogpθ(xt,t))x_{t-1} \leftarrow \text{ReverseStep}(x_t, \nabla_x \log p_\theta(x_t, t))

where ReverseStep implements the chosen sampler (DDPM, DDIM, or other).

The guidance schedule ωt\omega_t. The strength of the TAG correction is controlled by a timestep-dependent weight ωt\omega_t. In the theoretical framework (Section 3.1), TAG is defined as:

TAG(x,t)=xlogpt(x)+ωxlogpϕ(tx)\text{TAG}(x, t) = \nabla_x \log p_t(x) + \omega \cdot \nabla_x \log p_\phi(t | x)

where ω\omega is a hyperparameter controlling overall guidance strength. In practice, the paper uses a schedule:

ωt=ω01αˉt\omega_t = \omega_0 \cdot \sqrt{1 - \bar{\alpha}_t}

where ω0\omega_0 is a base guidance strength (tuned per task), and 1αˉt\sqrt{1 - \bar{\alpha}_t} scales the correction based on the noise level. Since αˉt\bar{\alpha}_t decreases monotonically from 1 (at t=0t = 0, clean data) to 0 (at t=Tt = T, pure noise), 1αˉt\sqrt{1 - \bar{\alpha}_t} increases from 0 to 1 over the course of generation. This means TAG applies weaker correction at early steps (when samples are noisy and temporal ambiguity is high) and stronger correction at later steps (when samples are more structured and temporal identity is clearer).

Why this schedule: Early in the reverse process, the sample is nearly pure noise and carries little temporal information—pt(x)p_t(x) and ps(x)p_s(x) are nearly identical for nearby tt, so the TLS gradient is small and potentially noisy. Applying strong correction here could introduce unnecessary variance. Later in the process, the sample has acquired structure and temporal identity becomes more distinct—the TLS gradient is more informative and the correction can be safely strengthened.

Integration with external guidance. When TAG is combined with training-free guidance like DPS or TFG, the combined score function becomes (Section 3.2, Equation 35):

ϵ~θ(xt,c,t)=ϵθ(xt,t)σtxtc(A(x^0),c)σtxtt(ϕ(xt,c),t)\tilde{\epsilon}_\theta(x_t, c, t) = \epsilon_\theta(x_t, t) - \sigma_t \nabla_{x_t} \ell_c(A(\hat{x}_0), c) - \sigma_t \nabla_{x_t} \ell_t(\phi(x'_t, c), t)

where ϵθ(xt,t)\epsilon_\theta(x_t, t) is the unconditional diffusion model output (related to the score by xlogpt(xt)=11αˉtϵθ(xt,t)\nabla_x \log p_t(x_t) = -\frac{1}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t)), c\ell_c is the property guidance loss, t\ell_t is the temporal alignment loss (cross-entropy between the time predictor output and the current timestep), and σt\sigma_t is a scaling factor related to the noise schedule.

What this computes: The first term is the standard unconditional denoising direction. The second term is the external guidance toward the target property (this is DPS/TFG as usual). The third term is the TAG correction—it adjusts the sample to improve its temporal alignment. The three terms are combined linearly with appropriate scaling.

Why this combination works: The key insight is Bayes' rule reinterpretation. The paper shows (Equation 11) that:

pt(xtc)pt(xt)p(cxt)p(txt,c)p_t(x_t | c) \propto p_t(x_t) \, p(c | x_t) \, p(t | x_t, c)

Taking the gradient of the log gives:

xtlogp(xtc)xtlogp(xt)+σtxtc(A(x^0),c)+ωtxtt(ϕ(xt,c),t)\nabla_{x_t} \log p(x_t|c) \approx \nabla_{x_t} \log p(x_t) + \sigma_t \nabla_{x_t} \ell_c(A(\hat{x}_0), c) + \omega_t \nabla_{x_t} \ell_t(\phi(x_t, c), t)

The temporal term acts as a regularizer—it penalizes samples that are temporally misaligned, encouraging the combined guidance to find solutions that both satisfy the target property and remain on the diffusion manifold. Without this term, the property guidance can push the sample arbitrarily far off-manifold in pursuit of the target property. With TAG, the optimization is constrained to stay near regions where the score function is reliable.

Multi-conditional TAG (Section 3.2, Appendix B.3). For multiple conditions c1,c2c_1, c_2, the paper factorizes the posterior as:

pt(xtc1,c2)pt(xt)p(c1xt)p(c2xt,c1)p(txt,c1,c2)p_t(x_t | c_1, c_2) \propto p_t(x_t) \, p(c_1 | x_t) \, p(c_2 | x_t, c_1) \, p(t | x_t, c_1, c_2)

The challenge is that training a separate time predictor for every combination of conditions is combinatorially infeasible. The paper proposes two practical approximations:

Single-condition time predictor (Proposition B.1): The sample is first updated to reflect condition c1c_1:

xtxtηt2xt1(A1(x^0),c1)x'_t \approx x_t - \eta^2_t \nabla_{x_t} \ell_1(A_1(\hat{x}_0), c_1)

Then the temporal alignment is computed using a time predictor conditioned only on c2c_2:

p(txt,c1,c2)p(txt,c2)p(t | x_t, c_1, c_2) \approx p(t | x'_t, c_2)

What this does: The first update xtxtx_t \to x'_t absorbs the effect of condition c1c_1 into the sample itself, similar to how classifier guidance modifies the sample before classification. This "reparameterization" means the time predictor sees a sample that already partially reflects c1c_1, so conditioning only on c2c_2 is sufficient to approximate the full multi-conditional temporal posterior. Proposition B.1 proves this is the posterior expectation under a Gaussian prior and a first-order Taylor expansion of the loss.

Unconditional time predictor (Proposition B.2): For even greater efficiency, an unconditional time predictor can be used by sequentially reparameterizing the sample:

xtxtηt2xt1(A1(x^0),c1)ηt2xt2(A2(x^0),c2)x''_t \approx x_t - \eta^2_t \nabla_{x_t} \ell_1(A_1(\hat{x}_0), c_1) - \eta^2_t \nabla_{x_t} \ell_2(A_2(\hat{x}'_0), c_2)

where x^0\hat{x}'_0 is the denoised estimate after the first update. The temporal alignment is then:

p(txt,c1,c2)p(txt)p(t | x_t, c_1, c_2) \approx p(t | x''_t)

What this does: Each condition is incorporated sequentially into the sample through gradient updates. The unconditional time predictor then evaluates temporal alignment on the fully condition-adjusted sample. This is the most computationally efficient approach—only one time predictor is needed regardless of the number of conditions—but relies on the sequential approximation being sufficiently accurate.

Why this sequential approach is reasonable: The conditions are applied in order of "strength," with each update moving the sample toward satisfying one condition while (approximately) maintaining the adjustments from previous conditions. Since each update is a small gradient step, the composition of multiple steps approximates the joint conditional gradient. The paper acknowledges this is an approximation but demonstrates empirically (Table 5) that both single-condition and unconditional time predictors match or exceed the performance of a dedicated multi-condition predictor, while being dramatically more practical.


Theoretical Analysis: Energy Barriers, Convergence, and Escape Times

The paper provides a three-part theoretical analysis that explains mechanistically why TAG helps: (1) it reshapes the energy landscape to accelerate escape from low-density regions, (2) it provides improved convergence guarantees in total variation distance, and (3) its effect can be understood through the JKO scheme for gradient flows.

Proposition 3.4 (Energy barrier reshaping). The standard energy barrier for a sample at timestep tkt_k is the negative log-density:

Uk(x)=logptk(x)U_k(x) = -\log p_{t_k}(x)

Langevin dynamics—which the diffusion reverse process approximates—must climb over or tunnel through these barriers to reach high-density regions. In low-density regions, Uk(x)U_k(x) is large and flat, meaning gradient-based dynamics become slow or stuck.

Applying TAG modifies this energy landscape to:

Φk(x)=Uk(x)iγiUi(x)\Phi_k(x) = U_k(x) - \sum_i \gamma_i U_i(x)

where γi=pi(x)ptot(x)\gamma_i = \frac{p_i(x)}{p_{tot}(x)} for iki \neq k and γk=1ikpi(x)ptot(x)\gamma_k = 1 - \sum_{i \neq k} \frac{p_i(x)}{p_{tot}(x)}.

What this computes: The modified potential Φk(x)\Phi_k(x) is the original potential Uk(x)U_k(x) minus a weighted sum of potentials at all other timesteps. Since γi0\gamma_i \geq 0 and iγi=1\sum_i \gamma_i = 1 (the coefficients form a probability distribution), the subtraction terms lower the effective barrier. In particular, if the sample is in a high-density region of some incorrect timestep jkj \neq k (meaning pj(x)p_j(x) is large and thus γj\gamma_j is large), then Uj(x)U_j(x) is small (low energy), but the subtraction γjUj(x)-\gamma_j U_j(x) further reduces Φk(x)\Phi_k(x), making the effective barrier even lower.

Why this matters: The modified Langevin dynamics under potential Φk\Phi_k evolves as:

dxt=Φk(xt)dt+2dWt=(sk(xt)ikγisi(xt))dt+2dWtdx_t = -\nabla \Phi_k(x_t) dt + \sqrt{2} dW_t = \left(s_k(x_t) - \sum_{i \neq k} \gamma_i s_i(x_t)\right) dt + \sqrt{2} dW_t

The drift term now includes attractive forces from the correct timestep's score sks_k and repulsive forces from all incorrect timesteps' scores si-s_i. This means the dynamics are not just climbing toward the tkt_k manifold—they are actively fleeing from tit_i manifolds for iki \neq k. This accelerates convergence because the sample doesn't need to slowly diffuse away from incorrect manifolds; it is directly pushed away.

JKO scheme analysis (Appendix C.7). The Jordan-Kinderlehrer-Otto scheme establishes that the Fokker-Planck equation for Langevin dynamics is the gradient flow of the KL divergence with respect to the Wasserstein-2 metric. The paper uses this to compare convergence rates.

Corollary C.6 (Gradient flow with TAG). For the original Langevin dynamics (without TAG), the KL divergence to the target distribution decays as:

ddtKL(qtpk)=Eqtrk2\frac{d}{dt} \text{KL}(q_t \| p_k) = -\mathbb{E}_{q_t} \|r_k\|^2

where rk(x,t)=xlogqt(x)pk(x)r_k(x, t) = \nabla_x \log \frac{q_t(x)}{p_k(x)} measures the score mismatch. For the TAG-modified dynamics, the decay is:

ddtKL(q~tpk)=Eq~t[r~k2+A(t)]\frac{d}{dt} \text{KL}(\tilde{q}_t \| p_k) = -\mathbb{E}_{\tilde{q}_t} \left[\|\tilde{r}_k\|^2 + A(t)\right]

where A(t)=iγiEq~t[r~k(x,t)si(x)]A(t) = \sum_i \gamma_i \mathbb{E}_{\tilde{q}_t} [\tilde{r}_k(x, t) \cdot s_i(x)] is the additional contribution from TAG.

What this means: The extra term A(t)A(t) captures the interaction between the corrected score mismatch r~k\tilde{r}_k and the score functions of incorrect timesteps. If this term is positive (which occurs when the sample is nearer to an incorrect manifold than the correct one), the KL divergence decays faster with TAG than without. The paper proves that when the sample is in a low-density region (Definition C.7: Dk,ϵ={pk(x)ϵ}D_{k,\epsilon} = \{p_k(x) \leq \epsilon\}) and the mixture score has sufficient magnitude (Equation 98), the escape time τ~\tilde{\tau} (time to leave the low-density region) satisfies:

E[τ~]KL(q0pk)β+β/η2\mathbb{E}[\tilde{\tau}] \leq \frac{\text{KL}(q_0 \| p_k)}{\beta + \beta/\eta^2}

where β\beta is a lower bound on the TAG correction strength and η\eta bounds the expected mixture score. The corresponding escape time for the original dynamics is E[τ]KL(q0pk)/β\mathbb{E}[\tau] \leq \text{KL}(q_0 \| p_k) / \beta.

Why TAG accelerates escape: The denominator for E[τ~]\mathbb{E}[\tilde{\tau}] contains the additional term β/η2\beta/\eta^2 compared to E[τ]\mathbb{E}[\tau]. This reduces the upper bound by a factor of 1+1/η21 + 1/\eta^2. When η\eta is large (meaning the mixture score of alternative timesteps is strong—i.e., the sample is clearly misaligned), the acceleration factor is modest. But when η\eta is small (the mixture score is weak—the sample is in an ambiguous region), the factor 1/η21/\eta^2 becomes large, dramatically reducing the escape time bound. This matches the intuition: TAG helps most when the sample is stuck in flat, low-density regions where the standard score gradient is too weak to provide direction.

Theorem 3.5 / C.12 (Improved convergence guarantee). The formal statement bounds the total variation distance between the TAG-corrected reverse process and the true data distribution:

dTV(p~0,qdata)dTV(p0,qdata)G4Fd_{TV}(\tilde{p}_0, q_{data}) \leq d_{TV}(p_0, q_{data}) - \frac{G}{4\sqrt{F}}

where FF is the original score approximation error (Equation 109), and G=mβ(1+1/η2)st0Tg(t)2dtG = m \beta (1 + 1/\eta^2) s \cdot \int_{t_0}^T g(t)^{-2} dt represents the improvement from TAG, with mm being the minimum gradient magnitude of the score approximation error function.

What this computes: A deterministic improvement in the convergence guarantee. The term G4F\frac{G}{4\sqrt{F}} is always positive (under the stated assumptions), so the upper bound on TAG's error is strictly lower than the bound on the standard process's error. The improvement depends on the TAG correction strength β\beta, the acceleration factor (1+1/η2)(1 + 1/\eta^2), the correction time ss, and the integrated inverse diffusion coefficient. Importantly, the improvement is independent of the specific external guidance being applied—it comes purely from the temporal alignment correction.

Why this form: The proof (Appendix C.7) uses a decomposition of the total error into truncation error, initial noise mismatch, and path-level score approximation error (as in Chen et al., 2023b; Oko et al., 2023). The path-level error term depends on the expected squared score approximation error at each timestep. Corollary C.10 shows that running TAG-modified Langevin dynamics for time ss reduces this expected error by at least mβ(1+1/η2)sm \beta (1 + 1/\eta^2) s. The algebra of subtracting this improvement from the original error bound (using ffgg/(2f)\sqrt{f} - \sqrt{f - g} \geq g / (2\sqrt{f})) yields the final result.


Corrupted Reverse Process: Isolating TAG's Corrective Mechanism

The experiments in Section 3.4 provide a controlled test of TAG's mechanism by deliberately corrupting the reverse process with artificial noise and measuring how effectively TAG restores generation quality.

Setup. At each reverse timestep tt, random Gaussian noise ztN(0,σ2I)z_t \sim \mathcal{N}(0, \sigma^2 I) is added to the sample xtx_t before the denoising step. This simulates an extreme off-manifold perturbation—the noise pushes the sample in a random direction, potentially far from the correct manifold. The experiment uses CIFAR-10 with 50 diffusion steps (DDIM sampling) and explores four noise levels: σ=0.05,0.1,0.2,0.3\sigma = 0.05, 0.1, 0.2, 0.3. Fifty thousand samples are generated, and quality is measured with FID (Fréchet Inception Distance, lower is better), IS (Inception Score, higher is better), and the Time-Gap metric.

The Time-Gap metric (Definition F.1). This is a diagnostic quantity that measures the degree of temporal misalignment during generation:

Time-Gap:=1Tt=1Targmaxϕ(xt)t\text{Time-Gap} := \frac{1}{T} \sum_{t=1}^T \left| \arg\max \phi(x_t) - t \right|

where ϕ(xt)\phi(x_t) is the time predictor's output distribution over timesteps, and argmax\arg\max returns the most likely timestep. The Time-Gap is the average absolute difference between the predicted timestep and the actual timestep, averaged over all TT steps of the reverse process. A Time-Gap of 0 means the time predictor perfectly identifies each sample's timestep—the sample is exactly on the expected manifold at every step. A Time-Gap of, say, 100 means the predictor thinks the sample is, on average, 100 timesteps away from where it should be—indicating severe off-manifold drift.

Why Time-Gap is a meaningful diagnostic: It directly operationalizes the concept of "temporal alignment" that TAG is designed to enforce. Unlike FID or IS, which measure only final output quality, Time-Gap can be measured at every intermediate timestep and reveals where in the process the off-manifold drift occurs (Appendix F.2, Figures 4-9). The paper demonstrates (Figure 10) that Time-Gap is well-correlated with standard quality metrics (FID and IS) across different numbers of function evaluations, validating it as a proxy for generation quality.

Results (Table 1, Figure 3). The effects are striking:

  • At σ=0.1\sigma = 0.1, no TAG: FID = 193.6, IS = 2.37, Time-Gap = 104.1. The generation quality is severely degraded compared to the uncorrupted baseline.

  • At σ=0.1\sigma = 0.1, TAG with ω=2.0\omega = 2.0: FID = 120.9, IS = 3.69, Time-Gap = 41.8. This represents an FID reduction of 37.5% and an IS improvement of 1.32 points—a dramatic recovery.

  • At σ=0.3\sigma = 0.3 (extreme corruption), no TAG: FID = 410.1, IS = 1.28—the output is essentially random.

  • At σ=0.3\sigma = 0.3, TAG with ω=200.0\omega = 200.0 (very strong correction): FID = 223.2, IS = 2.17—a 45.6% reduction in FID, recovering to a quality level comparable to the uncorrected model at σ=0.2\sigma = 0.2.

Key observations from the ω\omega sweep (Table 10):

  1. Optimal ω\omega increases with noise level. For σ=0.1\sigma = 0.1, the best FID (115.6) occurs at ω=1.0\omega = 1.0. For σ=0.2\sigma = 0.2, the best FID (230.9) occurs at ω=4.5\omega = 4.5. For σ=0.3\sigma = 0.3, the best FID (223.2) occurs at ω=200.0\omega = 200.0. Stronger corruption requires stronger correction—this makes intuitive sense: the further off-manifold the corruption pushes the sample, the stronger the TLS gradient needs to be to pull it back.

  2. Time-Gap decreases monotonically with ω\omega. As TAG strength increases from 0 to 200, the Time-Gap drops from 273.9 to 158.9. The Time-Gap drops most rapidly in the range ω=1.0\omega = 1.0 to ω=5.0\omega = 5.0 (from 250.9 to 185.1), then more slowly thereafter. This suggests that the most egregious temporal misalignments are corrected with moderate TAG strength, while finer corrections require progressively stronger guidance.

  3. IS (Inception Score) increases monotonically with ω\omega. From 1.27 at ω=0\omega = 0 to 2.17 at ω=200\omega = 200. This suggests that temporal alignment improves not just fidelity (FID) but also diversity and recognizability of generated classes (IS). Strong TAG correction does not collapse the generator into a single mode.

  4. Over-correction is possible. At ω=200.0\omega = 200.0, the FID (223.2) is slightly higher than at ω=100.0\omega = 100.0 (FID = 229.5), and the IS (2.17) is higher than at ω=100.0\omega = 100.0 (IS = 2.01). This suggests a mild trade-off: very strong TAG can improve class distinctiveness (IS) at a small cost to overall fidelity (FID).

Why TAG works under corruption but standard score correction fails. The corruption experiment highlights the fundamental difference between TAG and Langevin correction. When σ=0.3\sigma = 0.3 noise is added, the sample is pushed far from any manifold—the score function xlogpt(xt)\nabla_x \log p_t(x_t) becomes essentially random because the sample xtx_t is unlike anything seen during training. Running Langevin dynamics with this corrupted score function cannot recover because it has no correct signal to follow. TAG, however, relies on the time predictor, which was trained on the forward process and can still recognize that the corrupted sample "looks like" it came from some timestep (perhaps not the correct one). The TLS gradient xlogp(txt)\nabla_x \log p(t | x_t) provides a direction that points back toward regions where samples at timestep tt should reside, even if those regions are far away. This direction is not perfectly accurate—the time predictor itself becomes less reliable under extreme corruption (the Time-Gap at σ=0.3\sigma = 0.3 with TAG at ω=200.0\omega = 200.0 is still 158.9, meaning the predictor is off by 159 timesteps on average)—but it is directionally useful enough to recover significant quality.


Design Choices and Their Justifications

Why a separate time predictor rather than using the diffusion model's own internal representations? The diffusion model's internal representations are trained to denoise, not to recognize temporal identity. While they implicitly encode noise-level information, this information is entangled with content-specific features and becomes unreliable when the sample is off-manifold. A dedicated time predictor, trained purely on the timestep classification task with forward-process data, provides a signal that is statistically independent of the denoising objective. This independence is crucial: it means the TLS provides useful gradient information even when the score function is degraded, creating a complementary correction mechanism rather than a redundant one.

Why gradient-based correction rather than hard temporal reassignment? Hard reassignment (as in TCS) discards the current timestep information and replaces it with a predicted timestep, then evaluates the score function at this new timestep. This is brittle for two reasons. First, if the time predictor is wrong (which becomes more likely the further off-manifold the sample is), the reassignment compounds the error by evaluating the score function at an incorrect noise level. Second, even if the reassignment is correct, the diffusion model's score function at timestep t~\tilde{t} expects samples that are distributed according to pt~(x)p_{\tilde{t}}(x), not samples that have been artificially moved to resemble t~\tilde{t}. The score evaluation at the reassigned timestep may itself be unreliable. TAG's soft gradient-based correction avoids both problems: it makes small, incremental adjustments that keep the sample in regions where the score function is approximately valid, and it provides a continuous correction signal rather than a discrete jump.

Why classification (cross-entropy) rather than regression for the time predictor? The densities pt(x)p_t(x) and ps(x)p_s(x) overlap substantially for nearby tt, especially early in the diffusion process when samples are still highly noisy. A single noisy sample could plausibly have come from any of several timesteps. A regression objective would force the predictor to commit to a single tt, producing a sharp but potentially wrong estimate. The cross-entropy classification objective allows the predictor to express uncertainty by distributing probability mass across multiple timesteps, and the gradient xlogpϕ(tx)\nabla_x \log p_\phi(t | x) naturally accounts for this uncertainty through the softmax normalization. This is particularly important when the sample is off-manifold—the predictor may recognize that the sample is ambiguous between several timesteps, and the resulting gradient will reflect this ambiguity rather than forcing a potentially incorrect hard decision.

Why the ωt=ω01αˉt\omega_t = \omega_0 \sqrt{1 - \bar{\alpha}_t} schedule? This schedule is motivated by the observation that temporal information content varies with noise level. At high noise levels (αˉt1\bar{\alpha}_t \approx 1 at the start of reverse diffusion), the marginal distributions for nearby timesteps are nearly identical—a sample at t=900t = 900 looks almost indistinguishable from a sample at t=899t = 899. The TLS gradient in this regime is small and potentially noise-dominated, so weak correction is appropriate. At low noise levels (αˉt0\bar{\alpha}_t \approx 0 near the end of generation), the marginal distributions are sharply distinct—a sample at t=5t = 5 has clear structure that distinguishes it from t=4t = 4. The TLS gradient is informative and reliable, so stronger correction can be safely applied. The 1αˉt\sqrt{1 - \bar{\alpha}_t} factor provides a smooth interpolation between these regimes.

Why SimpleCNN for the time predictor? The timestep classification task depends primarily on global statistical properties (variance magnitude, spatial correlation structure) rather than fine spatial details. A SimpleCNN with aggressive pooling (2×2 average pooling after each convolutional block) is well-suited to capturing these global statistics while being computationally efficient. The paper's empirical validation (Table 18) confirms that the 1.48M-parameter SimpleCNN matches the time-gap performance of a 17.38M-parameter UNet encoder, validating that the task does not require complex architecture. The choice of 2×2 average pooling (rather than max pooling or strided convolutions) is appropriate for a statistics-estimation task—average pooling preserves distributional information (mean, variance) that max pooling might discard.

4. Key Insights and Innovations

Innovation 1: Temporal Misalignment as the Unifying Root Cause of Off-Manifold Degradation

The paper's most fundamental conceptual move is diagnostic rather than algorithmic: it reframes the diverse pathologies that plague guided diffusion—discretization errors in few-step sampling, score approximation errors under external guidance, combinatorial errors in multi-conditional generation—not as separate problems requiring separate fixes, but as manifestations of a single underlying phenomenon. When a sample deviates from the forward-process distribution at its current noise level, the learned score function becomes unreliable because it was trained under the assumption that the input would be a genuine forward-process sample. The sample loses its "temporal identity."

This is a genuine reframing, not an incremental observation. Prior work treated these failure modes as distinct. Exposure bias methods (Ning et al., 2023; 2024; Li et al., 2024a) addressed train-test mismatch as a statistical discrepancy in the sampler's input distribution. Time perturbation methods (Sadat et al., 2024; Li et al., 2024b) exploited local sensitivity to time inputs as a contrastive signal. TCS (Jung et al., 2024) treated temporal misalignment as a scheduling problem—correctable by hard reassignment to a predicted timestep. Each approach implicitly assumed a different mechanism of failure and designed a targeted patch.

TAG's re-framing is that all of these failures share the same root cause: the sample xtx_t at timestep tt has been pushed into a region where pt(xt)p_t(x_t) is small, and therefore the score function xlogpt(xt)\nabla_x \log p_t(x_t) is poorly estimated. The specific source of the push—whether it's an external guidance gradient, a multi-condition combination error, or a discretization error from skipped steps—is secondary. What matters is that the sample is temporally misaligned, and recovering alignment requires an independent signal that remains reliable when the score function degrades.

The significance of this reframing is that it unifies the solution space. Rather than developing separate corrections for guidance, for multi-condition combination, and for few-step sampling—which is what the prior literature largely did—TAG provides a single mechanism (the TLS gradient) that addresses all of them. The experimental results in Section 4 bear this out: the same time predictor and the same Algorithm 1, with only task-specific tuning of the guidance strength ω0\omega_0, works across CIFAR-10 label guidance, ImageNet label guidance, molecular property control, audio restoration, multi-conditional CelebA generation, few-step DDIM sampling, and large-scale text-to-image alignment. This universality is not an accident of architecture choice—it follows directly from the diagnostic insight that temporal misalignment is the common mechanism.

The paper operationalizes this diagnostic insight through the Time-Gap metric (Definition F.1), which gives the reframing empirical teeth. Time-Gap is not merely a heuristic quality score—it directly measures the quantity that TAG is designed to minimize (the average absolute deviation between predicted and actual timesteps), and the paper demonstrates (Figure 10, Appendix F.2) that it correlates with standard quality metrics (FID, IS) across different generation regimes. This transforms temporal alignment from an abstract concept into a measurable, optimizable quantity—a move that is conceptually analogous to how the Inception Score operationalized the intuition that good generated images should have recognizable class content.

Is this a fundamental shift or an incremental refinement? I would argue it is fundamental, because it changes what problem future work should try to solve. Before TAG, a researcher encountering poor sample quality under a new guidance method would reasonably ask: "What is wrong with my guidance design? Should I adjust the loss function? The step size? The noise schedule?" After TAG, the more productive question becomes: "Is my guidance causing temporal misalignment? Can I measure it with a time predictor? Can I correct it with a TLS gradient?" The diagnosis replaces ad-hoc debugging with a principled diagnostic procedure.

Innovation 2: The Time-Linked Score as a Competitive Attraction-Repulsion Mechanism

The second conceptually distinctive contribution is the structure of the TLS gradient revealed by Theorem 3.3. On its surface, TAG appears to be a straightforward application of Bayes' rule: treat time as a conditioning variable, train a classifier to estimate p(tx)p(t|x), take its gradient, and add it as a correction. This would be a reasonable but incremental idea—a "classifier guidance for time" in the spirit of Dhariwal & Nichol (2021).

What makes it more than incremental is the decomposition in Equation 12:

xlogp(tix)=kiptk(x)ptot(x)(xlogpti(x)xlogptk(x))\nabla_x \log p(t_i | x) = \sum_{k \neq i} \frac{p_{t_k}(x)}{p_{tot}(x)} \left(\nabla_x \log p_{t_i}(x) - \nabla_x \log p_{t_k}(x)\right)

This equation reveals that the TLS is not simply the score function of the correct timestep. It is a weighted combination of pairwise differences between the correct timestep's score and every other timestep's score. The weights ptk(x)ptot(x)\frac{p_{t_k}(x)}{p_{tot}(x)} are proportional to how strongly the sample "looks like" it came from timestep tkt_k. This means the TLS simultaneously attracts toward the correct manifold (through xlogpti(x)\nabla_x \log p_{t_i}(x) terms) and repels from incorrect manifolds (through xlogptk(x)-\nabla_x \log p_{t_k}(x) terms), with the repulsion being strongest from whichever incorrect timestep the sample currently most resembles.

This competitive structure is qualitatively different from what any prior time-based correction method achieves. TCS (Jung et al., 2024) performs a "hard" reassignment: if the predictor thinks the sample is at timestep t~\tilde{t}, it evaluates the score function at t~\tilde{t} and adjusts the noise schedule accordingly. This only provides an attractive force toward t~\tilde{t} (and, indirectly, away from other timesteps through the change in noise scaling). It provides no direct repulsive force from the manifold the sample is currently stuck on. Timestep Guidance (Sadat et al., 2024) and Self-Guidance (Li et al., 2024b) perturb the time input and use the difference in score outputs as guidance. This provides a local directional signal (how the score changes when time changes) but does not incorporate global information about which specific alternative timesteps the sample might be confused with. The TLS, by contrast, integrates information from all timesteps through the posterior distribution p(tx)p(t|x), producing a correction that is globally informed about the full temporal landscape.

The practical significance of this competitive structure is that it makes TAG adaptive to the nature of the misalignment. If the sample has drifted toward a manifold that is very close (in sample space) to the correct one, ptk(x)p_{t_k}(x) will be large for tktit_k \approx t_i, the weights will concentrate on nearby timesteps, and the TLS will provide a gentle nudge. If the sample has been pushed far away—for instance, to a manifold corresponding to a much earlier or later stage of diffusion—ptk(x)p_{t_k}(x) will be large for some kk far from ii, the repulsive force from that distant manifold will be strong, and the TLS will provide an aggressive correction. This adaptivity is automatic—it emerges from the structure of the posterior, not from any explicit difficulty-estimation module.

The paper's corrupted reverse process experiments provide indirect evidence for this adaptivity. At low corruption levels (σ=0.05\sigma = 0.05), the optimal TAG strength is low (ω=0.2\omega = 0.2), suggesting the TLS naturally provides mild corrections when misalignment is subtle. At high corruption levels (σ=0.3\sigma = 0.3), the optimal TAG strength is very high (ω=200.0\omega = 200.0), but critically, the correction remains effective (FID drops from 410.1 to 223.2) rather than overshooting or causing instability. A non-adaptive correction (e.g., a fixed-strength "push toward the correct timestep") would either be too weak at high corruption or too aggressive at low corruption. The TLS's competitive weighting provides a natural mechanism for scaling the correction to the severity of the misalignment.

Is this a fundamental insight or an emergent property of a straightforward application? I lean toward fundamental because the decomposition reveals that the seemingly simple idea of "gradient of a time classifier" contains a rich internal structure that was not obvious from the construction. This is analogous to how attention mechanisms were initially motivated as differentiable memory retrieval but later analyses revealed they implement more complex operations (content-based addressing, competition via softmax normalization). Similarly, the TLS appears to implement "classifier guidance for time" but actually implements a competitive attraction-repulsion dynamic that is far more robust than any single-timestep correction could be.

Innovation 3: Verifier-Independent Robustness Through Statistical Separation of Objectives

A third conceptual innovation—more subtle but equally important for practical deployment—is the independence between the TLS signal and the score function signal. TAG does not rely on the score function to correct itself (as Langevin dynamics does) or on perturbing the score function's inputs to extract additional information (as time perturbation methods do). Instead, it introduces a completely separate neural network (the time predictor) trained on a completely different objective (timestep classification via cross-entropy on forward-process data) that happens to produce gradients aligned with the correction direction.

This statistical separation is what makes TAG robust in regimes where score-based correction fails. The paper explicitly demonstrates this failure mode: Langevin dynamics applied on top of DPS on CIFAR-10 increases FID from 217.1 to 226.8 (Table 8). The reason is straightforward: in regions where the score function is unreliable, following its gradient more carefully (which is what Langevin dynamics does, by taking multiple small steps) only compounds the error. The score function cannot bootstrap itself out of a low-density region because it has no training signal there. TAG's time predictor, by contrast, was trained on the forward process—it has seen noisy samples at every timestep and every region of sample space that the forward process can reach. Even when a sample is pushed to a location that the forward process would never produce at its current timestep, the time predictor can recognize that this location is statistically similar to what the forward process produces at some other timestep, and it can provide a gradient that points back toward the high-probability region of the correct timestep.

This independence is not merely an architectural convenience—it is a design principle. The paper makes it explicit by choosing a time predictor architecture (SimpleCNN) that is architecturally different from the diffusion model backbone (UNet) and by training it with a fundamentally different loss function (cross-entropy classification vs. denoising score matching). If the time predictor shared architecture or training objectives with the score model, their errors would likely be correlated, and the TLS would degrade in the same regions where the score function degrades.

The practical significance is that TAG can be deployed as a bolt-on correction to any existing diffusion model without modifying the base model's weights, retraining it, or even accessing its internal representations. The time predictor is a separate network that only needs the sample xtx_t (and optionally the condition cc) as input. This makes TAG applicable in black-box or API-based settings where the diffusion model is provided as a service and only inputs/outputs are accessible. It also means that improvements to the time predictor (better architecture, more training data, adversarial training) automatically improve TAG's correction quality without any changes to the base model—the two systems evolve independently.

Is this incremental relative to prior work that also used auxiliary networks (e.g., TCS's time predictor, CDM's time classifier)? Yes and no. The use of an auxiliary time predictor is not novel—Jung et al. (2024) and Yadin et al. (2024) both train time classifiers. What is novel is the recognition that the gradient of this classifier, rather than its argmax (as in TCS) or its probability ratios (as in CDM), is the quantity that provides robust corrective force. TCS uses t~=argmaxϕ(xt)\tilde{t} = \arg\max \phi(x_t) to reassign timesteps—effectively throwing away all the gradient information except a single hard decision. CDM uses the classifier probabilities to estimate density ratios for score approximation—using the classifier's output values but not its input gradient. TAG uses xlogpϕ(txt)\nabla_x \log p_\phi(t | x_t)—the rate at which the classifier's confidence in the correct timestep changes as the sample moves. This is a fundamentally different way of extracting value from the same auxiliary model, and the empirical results (TCS achieving FID 213.4 and accuracy 29.4% vs. TAG achieving FID 190.4 and accuracy 63.2% on CIFAR-10 DPS, Table 2) suggest the gradient-based approach is substantially more effective.

Innovation 4: Energy Barrier Reshaping as a Mechanistic Explanation for Accelerated Correction

The fourth conceptual contribution is the theoretical framework that explains why TAG's correction accelerates convergence out of low-density regions. Proposition 3.4 and the subsequent JKO scheme analysis (Appendix C.7) are not merely formal justifications—they provide a mechanistic model that yields specific, testable predictions about TAG's behavior.

The key insight is that the effective potential governing the reverse diffusion dynamics is reshaped by TAG from Uk(x)=logptk(x)U_k(x) = -\log p_{t_k}(x) to Φk(x)=Uk(x)iγiUi(x)\Phi_k(x) = U_k(x) - \sum_i \gamma_i U_i(x). In flat, low-density regions where Uk(x)U_k(x) is large and its gradient is small (making standard Langevin dynamics slow), the subtraction terms γiUi(x)-\gamma_i U_i(x) can reduce the effective barrier. Specifically, if the sample has drifted into a region where it strongly resembles some alternative timestep jj, then γj\gamma_j will be large, and even though Uj(x)U_j(x) might be small (the sample is in a high-density region for tjt_j), the subtraction γjUj(x)-\gamma_j U_j(x) still lowers Φk(x)\Phi_k(x), making the effective landscape less flat and the gradient flow faster.

This is a genuinely novel mechanistic insight, not just a post-hoc rationalization. It explains the counterintuitive empirical finding that TAG helps even when the diffusion model's score function is accurate but the landscape is flat. Prior work attributed slow convergence in low-density regions to score approximation error (Oko et al., 2023; Chen et al., 2023a)—the model hasn't seen enough training data there, so its score estimates are noisy. But the energy barrier analysis reveals a second, independent mechanism: even with a perfectly accurate score function, a flat potential landscape means gradient-based dynamics are slow. The score function can accurately point "uphill" toward higher density, but if the slope is very shallow, progress is slow. TAG's repulsive terms from alternative manifolds steepen the effective landscape, accelerating convergence even when score accuracy is not the bottleneck.

This has implications beyond TAG itself. It suggests that any correction method that introduces competitive forces from alternative modes or manifolds can accelerate sampling, even if those forces are derived from imperfect models. The specific mechanism—subtracting weighted potential functions of competing modes—is related to techniques in energy-based modeling and MCMC (e.g., parallel tempering, which also uses information from alternative distributions to accelerate mixing), but its application to diffusion model correction is new.

The convergence theorem (Theorem 3.5/C.12) quantifies this improvement: the upper bound on total variation distance is reduced by a term proportional to β(1+1/η2)s\beta(1 + 1/\eta^2)s, where β\beta measures the TAG correction strength, η\eta is related to the expected mixture score of alternative timesteps, and ss is the time spent applying correction. The factor (1+1/η2)(1 + 1/\eta^2) captures the acceleration: when η\eta is small (the mixture score is weak, meaning the sample is in a flat region with no strong pull toward any alternative timestep), 1/η21/\eta^2 is large, and the improvement is dramatic. When η\eta is large (the sample is strongly pulled toward some specific alternative timestep), the improvement is more modest—but in that case, the sample is already near some manifold, so less correction is needed.

Is this a fundamental theoretical advance or an incremental application of known results? The JKO scheme and the Girsanov-based analysis of diffusion model convergence (Chen et al., 2023b; Oko et al., 2023) are established tools. The novelty is in applying them to analyze an auxiliary correction term that reshapes the potential landscape, and in deriving concrete, interpretable improvement bounds that explain when and why the correction helps. The result that escape time bounds are reduced by a factor of 1+1/η21 + 1/\eta^2 is simple enough to be memorable and specific enough to generate testable predictions (e.g., TAG should help most in scenarios where the sample is in a flat region with no strong directional signal, which is exactly the corrupted reverse process regime).

Summary of the Intellectual Architecture

Taken together, these four innovations form a coherent intellectual structure. Innovation 1 (temporal misalignment as root cause) provides the diagnostic insight—what is the problem? Innovation 2 (competitive TLS structure) provides the mechanism—how does the correction work internally? Innovation 3 (verifier-independent robustness) provides the practical principle—why can the correction succeed where score-based methods fail? Innovation 4 (energy barrier reshaping) provides the theoretical explanation—why does the mechanism accelerate convergence?

The result is not just a method that works, but a framework that explains why it works and, in doing so, provides guidance for future improvements: train better time predictors (Innovation 3), exploit the competitive structure for adaptive correction schedules (Innovation 2), and target the correction toward flat regions of the energy landscape (Innovation 4). The empirical results—4× efficiency gains in compute-optimal scaling, 47.9% FAD reduction on audio declipping, recovery from extreme corruption—validate the framework, but the framework itself is the lasting contribution.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The primary benchmark is the TFG benchmark (Ye et al., 2024), a standardized set of training-free guidance tasks spanning image, molecular, and audio domains. Image tasks use CIFAR-10 (Krizhevsky et al., 2009; 60,000 32×32 images, 10 classes), ImageNet (Russakovsky et al., 2015; 256×256 images, class-conditional), Cat (Elson et al., 2007; 256×256 cat faces), and CelebA-HQ (Karras et al., 2018; 256×256 face images with attribute labels). Molecular generation uses the QM9 dataset (Ramakrishnan et al., 2014; 134k molecules with up to 9 heavy atoms, 12 quantum chemical properties, split 130k/18k/13k for train/valid/test). Audio tasks use open-source mel-spectrogram data from Audio-DDPM at 256×256 resolution. For large-scale text-to-image experiments, prompts are drawn from HPSv2 (Wu et al., 2023; 256 prompts) and Partiprompts (Yu et al., 2022; 64 prompts), with style images from WikiArt (Yu et al., 2022). Sample sizes vary by experiment: 512 for CIFAR-10 label guidance (Table 2), 256 for ImageNet and CelebA (Table 2), 4096 for molecular generation (Table 2), 256 for audio (Table 2), and 256 for text-to-image alignment (Table 6). For the larger-scale CIFAR-10 evaluation in Table 13, 50,000 samples are used.

  • Base model(s). Experiments span six pretrained diffusion models: CIFAR10-DDPM (Nichol & Dhariwal, 2021; unconditional, VP-SDE), ImageNet-DDPM (Dhariwal & Nichol, 2021; class-conditional, 256×256), Cat-DDPM (Elson et al., 2007; unconditional, 256×256), CelebA-DDPM (Karras et al., 2018; unconditional, 256×256), Molecule-EDM (Hoogeboom et al., 2022; unconditional, equivariant diffusion on 3D molecular graphs), and Audio-Diffusion (Kong et al., 2021; Popov et al., 2021; unconditional, mel-spectrogram generation). All are pretrained and unmodified by TAG—the framework is applied as an inference-time correction. For large-scale text-to-image experiments, Stable Diffusion v1.5 (Rombach et al., 2022) is used as the base model. The diversity of architectures (UNet-based for images, EGNN-based for molecules, WaveNet-style for audio, latent diffusion for text-to-image) and data modalities (pixel images, 3D molecular coordinates, mel-spectrograms, latent representations) is deliberately chosen to test TAG's generality. No model is fine-tuned for TAG; only the lightweight time predictor is trained separately for each dataset.

  • Metrics. Generation fidelity is measured by FID (Fréchet Inception Distance; Heusel et al., 2017) for images, FAD (Fréchet Audio Distance; Kilgour et al., 2018) for audio, Atom Stability (percentage of atoms with valid valencies; Hoogeboom et al., 2022) for molecules, and KID (Kernel Inception Distance; Bińkowski et al., 2018) for multi-conditional CelebA generation (reported in log scale). Generation validity (conditioning accuracy) is measured by classification accuracy (for label-guided and multi-attribute generation), LPIPS (Zhang et al., 2018) for image restoration, MAE (Mean Absolute Error) between predicted and target molecular properties, DTW (Dynamic Time Warping; Müller, 2007) for audio restoration, and reward scores (Aesthetic predictor from Schuhmann et al., 2022; CLIPScore from Radford et al., 2021) for text-to-image alignment. The paper also introduces the Time-Gap metric (Definition F.1), computed as the average absolute difference between the time predictor's argmax prediction and the current timestep, averaged over all T reverse steps: Time-Gap:=1Tt=1Targmaxϕ(xt)t\text{Time-Gap} := \frac{1}{T} \sum_{t=1}^T |\arg\max \phi(x_t) - t|. Lower Time-Gap indicates better temporal alignment. For the few-step generation experiments, IS (Inception Score; Salimans et al., 2016) is additionally reported. For style transfer, Style Score (matching Gram matrices of CLIP features) is used. Metrics are averaged across target values within each task, with the best-performing guidance strength ω0\omega_0 selected per target via grid search over [0.01,0.05,0.15,0.3,0.5,1.0,1.5,2.0,3.0,4.0,5.0][0.01, 0.05, 0.15, 0.3, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 5.0] (exact ranges vary per task; Appendix E.3.7, Table 14 provides hyperparameter details).

  • Baselines. The paper compares TAG against multiple families of prior correction methods, all applied on top of DPS (Diffusion Posterior Sampling; Chung et al., 2023) or TFG (Training-Free Guidance; Ye et al., 2024):

    • DPS alone and TFG alone — the base guidance methods with their reported optimal hyperparameters from Ye et al. (2024).
    • TCS (Time Correction Sampler; Jung et al., 2024) — uses a time predictor to hard-reassign timesteps, evaluated with DPS.
    • Timestep Guidance (TSG; Sadat et al., 2024) — perturbs the time input to the score model to derive contrastive signals.
    • Self-Guidance (SG; Li et al., 2024b) — similar time-perturbation approach for contrastive guidance.
    • Epsilon Scaling (Ning et al., 2024) — scales diffusion model output vector norms to mitigate exposure bias.
    • Time-Shift Sampler (Li et al., 2024a) — adjusts time information based on batch variance.
    • Input Perturbation (Ning et al., 2023) — adds noise during training to reduce exposure bias (requires retraining models from scratch).
    • Langevin Dynamics (corrector sampling; Song et al., 2021b; Song & Ermon, 2019) — score-based refinement steps between predictor steps.
    • For text-to-image: DAS (Kim et al., 2025) — a test-time sampler optimizing reward alignment; TAG is applied on top of DAS.
    • For multi-conditional generation: the baseline is naive score summation (combining independent guidance terms without correction), compared against TAG with multi-condition, single-condition, and unconditional time predictors.
  • Generation budget / compute accounting. The primary unit of computation is number of diffusion model function evaluations (NFE) — one NFE equals one call to the score model ϵθ(xt,t)\epsilon_\theta(x_t, t). TAG adds a constant per-step overhead of one forward pass through the time predictor plus one backward pass to compute the TLS gradient. For the SimpleCNN architecture (1.48M parameters), this overhead is approximately 8.5% of a UNet encoder forward pass and a negligible fraction of the full diffusion backbone cost. The paper does not explicitly report wall-clock time, but the architectural comparisons in Table 18 and Appendix E.4 demonstrate that the time predictor is intentionally lightweight. For the corrupted reverse process experiments (Section 3.4), "budget" is not varied—all experiments use 50 diffusion steps (DDIM sampling) with fixed 50,000-sample evaluation. For few-step generation (Section 4.3, Table 4), NFE is explicitly varied from 1 to 100, measuring quality degradation as steps decrease. No FLOPs-matched pretraining comparisons are performed (the paper does not study the pretraining-vs-inference tradeoff).

  • Cross-validation / statistical protocol. No cross-validation is used for strategy selection (TAG has a single main hyperparameter ω0\omega_0, tuned via grid search per task and target). For the TFG benchmark (Table 2), results are averaged across all target values within each task (e.g., all 10 CIFAR-10 classes, all selected ImageNet classes, all 6 molecular properties). The paper reports that "the final results are averaged over the best-performing guidance strength ω0\omega_0 according to the grid search for all target values in each task" (Section 4.1). For the multi-conditional experiments (Table 5), results are reported per target condition combination. For the few-step experiments (Table 4, Table 17), three independent runs are not reported—results are single evaluations. For the corrupted reverse process (Section 3.4, Table 1), the grid search over ω{0,0.5,1.0,2.0,4.0}\omega \in \{0, 0.5, 1.0, 2.0, 4.0\} at each noise level is exhaustive, and the best FID is reported. For large-scale text-to-image (Table 6), "three independent runs" are mentioned but only mean values are reported without standard deviations or confidence intervals. The primary robustness check is consistency across diverse tasks and models rather than within-task statistical significance.


Main Quantitative Results

Training-Free Guidance Benchmark (Table 2): TAG Consistently Improves Fidelity While Maintaining or Enhancing Validity

The central quantitative claim of the paper is that TAG improves generation quality across the TFG benchmark without sacrificing conditioning accuracy—and in many cases, improving both simultaneously. Table 2 presents the comprehensive comparison across six datasets, eight tasks, two base guidance methods (DPS and TFG), and four baseline correction methods (TCS, Timestep Guidance, Self-Guidance, plus the base methods alone).

Image restoration tasks (deblur, super-resolution on Cat-DDPM):

  • DPS + TAG vs. DPS alone (Gaussian deblur): FID improves from 139.7 to 128.9 (7.7% relative improvement); LPIPS improves from 0.613 to 0.570 (7.0% relative improvement).
  • TFG + TAG vs. TFG alone (Gaussian deblur): FID improves from 64.2 to 62.7 (2.3% relative improvement); LPIPS from 0.154 to 0.151 (1.9% relative improvement).
  • DPS + TAG vs. DPS alone (super-resolution): FID improves from 139.0 to 128.3 (7.7%); LPIPS from 0.614 to 0.572 (6.8%).
  • TFG + TAG vs. TFG alone (super-resolution): FID improves from 65.5 to 64.7 (1.2%); LPIPS from 0.187 to 0.175 (6.4%).

The relative improvements are larger for DPS (7-8%) than for TFG (1-6%), consistent with TAG's design principle: TFG is a more carefully optimized method that produces less off-manifold drift (its FID scores are already substantially lower than DPS's), so TAG has less room to help. DPS, being a simpler and more aggressive guidance method, pushes samples further off-manifold, creating larger corrections for TAG to address.

Label-guided generation (CIFAR-10, ImageNet):

  • CIFAR-10, DPS + TAG vs. DPS: FID improves from 217.1 to 190.4 (12.3% relative improvement); accuracy improves from 57.5% to 63.2% (9.9% relative improvement).
  • CIFAR-10, TFG + TAG vs. TFG: FID improves from 114.1 to 102.7 (10.0% relative improvement); accuracy improves from 55.8% to 61.5% (10.2% relative improvement).
  • ImageNet, DPS + TAG vs. DPS: FID improves from 196.9 to 192.2 (2.4%); accuracy decreases from 24.5% to 22.9% (−6.5%).
  • ImageNet, TFG + TAG vs. TFG: FID improves from 231.0 to 219.4 (5.0%); accuracy improves from 14.3% to 17.8% (24.5%).

The ImageNet DPS case is the only instance in Table 2 where TAG reduces conditioning accuracy (by 6.5%). This is notable: DPS on ImageNet is already operating at low accuracy (24.5%), and the TAG correction—which prioritizes temporal alignment—slightly reduces the effectiveness of the guidance toward the target label. However, the FID improvement (2.4%) suggests the quality-fidelity tradeoff still favors TAG. For TFG on ImageNet, both metrics improve substantially (FID +5.0%, accuracy +24.5%), suggesting the combination of better base guidance (TFG) with TAG's temporal correction is particularly synergistic.

Molecular property control (QM9, six properties):

  • DPS + TAG vs. DPS (polarizability α): MAE improves from 13.33 to 7.96 (40.3% relative improvement); Atom Stability improves from 28.4% to 96.4% (239.7% relative improvement — a truly dramatic gain).
  • DPS + TAG vs. DPS (dipole moment µ): MAE from 4779.92 to 1.48 (99.9% improvement); Stability from 34.4% to 97.2%.
  • TFG + TAG vs. TFG (α): MAE from 8.91 to 4.46 (49.9% improvement); Stability from 19.2% to 43.6%.
  • TFG + TAG vs. TFG (µ): MAE from 2.41 to 1.28 (46.9% improvement); Stability from 26.3% to 94.3%.

The molecular generation results are where TAG shows its most dramatic improvements, particularly in Atom Stability. For DPS on polarizability, stability jumps from 28.4% (meaning 71.6% of atoms have incorrect valencies—the molecules are chemically invalid) to 96.4% (nearly all atoms are valid). This is a qualitative transformation: without TAG, DPS produces molecules that satisfy the target property but are chemically nonsensical; with TAG, the molecules are both property-satisfying and chemically valid. The MAE improvements are also substantial (40-50% for most properties), showing that TAG does not merely improve validity at the cost of property control—it improves both simultaneously. The paper's interpretation (implicit in the off-manifold framework) is that temporal misalignment causes the external guidance to optimize properties in ways that break molecular structure; TAG's correction constrains the optimization to remain on the manifold of valid molecules.

Audio restoration (declipping, inpainting):

  • DPS + TAG vs. DPS (declipping): FAD improves from 2.41 to 2.33 (3.3%); DTW from 191 to 189 (1.0%).
  • TFG + TAG vs. TFG (declipping): FAD from 1.42 to 0.74 (47.9% relative improvement); DTW from 256 to 120 (53.1% relative improvement). These are the largest relative FAD/DTW improvements in the table.
  • DPS + TAG vs. DPS (inpainting): FAD from 2.26 to 2.25 (0.4%); DTW from 176 to 157 (10.8%).
  • TFG + TAG vs. TFG (inpainting): FAD from 0.52 to 0.42 (19.3%); DTW from 74 to 51 (31.1%).

The audio results reveal a stark asymmetry: TAG provides minimal improvements on DPS (3.3% FAD on declipping, 0.4% on inpainting) but massive improvements on TFG (47.9% FAD on declipping, 19.3% on inpainting). This is the opposite of the image restoration pattern (where DPS benefited more from TAG). The paper does not directly discuss this reversal, but it is consistent with the off-manifold framework: for these particular audio tasks and models, TFG apparently produces more severe temporal misalignment than DPS does (or the time predictor is better at recognizing TFG-induced drift), making TAG's correction more impactful. The DTW improvements (which measure conditioning accuracy) are also dramatically larger for TFG+TAG (53.1% on declipping, 31.1% on inpainting), showing that temporal alignment significantly improves the model's ability to satisfy the restoration objective.

Comparison with baseline correction methods (bottom rows of Table 2):

  • TCS (Jung et al., 2024): On CIFAR-10 DPS, TCS achieves FID 213.4 and accuracy 29.4% — worse than DPS alone (217.1 FID, 57.5% accuracy) on accuracy by a large margin. On molecular α, TCS achieves MAE 11.44 (vs. DPS 13.33) but Stability 15.3% (vs. DPS 28.4%) — better property control but catastrophic validity loss.
  • Timestep Guidance (Sadat et al., 2024): On CIFAR-10 DPS, FID 393.2, accuracy 11.3% — substantially worse than DPS alone on both metrics.
  • Self-Guidance (Li et al., 2024b): On CIFAR-10 DPS, FID 205.4 (better than DPS's 217.1) but accuracy 51.6% (worse than DPS's 57.5%).

These comparisons demonstrate that prior time-based and score-based correction methods fail under external guidance — they either degrade fidelity, conditioning accuracy, or both. TAG is the only method that consistently improves over the base guidance across all tasks, and often does so substantially. The contrast with TCS is particularly instructive: TCS and TAG both use a time predictor, but TCS's hard reassignment strategy fails where TAG's soft gradient-based correction succeeds. This directly validates the paper's claim that the gradient of the time predictor (TLS), rather than its argmax, is the operative mechanism.

Increasing guidance strength (Table 3): robustness under aggressive guidance.

Table 3 probes how TAG performs as the external guidance becomes stronger—a regime where off-manifold drift should be more severe. DPS guidance strength is varied from 1.0 (the standard setting) to 5.0.

  • At strength 1.0, CIFAR-10: DPS achieves FID 217.1, accuracy 57.5%. DPS+TAG achieves FID 190.4, accuracy 63.2%. TAG improves both metrics.
  • At strength 2.5, CIFAR-10: DPS degrades to FID 334.1, accuracy 41.9%. DPS+TAG achieves FID 289.7, accuracy 51.9% — FID is 44.4 points better and accuracy is 10 points higher.
  • At strength 5.0, CIFAR-10: DPS further degrades to FID 384.8, accuracy 29.4%. DPS+TAG achieves FID 347.8, accuracy 41.0%. TAG still improves over DPS, but the gap narrows.
  • For molecular polarizability α, at strength 1.0: DPS achieves MAE 103.7, Stability 1.1%. DPS+TAG: MAE 48.5, Stability 32.2%. TAG improves stability by 31.1 percentage points.
  • At strength 5.0 (α): DPS: MAE 112.7, Stability 1.1%. DPS+TAG: MAE 51.7, Stability 30.4%. TAG maintains roughly the same improvement despite the much stronger (and more off-manifold-inducing) guidance.

Key observations: (1) The absolute FID of DPS+TAG increases with guidance strength (from 190.4 at strength 1.0 to 347.8 at strength 5.0), indicating that stronger guidance still pushes samples further off-manifold despite TAG's correction — TAG mitigates but does not eliminate the effect. (2) The relative improvement of TAG over DPS shrinks as strength increases (from 12.3% FID reduction at 1.0 to 9.6% at 5.0), suggesting there is a limit to how much temporal correction can compensate for extreme guidance. (3) On α, TAG maintains a roughly constant stability improvement (~30 percentage points) across all guidance strengths, while DPS alone produces essentially no stable molecules (stability 1.1%) regardless of guidance strength. This suggests that for this task, TAG's correction addresses a fundamental structural issue (molecular validity) that DPS alone cannot resolve at any guidance strength.

Multi-Conditional Guidance (Table 5): TAG Scales to Multiple Conditions Without Combinatorial Predictor Training

Table 5 reports results for generating CelebA faces conditioned on two attributes (Gender+Age, Gender+Hair) and molecules conditioned on multiple properties (α + µ; Cv + µ; all six properties simultaneously). The baseline is naive score summation without any temporal correction. TAG is tested with three time predictor variants: Multi. (a separate time predictor trained for each condition combination), Single. (a single-condition time predictor using the reparameterization in Proposition B.1), and Uncon. (an unconditional time predictor using iterative reparameterization in Proposition B.2).

CelebA Gender + Age (KID↓ / Acc↑):

  • Baseline (naive summation): KID −2.75, Accuracy 80.5%.
  • TAG Multi.: KID −2.85, Accuracy 87.1% (+6.6 percentage points accuracy).
  • TAG Single.: KID −2.86, Accuracy 91.0% (+10.5 pp accuracy).
  • TAG Uncon.: KID −2.87, Accuracy 89.1% (+8.6 pp accuracy).

CelebA Gender + Hair:

  • Baseline: KID −3.16, Accuracy 92.1%.
  • TAG Multi.: KID −3.19, Accuracy 94.9% (+2.8 pp).
  • TAG Single.: KID −3.27, Accuracy 96.1% (+4.0 pp).
  • TAG Uncon.: KID −3.08, Accuracy 96.0% (+3.9 pp).

Molecule α + µ (MAE↓ / Stability↑ for each property):

  • Baseline: MAE 13.7 / 1782.8 / 68.9, Stability 4.97% / 70.9%.
  • TAG Multi.: MAE 4.56 / 1.31 / 84.7, Stability 84.7%.
  • TAG Single.: MAE 4.65 / 1.33 / 83.9, Stability 83.9%.
  • TAG Uncon.: MAE 4.56 / 1.35 / 84.9, Stability 84.9%.

Molecule Cv + µ:

  • Baseline: MAE 1425.2 / 70.9 / 10.1, Stability 31.9% / 4.33%.
  • TAG Multi.: MAE 2.72 / 1.33 / 84.2, Stability 84.2%.
  • TAG Single.: MAE 2.63 / 1.40 / 82.9, Stability 82.9%.
  • TAG Uncon.: MAE 2.74 / 1.36 / 84.2, Stability 84.2%.

Molecule all six properties (α, µ, Cv, εHOMO, εLUMO, Δε):

  • Baseline: MAE 0.635 / 1.14 / 1.18 / 56.0 / ..., Stability 91.2% (reported as a single number across properties).
  • TAG Multi.: MAE 0.610 / 1.13 / 1.15 / 91.2%.
  • TAG Single.: MAE 0.577 / 1.05 / 1.11 / 85.9%.
  • TAG Uncon.: MAE 0.530 / 1.07 / 1.15 / 85.9%.

The headline finding is clear: all three TAG variants substantially outperform the naive multi-condition baseline across all tasks. On molecular α + µ, the baseline produces catastrophic dipole MAE of 1782.8 and stability of 4.97% — the molecules are essentially random. TAG reduces α MAE by ~67% (from 13.7 to 4.56), µ MAE by ~99.9% (from 1782.8 to ~1.3), and improves stability from 4.97% to ~84%. This is not a marginal improvement; it is the difference between useless and useful multi-conditional generation.

The second key finding is that Single. and Uncon. time predictors match or exceed Multi. performance. On CelebA Gender+Age, Single. achieves the highest accuracy (91.0% vs. Multi.'s 87.1%). On molecules α + µ, all three variants achieve essentially identical MAE and stability. On the six-property task, Uncon. achieves the best MAE for α (0.530 vs. Multi.'s 0.610) and comparable performance on other properties. This is a practically crucial result: it means practitioners do not need to train combinatorial numbers of time predictors for different condition combinations. A single unconditional time predictor, applied with iterative reparameterization (Algorithm 2), provides effective multi-conditional temporal alignment without any condition-specific training.

The slight variation across predictor types (e.g., Multi. achieving 91.2% stability on six-property while Single. achieves 85.9%) suggests that for very complex multi-condition scenarios (six simultaneous properties), the dedicated multi-condition predictor has an edge in some metrics, but the gap is surprisingly small given the dramatic difference in training requirements. The paper does not explore whether this gap widens with more conditions or different condition types.

Few-Step Generation (Tables 4 and 17): TAG Mitigates Discretization Error in Accelerated Sampling

Tables 4 and 17 evaluate TAG in an unconditional setting (no external guidance), where the only source of off-manifold drift is discretization error from reducing the number of sampling steps.

Table 4 (DDIM sampling, 1 to 100 steps):

  • CIFAR-10, 1 step: Without TAG: FID 460.0. With TAG: FID 271.1 — a 41.1% reduction. At 3 steps: 234.1 → 160.5 (31.4% reduction). At 5 steps: 158.6 → 118.8 (25.1% reduction). At 10 steps: 106.3 → 93.1 (12.4% reduction). At 50 steps: 71.8 → 70.9 (1.3% reduction). At 100 steps: 67.6 → 66.5 (1.6% reduction).
  • ImageNet, 1 step: 430.3 → 352.8 (18.0% reduction). At 3 steps: 297.6 → 265.1 (10.9% reduction). At higher step counts, the improvement narrows to ~5% or less.
  • Cat, 1 step: 433.7 → 314.8 (27.4% reduction). At 3 steps: 313.5 → 178.8 (43.0% reduction — the largest relative improvement in the table). At 5 steps: 243.9 → 199.5 (18.2% reduction).

Table 17 (DDPM sampling, 50k samples, CIFAR-10):

  • NFE 1 (single-step generation): Without TAG: FID 449.8, IS 1.26. With TAG: FID 232.9, IS 2.26 — FID reduced by 48.2%, IS nearly doubled.
  • NFE 3: FID 194.5 → 124.2, IS 2.04 → 3.55.
  • NFE 5: FID 116.5 → 97.4, IS 3.08 → 3.66.

The pattern is consistent and strong: TAG's benefit is inversely proportional to the number of sampling steps. At 100 steps (near the full schedule), TAG provides negligible improvement (~1-2% FID reduction). At 1 step (extreme discretization), TAG provides 18-48% FID reduction depending on the dataset. This directly validates the paper's claim that TAG corrects off-manifold drift: when steps are plentiful and discretization error is small, samples naturally remain near the correct manifold, and TAG has little to do. When steps are scarce and discretization error is large, samples drift significantly off-manifold, and TAG's correction becomes crucial.

The Inception Score improvements in Table 17 are particularly telling: at NFE 1, IS increases from 1.26 (near-random — indistinguishable classes) to 2.26 with TAG, despite the extreme single-step setting. This suggests TAG is not merely improving overall image quality (FID) but specifically improving class distinctiveness—the images are more recognizable as belonging to specific classes.

The Cat dataset's anomalous behavior (largest improvement at 3 steps rather than 1 step) may reflect dataset-specific properties of the discretization error—perhaps at 1 step, even TAG cannot recover meaningful structure, while at 3 steps there is just enough signal for temporal alignment to provide substantial benefit.

Large-Scale Text-to-Image Generation (Table 6): TAG Scales to Production Models

Table 6 evaluates TAG on Stable Diffusion v1.5 in two practical scenarios: reward alignment (optimizing generated images to maximize Aesthetic or CLIP scores) and style transfer (matching both text prompts and reference style images).

Single-objective DAS (Aesthetic score ↑, Time-Gap ↓):

  • DAS alone: Aesthetic 7.948, Time-Gap 90.04.
  • DAS + TAG: Aesthetic 9.087 (14.3% improvement), Time-Gap 28.84 (68.0% reduction).

Single-objective DAS (CLIP score ↑):

  • DAS alone: CLIP 0.389, Time-Gap 20.73.
  • DAS + TAG: CLIP 0.439 (12.9% improvement), Time-Gap 11.62 (43.9% reduction).

Multi-objective DAS (Aesthetic + CLIP):

  • DAS alone: Aesthetic 8.107, CLIP 0.439, Time-Gap 20.73.
  • DAS + TAG: Aesthetic 8.572 (5.7% improvement), CLIP 0.463 (5.5% improvement), Time-Gap 9.765 (52.9% reduction).

Style Transfer (Style Score ↓, Time-Gap ↓):

  • TFG alone: Style Score 4.82, Time-Gap 80.6.
  • TFG + TAG: Style Score 3.03 (37.1% reduction of style error), Time-Gap 23.6 (70.7% reduction).

The results demonstrate that TAG's mechanism scales to latent diffusion models (Stable Diffusion operates in a compressed latent space, not pixel space) and to reward-based optimization (DAS uses importance sampling and resampling, not just gradient-based guidance). The Time-Gap reductions are particularly dramatic (44-71% across all settings), providing direct evidence that the reward optimization process causes temporal misalignment and that TAG effectively corrects it.

Critically, TAG improves both the reward (Aesthetic/CLIP scores increase) and reduces Time-Gap simultaneously. This is the non-trivial claim: one might expect that constraining the optimization to stay on the temporal manifold would reduce the achievable reward (a "quality-validity tradeoff"), but the opposite occurs. The paper's interpretation—implicit in the Bayesian factorization of Section 3.2—is that temporal alignment and property satisfaction are complementary: samples that remain on the diffusion manifold are more amenable to accurate property optimization because the score function remains reliable, enabling the guidance to find genuinely better solutions rather than exploiting score approximation errors.

The style transfer results are particularly strong: Style Score drops from 4.82 to 3.03, while Time-Gap drops from 80.6 to 23.6. This means TAG simultaneously improves style matching (the primary objective) and temporal alignment (the diagnostic of off-manifold drift). The paper does not provide qualitative examples for the style transfer experiments (the visualizations in Appendix G are limited to CIFAR-10, ImageNet, molecules, and CelebA), which is a gap—style transfer is an inherently visual task, and the Style Score metric may not fully capture perceptual quality.


Ablation Studies and Robustness Checks

Time predictor architecture comparison (Table 18, Appendix E.4): The SimpleCNN time predictor (1.48M parameters, four convolutional layers with average pooling) is compared against a UNet encoder (17.38M parameters) on CIFAR-10 for both unconditional and conditional time predictors. Across training checkpoints from 50K to 300K iterations, both architectures achieve comparable time-gap performance—for example, at 300K iterations, SimpleCNN unconditional achieves Time-Gap 22.93 vs. UNet's 24.40; SimpleCNN conditional achieves 21.11 vs. UNet's 22.49. The SimpleCNN is not merely "good enough for its size"—it is competitive with a model 11.7× larger. This validates the paper's claim that timestep classification is an inherently simpler task than denoising and does not require the representational capacity of the full diffusion backbone.

Effect of time predictor training steps (Table 16): On CIFAR-10 with TFG+TAG, a time predictor trained for 10K iterations achieves FID 116.0 and accuracy 55.3%; training for 30K iterations improves to FID 102.7 and accuracy 61.5%. This monotonic improvement confirms that TAG's performance is limited by time predictor quality—better temporal classification directly translates to more effective correction. The improvement is substantial (11.5% FID reduction, 6.2 percentage points accuracy gain from 3× more training), suggesting that further gains might be achievable with even more training or better predictor architectures.

Guidance strength ω sweep under corrupted reverse process (Tables 1, 10, 11): At fixed noise σ = 0.2, sweeping ω from 0 to 200 reveals several patterns:

  • Time-Gap decreases monotonically: from 273.9 at ω = 0 to 158.9 at ω = 200 (Table 10). The most rapid decline occurs between ω = 1.0 and ω = 5.0 (from 250.9 to 185.1), after which improvements are more gradual. This suggests there are "easy" temporal misalignments that moderate TAG corrects quickly, and "hard" misalignments that require progressively stronger correction.
  • FID follows a U-shaped curve: FID decreases from 410.1 at ω = 0 to a minimum of 223.2 at ω = 200, with monotonic improvement across the full sweep. There is no observed over-correction regime where FID starts increasing—even ω = 200 (extremely strong correction) maintains the lowest FID. However, the paper notes (Table 1) that for σ = 0.1, the optimal ω is ~2.0, with FID actually increasing slightly at ω = 4.0 (159.8 vs. 120.9 at ω = 2.0), suggesting that over-correction can occur when the noise level is lower.
  • IS increases monotonically with ω: from 1.27 at ω = 0 to 2.17 at ω = 200. This suggests stronger temporal alignment improves class distinctiveness without collapsing mode coverage.
  • Optimal ω scales with noise level: For σ = 0.05, best FID at ω = 0.2 (FID 62.5). For σ = 0.1, ω = 1.0 (FID 115.6). For σ = 0.2, ω = 4.5 (FID 230.9). For σ = 0.3, ω = 200 (FID 223.2). The relationship is super-linear: doubling noise from 0.1 to 0.2 requires ~4.5× stronger correction, while going from 0.2 to 0.3 requires ~44× stronger correction. This suggests that at very high noise, the TLS gradient magnitude becomes small (the time predictor becomes uncertain about which manifold to target), requiring much larger ω to achieve meaningful correction.

Time-Gap correlation with standard quality metrics (Figure 10, Appendix F.2): On CIFAR-10 unconditional generation, Time-Gap, FID, and IS are measured as the number of function evaluations (NFE) varies from 1 to 50. As NFE increases, Time-Gap decreases from ~275 at NFE 1 to ~17 at NFE 50, while FID decreases from ~440 to ~70 and IS increases from ~1.3 to ~6.8. The correlation is strong and monotonic, validating Time-Gap as a diagnostic proxy for generation quality. The paper notes a limitation: "Once the reverse diffusion process is good enough (i.e., time gap is already small), it often loses correlation with FID measure" — at high NFE, Time-Gap continues to decrease slightly while FID plateaus, suggesting that temporal alignment is necessary but not sufficient for optimal quality.

Varying DPS guidance strength (Table 3, discussed above under Main Results): This ablation demonstrates that TAG's benefit is robust to guidance strength but not unlimited—at extreme strength (5.0 on CIFAR-10), TAG still helps (FID 347.8 vs. 384.8) but the absolute quality is poor.

Multi-conditional time predictor variants (Table 5, discussed above): Single-condition and unconditional predictors match multi-condition performance, validating the reparameterization approximations in Propositions B.1 and B.2.

Input perturbation during training (Table 7): Ning et al. (2023) propose adding noise during diffusion model training (controlled by η) to reduce exposure bias. The paper retrains CIFAR-10 diffusion models with η ∈ {0, 0.05, 0.10, 0.15} from scratch and evaluates with DPS. No input perturbation level improves over η = 0 (no perturbation): at η = 0, FID 332.0, accuracy 28.5%; at η = 0.15, FID 326.7, accuracy 29.2% (slightly better accuracy but worse FID). This demonstrates that training-time exposure bias reduction does not generalize to the off-manifold regime caused by external guidance—the distribution shift during guided sampling is qualitatively different from the train-test mismatch that input perturbation addresses.

Comparison with additional baselines on CIFAR-10 DPS (Table 8): This consolidates the comparison from Table 2 and adds Epsilon Scaling, Time-Shift Sampler, and Langevin Dynamics (all applied on top of DPS). Epsilon Scaling achieves the best FID among non-TAG methods (186.0) but at reduced accuracy (53.0% vs. DPS's 57.5%). Time-Shift Sampler achieves FID 237.0 (worse than DPS) with accuracy 60.8% (better than DPS). Langevin Dynamics achieves FID 226.8 (worse than DPS) with accuracy 58.2% (marginally better). Only TAG improves both metrics substantially: FID 190.4 (better than all others), accuracy 63.2% (better than all others). This demonstrates that TAG is not merely trading off fidelity for accuracy or vice versa—it provides a genuine Pareto improvement.

Time-Gap analysis across datasets (Figures 4-9, Appendix F.2): The paper measures per-timestep Time-Gap for unconditional and conditional generation across CIFAR-10, ImageNet, Cat, Audio, CelebA, and Molecule. In unconditional settings, Time-Gap is small across all timesteps for most datasets (the time predictor accurately identifies the correct timestep). Under conditional guidance, Time-Gap increases substantially, particularly at intermediate timesteps (roughly t = 200-800), confirming that external guidance causes temporal misalignment and that TAG's correction is most needed in these regimes. For molecules, the Time-Gap is notably larger even in unconditional generation (Figure 9a), suggesting that the EGNN-based diffusion model produces inherently more temporal ambiguity than image-based models—and indeed, TAG shows its largest relative improvements on molecular tasks.

50k-sample evaluation on CIFAR-10 (Table 13): The main Table 2 uses only 512 samples for CIFAR-10 to enable rapid experimentation across many targets and methods. To verify that results hold at standard benchmark scales, the paper evaluates TFG+TAG vs. TFG on 50,000 samples with 100 inference steps. TFG alone: FID 77.5, accuracy 54.3%. TFG+TAG: FID 47.1, accuracy 84.4% — FID improves by 39.2% and accuracy by 30.1 percentage points. Notably, the relative improvement is actually larger at 50k samples than at 512 samples (where FID improved from 114.1 to 102.7, a 10.0% reduction). This may reflect that the time predictor benefits from the more stable FID estimation at larger sample sizes, or that the smaller 512-sample evaluation introduced noise that obscured some of TAG's benefit. Either way, the large-scale evaluation confirms and strengthens the paper's claims.


Critical Assessment

Claim 1: TAG Consistently Improves Fidelity and Validity Across the TFG Benchmark

The experimental evidence in Table 2 supports this claim with significant breadth (six datasets, eight tasks, two base methods, four baseline comparisons) but variable depth of verification. The claim holds most strongly for molecular generation (where TAG produces dramatic, qualitative improvements—stability jumping from 28.4% to 96.4% for DPS on polarizability) and for TFG on audio (47.9% FAD reduction on declipping). It holds moderately for image label guidance (10-12% FID improvement on CIFAR-10, 2.4-5.0% on ImageNet) and weakly for DPS on audio (0.4-3.3% FAD improvement). This variance is not explained by the paper—why does TAG provide massive gains on TFG audio but minimal gains on DPS audio? The paper's off-manifold framework would predict that TAG helps more when off-manifold drift is more severe, but then one would need to measure or estimate the severity of drift for each task-method combination to verify this explanation. The paper does not provide such measurements.

A genuine weakness is the single sample size per evaluation. For most experiments in Table 2, only 128-512 samples are generated per target condition. At these sample sizes, FID estimates have non-trivial variance—a difference of a few FID points may not be statistically reliable. The paper provides no confidence intervals, no standard deviations, and no multiple-run error bars. The 50k-sample evaluation in Table 13 partially addresses this for CIFAR-10 (and reassuringly shows even larger improvements at scale), but equivalent large-scale evaluations are absent for ImageNet, Cat, molecules, and audio.

A second weakness is the single model per dataset design. All CIFAR-10 experiments use the same DDPM model (Nichol & Dhariwal, 2021); all ImageNet experiments use the same DDPM model (Dhariwal & Nichol, 2021). If these particular models happen to produce score functions that are unusually sensitive to off-manifold drift (or unusually amenable to temporal correction), the results would not generalize. The paper acknowledges this indirectly by testing across diverse architectures (UNet, EGNN, WaveNet-style) and modalities, which provides cross-architecture evidence but not within-task replication.

The comparison with baseline correction methods (TCS, TSG, SG) is comprehensive and fair—these baselines represent the most directly comparable prior work, and TAG outperforms them in essentially all settings. However, the baselines are applied with their reported hyperparameters rather than being re-tuned for the specific DPS/TFG configurations used in this paper. It is possible that these baselines could perform better with task-specific hyperparameter optimization, though the magnitude of TAG's advantage (e.g., TCS accuracy 29.4% vs. DPS+TAG accuracy 63.2% on CIFAR-10) makes it unlikely that tuning alone would close the gap.

A missing experiment is the direct measurement of off-manifold drift for each task-method combination using the Time-Gap metric. Table 2 reports FID, accuracy, FAD, etc., but does not report Time-Gap for these experiments (Time-Gap data is only shown for the corrupted reverse process in Section 3.4 and partially for text-to-image in Table 6). If the paper's central claim is that TAG works by reducing temporal misalignment, then showing that TAG reduces Time-Gap in proportion to its improvement in FID/accuracy across all tasks would substantially strengthen the mechanistic argument. The correlation analysis in Appendix F.2 (Figures 4-9) shows per-timestep Time-Gap for some settings, but these are not systematically matched to the quantitative results in Table 2.

Claim 2: TAG's Benefit Is Largest When Off-Manifold Drift Is Most Severe

This claim is supported by two lines of evidence: the corrupted reverse process experiments (Section 3.4) and the DPS guidance strength sweep (Table 3).

The corrupted reverse process experiments are clean and convincing for the specific noise perturbation studied. As σ increases from 0.05 to 0.3, FID without TAG degrades from 78.9 to 410.1, and TAG's relative improvement grows from 20.8% FID reduction (78.9 → 62.5) to 45.6% (410.1 → 223.2). The monotonic relationship between corruption level and TAG benefit is exactly what the off-manifold framework predicts. However, this experiment studies only one type of "drift"—isotropic Gaussian noise added uniformly at every step. This is not representative of the structured drift produced by external guidance (which pushes samples in specific, content-dependent directions) or discretization error (which depends on the local curvature of the probability flow ODE). The paper implicitly acknowledges this by treating the corrupted reverse process as a separate analysis (Section 3.4) rather than integrating it into the main benchmark, but does not discuss how the noise perturbation relates quantitatively to the drift produced by DPS or TFG.

The DPS guidance strength sweep (Table 3) provides evidence in a more realistic setting. As DPS strength increases from 1.0 to 5.0 on CIFAR-10, DPS-only FID degrades from 217.1 to 384.8, and TAG's absolute improvement grows from 26.7 FID points to 37.0 FID points. However, the relative improvement shrinks (12.3% at strength 1.0, 9.6% at strength 5.0), which complicates the narrative. If TAG's benefit were strictly proportional to drift severity, relative improvement should grow, not shrink. One interpretation is that at very high guidance strengths, the drift is so severe that even TAG's temporal correction cannot fully compensate—the sample is pushed so far off-manifold that the time predictor itself becomes unreliable. This interpretation is consistent with the Time-Gap analysis showing that even with TAG, Time-Gap remains high under extreme corruption (158.9 at σ = 0.3 with ω = 200, Table 10), but the paper does not directly address this limit.

A missing experiment is a systematic study of TAG's failure mode at extreme drift. At what noise level or guidance strength does TAG provide zero benefit? Does TAG ever hurt (increase FID relative to no correction) in any regime? The corrupted reverse process experiment shows TAG always helps (FID is lower with TAG for all σ and all ω > 0 in Table 1), but the range of ω tested at each σ was limited (0 to 4.0 for σ ≤ 0.2; the extended sweep in Table 10 only covers σ = 0.2). A full grid of (σ, ω) with very high values of both would reveal whether there is a regime where TAG over-corrects and degrades quality—the existence of such a regime is suggested by the slight FID increase at ω = 4.0 vs. ω = 2.0 for σ = 0.1 (Table 1: FID 159.8 at ω = 4.0 vs. 120.9 at ω = 2.0).

Claim 3: TAG Is a General Solution Applicable Across Diverse Domains and Tasks

The breadth of evaluation—six datasets spanning images, molecules, and audio; tasks including restoration, conditional generation, multi-conditional generation, few-step sampling, and text-to-image alignment—provides strong prima facie evidence for generality. The consistent improvement pattern (TAG never substantially degrades performance, almost always improves it) across this diversity is impressive.

However, there are important caveats:

Single model architecture per domain. Within each domain, only one diffusion model architecture is tested. For images, this is DDPM (UNet-based); for molecules, EDM (EGNN-based); for audio, Audio-Diffusion (WaveNet-style). It is possible that these specific architectures have properties (e.g., noise schedule, score parameterization) that make them particularly amenable to temporal correction. The text-to-image experiments with Stable Diffusion (latent diffusion, different architecture family) provide some cross-architecture evidence within the image domain, but equivalent diversity is not shown for molecules or audio.

Small sample sizes for non-image tasks. The molecular evaluations use 1024-4096 samples; audio uses 256 samples. For molecular property prediction, MAE computed on 1024 molecules has substantial estimation variance—a difference of a few MAE units may not be statistically significant. The paper does not report error bars.

Hyperparameter tuning is per-task and per-target. The guidance strength ω₀ is tuned via grid search for each task and each target value individually. This means the reported results represent the best achievable performance given oracle hyperparameter knowledge, not what a practitioner would achieve on a new task without tuning. The paper does not study whether the optimal ω₀ transfers across targets within a task (e.g., does the ω₀ that works best for CIFAR-10 class "airplane" also work well for "automobile"?), nor does it propose a method for automatic ω₀ selection. In fairness, the same tuning protocol applies to the baseline methods (DPS, TFG), so the comparison is fair given oracle tuning for all methods. But the claim of "generality" should be qualified: TAG generalizes across domains and tasks given task-specific hyperparameter tuning, not in a zero-shot plug-and-play manner.

Missing domains. The paper does not evaluate TAG on video generation, language modeling, or reinforcement learning (all mentioned as future work in Section 5). Given that the off-manifold phenomenon is framed as universal to all diffusion model applications, the absence of any non-image/non-audio/non-molecule domain limits the generality claim. The mechanism (temporal misalignment → score approximation error) should apply to any diffusion model regardless of modality, but empirical verification is absent.

Claim 4: The Time Predictor's Gradient (TLS) Is the Operative Mechanism, Not Its Argmax

This claim is supported by the comparison with TCS (Table 2). TCS and TAG both use a time predictor, but TCS uses arg max φ(x_t) to hard-reassign timesteps, while TAG uses ∇_x log p_φ(t | x_t) as a soft gradient correction. On CIFAR-10 DPS, TCS achieves FID 213.4 and accuracy 29.4% (dramatically worse than DPS alone at 217.1 and 57.5%), while TAG achieves FID 190.4 and accuracy 63.2%. This is a clean ablation: same auxiliary model (time predictor), different usage of its output, dramatically different performance. The gradient-based approach is clearly superior.

However, the TCS comparison is limited to CIFAR-10 DPS (and partially to other tasks in the bottom rows of Table 2, where TCS uniformly performs poorly). A more thorough ablation would compare TAG against TCS across all tasks and guidance methods in the full benchmark. It is possible that TCS performs better on some tasks (e.g., where the off-manifold drift is small and hard reassignment is less risky), and such cases would refine the understanding of when soft gradient correction is necessary vs. when hard reassignment suffices.

A second missing ablation is the comparison against no time predictor at all—i.e., what happens if the TLS is replaced with a simpler heuristic for temporal alignment? The paper argues that the learned time predictor is necessary because temporal identity is complex and cannot be captured analytically, but it does not test any non-learned alternatives (e.g., using the diffusion model's own internal uncertainty estimates, or a simple variance-based heuristic). Including such a baseline would strengthen the claim that the learned time predictor specifically—rather than any temporal correction signal—is the operative mechanism.

A third missing ablation is the randomized time predictor: replacing the time predictor's gradient with random noise scaled to have similar magnitude as the TLS. This would test whether the performance improvement comes from the specific direction of the TLS gradient or simply from adding any perturbation that breaks the accumulation of deterministic errors (a known phenomenon in some iterative algorithms). If random noise also helps (even if less effectively), it would suggest that part of TAG's benefit comes from stochasticity rather than directed temporal correction.

Claim 5: Multi-Conditional TAG Works Without Combinatorial Time Predictor Training

The evidence in Table 5 strongly supports this claim for the specific condition combinations and models tested. Single-condition and unconditional time predictors match or exceed dedicated multi-condition predictors on CelebA (two conditions) and molecules (up to six conditions). This is a practically important result that makes multi-conditional TAG feasible without exponential training cost.

However, the evaluation is limited to two specific multi-condition tasks (CelebA attributes and molecular properties). The reparameterization approximations in Propositions B.1 and B.2 rely on the conditions being "compatible" in the sense that sequential gradient updates approximate the joint conditional gradient. This may break down for conditions that interact non-additively—for instance, conditions that are mutually exclusive or that require trade-offs (e.g., "generate a cat that is both very large and very small"). The paper does not test adversarial or conflicting conditions, nor does it characterize when the single-condition/unconditional approximation fails.

Additionally, the multi-condition experiments use the same base models and time predictor architectures as the single-condition experiments. The fact that an unconditional time predictor trained only on noisy samples (without any condition information) can provide effective multi-condition temporal alignment is surprising and deserves more investigation. The paper's explanation—that iterative reparameterization absorbs condition information into the sample before time prediction—is plausible but not empirically validated (e.g., by showing that the reparameterized sample x″_t is closer to the correct manifold than the original x_t). A diagnostic experiment measuring Time-Gap before and after reparameterization would clarify whether the mechanism works as theorized.

Overall Assessment of Strengths and Weaknesses

Strengths:

  1. Breadth of evaluation across 6 datasets, 4 modalities, 8+ task types is exceptional for a methods paper and provides strong evidence for generality.
  2. Clean ablation against TCS isolates the TLS gradient mechanism and demonstrates its superiority over hard temporal reassignment.
  3. Corrupted reverse process experiments provide a controlled, quantitative demonstration that TAG's benefit scales with off-manifold drift severity.
  4. Multi-condition analysis shows that practical approximations (single-condition and unconditional time predictors) work nearly as well as dedicated multi-condition models, removing a major barrier to adoption.
  5. The 50k-sample CIFAR-10 evaluation confirms that results hold at standard benchmark scales.
  6. Baseline comparisons are extensive (TCS, TSG, SG, Epsilon Scaling, Time-Shift Sampler, Input Perturbation, Langevin Dynamics) and fairly implemented.

Weaknesses:

  1. No confidence intervals, standard deviations, or multiple-run error bars for any experiment. Given small sample sizes for several tasks (128-256 samples), the statistical reliability of FID/accuracy differences is uncertain.
  2. Time-Gap is not systematically reported for the main benchmark experiments (Table 2), making it impossible to verify the mechanistic claim that TAG works by reducing temporal misalignment in these specific settings.
  3. Hyperparameter tuning is per-target and per-task with oracle grid search; no method for automatic ω₀ selection is proposed or evaluated. The reported performance is an upper bound on what a practitioner would achieve without extensive tuning.
  4. Missing failure mode analysis. The paper does not systematically explore when TAG provides zero or negative benefit. The ImageNet DPS case (accuracy -6.5%) and the slight FID increase at high ω for σ = 0.1 suggest failure modes exist but are not characterized.
  5. Time predictor training cost is not systematically compared to the cost of alternative approaches (e.g., TCS, which also requires a time predictor; or fine-tuning methods, which the paper argues against). The claim that the time predictor trains "in minutes" is qualitative without wall-clock comparisons.
  6. No comparison against fine-tuning-based approaches (ControlNet, IP-Adapter, DDPO) for the tasks where those methods are applicable. The paper argues that fine-tuning is too costly for real-time varying conditions, but for fixed-condition tasks (e.g., standard label guidance), fine-tuning is a viable baseline that is not compared.
  7. Single model per domain limits within-domain generality claims. No alternative CIFAR-10 diffusion models, no alternative molecular diffusion models.
  8. The correlated reverse process experiment uses Gaussian noise, which may not be representative of the structured off-manifold drift produced by real guidance methods.
  9. No evaluation on video, language, or RL tasks despite the paper's claim of generality and explicit mention of these as future work.

The experiments do support the paper's central claims—TAG improves generation quality across diverse settings, the improvement correlates with off-manifold drift severity, the TLS gradient mechanism is superior to alternatives, and multi-conditional TAG is practical without combinatorial training. However, the support is qualitative (consistent improvement pattern) rather than quantitative (precise characterization of when and by how much TAG helps), and several practically important questions (how to set ω₀ automatically, when TAG fails, statistical reliability of the reported numbers) remain unanswered. The paper's contribution is primarily the method and its demonstrated effectiveness, not a complete characterization of its operating envelope.

6. Limitations and Trade-offs

The Cost of Difficulty Estimation Remains Unaccounted for in the Headline Efficiency Claims

The TAG framework requires a trained time predictor to compute the TLS gradient at each reverse step. The paper describes this as a lightweight auxiliary model—for images, a SimpleCNN with approximately 1.48M parameters (Appendix E.4), roughly 8.5% the size of a UNet encoder. Training completes in "minutes on a minimal computational resources" (Appendix D.2) and the per-step inference overhead is one forward and one backward pass through this small network.

However, the total cost is not amortized in any efficiency comparison. The paper never reports wall-clock time, total FLOPs for the time predictor training, or the per-step inference cost relative to the diffusion model forward pass in a way that allows direct comparison. For the TFG benchmark (Table 2), the headline numbers show TAG improving FID and accuracy—but do not report the additional computation required to achieve those gains. A practitioner comparing TAG against simply increasing the number of diffusion steps (which also improves quality without any auxiliary model) has no basis for deciding which is more compute-efficient.

The paper acknowledges this indirectly by emphasizing the predictor's lightweight architecture, but the claim that training is "minutes" is qualitative and dataset-dependent—Table 15 shows 600K iterations on 4 A100 GPUs for ImageNet, which is non-trivial. The paper does not ablate whether a simpler, cheaper time predictor (fewer training iterations, smaller architecture) would suffice, nor does it measure how predictor quality trades off against TAG performance beyond the single 10K vs. 30K comparison in Table 16 (which shows a 11.5% FID reduction from 3× more training, suggesting performance is indeed sensitive to predictor quality).

Consequence: The reported improvements (e.g., 12.3% FID reduction on CIFAR-10 DPS) are achieved with an unquantified additional computational cost. The "8.5% of UNet encoder" figure refers to parameter count, not FLOPs—the backward pass for computing the TLS gradient adds additional cost beyond the forward pass. Without a FLOPs-matched or wall-clock-matched comparison against alternatives (more diffusion steps, Langevin correction, or simply stronger guidance), the practitioner cannot determine whether TAG is the most compute-efficient way to improve quality for their specific budget.

Mitigation status: Not addressed. The paper treats the time predictor as "lightweight" and moves on. No ablation varies time predictor capacity or training budget to characterize the cost-quality tradeoff. No comparison reports wall-clock time for any experiment.


Hyperparameter Sensitivity Requiring Per-Task Oracle Tuning

TAG introduces at least one critical hyperparameter: the guidance strength ω₀ (which scales the TLS correction via the schedule ω_t = ω₀ √(1 − ᾱ_t)). The paper tunes ω₀ via grid search for each task and each target value independently. For the TFG benchmark (Section 4.1), the authors state:

"The final results are averaged over the best-performing guidance strength ω₀ according to the grid search for all target values in each task."

The search range varies per task (Appendix E.3.7, Table 14), and the optimal value differs dramatically across settings. The corrupted reverse process experiments (Table 10) reveal that the optimal ω is highly sensitive to the severity of off-manifold drift: for noise level σ = 0.05, the best FID occurs at ω = 0.2; for σ = 0.1, at ω = 1.0; for σ = 0.2, at ω = 4.5; for σ = 0.3, at ω = 200.0—a three-order-of-magnitude range. In the DPS strength sweep (Table 3), the optimal TAG strength varies across guidance strength levels, tasks (CIFAR-10 vs. ImageNet vs. molecules), and metrics (fidelity vs. validity).

Consequence: The reported results represent an oracle upper bound on TAG's performance—what is achievable when the best ω₀ for each task and target is known in advance. A practitioner deploying TAG on a new task, with a new diffusion model, or with a new guidance method has no principled way to select ω₀ without running their own expensive grid search. If the optimal ω₀ varies across targets within a task (as is likely given the per-target tuning), using a single ω₀ averaged across targets will degrade performance relative to the reported numbers.

Worse, the paper provides no diagnostic or heuristic for setting ω₀. The Time-Gap metric (Definition F.1) could in principle serve as a tuning signal—one could sweep ω₀ and select the value that minimizes Time-Gap without requiring ground-truth quality metrics—but this is neither proposed nor evaluated. The guidance schedule ω_t = ω₀ √(1 − ᾱ_t) is motivated by the observation that temporal information increases as noise decreases (Section 3.4), but the square-root scaling is heuristic; the paper does not ablate alternative schedules (linear, constant, learned) to determine whether this specific choice matters.

Mitigation status: Not addressed. The paper does not propose an automatic ω₀ selection method, does not evaluate how performance degrades under suboptimal ω₀, and does not study whether the optimal ω₀ transfers across similar tasks or targets. This is a substantial gap between the paper's demonstration of TAG's potential (with oracle tuning) and its practical deployability (where tuning is expensive or impossible).


Failure Mode Under Extreme Off-Manifold Drift: The Time Predictor Itself Becomes Unreliable

TAG's central premise is that the time predictor remains reliable when the diffusion model's score function becomes unreliable—that the predictor, trained on forward-process data, can recognize temporal identity even for samples far from the correct manifold. The paper's own data partially contradicts this premise.

The corrupted reverse process experiments (Table 10) show that at σ = 0.3 (extreme noise), even with the strongest TAG correction tested (ω = 200), the Time-Gap remains 158.9—meaning the time predictor's argmax is, on average, 159 timesteps away from the correct timestep. At this level of temporal misalignment, the TLS gradient is being computed for a timestep that the predictor cannot reliably identify. The achieved FID (223.2) is dramatically better than without TAG (410.1), so the correction is still directionally useful, but it is far from the clean-generation baseline (FID ~12). The paper does not explore whether further increasing ω would eventually cause the TLS gradient to become actively harmful (pointing toward the wrong manifold because the predictor is too uncertain) or whether there is a hard ceiling on recoverable quality determined by time predictor accuracy.

The DPS guidance strength sweep (Table 3) shows a related pattern: TAG's relative improvement over DPS alone shrinks as guidance strength increases. On CIFAR-10, TAG reduces FID by 12.3% at DPS strength 1.0, but only by 9.6% at strength 5.0. This suggests that at very high drift severity, TAG's correction becomes less effective—consistent with the time predictor's accuracy degrading as samples are pushed further from the forward-process distribution.

Consequence: TAG cannot fully compensate for arbitrary off-manifold drift. There exists a regime—not precisely characterized in the paper—where the time predictor's accuracy degrades to the point that the TLS provides only weak or noisy correction, and further increasing ω does not recover quality. This sets a hard ceiling on the method's applicability: for guidance methods or corruption levels that push samples into regions where even the time predictor cannot recognize temporal identity, TAG offers diminishing returns. The paper does not identify where this ceiling lies for the standard TFG benchmark tasks, making it impossible for a practitioner to predict whether TAG will help substantially or marginally for their specific use case.

The paper acknowledges a related limitation in Appendix F.2: "Once the reverse diffusion process is good enough (i.e., time gap is already small), it often loses correlation with FID measure" and "improving the performance of the time predictor network will reduce this problem and thereby further boost the effect of the TAG." But this is framed as a saturation effect at the high-quality end, not as a failure mode at the low-quality (high-drift) end where the predictor breaks down. The breakdown regime is empirically present in the paper's own data but is not analyzed as a limitation.

Mitigation status: Partially acknowledged in qualitative terms (Appendix F.2 notes that predictor performance degrades for low-dimensional data and near the final timestep T), but the failure mode under extreme drift is not systematically characterized. The paper suggests improving the time predictor as future work but does not propose specific directions (adversarial training? ensemble predictors? confidence-based gating of the TLS?).


Single Model per Domain and Absence of Within-Task Replication

All experiments within each domain use exactly one pretrained diffusion model: CIFAR10-DDPM (Nichol & Dhariwal, 2021) for CIFAR-10, ImageNet-DDPM (Dhariwal & Nichol, 2021) for ImageNet, Cat-DDPM for Cat, CelebA-DDPM for CelebA, Molecule-EDM (Hoogeboom et al., 2022) for molecular generation, Audio-Diffusion for audio, and Stable Diffusion v1.5 for text-to-image. For any given dataset and task, TAG is evaluated against baselines on that single model. There is no within-domain replication across different model architectures, training recipes, or noise schedules.

The paper argues for generality by testing across diverse architectures (UNet, EGNN, WaveNet-style, latent diffusion) and modalities (pixels, 3D coordinates, mel-spectrograms, latent codes). This provides cross-domain evidence but no within-domain evidence. If the specific CIFAR-10 DDPM model happens to have properties that make temporal misalignment particularly detectable by a SimpleCNN (e.g., its noise schedule produces clearly separable temporal signatures), then TAG's performance on CIFAR-10 may not transfer to a different CIFAR-10 diffusion model with a different architecture or training procedure.

Consequence: The reported results are model-specific within each domain. A practitioner using a different diffusion model for CIFAR-10 (e.g., a score-based model with VE-SDE rather than VP-SDE, or a latent diffusion model at a different resolution) has no direct evidence that TAG will provide similar improvements. The paper's off-manifold framework predicts that TAG should help for any diffusion model where external guidance causes temporal misalignment, but the magnitude of the improvement depends on properties of the specific model (how distinguishable its temporal manifolds are, how its score function degrades off-manifold, how its noise schedule interacts with the time predictor's classification accuracy). These properties are not characterized or compared across models.

Additionally, the molecular generation experiments use an EGNN-based model (Hoogeboom et al., 2022), which is the dominant architecture for 3D molecular diffusion but not the only one. Alternative molecular diffusion models (e.g., GeoLDM, Xu et al., 2023; or EDM with different equivariance constraints) might exhibit different temporal alignment characteristics. The paper's dramatic stability improvements (28.4% → 96.4% for DPS on polarizability, Table 2) are the strongest results in the paper; if these are partly specific to the EGNN architecture's tendency to produce severe temporal misalignment under guidance, the results would not generalize to other molecular diffusion models.

Mitigation status: Not addressed. The paper treats diversity of domains as sufficient evidence for generality, without acknowledging the single-model-per-domain limitation or calling for within-domain replication.


The Time Predictor's Design and Training Are Insufficiently Validated as Optimal

The time predictor is the single most critical component of TAG—if it provides inaccurate TLS gradients, the entire correction mechanism fails. The paper makes specific design choices: SimpleCNN architecture with 2×2 average pooling, cross-entropy classification objective rather than regression, training on forward-process samples with uniform timestep sampling, and the guidance schedule ω_t = ω₀ √(1 − ᾱ_t). While the paper provides justifications for these choices (Section 3, Appendix E.4), it does not systematically ablate them.

Several potential issues are raised by the paper's own data:

Predictor accuracy degrades at high t (near the end of diffusion, t close to T = 1000). Figures 4-9 (Appendix F.2) show per-timestep Time-Gap for various datasets. For CIFAR-10 and molecules, Time-Gap increases substantially as t → T, indicating the predictor cannot reliably distinguish late-stage timesteps. The paper notes this as consistent with Kahouli et al. (2024)'s finding that "overlapping distributions near T impede accurate predictions." However, this is precisely where TAG applies the strongest correction (because √(1 − ᾱ_t) is largest at large t). If the TLS gradient is least reliable when the correction is strongest, this schedule may be suboptimal—the paper does not test whether reversing the schedule (stronger correction at low t where prediction is more accurate) would perform better.

Cross-entropy classification may not be optimal for gradient quality. The time predictor is trained to maximize the probability of the correct timestep, which encourages sharp posteriors. But the TLS uses the gradient of the log-probability, not the probability itself. A predictor that is overconfident (assigning near-zero probability to all incorrect timesteps) will produce TLS gradients that point strongly toward the correct manifold but may lack the repulsive components from alternative manifolds that Theorem 3.3 shows are essential to TAG's mechanism. Conversely, a predictor with high entropy (spreading probability across many timesteps) will produce weaker but potentially more robust gradients. The paper does not explore how training objectives (label smoothing, temperature scaling, confidence calibration) affect TLS quality. The comparison between different training checkpoints (Table 16, 10K vs. 30K iterations) shows that more training improves TAG performance, but this only demonstrates that accuracy matters—not that the chosen objective is optimal for gradient-based correction.

The time predictor architecture was chosen by comparison against a UNet encoder (Table 18), but no ablations of the SimpleCNN's internal design are reported. The number of layers, channel widths, pooling strategy, and activation functions are stated but not varied. A practitioner wanting to deploy TAG on a new dataset has no guidance on how to scale the predictor (e.g., does a larger predictor always improve TAG performance, or is there a saturation point?).

Consequence: TAG's performance may be improvable with better time predictor design, but the current design choices are not validated as sufficient or near-optimal. More importantly, a practitioner who trains a time predictor following the paper's recipe but achieves a slightly different accuracy profile (e.g., better at high t, worse at low t) may observe qualitatively different TAG behavior than reported—perhaps weaker correction overall, or correction that helps on some tasks but not others. The paper provides no diagnostic for determining whether a given time predictor is "good enough" for effective TAG beyond the qualitative Time-Gap correlation in Figure 10.

Mitigation status: Partially addressed. The paper acknowledges in Appendix A (Limitations) that "more sophisticated predictor architectures could unlock additional gains" and leaves this to future work. The UNet comparison (Table 18) shows that a much larger model does not help, suggesting the SimpleCNN is sufficient for current performance levels—but this does not validate that it is optimal, nor does it rule out that a different small architecture (e.g., with attention, with residual connections) might perform substantially better.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a diagnostic reframing rather than a paradigm shift: it identifies temporal misalignment—a sample losing its correct noise-level identity during the reverse process—as the common root cause of the diverse failure modes that plague guided diffusion models. This is not a new architectural paradigm like the introduction of diffusion models themselves, nor a new training objective like classifier-free guidance. It is a reclassification of symptoms that pushes the field toward a specific, measurable, and correctable mechanism.

The landscape change is methodological rather than conceptual. Before TAG, a researcher encountering poor sample quality under a new guidance method faced a diffuse set of possible culprits: exposure bias, score approximation error, discretization error, or an ill-chosen loss function. The toolbox of fixes was correspondingly scattered—add noise during training (Ning et al., 2023), scale model outputs (Ning et al., 2024), perturb time inputs (Sadat et al., 2024; Li et al., 2024b), or hard-reassign timesteps (Jung et al., 2024). Each fix addressed a different hypothesized mechanism, and their inconsistent performance across tasks (Tables 2, 8) reflected the lack of a unifying diagnosis.

TAG consolidates this fractured landscape around a single measurable quantity: Time-Gap (Definition F.1). The claim is not that temporal misalignment is the only source of error in guided diffusion—score approximation error, for instance, can occur even for perfectly aligned samples if the score model is undertrained. Rather, the claim is that temporal misalignment is the dominant source when external perturbations are applied, and that correcting it yields consistent, substantial improvements across a wide range of scenarios. The evidence for this consolidation is the breadth of the experimental results: the same mechanism (TLS gradient), applied with the same algorithmic structure (Algorithm 1), and the same auxiliary model (time predictor), improves quality across image restoration, label-guided generation, molecular property control, audio restoration, multi-conditional generation, few-step sampling, and large-scale text-to-image alignment (Tables 2, 4, 5, 6).

The paper reconciles several prior contradictory findings, though it does so implicitly rather than through direct replication of prior work:

  • Why Langevin correction sometimes helps and sometimes hurts. Song & Ermon (2019) and Song et al. (2021b) proposed corrector steps that follow the learned score function for additional refinement. The paper's Table 8 shows Langevin dynamics applied to DPS on CIFAR-10 increases FID from 217.1 to 226.8. TAG's framework explains this: when a sample is off-manifold, the score function itself is unreliable (it was trained on forward-process samples that assume temporal alignment), so following it more carefully only compounds the error. Langevin correction works when samples are already near the manifold (the standard unconditional setting); it fails when they have been pushed away. This reconciles the positive results in the original score-based generative modeling literature with the negative results observed under external guidance.

  • Why time perturbation methods (TSG, SG) underperform. Sadat et al. (2024) and Li et al. (2024b) derive contrastive guidance by perturbing the time input to the score model and measuring the output change. Table 2 shows these methods achieve FID 393.2 and 205.4 respectively on CIFAR-10 DPS (vs. 217.1 for DPS alone). TAG's decomposition (Theorem 3.3) explains why: the TLS is a weighted combination of score functions from all timesteps, providing both attraction toward the correct manifold and repulsion from incorrect ones. Time perturbation methods, by contrast, capture only local sensitivity around the current timestep—a first-order approximation that becomes inaccurate when the sample is far from the correct manifold. The competitive structure of the TLS (repulsion from wherever the sample currently "looks like" it belongs) provides a qualitatively stronger correction signal.

  • Why TCS's hard temporal reassignment degrades so severely. TCS (Jung et al., 2024) uses the same time predictor architecture but reassigns timesteps rather than using the gradient. Table 2 shows TCS achieves FID 213.4 and accuracy 29.4% on CIFAR-10 DPS—accuracy drops by nearly half relative to DPS alone (57.5%). Theorem 3.3 explains this failure: hard reassignment discards the competitive information in the posterior distribution p(tx)p(t|x) and commits to a single timestep. If the predictor is slightly wrong (which becomes increasingly likely as off-manifold drift increases), the hard reassignment evaluates the score function at an incorrect noise level, compounding the error. TAG's soft gradient correction avoids this commitment—it moves the sample incrementally, maintaining approximate validity at each step.

These reconciliations shift the research agenda. The question "which correction method works best?" is replaced by "how can we best measure and reduce temporal misalignment?" This is a more productive framing because it points toward a specific optimization target (the Time-Gap) rather than an empirical horse race among heuristics.

The paper also redirects investment in auxiliary models. Prior work that used time predictors (TCS, CDM) treated them as means to an end—a component whose output (argmax, density ratio) feeds into the main algorithm. TAG identifies the gradient of the time predictor as the valuable quantity, not its forward-pass output. This changes how future work should design and evaluate time predictors: the relevant metric is not classification accuracy but gradient quality. A predictor with 95% timestep classification accuracy but noisy gradients may produce worse TLS correction than a predictor with 90% accuracy but smooth, well-behaved gradients. This insight is not explicitly stated in the paper but follows directly from Theorem 3.3 and the empirical comparison with TCS.

Finally, the paper establishes Time-Gap as a diagnostic tool for the field. Figure 10 demonstrates that Time-Gap correlates with standard quality metrics (FID, IS) across different generation regimes. Figures 4-9 provide per-timestep Time-Gap profiles for various datasets and tasks, revealing where in the diffusion process off-manifold drift is most severe. This is comparable to how the Inception Score operationalized the previously vague notion of "generated image quality" for GANs—Time-Gap operationalizes the previously vague notion of "off-manifold drift" for guided diffusion. Future papers proposing new guidance methods can (and should) report Time-Gap alongside FID and accuracy to demonstrate that their method does not inadvertently cause temporal misalignment.

What does not change: The paper does not challenge the fundamental structure of diffusion models, the necessity of guidance for controlled generation, or the basic tradeoff between fidelity and conditioning accuracy. TAG is a corrective mechanism that makes existing methods more robust, not a replacement for them. The paper does not claim that temporal alignment is sufficient for optimal quality—the Appendix F.2 limitation acknowledges that at high quality levels (small Time-Gap), further reducing Time-Gap yields diminishing returns. Score approximation error, architectural limitations, and guidance design choices remain important factors that TAG does not eliminate.

Follow-Up Research This Work Enables

1. Robust time predictors trained for gradient quality rather than classification accuracy. The paper trains time predictors with standard cross-entropy loss, which optimizes classification accuracy—the probability of the correct timestep. But TAG uses the gradient of the log-probability, and classification accuracy may not be the right proxy for gradient quality. A predictor that is overconfident (assigning probability near 1.0 to the correct timestep and near 0.0 to all others) will produce TLS gradients that lack the repulsive components from alternative manifolds that Theorem 3.3 shows are essential. Conversely, a predictor that is well-calibrated but produces noisy gradients (large Hessian, sharp curvature) may inject harmful variance into the reverse process. A natural follow-up would train time predictors with objectives that directly optimize gradient quality—for instance, adding a gradient smoothness regularizer (penalizing large second derivatives of the log-probability with respect to the input), using adversarial training to improve robustness in low-density regions, or employing score matching on the time predictor's own output to ensure the TLS aligns with the true xlogp(tx)\nabla_x \log p(t|x). The evaluation would measure both Time-Gap and downstream TAG performance (FID, accuracy) on the TFG benchmark, with the hypothesis that gradient-optimized predictors enable stronger TAG correction (higher ω without over-correction artifacts) and achieve lower minimum FID than classification-optimized predictors.

2. The over-correction failure mode: characterizing when and why the TLS becomes harmful. The paper's data contain hints of a failure regime that is not systematically explored. Table 1 shows that for noise level σ = 0.1, FID improves from 120.9 at ω = 1.0 to 132.6 at ω = 2.0—the correction overshoots. Table 3 shows that TAG's relative improvement over DPS shrinks from 12.3% at DPS strength 1.0 to 9.6% at strength 5.0 on CIFAR-10, suggesting correction becomes less effective at extreme drift. The Time-Gap at σ = 0.3 with ω = 200 remains 158.9 (Table 10)—the predictor is badly wrong about temporal identity, yet the TLS is still being used for correction. A targeted follow-up would systematically map the (drift severity, TAG strength) space to identify where TAG transitions from helpful to harmful. The experiment would use the corrupted reverse process setup with a dense grid of (σ, ω) values, extended to very high ω (up to 1000 or until quality collapses), measuring FID, IS, and Time-Gap. The analysis would examine whether failure correlates with Time-Gap exceeding some threshold (suggesting the predictor's accuracy has degraded below a critical level) or with the TLS gradient norm becoming too large relative to the score function norm (suggesting the correction dominates the denoising). The practical output would be a diagnostic—a simple rule based on the time predictor's confidence or the TLS gradient magnitude—that practitioners can use to detect when TAG is about to become counterproductive and should be weakened or disabled.

3. Combining TAG with on-policy time predictor training. The paper's time predictor is trained once on forward-process data and remains frozen during inference. But as TAG corrects samples during generation, the distribution of samples seen by the time predictor shifts away from the forward-process distribution—the predictor is evaluated on samples that have been modified by both external guidance and previous TAG corrections. This is the same distribution-shift problem that motivates on-policy training in reinforcement learning. A natural extension would train the time predictor iteratively: generate samples with TAG using the current predictor, measure the Time-Gap, and fine-tune the predictor on the off-policy samples to improve its accuracy in the regions actually visited during guided generation. This could be done in a lightweight online fashion (updating the predictor every few generation steps) or in an offline bootstrap (generate a dataset of TAG trajectories, retrain, repeat). The evaluation would compare standard TAG against iteratively retrained TAG on the TFG benchmark, with the hypothesis that on-policy training reduces Time-Gap (particularly at high noise levels where the paper shows the predictor degrades) and enables stronger, more effective correction. The danger—acknowledged by the paper's negative result with ReST^EM-trained revision models—is that on-policy training could amplify spurious correlations or overfit to specific guidance trajectories, so the evaluation should include out-of-distribution tests on novel guidance strengths or condition combinations not seen during retraining.

4. Tagless temporal alignment: can the diffusion model's internal representations substitute for a separate time predictor? TAG's key innovation is the TLS—a corrective gradient derived from a dedicated time predictor. But diffusion models already contain temporal information in their internal representations: the timestep embedding modulates every layer, and intermediate feature maps presumably encode noise-level statistics. A compelling follow-up would attempt to extract a TLS-like gradient directly from the diffusion model itself, eliminating the need for a separate time predictor. The approach would probe the diffusion model's internal activations at various layers, train a lightweight linear readout to predict the timestep (analogous to how classifier guidance reuses a pretrained classifier rather than training from scratch), and compute the gradient of this readout with respect to the input. This would test whether the TLS signal is already latent in the diffusion model and merely needs to be extracted, or whether (as the paper's off-manifold argument suggests) the diffusion model's temporal representations degrade in exactly the same regions where its score function degrades, making the separate predictor essential. The experiment would compare TAG with a dedicated time predictor against TAG using internal readouts from the diffusion model at various depths, on the corrupted reverse process benchmark. If internal readouts work, TAG becomes truly zero-overhead (no auxiliary model training). If they fail, it validates the paper's implicit claim that statistical independence between the score function and the TLS is necessary for robustness.

5. Continuous adaptive guidance schedules informed by real-time Time-Gap monitoring. The paper uses a fixed guidance schedule ω_t = ω₀ √(1 − ᾱ_t) that depends only on the timestep, not on the actual degree of misalignment. But the Time-Gap can be computed at every step using the time predictor's forward pass (before computing the TLS gradient). This opens the possibility of closed-loop correction: at each step, measure the current Time-Gap, and set ω_t proportional to the measured misalignment. When Time-Gap is small (the sample is well-aligned), apply weak correction to avoid unnecessary perturbation. When Time-Gap spikes (the sample has been pushed off-manifold by external guidance), apply strong correction to recover. This would replace the current oracle grid search over ω₀ with an adaptive mechanism that sets correction strength based on measured need. The experiment would implement a proportional controller: ω_t = κ · Time-Gap(x_t), where κ is a global gain (much less sensitive than ω₀ since it multiplies a measured quantity rather than setting absolute scale). The evaluation would compare the adaptive schedule against the fixed √(1 − ᾱ_t) schedule on the DPS strength sweep (Table 3) and corrupted reverse process (Table 1), with the hypothesis that adaptive correction maintains near-optimal performance across a wide range of drift severities without per-setting tuning. The key metric is whether a single κ value works across all σ levels or DPS strengths—if so, it largely solves the hyperparameter sensitivity limitation.

6. Stress-testing multi-conditional TAG with adversarial or conflicting conditions. The multi-conditional results (Table 5) show that single-condition and unconditional time predictors match dedicated multi-condition predictors for compatible conditions (CelebA attributes, molecular properties). But the reparameterization approximations in Propositions B.1 and B.2 assume conditions can be sequentially absorbed into the sample without destructive interference. This assumption likely breaks down for conflicting conditions—for instance, generating an image that is simultaneously "bright" and "dark," or a molecule with both high and low polarizability. A stress-test follow-up would construct synthetic conflicting condition pairs (using the same datasets and models as the paper) and measure where the sequential reparameterization fails. The evaluation would use increasing degrees of conflict (from orthogonal to mildly opposed to strictly contradictory conditions) and measure both TAG's ability to maintain temporal alignment (Time-Gap) and the multi-condition guidance's ability to find acceptable trade-offs (Pareto frontier analysis). The hypothesis is that sequential reparameterization introduces order-dependence—applying condition A then condition B produces different results than B then A—and that this order-dependence correlates with condition conflict severity. The practical output would be guidance on when the efficient single-condition/unconditional predictor approximation is safe and when practitioners should invest in dedicated multi-condition time predictors or alternative multi-condition combination strategies.

Practical Applications and Downstream Use Cases

1. Robust molecular design with property constraints. The paper's molecular generation results (Table 2) show that TAG can transform a property-guided molecular generator from producing chemically invalid outputs to producing valid, property-satisfying molecules. For DPS on polarizability α, TAG improves Atom Stability from 28.4% to 96.4% while simultaneously reducing MAE from 13.33 to 7.96. This is not a marginal quality improvement—it is the difference between a system that produces mostly nonsense (71.6% of atoms have incorrect valencies) and one that produces chemically sensible candidates. In drug discovery and materials design, where generated molecules must be both functional (satisfying target property constraints) and synthesizable (satisfying chemical validity constraints), TAG addresses the core failure mode of training-free guided generation: the guidance term optimizes the property at the expense of molecular structure. A deployment pipeline would use an unconditional molecular diffusion model (e.g., EDM) with property-guided sampling (e.g., TFG) augmented with TAG's temporal correction, enabling zero-shot conditional generation for novel property targets without retraining the diffusion model for each new property or property combination. The multi-conditional results (Table 5) further show that a single unconditional time predictor supports simultaneous control of up to six molecular properties without combinatorial predictor training, making the approach scalable to the dozens of properties typically considered in lead optimization.

2. Accelerated sampling for interactive image editing. The few-step generation results (Table 4, Table 17) show TAG's largest relative improvements at the smallest number of function evaluations. On CIFAR-10 with 1 NFE, TAG reduces FID from 460.0 to 271.1 (41.1% reduction) and nearly doubles IS from 1.26 to 2.26 (Table 17). In interactive applications—image inpainting where a user draws a mask and expects real-time feedback, style transfer where sliders control the strength of different artistic effects, or personalization where text prompts are iteratively refined—latency is critical, and single-digit NFE is the target regime. TAG's mechanism is particularly well-suited here: aggressive timestep skipping causes severe discretization error (off-manifold drift), which is exactly the regime where TAG provides the largest correction. A practical system would integrate TAG into an accelerated DDIM sampler (or a distilled few-step model) and apply correction with a fixed ω₀ tuned once on a validation set of editing operations. The lightweight time predictor (1.48M parameters, ~8.5% of a UNet encoder) adds negligible latency per step relative to the diffusion model forward pass, making the quality improvement essentially free in wall-clock terms for GPU-bound inference.

3. Test-time reward alignment for text-to-image models without reward over-optimization. The DAS+TAG results (Table 6) demonstrate that TAG can improve reward alignment (Aesthetic score, CLIP score) while simultaneously reducing temporal misalignment. On single-objective aesthetic optimization, DAS+TAG achieves Aesthetic 9.087 vs. 7.948 for DAS alone (14.3% improvement) while reducing Time-Gap from 90.04 to 28.84 (68.0% reduction). This addresses a known failure mode in reward-guided generation: the optimization exploits the reward model (producing high-scoring but unrealistic images) by pushing samples into regions where the diffusion model's score function is unreliable. TAG's temporal correction acts as an implicit regularizer—it penalizes samples that drift off the diffusion manifold, constraining the reward optimization to stay in regions where the model can produce realistic outputs. In production text-to-image systems where user satisfaction depends on both prompt alignment and image quality, TAG provides a lightweight (inference-time only, no fine-tuning) mechanism to increase reward while maintaining or improving realism. The multi-objective results (Aesthetic + CLIP, Table 6) show that this regularization does not prevent satisfying multiple objectives simultaneously—both scores improve with TAG.

4. Black-box deployment as a bolt-on correction for API-accessed diffusion models. TAG's most distinctive practical advantage is that it requires no access to the diffusion model's weights, architecture, or training procedure. The time predictor is a separate network that only needs the noisy sample (and optionally the condition) as input. This means TAG can be deployed as a wrapper around a diffusion model API: the client sends the current noisy sample to the API, receives the denoised sample, computes the TLS gradient locally using a lightweight time predictor, applies the TAG correction, and proceeds to the next step. This is not possible with methods that modify the score function's inputs (time perturbation), adjust its outputs (epsilon scaling), or require access to internal representations. For organizations that use third-party diffusion models (e.g., Stability AI's API, OpenAI's DALL-E API) and want to apply custom guidance for their specific use case, TAG provides a model-agnostic correction that works regardless of the underlying architecture. The primary cost is training a time predictor for the API model's noise schedule—which requires only the ability to run the forward diffusion process (adding noise to clean images), not access to the model itself. This application is enabled directly by the paper's design principle of statistical separation between the time predictor and the diffusion model (Innovation 3).

When to Prefer This Method

The paper does not explicitly position TAG against a clearly named set of alternatives with defined tradeoff conditions. TAG is presented as a universal corrective mechanism that can be applied on top of any guidance method, not as a replacement for specific existing approaches. The experimental comparisons against TCS, Timestep Guidance, Self-Guidance, Epsilon Scaling, Time-Shift Sampler, Input Perturbation, and Langevin Dynamics (Tables 2, 7, 8) demonstrate TAG's superiority in the specific settings tested, but the paper does not articulate conditions under which a practitioner should prefer one of these alternatives—in part because TAG outperforms all of them in essentially every comparison.

The closest the paper comes to a tradeoff discussion is in Appendix D.2, where it contrasts TAG against fine-tuning approaches (ControlNet, IP-Adapter, RL-based reward tuning). The argument is not that TAG is universally better than fine-tuning, but that TAG is preferable when:

  • The target condition changes at inference time (e.g., molecular property targets specified per query, dynamic text prompts in interactive editing), making per-condition fine-tuning infeasible.
  • The practitioner lacks task-specific labeled data for fine-tuning (TAG's time predictor is trained purely on forward-process samples with timestep labels, which are generated automatically).
  • Computational budget for adaptation is limited (time predictor training in "minutes" vs. hours of gradient-based fine-tuning).

However, the paper does not empirically compare TAG against any fine-tuning method. The text-to-image experiments use DAS (a test-time sampler, not a fine-tuning method) as the base, and the TFG benchmark comparisons are against inference-time correction methods only. There is no experiment showing, for instance, that TAG on top of an unconditional model matches or exceeds a task-specific fine-tuned conditional model on label guidance. The tradeoff between "train a separate lightweight auxiliary model and apply inference-time correction" vs. "fine-tune the base model for the specific task" is stated qualitatively but never quantified.

Given the paper's experimental scope, the only empirically justified decision rule is:

  • When applying training-free guidance (DPS, TFG) to an unconditional diffusion model for any task and you observe quality degradation: add TAG with a time predictor trained on the same base model's forward process, and tune ω₀ via grid search over approximately [0.01, 5.0] on a validation set. Expect FID improvements of 2-48% depending on task and guidance strength, with larger gains when guidance is stronger or sampling steps are fewer.

This is a narrow (but practically useful) prescription. The paper does not provide guidance for settings outside the TFG benchmark paradigm—for instance, when using classifier-free guidance (which already conditions the score function during training), when the base model is conditional rather than unconditional, or when the primary quality bottleneck is not off-manifold drift but something else (e.g., training data diversity, model capacity). These remain open questions.