class

FlowMatchingModel

extendsPretrainedModel

FlowMatchingModel(config: FlowMatchingConfig)

source edit

Velocity field trained by Conditional Flow Matching.

forward(sample, t) evaluates $v_t(x; \theta)$ . The pieces that make it a generative model are separate and each maps onto one part of the paper:

path_sample — draw $x_t$ from the conditional path.
conditional_target — the field Theorem 3 says generates it.
flow_matching_loss — the objective, with no solve in it.
sample — integrate $t: 0 \to 1$ .
log_prob / bits_per_dim — integrate back, tracking the divergence.

Parameters

configFlowMatchingConfig

Architecture, path and solver settings.

Attributes

fieldnn.Module

The velocity field itself.

Notes

Reference: Lipman, Chen, Ben-Hamu, Nickel, and Le, "Flow Matching for Generative Modeling", ICLR, 2023 (arXiv:2210.02747). Theorem 2 gives the objective's equivalence to the intractable one; Theorem 3 gives the conditional field; §4.2 gives the optimal-transport path.

Where to run it. Training is a plain U-Net forward and backward with no solve, so it behaves like any convolutional model and belongs on "metal". Sampling and likelihood are solves, and inherit the same guidance as any continuous flow: the GPU wins once the batch is real.

Examples

>>> import lucid
>>> from lucid.models.generative.flow_matching import (
...     FlowMatchingConfig, FlowMatchingModel,
... )
>>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16,
...                          channel_mult=(1, 2), num_res_blocks=1,
...                          attention_resolutions=(), resnet_groups=8)
>>> model = FlowMatchingModel(cfg).eval()
>>> loss, _, _ = model.flow_matching_loss(lucid.randn((2, 3, 8, 8)))
>>> loss.shape
()
>>> model.sample(n_samples=2, steps=4).shape
(2, 3, 8, 8)

Used by 2

Constructors

dunder

init

→None

__init__(config: FlowMatchingConfig)

source edit

Properties

prop

input_dim

→int

input_dim: int

source edit

int: Flattened width D of one sample.

prop

nfe

→int

nfe: int

source edit

int: Field evaluations spent by the most recent solve.

A training step solves nothing and leaves the count where the last solve put it. The number the paper cares about: straighter paths need fewer.

prop

path

→str

path: str

source edit

str: Which conditional probability path the target comes from.

prop

trace_method

→str

trace_method: str

source edit

str: How the divergence is obtained when scoring likelihood.

Instance methods

bits_per_dim

→Tensor

bits_per_dim(x: Tensor)

source edit

Per-sample negative log-likelihood in bits per dimension.

coefficients

→tuple[Tensor, Tensor, Tensor, Tensor]

coefficients(t: Tensor)

source edit

(a_t, sigma_t, a_t', sigma_t') for the configured path.

Exposed because it is the only place the two paths differ, and because checking a schedule's derivative against a difference quotient is the cheapest way to know a new one is right.

conditional_target

→Tensor

conditional_target(x1: Tensor, x0: Tensor, t: Tensor)

source edit

The field $u_t(x \mid x_1)$ that generates the path.

Theorem 3 states it as $(\sigma_t'/\sigma_t)(x - a_t x_1) + a_t' x_1$ . Since $x - a_t x_1 = \sigma_t x_0$ for the sample actually drawn, that reduces to $\sigma_t' x_0 + a_t' x_1$ — the same quantity with no division by $\sigma_t$ , which matters because $\sigma_t \to 0$ at t = 1 on the diffusion path. For the optimal-transport path it collapses further, to the $t$ -independent $x_1 - (1 - \sigma_{\min}) x_0$ of paper eq. (23).

Parameters

x1Tensor

(B, C, H, W) data and noise samples.

x0Tensor

(B, C, H, W) data and noise samples.

tTensor

(B,) times in [0, 1].

Returns

Tensor

(B, C, H, W) regression target.

Examples

>>> import lucid
>>> from lucid.models.generative.flow_matching import (
...     FlowMatchingConfig, FlowMatchingModel,
... )
>>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16,
...                          channel_mult=(1, 2), num_res_blocks=1,
...                          attention_resolutions=(), resnet_groups=8)
>>> model = FlowMatchingModel(cfg).eval()
>>> x1, x0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8))
>>> t = lucid.tensor([0.3, 0.6])
>>> target = model.conditional_target(x1, x0, t)
>>> target.shape
(2, 3, 8, 8)
>>> step = 1e-2  # the target is the time derivative of the path
>>> ahead = model.path_sample(x1, x0, t + step)
>>> slope = (ahead - model.path_sample(x1, x0, t)) / step
>>> bool(lucid.allclose(slope, target, atol=1e-3))
True
>>> bool(lucid.allclose(target, x1 - (1 - cfg.sigma_min) * x0))  # eq. (23)
True

