flow_matching#

Rectified-flow (RF) helpers: forward process, inversions, timestep sampling.

This module intentionally does not define a NoiseScheduler class. It exposes the handful of primitives that DMD2 actually needs as plain functions, so callers can plug fastgen into any training stack without adopting a new scheduler object.

RF convention used throughout: alpha_t = 1 - t and sigma_t = t, so x_t = (1 - t) * x_0 + t * eps with t in [0, 1]. Existing RF conversions use float64 intermediates and cast back to the input dtype. PDD grid and interval helpers instead preserve float32-or-higher outputs because their result remains in the decoding path.

Functions

add_noise

Forward process under rectified flow: x_t = (1 - t) * x_0 + t * eps.

fusion_coefficients

Return step-width-normalized coefficients for a contiguous block [start, end).

integrate_interval_velocities

Integrate per-sample velocity heads over half-open blocks [start, end).

make_shifted_flow_grid

Construct the fixed decreasing shifted rectified-flow grid.

pred_noise_to_pred_x0

Convert an eps-parameterized prediction to an x_0 prediction under RF.

pred_x0_from_flow

Convert a flow-parameterized prediction (v = eps - x_0) to an x_0 prediction.

rf_alpha

Rectified-flow data coefficient alpha_t = 1 - t.

rf_sigma

Rectified-flow noise coefficient sigma_t = t.

sample_from_t_list

Sample n starting timesteps uniformly from t_list[:-1].

sample_timesteps

Sample n training timesteps according to cfg.

x0_to_eps

Invert the RF forward process: eps = (x_t - (1 - t) * x_0) / t.

x0_to_flow

Convert an x_0 prediction back into a flow-parameterized prediction under RF.
add_noise(x0, eps, t)#

Forward process under rectified flow: x_t = (1 - t) * x_0 + t * eps.

t is broadcast across the spatial axes of x_0 via expand_like(). Computation is performed in float64 for numerical stability and the output is cast back to x_0’s dtype.

Parameters:

  • x0 (Tensor)

  • eps (Tensor)

  • t (Tensor)

Return type:

Tensor

fusion_coefficients(grid, start, end)#

Return step-width-normalized coefficients for a contiguous block [start, end).

Parameters:

  • grid (Tensor)

  • start (int)

  • end (int)

Return type:

Tensor

integrate_interval_velocities(state, velocities, grid, start, end)#

Integrate per-sample velocity heads over half-open blocks [start, end).

state is batch-first with shape [B, ...] and velocities has shape [B, N, ...]. start and end may be scalar integers or integer tensors of shape [B]. The empty block start == end returns state promoted to the PDD math dtype.

Parameters:

  • state (Tensor)

  • velocities (Tensor)

  • grid (Tensor)

  • start (int | Tensor)

  • end (int | Tensor)

Return type:

Tensor

make_shifted_flow_grid(grid_size, shift, *, max_t, device=None, dtype=torch.float32)#

Construct the fixed decreasing shifted rectified-flow grid.

Schedule construction and upper clamps use float64, matching FastGen’s RF scheduler. The completed grid is cast to the requested PDD math dtype. Low precision requests are promoted to float32 so distinct early intervals do not collapse for the canonical 128-node, shift-5 schedule.

Parameters:

  • grid_size (int)

  • shift (float)

  • max_t (float)

  • device (device | str | None)

  • dtype (dtype)

Return type:

Tensor

pred_noise_to_pred_x0(pred_noise, noisy_latents, t)#

Convert an eps-parameterized prediction to an x_0 prediction under RF.

Solves x_t = (1 - t) * x_0 + t * eps for x_0: x_0 = (x_t - t * eps) / (1 - t).

Parameters:

  • pred_noise (Tensor)

  • noisy_latents (Tensor)

  • t (Tensor)

Return type:

Tensor

pred_x0_from_flow(pred_flow, noisy_latents, t)#

Convert a flow-parameterized prediction (v = eps - x_0) to an x_0 prediction.

Under RF x_t = (1 - t) * x_0 + t * eps and v = eps - x_0 combine to x_t = x_0 + t * v, so x_0 = x_t - t * v.

Parameters:

  • pred_flow (Tensor)

  • noisy_latents (Tensor)

  • t (Tensor)

Return type:

Tensor

rf_alpha(t)#

Rectified-flow data coefficient alpha_t = 1 - t.

Parameters:

t (Tensor)

Return type:

Tensor

rf_sigma(t)#

Rectified-flow noise coefficient sigma_t = t.

Parameters:

t (Tensor)

Return type:

Tensor

sample_from_t_list(n, t_list, *, device, dtype=torch.float32)#

Sample n starting timesteps uniformly from t_list[:-1].

Used for multi-step student training: t_list encodes the inference trajectory (t_list[-1] must be 0), and a random intermediate timestep is sampled so the student is trained at every rung of the trajectory.

Parameters:

  • n (int)

  • t_list (list[float])

  • device (device)

  • dtype (dtype)

Return type:

Tensor

sample_timesteps(n, cfg, *, device, dtype=torch.float32)#

Sample n training timesteps according to cfg.

Supports uniform, logitnormal, lognormal, shifted, and polynomial distributions. polynomial for RF degenerates to discrete uniform sampling from a linspace(min_t, max_t, 1000) grid (EDM’s polynomial-spaced _sigmas is EDM-specific and not applicable under RF).

Parameters:
Return type:

torch.Tensor

x0_to_eps(x0, x_t, t)#

Invert the RF forward process: eps = (x_t - (1 - t) * x_0) / t.

Used when unrolling the student in ODE mode — given the current x_t and the student’s x_0 prediction, we can recover the implied eps deterministically.

Parameters:

  • x0 (Tensor)

  • x_t (Tensor)

  • t (Tensor)

Return type:

Tensor

x0_to_flow(x0, x_t, t)#

Convert an x_0 prediction back into a flow-parameterized prediction under RF.

Under RF x_t = (1 - t) * x_0 + t * eps and v = eps - x_0, so x_t - x_0 = t * (eps - x_0) = t * v and therefore v = (x_t - x_0) / t.

Used when the fake score is flow-native but the DSM loss is computed in a different target parameterization: convert raw flow → x_0 via pred_x0_from_flow(), then back to the loss space (which may coincide with flow, in which case the round-trip is identity up to fp64 round-off).

Parameters:

  • x0 (Tensor)

  • x_t (Tensor)

  • t (Tensor)

Return type:

Tensor