Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Diffusion and Flow World Models › Conditional Denoising Dynamics

Michael BrenndoerferJune 28, 20268 min read

World Models HandbookAll chapters

Part 5 of 12World-Model Architectures

Chapter 28Diffusion and Flow World Models › Conditional Denoising Dynamics

Reading section 2 of 9 · Conditional Denoising Dynamics · 8 minSection 2 of 9 · 8 min

Section 2 of 9 · 8 min read

Diffusion and Flow World Models › Conditional Denoising DynamicsOne section · Sections in this chapter2 / 9 ↓

Sections in this chapter9

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The object we want to model is a conditional distribution $p(x_0 \mid c)$ , where $x_0$ is the target (a next state, a latent, or an entire trajectory) and $c$ is the conditioning signal (the current state and action, plus anything else we choose to observe). In a world-modeling context, $c = (s_t, a_t)$ and $x_0 = s_{t+1}$ , but we will keep the notation general. The Gaussian construction below applies directly to suitably scaled continuous Euclidean targets, including scalar states, vector encodings, latents, and flattened horizons. Periodic, discrete, manifold-valued, or hard-constrained targets may require a different representation or corruption process.

A standard Gaussian DDPM defines a forward corruption process that progressively corrupts $x_0$ . Let $\beta_t\in(0,1)$ be the scalar noise variance at step $t$ , define $\alpha_t=1-\beta_t$ and $\bar\alpha_t=\prod_{s=1}^t\alpha_s$ , and let $\mathbf I$ be the identity matrix matching the target dimension. With suitable data scaling and a schedule for which $\bar\alpha_T$ is near zero, the terminal marginal is close to a Gaussian base. A network then learns approximate reverse conditionals rather than exactly inverting each realized corruption path. The forward process is Markovian, meaning each step depends only on the immediately preceding state:

q(x_t \mid x_{t-1}) = \mathcal{N}\!\left(x_t;\ \sqrt{1 - \beta_t}\,x_{t-1},\ \beta_t \mathbf{I}\right).

The schedule typically starts with small $\beta_t$ values and grows until $\bar\alpha_T$ is nearly zero. Each step multiplies the preceding state by $\sqrt{1-\beta_t}$ and adds noise with amplitude $\sqrt{\beta_t}$ , equivalently covariance $\beta_t\mathbf I$ . For $0<\beta_t<1$ , the marginal signal-to-noise ratio $\bar\alpha_t/(1-\bar\alpha_t)$ decreases with $t$ . Low-noise levels preserve fine detail and are traversed near the end of reverse sampling; sufficiently high cumulative noise makes the terminal marginal close to a simple Gaussian base. A useful bookkeeping identity, derived by composing Gaussians, gives the marginal distribution of $x_t$ directly from $x_0$ :

q(x_t \mid x_0) = \mathcal{N}\!\left(x_t;\ \sqrt{\bar\alpha_t}\,x_0,\ (1 - \bar\alpha_t)\mathbf{I}\right).

Equivalently, $x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1 - \bar\alpha_t}\,\epsilon$ with $\epsilon \sim \mathcal{N}(0, \mathbf{I})$ . This closed form makes training tractable by jumping directly from $x_0$ to $x_t$ . Without it, producing $x_t$ would require simulating $t$ forward transitions, about $T/2$ on average under uniform timestep sampling, which would be prohibitively expensive for the thousand-step schedules used in early DDPM experiments.

The closed-form marginal above lets us visualize the forward corruption process directly. The plot below shows several noisy trajectories along a linear schedule, together with the signal amplitude $\sqrt{\bar\alpha_t}$ .

In[3]:

Code

import numpy as np

# Small linear schedule for visualization.
betas_viz = np.linspace(1e-4, 0.05, 200)
alphas_viz = 1.0 - betas_viz
alpha_bars_viz = np.cumprod(alphas_viz)

x0_viz = 1.0
viz_rng = np.random.default_rng(42)
viz_trajectories = []
for _ in range(5):
    x_t_viz = x0_viz
    path_viz = []
    for alpha_t_viz, beta_t_viz in zip(alphas_viz, betas_viz):
        x_t_viz = (
            np.sqrt(alpha_t_viz) * x_t_viz
            + np.sqrt(beta_t_viz) * viz_rng.normal()
        )
        path_viz.append(x_t_viz)
    viz_trajectories.append(path_viz)
viz_trajectories = np.stack(viz_trajectories, axis=0)
signal_viz = np.sqrt(alpha_bars_viz)
timesteps_viz = np.arange(1, len(betas_viz) + 1)

Out[4]:

Visualization

Five forward-corruption paths start near x equals 1, spread out, and increasingly fluctuate around zero as the signal-amplitude curve falls across 200 timesteps. — Five recursively sampled trajectories of the forward Markov chain from x0 = 1 across 200 linear-schedule timesteps. The signal amplitude sqrt(alpha_bar_t) shrinks toward zero while path variability grows, so the terminal marginal is close to Gaussian noise.

The noise schedule

Many canonical diffusion implementations fix a hand-designed schedule $\{\beta_t\}$ , although learned and optimized schedules also exist. A linear schedule from $10^{-4}$ to $0.02$ was the original DDPM experimental default from Ho, Jain, and Abbeel (2020). Nichol and Dhariwal (2021) found a cosine schedule helpful in their image experiments because it preserved information longer than the linear schedule they tested. For world models, the schedule must be evaluated against the scale and distribution of the target. The terminal cumulative signal $\bar\alpha_T$ should be small enough that the chosen Gaussian base is a good approximation, while early noise levels should still expose fine state differences. There is no scale-only rule saying that values in $[-5,5]$ require a particular terminal schedule.

