Diffusion Spacetime, Interactive

← text notes · Karczewski, Heinonen, Pouplin, Hauberg & Garg, ICLR 2026. Five widgets for the paper's constructs. Everything is computed in your browser, and nothing is learned. The data are Gaussian mixtures, so every denoising distribution p(x₀ | xt) is a Gaussian mixture in closed form. That makes η, μ, ψ, the KL divergences and the Fisher–Rao metric exact. The geodesic solver is the one in code/spacetime_checks.py: it minimises the paper's Eq 14 with a natural-gradient step.

1. The spacetime explorer

The data are the paper's toy (App G.1): three Gaussians at −2.5, 0.5 and 2.5 with standard deviation 0.75, under a variance-preserving schedule with log-SNR falling linearly from 10 at t = 0 to −10 at t = 1. The shading is the noisy marginal pt(x), normalised per column as in the paper's Fig. 1. Each point (xt, t) of the plane is a distribution: the denoising posterior p(x₀ | xt). Drag the two white handles. The blue curve is the spacetime geodesic between them, found by minimising the discrete energy, Eq 14. The slider moves a point along it, and the right panel shows the posterior at that point.

spacetime geodesicstraight line in (xt, t)up to t = 1, across, downPF-ODEdata q(x₀)p(x₀ | γs)
Things to try. (1) With the Fig. 1 endpoints, move s from 0 to 1. The posterior starts narrow on the left mode, widens until it covers all three modes (the "forget" half), then narrows onto the right mode (the "denoise" half). (2) Compare the lengths. The geodesic climbs only to t ≈ 0.62, yet the path that goes all the way to t = 1 is barely longer: about 1% here. Everything above t ≈ 0.6 is almost free, because the posteriors there barely depend on xt. Press same mode and the ratio drops to about 0.2. (3) Turn on metric ellipses. Each one is the set of steps of equal cost, √(dz⊤G dz) = 0.3. Near t = 1 they are clipped to their cell: moving xt there costs nothing, which is the memorylessness of Eq 11. Near t = 0 they shrink to dots: the posteriors are nearly Dirac, and every step is expensive. (4) Press the Fig. 3 button. The geodesic and the PF-ODE path share their endpoints, and their lengths agree to within a few tenths of a percent.

2. Two points, four numbers

For an exponential family, Lemma D.2 says (η₁ − η₂)⊤(μ₁ − μ₂) = KL(1‖2) + KL(2‖1). This is an identity for any two points, not an approximation for nearby ones. Eq 14 is therefore (N − 1) × a sum of symmetrised KLs between consecutive points of the curve. That is exactly the local-KL energy of Eq 13 with KL replaced by its symmetrised version. The only approximation left is the one every discretised energy makes. The usual local approximation, ½Δz⊤G Δz per KL, is shown for contrast. It is only right for nearby points.

q(x₀)p(x₀ | z₁)p(x₀ | z₂)Eq 14 for the straight segment z₁ → z₂, against N
Put the two points close together and all rows agree. Pull them apart and the quadratic approximation ½Δz⊤G(z₁)Δz drifts, by tens of percent within a mode, while the chord row still matches the sum of the two KLs to every printed digit. The right panel shows why Eq 14 needs only modest N: each chord is exact, so the energy of a fixed curve converges quickly as N grows. In the code, the straight Fig. 1 segment gives 129.3 at N = 4, 127.0 at N = 16 and 126.96 in the limit.

3. Gaussian data: the spacetime is a hyperbolic half-plane

Take q = N(0, 1) and the VP schedule. Then p(x₀ | xt) = N(m, s²), with m = αtxt and s = σt. So spacetime is the family of 1-D Gaussians with s < 1. Its Fisher–Rao metric is (dm² + 2ds²)/s², the hyperbolic half-plane stretched by √2 along m. Geodesics are semi-ellipses centred on the s = 0 axis, (m − m₀)²/2 + s² = R², and the distance has a closed form. Drag the handles. Both are data points anchored at a small noise level s, the way DiffED anchors at tmin.

