One-Step Generative Surrogate Models via Block-Triangular Joint Drifting
Nicholas Geissler, Shreya Jha, Ricardo Baptista, Benjamin Peherstorfer
cs.LG, math.NA
2026-09-22
BTJD learns a one-NFE transition sampler by drifting the joint law of consecutive states. On 2D turbulence, energy error is 2.14e-2, about 4× below one-step ReFlow distillation.
Cheap surrogates for stochastic systems need to sample the next-state conditional, then roll out trajectories autoregressively. Diffusion and flow models can represent those transition laws, but each physical step often costs K denoising or integration evaluations. A T-step ensemble then costs KT network calls.
Distillation (MeanFlow, ReFlow, consistency models) can shrink K to one, at the price of training a multi-step teacher first and trading accuracy for speed. Methods that only match time marginals (DICE, marginal diffusion) are cheaper still, yet marginals do not determine path-dependent statistics. A deterministic stepper from a fixed initial state always produces the same path, so it cannot express intrinsic noise.
Drifting was designed for one-step generation: during training a distribution-dependent field pushes model samples toward a target; after convergence, noise maps to a sample in one forward pass. The field is built from an empirical sample of the target. Trajectory logs give one realized successor per observed state, so the empirical conditional is a point mass with no width or shape. Score and flow-matching regressions do not need repeated successors at the exact same state. Drifting does.
BTJD changes the target. It drifts the joint law of consecutive states, for which transition pairs are actual samples. The conditional those pairs never fully populate is recovered afterwards from structure.
The generator is constrained to block-triangular form: the first block copies the current state and ignores reference noise, so the current-state marginal stays the true marginal for every parameter value. Only the second block g is learned. Once the joint matches, g pushes noise to the conditional transition law, by a standard result on block-triangular transport. At inference, draw z, evaluate g once, get the next state.
The drift field is the first variation of Sinkhorn divergence on the joint, then orthogonally projected onto the second block. The current state does not move. The loss collapses to a squared regression on g. Parameters are shared across time; each batch draws several times and several pairs, and estimates the field with empirical Sinkhorn barycentric projections.
A theorem covers the unregularized case ε=0: compact support, densities bounded away from zero and infinity, identical first marginals. Then the projected field vanishes almost everywhere if and only if the model joint equals the target. Training uses ε>0 throughout (0.01 with 50 Sinkhorn steps on Duffing; about 0.07–0.1 with 20 steps on the PDEs). High-dimensional fields are encoded into a latent space before drifting.
Four problems. Low-dimensional tests report sliced Wasserstein-2 on time marginals and a path-dependent quantity of interest.
Duffing oscillator, σ=0.5, 1200 steps, 5000 training trajectories, random initial conditions:
| Method | sliced-W2 | traj. QoI err. |
| Deterministic stepper | 3.48e-1 | 3.68e-2 |
| DICE | 5.80e-1 | 8.30e-2 |
| Marginal diffusion | 8.10e-2 | 1.07e-1 |
| SDE learning | 8.30e-2 | 5.70e-3 |
| SDE matching | 5.49e-1 | 4.47e-2 |
| BTJD | 4.80e-2 | 2.75e-3 |
With a fixed initial state, BTJD still leads: W2 5.80e-2 and QoI 2.49e-3, against SDE learning at 1.08e-1 and 6.80e-3. Rollouts from one point fan out. The randomness sits in the transitions.
Nine-dimensional Rayleigh–Bénard, trained on μ in {13.5, …, 14.2}, tested at unseen μ=13.65: BTJD W2 4.60e-2, rotational-current QoI 3.00e-4. Next-best SDE learning: 5.40e-2 and 2.20e-2. Matching the histogram is not enough; the sense of rotation, a trajectory-level current, matches too.
PDE tests add multi-step generative baselines. Stochastic Burgers on 64 grid points, 16-dimensional latent:
| Method | energy err. | enstrophy err. | NFE/step |
| Operator learning | 2.21e-2 | 2.55e-1 | 1 |
| ARDM 100 steps | 1.24e-2 | 2.11e-1 | 100 |
| CFM 20 steps | 2.71e-3 | 1.53e-1 | 20 |
| MeanFlow 1 step | 8.22e-1 | 3.59e+2 | 1 |
| ReFlow+distill 1 step | 2.78e-3 | 1.40e-1 | 1 |
| BTJD | 2.22e-3 | 5.87e-2 | 1 |
Enstrophy error is less than half the next-best. One-step MeanFlow collapsed on this problem.
Two-dimensional stochastically forced turbulence, trained at 64×64:
| Method | energy err. | enstrophy err. | NFE/step |
| ARDM 100 steps | 1.16e-1 | 1.34e-1 | 100 |
| CFM 20 steps | 1.39e-1 | 1.04e-1 | 20 |
| MeanFlow 4 steps | 1.22e-1 | 6.74e-2 | 4 |
| ReFlow+distill 1 step | 8.60e-2 | 7.16e-2 | 1 |
| BTJD | 2.14e-2 | 2.52e-2 | 1 |
Energy error is about a quarter of one-step ReFlow distillation; enstrophy about a third. Repeated rollouts from the same initial vorticity diverge into distinct turbulent realizations. Operator learning from the same initial condition always returns the same trajectory.
Ensemble forecasts, UQ, and outer-loop design need stochastic trajectories that are both accurate and cheap. BTJD makes one-step sampling the training objective, with no teacher. Laboratory and simulation logs already look like "one successor per state."
Use it when a bank of stochastic trajectories exists and the quantities that matter are path-level (energy, enstrophy, barrier crossings), not a histogram at a single time. Training still pays for Sinkhorn every step. Inference is the cheap part.
The Burgers MeanFlow collapse is a reminder that compressing a CFM teacher to one NFE is not free. This is a scientific reduced-order generator, not a general image model.
The fixed-point result is for ε=0, compact support, and bounded densities. Every experiment uses ε>0. Whether projection adds extra zeros is settled only in that idealized setting.
Training cost is barely reported. Each step runs 20–50 Sinkhorn iterations, batches go up to 8192 particles, low-dimensional runs take 100k steps and turbulence 65k. One NFE is an inference number.
Burgers and turbulence both go through autoencoders. Table errors are on decoded fields; reconstruction error is not isolated. Low-dimensional experiments omit CFM, ARDM, and distillation, so a blanket "best everywhere" claim is stronger than the grid. How hard MeanFlow was tuned on Burgers is unclear.
Conditioning inputs are jittered with 1–2% Gaussian noise. Long-horizon accumulation and extrapolation past the training window are not stress-tested.