Score Accuracy Along the Forward Diffusion Does Not Certify Numerical Stability in Diffusion Sampling
By Yiwei Zhou
"Small forward-marginal error in score matching does not guarantee numerical stability; discretized reverse samplers can diverge in Wasserstein distances while converging weakly."
Abstract
Score matching controls average error under the forward marginals, but a discretized reverse-time sampler evaluates the learned score along its own trajectory. We show that small forward-marginal error does not guarantee numerical stability. We construct a single smooth score field with arbitrarily small forward-marginal $L^2$ error. The learned reverse-time process is nonexplosive, has moments of every order, and can be arbitrarily close to the exact reverse-time process in path-space total variation. Yet its Euler--Maruyama discretizations converge in probability while every positive moment diverges. Thus weak convergence can hold even though every Wasserstein distance $W_p$, $p\ge1$, diverges. The same failure can occur within one fixed finite neural architecture. We construct a family of bounded, globally Lipschitz denoisers for which both the forward-marginal error and the path-space total variation distance tend to zero, while their Euler--Maruyama endpoints diverge in every $W_p$. For compactly supported data, we also give a simple positive result. Projecting the learned denoiser onto a known bounded closed convex set containing the support preserves pointwise accuracy, gives grid-uniform moment bounds, and yields Wasserstein convergence under mild local regularity. Experiments with a small fixed DiT-style network show large growth along rare numerical trajectories and its suppression by denoiser projection, while overall trajectory errors remain small.
Technical Analysis & Implementation
Score Accuracy Along the Forward Diffusion Does Not Certify Numerical Stability in Diffusion Sampling§
Core Problem§
Score matching minimizes average error under forward marginals, but reverse-time samplers evaluate the learned score along their own trajectories. The paper shows that small forward-marginal $L^2$ error does not imply numerical stability of Euler–Maruyama discretizations. In fact, even when the learned score is arbitrarily close to the true score in path-space TV, the discretized process can have diverging moments and Wasserstein distances $W_p$ for any $p \ge 1$.
Main Results§
1. Counterexample construction: A smooth score field with arbitrarily small forward-marginal $L^2$ error but whose Euler–Maruyama discretizations converge in probability while every positive moment diverges. 2. Same failure within fixed architecture: A family of bounded, globally Lipschitz denoisers (e.g., small DiT-style networks) with forward-marginal error and path-space TV tending to zero, yet Euler–Maruyama endpoints diverge in every $W_p$. 3. Positive result for compactly supported data: Projecting the learned denoiser onto a known bounded closed convex set containing the data support preserves accuracy, yields uniform moment bounds, and ensures Wasserstein convergence.
Methodology Details§
Let $p_t$ be the forward marginal density and $\nabla \log p_t$ the true score. A learned score $s_\theta(t,x)$ has forward-marginal error $\mathbb{E}_{p_t}[\|s_\theta(t,x) - \nabla \log p_t(x)\|^2]$. The reverse-time SDE is $$ dX_t = -2 \nabla \log p_{T-t}(X_t) dt + \sqrt{2} dW_t, $$ which is discretized via Euler–Maruyama with step size $h$. The constructed counterexample uses a score that oscillates rapidly along the reverse trajectory, causing large errors in the drift that accumulate but vanish in the continuous limit.
Implementation Sketch§
A minimal PyTorch snippet for Euler–Maruyama sampling with a learned score model:
import torch
def euler_maruyama_sample(score_model, num_steps, device='cpu'):
# score_model: (t, x) -> score estimate
dt = 1.0 / num_steps
x = torch.randn(1, 784, device=device) # initial noise
for i in range(num_steps):
t = torch.full((1,), 1.0 - i*dt, device=device)
drift = -2.0 * score_model(t, x)
x = x + drift * dt + torch.sqrt(2*dt) * torch.randn_like(x)
return xThe paper shows that even with a perfect score in the forward sense, this discretization can cause moments to blow up.
Positive Result§
For data supported on a bounded convex set $K$, projecting the denoiser $D_\theta$ (related to score via $s_\theta = (D_\theta(x) - x)/\sigma_t^2$) onto $K$ yields $$\tilde D_\theta = \text{proj}_K \circ D_\theta,$$ which maintains $L^2$ accuracy and provides moment bounds. The resulting Euler–Maruyama sampler converges in Wasserstein distance under mild regularity.
Implications§
Practitioners should not trust accurate score matching alone; numerical stability requires additional regularity (e.g., denoiser projection). The results highlight the gap between theory and practice in diffusion models.