geodesic (closed form)PF-ODE paths (the one with x = √2 is a geodesic)DiffED(d), both points at noise slog law 2√2 log(d/√2s)"noise everything" upper bound
What the formula says about DiffED. (1) Once d ≫ s it grows like the logarithm of the Euclidean distance d between the two points, 2√2 log(d/√2 s). So it is a compressed Euclidean distance, not a new notion of similarity. (2) It depends on the anchor noise s as −2√2 log s. Halving tmin changes every DiffED by roughly the same additive amount. (3) The apex of the geodesic sits at s ≈ d/(2√2). That is the paper's "the less similar the endpoints, the noisier the midpoint" (Fig. 4), now with a number. (4) When the apex would pass s = 1, the geodesic has to leave the family. The true path then runs along the top edge, where every point is essentially the prior, and DiffED is squeezed between the purple curve and the orange bound, the length of the path that goes via q itself. Turn on the PF-ODE paths: they are quarter-ellipses with semi-axes (x, 1), and only the one with x = √2 has the geodesic shape.

4. Pullback versus spacetime on a ring of data

The data lie near a circle of radius 2: a mixture of 24 small Gaussians, standard deviation 0.12. Pullback geometry (Section 4) uses the PF-ODE as a deterministic, bijective decoder. Whatever that decoder is, Prop B.1 says its shortest paths decode to the straight chord, the dashed grey line, straight through the empty middle. The spacetime geodesic between the same two points, anchored at log-SNR 4, is a path of distributions. The purple dots are 300 samples from the posterior at position s along it. Their mean (blue) does cut inside the ring, because the mean of a spread-out posterior is off the data. The samples themselves stay on the ring.

pullback geodesic, decoded (the chord)E[x₀ | γs] along the spacetime geodesicsamples of p(x₀ | γs)noise level σ/α = SNR−1/2 along the geodesic
Move the endpoints close together. The geodesic hardly rises in noise, and its posteriors are small blobs sliding along the ring. Move them to opposite sides. The geodesic climbs until one posterior covers most of the ring, then descends. It never tries to squeeze a sharp distribution through the middle, where there is no data. That is the sense in which spacetime geometry "sees" the data manifold and pullback geometry cannot: the chord is the same for every bijective decoder, and so for every data distribution.

5. Transition paths along a spacetime geodesic

A Boltzmann density q ∝ e−U with two wells, A and B. The straight route between them crosses a high barrier. A curved channel over the top is much lower. The denoising posterior is again Boltzmann, with energy U(x₀) + ½SNR‖x₀ − xt/α‖² (Eq 19). Algorithm 1 runs 120 Langevin steps on the posterior at each of 40 points on the geodesic, carrying the state along. The result is a chain from A to B that is always near-typical for a slowly moving target. Baseline: the same annealed Langevin along a straight line at fixed noise, log-SNR 4 all the way. The optional penalty is the paper's low-variance constraint, h = max(−log SNR, ρ).

Langevin along the spacetime geodesicLangevin along a fixed-noise straight linestraight line with the geodesic's noise scheduleE[x₀ | γs]samples of p(x₀ | γs)SNR−1/2 along the geodesic
Look at the blue curve first. The geodesic's posterior mean runs straight along the barrier, not over the channel. Slide s to 0.5: the posterior there is the two wells, reweighted, with almost no mass in the channel. That is how Fisher–Rao geometry moves probability. It has no notion of transport, so mass moves from A to B by reweighting the two modes, and nothing ever has to pass through the middle. The geodesic even pushes xt/α far below the wells, to y ≈ −4.4, to keep the channel out of the posterior. Even so, the transition paths are much better than the baseline's, and the reason is the Langevin chain. Once the noise is high, a chain stranded in A is out of equilibrium and has to move to B. A slow random walk gets there by the lowest route it can find: the channel, or the flank of the barrier just below it. In spacetime_checks.py, 44% of 32 chains cross x = 0 inside the channel, and none of the baseline's do. So the geodesic's job is to supply a schedule of broad, then narrow, targets. Routing the path is done by the dynamics. The fixed-noise baseline never gets broad enough, so every one of its chains is dragged through the barrier. The ablation (green) tells you which part matters. Keep xt/α on the straight line and borrow only the geodesic's noise schedule. Its chains do better than the geodesic's: in the code, MaxEnergy 8.6 against 11.9, and 94% against 44% crossing in the channel. The geodesic's dip away from the channel pulls chains toward the barrier. In this toy, what helps is noising up and back down; the geodesic's shape does not. Compare both MaxEnergy numbers with the flood-fill lower bound, the lowest barrier any continuous path can have. Note what MaxEnergy measures: the chain runs on a moving target, so a path's timing has no physical meaning, and only the set of states it visits is scored.