Struct FlowMatchingModel 

pub struct FlowMatchingModel {
    pub velocity_net: CnfMlp,
    pub config: FlowMatchingConfig,
}

Conditional Flow Matching model (Lipman et al. 2022).

Learns a velocity field v_θ(x, t) such that Euler integration from t=0 to t=1 transforms noise p_0 = N(0,I) into data p_1.

Loss: E[||v_θ(x_t, t) - u_t(x_t|x_0,x_1)||^2] where x_t = (1-(1-σ)t)·x_0 + t·x_1 and u_t = x_1 - (1-σ)·x_0.

Fields

velocity_net: CnfMlp

Velocity network: input [z; t] (dim z_dim + 1) → output z_dim.

config: FlowMatchingConfig

Configuration.

Implementations

impl FlowMatchingModel

pub fn new(config: FlowMatchingConfig) -> Self

Create a new FlowMatchingModel.

pub fn velocity(&self, z: &[f64], t: f64) -> Vec<f64>

Evaluate the velocity field v_θ(z, t).

pub fn cfm_loss( &self, x0_batch: &[Vec<f64>], x1_batch: &[Vec<f64>], rng: &mut StdRng, ) -> f64

Conditional Flow Matching loss.

For each pair (x0, x1):

• Sample t ~ U[0,1]
• Compute x_t = (1 - (1 - σ_min) * t) * x0 + t * x1
• Conditional vector field: u_t = x1 - (1 - σ_min) * x0
• Loss term: ||v_θ(x_t, t) - u_t||^2

pub fn train_step( &mut self, x1_batch: &[Vec<f64>], lr: f64, rng: &mut StdRng, ) -> f64

One training step: compute CFM loss and update via finite-differences.

Returns the CFM loss.

pub fn sample(&self, n_steps: usize, rng: &mut StdRng) -> Vec<f64>

Sample a single data point by Euler-integrating the velocity field from t=0 to t=1.

pub fn sample_batch( &self, n_samples: usize, n_steps: usize, rng: &mut StdRng, ) -> Vec<Vec<f64>>

Sample a batch of data points.