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]. Internally all arithmetic is in float64 for numerical stability, and the result is cast back to the input dtype.

Functions

add_noise

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

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

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