flow_matching_loss

→tuple[Tensor, Tensor, Tensor]

flow_matching_loss(x1: Tensor)

source edit

One Conditional Flow Matching step — no ODE is solved.

Parameters

x1Tensor

(B, C, H, W) data batch.

Returns

tuple[Tensor, Tensor, Tensor]

(loss, prediction, target) — the scalar mean squared error, the field's output at the sampled points, and what it was regressed against.

Notes

t is drawn uniformly per sample, which is the estimator the objective is written as an expectation over. Sharing one t across the batch would still be unbiased but noisier.

Examples

>>> import lucid
>>> from lucid.models.generative.flow_matching import (
...     FlowMatchingConfig, FlowMatchingModel,
... )
>>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16,
...                          channel_mult=(1, 2), num_res_blocks=1,
...                          attention_resolutions=(), resnet_groups=8)
>>> model = FlowMatchingModel(cfg).eval()
>>> loss, prediction, target = model.flow_matching_loss(
...     lucid.randn((2, 3, 8, 8)))
>>> loss.shape, prediction.shape, target.shape
((), (2, 3, 8, 8), (2, 3, 8, 8))
>>> bool(lucid.allclose(loss, ((prediction - target) ** 2).mean()))
True

forward

→DiffusionModelOutput

forward(sample: Tensor, t: Tensor)

source edit

Evaluate the velocity field at (t, sample).

log_prob

→Tensor

log_prob(x: Tensor)

source edit

Exact per-sample $\log p_1(x)$ in nats.

Integrates the field from t = 1 back to t = 0 alongside the accumulated divergence, then scores the arrival against the standard normal the path starts from.

Parameters

xTensor

(B, C, H, W) data samples.

Returns

Tensor

(B,) log-likelihood.

Notes

With trace_method resolving to "hutchinson" — which it does at any image size — this is an unbiased estimate, and two calls on the same input will not agree. That is the same estimator the paper reports bits/dim with.

Examples

>>> import math
>>> import lucid
>>> from lucid.models.generative.flow_matching import (
...     FlowMatchingConfig, FlowMatchingModel,
... )
>>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16,
...                          channel_mult=(1, 2), num_res_blocks=1,
...                          attention_resolutions=(), resnet_groups=8)
>>> model = FlowMatchingModel(cfg).eval()
>>> x = lucid.randn((2, 3, 8, 8))
>>> log_p = model.log_prob(x)
>>> log_p.shape
(2,)
>>> # A fresh field is exactly zero (its output layer starts at zero),
>>> # so the flow is the identity and x is scored as N(0, I) itself.
>>> const = 0.5 * model.input_dim * math.log(2 * math.pi)
>>> gauss = -0.5 * (x.reshape(2, -1) ** 2).sum(dim=-1) - const
>>> bool(lucid.allclose(log_p, gauss, atol=1e-3))
True

path_sample

→Tensor

path_sample(x1: Tensor, x0: Tensor, t: Tensor)

source edit

Draw $x_t \sim p_t(\cdot \mid x_1)$ , given the noise x0.

Parameters

x1Tensor

(B, C, H, W) data samples — the conditioning point.

x0Tensor

(B, C, H, W) standard-normal samples.

tTensor

(B,) times in [0, 1].

Returns

Tensor

(B, C, H, W) point on the conditional path. For the optimal-transport path this is the straight-line interpolation of paper eq. (22).

Examples

>>> import lucid
>>> from lucid.models.generative.flow_matching import (
...     FlowMatchingConfig, FlowMatchingModel,
... )
>>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16,
...                          channel_mult=(1, 2), num_res_blocks=1,
...                          attention_resolutions=(), resnet_groups=8)
>>> model = FlowMatchingModel(cfg).eval()
>>> x1, x0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8))
>>> model.path_sample(x1, x0, lucid.tensor([0.25, 0.75])).shape
(2, 3, 8, 8)
>>> bool(lucid.allclose(model.path_sample(x1, x0, lucid.zeros(2)), x0))
True
>>> end = model.path_sample(x1, x0, lucid.ones(2))
>>> bool(lucid.allclose(end, x1 + cfg.sigma_min * x0))  # still sigma_min wide
True

sample

→Tensor

sample(n_samples: int = 1, steps: int | None = None, device: str | None = None, noise: Tensor | None = None)

source edit

Generate by integrating the field from noise to data.

Parameters

n_samplesint= 1

How many samples to draw. Ignored when noise is given.

stepsint= None