Training consists of one pleasantly simple objective. Sample a timestep $t \sim \mathrm{Uniform}\{1, \dots, T\}$ , sample $x_0$ from the data, sample $\epsilon \sim \mathcal{N}(0, \mathbf{I})$ , and ask the network to predict the noise:

\mathcal{L}(\theta) = \mathbb{E}_{t, (x_0,c)\sim p_{\mathrm{data}}, \epsilon}\!\left[\left\|\epsilon - \epsilon_\theta(x_t, t, c)\right\|_2^2\right].

Despite its simplicity, this objective is a denoising score matching loss in disguise. The optimal predictor satisfies

\epsilon_\theta(x_t, t, c) = -\sqrt{1 - \bar\alpha_t}\,\nabla_{x_t} \log q(x_t \mid c).

At each fixed timestep, the population-optimal score is obtained from the noise predictor as $\nabla_{x_t}\log q(x_t\mid c)=-\epsilon_\theta(x_t,t,c)/\sqrt{1-\bar\alpha_t}$ . Across timesteps, the unweighted noise loss corresponds to a timestep-weighted score-matching objective. The connection, developed in denoising score matching and used by DDPMs, avoids asking the network to evaluate or normalize a density explicitly. An exact score still has the integrability structure of a log-density gradient; the practical convenience is that training does not require computing a partition function.

A score field gives a complementary local view of a density: it shows the direction and strength of increasing log density. For the well-separated symmetric mixture below, the score is zero at the midpoint and at the two mixture modes, whose locations are numerically indistinguishable from the component means, and points outward toward the modes between those zeros. The stabilized finite-difference curve preserves these central directions and zero crossings; adding $10^{-12}$ to the density before taking its logarithm attenuates the far-tail score magnitude.

In[5]:

Code

import numpy as np

# Score field for a symmetric two-component Gaussian mixture.
x_grid = np.linspace(-3.0, 3.0, 601)
mu_left, mu_right, sigma_mix = -1.2, 1.2, 0.25
density_left = np.exp(-0.5 * ((x_grid - mu_left) / sigma_mix) ** 2) / (
    sigma_mix * np.sqrt(2.0 * np.pi)
)
density_right = np.exp(-0.5 * ((x_grid - mu_right) / sigma_mix) ** 2) / (
    sigma_mix * np.sqrt(2.0 * np.pi)
)
density_mix = 0.5 * density_left + 0.5 * density_right
log_density_mix = np.log(density_mix + 1e-12)
score_mix = np.gradient(log_density_mix, x_grid)

Out[6]:

Visualization

Score curve for a symmetric bimodal density: it crosses zero at minus 1.2, zero, and plus 1.2; points outward from the midpoint between the modes; and points back toward the modes outside them. — Score field of a symmetric two-component Gaussian mixture whose modes are numerically indistinguishable from x = -1.2 and x = 1.2. The score crosses zero at both mixture modes and at the low-density midpoint, and points away from the middle between those zeros.

At sampling time we start from a new Gaussian draw and apply learned reverse-time transitions; this approximates the time reversal of the forward distributional dynamics rather than undoing a realized noise path. For sufficiently small $\beta_t$ and a regular noisy marginal density, the local conditional $q(x_{t-1}\mid x_t,c)$ is approximately Gaussian; by contrast, $q(x_{t-1}\mid x_t,x_0)$ is exactly Gaussian for this forward process. Define $\alpha_t=1-\beta_t$ . In the common epsilon parameterization, the reverse-model mean is

\mu_\theta(x_t, t, c) = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1 - \bar\alpha_t}}\,\epsilon_\theta(x_t, t, c)\right).

With the fixed-variance choice used in this chapter's code, sampling iterates $x_{t-1} = \mu_\theta(x_t, t, c) + \sqrt{\beta_t}\,\eta$ with $\eta \sim \mathcal{N}(0, \mathbf{I})$ for $t > 1$ , then uses the mean at $t=1$ . Other DDPM implementations use the forward-posterior variance $\tilde\beta_t$ or a learned variance. Fresh noise makes ancestral paths stochastic; deterministic samplers instead obtain diversity solely from their initial draw. A useful denoised estimate is

\hat{x}_0(x_t, t, c) = \frac{x_t - \sqrt{1 - \bar\alpha_t}\,\epsilon_\theta(x_t, t, c)}{\sqrt{\bar\alpha_t}}.

This quantity estimates a clean target compatible with the noisy sample; it does not recover a unique originating example. Samplers can be written using data, score, noise, or velocity predictions. These forms are algebraically convertible when they represent the same exact field under one schedule, but separately trained finite-capacity implementations need not behave identically. Some updates are deterministic, some add noise, and multistage solvers may spend several network evaluations per accepted step. We revisit these choices in the sampling-tradeoff section and implement DDPM and DDIM in the worked example.

To model $p(s_{t+1}\mid s_t,a_t)$ rather than the marginal next-state distribution, the network must receive the relevant conditioning. The simple formulation here uses the current state and action, but a one-step model may also condition on histories, beliefs, observations, or goals. Implementations concatenate these features or inject them through mechanisms such as cross-attention. When a learned score accurately represents a separated bimodal conditional, its reverse-time dynamics can transport different base samples toward different modes instead of returning only their conditional mean.