Skip to main content

Crate flowmatch

Crate flowmatch 

Source
Expand description

§flowmatch

Flow matching as a library primitive.

This crate is intentionally small:

  • it implements training loops and sampling for flow-matching style models,
  • it depends on wass for OT-shaped coupling primitives (e.g. semidiscrete assignments),
  • it does not provide a CLI or experiment runner.

§Public invariants (must not change)

  • Determinism knobs are explicit: training/sampling functions take seed (or configs do).
  • No hidden normalization: if inputs are normalized, it is stated in the doc comment.
  • Backend-agnostic by default: this crate uses ndarray and simple SGD; no GPU framework types leak through the public API in the default feature set.
    • Optional training backends (e.g. burn) are feature-gated.

§How this maps to “Flow Matching” (papers)

The core training objective used here is the standard conditional flow matching regression (sample t, sample a point on a path x_t, regress a vector field v_theta(x_t, t) toward a target velocity u_t). Concretely:

  • sd_fm::train_sd_fm_semidiscrete_linear uses a linear interpolation path x_t = (1-t) x_0 + t y_j and target u_t = y_j - x_0.
  • A semidiscrete “pick j” step is provided by wass::semidiscrete (potentials + hard assignment), which acts like a simple coupling / conditioning mechanism.

§References (conceptual anchors; not “implemented fully”)

  • Lipman et al., Flow Matching for Generative Modeling (arXiv:2210.02747): the canonical FM objective and linear-path baselines.
  • Lipman et al., Flow Matching Guide and Code (arXiv:2412.06264): a comprehensive reference covering the full design space.
  • Li et al., Flow Matching Meets Biology and Life Science: A Survey (arXiv:2507.17731, 2025): a taxonomy of variants (CFM/RFM, non-Euclidean, discrete) and a map of applications/tooling.
  • Gat et al., Discrete Flow Matching (NeurIPS 2024): extending the FM paradigm to discrete data (language, graphs).
  • Chen & Lipman, Riemannian Flow Matching on General Geometries (arXiv:2302.03660): the foundation for FM on manifolds (like the Poincaré ball in hyperball).

Related variants that are not implemented here (yet):

  • Dao et al., Flow Matching in Latent Space (arXiv:2307.08698) — latent FM + guidance details
  • Klein et al., Equivariant Flow Matching (NeurIPS 2023) — symmetry constraints
  • Zaghen et al., Towards Variational Flow Matching on General Geometries (arXiv:2502.12981, 2025) — variational objectives with Riemannian Gaussians (RG-VFM).

Applications & Extensions:

  • Qin et al., DeFoG: Discrete Flow Matching for Graph Generation (arXiv:2410.04263, 2025).
  • FlowMM: Generating Materials with Riemannian Flow Matching (2024/2025).

§What can change later

  • The parameterization of vector fields (linear vs MLP vs backend-specific).
  • ODE integrators (Euler → Heun/RK).
  • Adding optional Tweedie correction utilities (diffusion-specific).

§Module map

  • flow: VectorField trait and flow_drift primitive (canonical home; was wass::flow)
  • sd_fm: semidiscrete conditional FM training and sampling
  • rfm: rectified-flow coupling helpers (minibatch OT pairing)
  • linear: simple linear vector-field parameterizations
  • ode: fixed-step ODE integrators (Euler, Heun)
  • ot_cfm: OT-conditional flow matching (mini-batch OT coupling for CFM training)
  • energy: energy matching (scalar energy whose gradient gives the velocity field)
  • metrics: evaluation metrics (JS divergence, entropic OT cost)
  • discrete_ctmc: CTMC generator scaffolding for discrete FM
  • simplex: simplex utilities for discrete FM variants
  • riemannian: Riemannian FM training (feature-gated: riemannian)
  • riemannian_ode: manifold ODE integrators (feature-gated: riemannian)
  • burn_euclidean: Burn-backed Euclidean FM training (feature-gated: burn)
  • burn_sd_fm: Burn-backed SD-FM/RFM training (feature-gated: burn)

Re-exports§

pub use flow::flow_drift;
pub use flow::VectorField;

Modules§

discrete_ctmc
Discrete FM (CTMC-based) scaffolding.
energy
Energy matching for flow matching.
flow
Flow / drift primitives.
linear
Simple vector-field parameterizations for flow matching.
metrics
Evaluation utilities for flow models (small + explicit).
ode
Minimal ODE integrators for sampling flow models.
ot_cfm
OT-conditional flow matching (OT-CFM).
rfm
Rectified Flow Matching (RFM) helpers.
sd_fm
Semidiscrete flow matching (SD-FM), minimal version.

Enums§

Error
flowmatch error variants.

Type Aliases§

Result
Convenience alias for std::result::Result<T, flowmatch::Error>.