Force a fixed-step solve with this many uniform steps instead of the adaptive method. The honest way to report the few-evaluation sampling the straight paths are supposed to allow: fix the budget rather than the tolerance.

devicestr= None

Where to allocate the prior draw. Defaults to the device the model's parameters are on.

noiseTensor= None

(N, C, H, W) starting point, in place of a fresh standard-normal draw — pass the same tensor to two models to compare them on identical initial conditions.

Returns

Tensor

(N, C, H, W) samples at t = 1.

Examples

>>> import lucid
>>> from lucid.models.generative.flow_matching import (
...     FlowMatchingConfig, FlowMatchingModel,
... )
>>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16,
...                          channel_mult=(1, 2), num_res_blocks=1,
...                          attention_resolutions=(), resnet_groups=8)
>>> model = FlowMatchingModel(cfg).eval()
>>> model.sample(n_samples=2, steps=4).shape
(2, 3, 8, 8)
>>> model.nfe  # the adaptive default falls back to Euler: one call a step
4
>>> noise = lucid.randn((3, 3, 8, 8))
>>> model.sample(noise=noise, steps=4).shape  # the batch follows noise
(3, 3, 8, 8)

>>> import lucid >>> from lucid.models.generative.flow_matching import ( ... FlowMatchingConfig, FlowMatchingModel, ... ) >>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=(), resnet_groups=8) >>> model = FlowMatchingModel(cfg).eval() >>> loss, _, _ = model.flow_matching_loss(lucid.randn((2, 3, 8, 8))) >>> loss.shape () >>> model.sample(n_samples=2, steps=4).shape (2, 3, 8, 8)

>>> import lucid >>> from lucid.models.generative.flow_matching import ( ... FlowMatchingConfig, FlowMatchingModel, ... ) >>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=(), resnet_groups=8) >>> model = FlowMatchingModel(cfg).eval() >>> x1, x0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8)) >>> t = lucid.tensor([0.3, 0.6]) >>> target = model.conditional_target(x1, x0, t) >>> target.shape (2, 3, 8, 8) >>> step = 1e-2 # the target is the time derivative of the path >>> ahead = model.path_sample(x1, x0, t + step) >>> slope = (ahead - model.path_sample(x1, x0, t)) / step >>> bool(lucid.allclose(slope, target, atol=1e-3)) True >>> bool(lucid.allclose(target, x1 - (1 - cfg.sigma_min) * x0)) # eq. (23) True

>>> import math >>> import lucid >>> from lucid.models.generative.flow_matching import ( ... FlowMatchingConfig, FlowMatchingModel, ... ) >>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=(), resnet_groups=8) >>> model = FlowMatchingModel(cfg).eval() >>> x = lucid.randn((2, 3, 8, 8)) >>> log_p = model.log_prob(x) >>> log_p.shape (2,) >>> # A fresh field is exactly zero (its output layer starts at zero), >>> # so the flow is the identity and x is scored as N(0, I) itself. >>> const = 0.5 * model.input_dim * math.log(2 * math.pi) >>> gauss = -0.5 * (x.reshape(2, -1) ** 2).sum(dim=-1) - const >>> bool(lucid.allclose(log_p, gauss, atol=1e-3)) True

>>> import lucid >>> from lucid.models.generative.flow_matching import ( ... FlowMatchingConfig, FlowMatchingModel, ... ) >>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=(), resnet_groups=8) >>> model = FlowMatchingModel(cfg).eval() >>> x1, x0 = lucid.randn((2, 3, 8, 8)), lucid.randn((2, 3, 8, 8)) >>> model.path_sample(x1, x0, lucid.tensor([0.25, 0.75])).shape (2, 3, 8, 8) >>> bool(lucid.allclose(model.path_sample(x1, x0, lucid.zeros(2)), x0)) True >>> end = model.path_sample(x1, x0, lucid.ones(2)) >>> bool(lucid.allclose(end, x1 + cfg.sigma_min * x0)) # still sigma_min wide True

>>> import lucid >>> from lucid.models.generative.flow_matching import ( ... FlowMatchingConfig, FlowMatchingModel, ... ) >>> cfg = FlowMatchingConfig(sample_size=8, base_channels=16, ... channel_mult=(1, 2), num_res_blocks=1, ... attention_resolutions=(), resnet_groups=8) >>> model = FlowMatchingModel(cfg).eval() >>> model.sample(n_samples=2, steps=4).shape (2, 3, 8, 8) >>> model.nfe # the adaptive default falls back to Euler: one call a step 4 >>> noise = lucid.randn((3, 3, 8, 8)) >>> model.sample(noise=noise, steps=4).shape # the batch follows noise (3, 3, 8, 8)