Skip to content

conditional_flow_matching

conditional_flow_matching

Conditional flow matching utilities.

ConditionalFlowMatcher 

ConditionalFlowMatcher(sigma=0.0)

Implements the conditional flow matching interpolation API used by VersatIL.

The path convention is t=0 at source noise x0 and t=1 at target data x1. For a batch time vector t and samples with arbitrary trailing dimensions, this class builds the Gaussian bridge x_t = t * x1 + (1 - t) * x0 + sigma * epsilon and the target velocity u_t = x1 - x0.

Initialize the matcher with a constant bridge noise standard deviation.

Source code in src/versatil/models/layers/denoising/conditional_flow_matching.py 
26
27
28
def __init__(self, sigma: float = 0.0):
    """Initialize the matcher with a constant bridge noise standard deviation."""
    self.sigma = sigma

compute_mu_t 

compute_mu_t(x0, x1, t)

Compute the bridge mean mu_t = t * x1 + (1 - t) * x0.

Parameters:

Name Type Description Default
x0 Tensor

Source samples, usually Gaussian noise, with shape (B, ...).

 required
x1 Tensor

Target samples with the same shape as x0.

 required
t Tensor

Per-sample times with shape (B,).

 required
Source code in src/versatil/models/layers/denoising/conditional_flow_matching.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def compute_mu_t(
    self,
    x0: torch.Tensor,
    x1: torch.Tensor,
    t: torch.Tensor,
) -> torch.Tensor:
    """Compute the bridge mean ``mu_t = t * x1 + (1 - t) * x0``.

    Args:
        x0: Source samples, usually Gaussian noise, with shape ``(B, ...)``.
        x1: Target samples with the same shape as ``x0``.
        t: Per-sample times with shape ``(B,)``.
    """
    padded_time = _pad_time_like_tensor(time=t, tensor=x0)
    return padded_time * x1 + (1 - padded_time) * x0

compute_sigma_t 

compute_sigma_t(t)

Return the constant path standard deviation at time t.

Source code in src/versatil/models/layers/denoising/conditional_flow_matching.py 
46
47
48
49
def compute_sigma_t(self, t: torch.Tensor) -> float:
    """Return the constant path standard deviation at time ``t``."""
    del t
    return self.sigma

compute_lambda 

compute_lambda(t)

Compute the score weighting coefficient used by score-CFM variants.

VersatIL currently trains the velocity target directly, but this method preserves the standard matcher API for callers that expect the 2 * sigma_t / sigma**2 score weight.

Source code in src/versatil/models/layers/denoising/conditional_flow_matching.py 
51
52
53
54
55
56
57
58
59
def compute_lambda(self, t: torch.Tensor) -> float:
    """Compute the score weighting coefficient used by score-CFM variants.

    VersatIL currently trains the velocity target directly, but this method
    preserves the standard matcher API for callers that expect the
    ``2 * sigma_t / sigma**2`` score weight.
    """
    sigma_t = self.compute_sigma_t(t)
    return 2 * sigma_t / (self.sigma**2 + 1e-8)

sample_xt 

sample_xt(x0, x1, t, epsilon)

Sample x_t from the conditional Gaussian bridge.

Parameters:

Name Type Description Default
x0 Tensor

Source samples, usually Gaussian noise.

 required
x1 Tensor

Target samples.

 required
t Tensor

Per-sample times.

 required
epsilon Tensor

Standard Gaussian noise with the same shape as x0.

 required
Source code in src/versatil/models/layers/denoising/conditional_flow_matching.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def sample_xt(
    self,
    x0: torch.Tensor,
    x1: torch.Tensor,
    t: torch.Tensor,
    epsilon: torch.Tensor,
) -> torch.Tensor:
    """Sample ``x_t`` from the conditional Gaussian bridge.

    Args:
        x0: Source samples, usually Gaussian noise.
        x1: Target samples.
        t: Per-sample times.
        epsilon: Standard Gaussian noise with the same shape as ``x0``.
    """
    mu_t = self.compute_mu_t(x0=x0, x1=x1, t=t)
    sigma_t = _pad_time_like_tensor(time=self.compute_sigma_t(t), tensor=x0)
    return mu_t + sigma_t * epsilon

compute_conditional_flow 

compute_conditional_flow(x0, x1, t, xt)

Compute the conditional velocity target u_t = x1 - x0.

The independent CFM path has a constant velocity along each source-target pair, so t and xt are accepted for API symmetry but do not affect the returned target.

Source code in src/versatil/models/layers/denoising/conditional_flow_matching.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def compute_conditional_flow(
    self,
    x0: torch.Tensor,
    x1: torch.Tensor,
    t: torch.Tensor,
    xt: torch.Tensor,
) -> torch.Tensor:
    """Compute the conditional velocity target ``u_t = x1 - x0``.

    The independent CFM path has a constant velocity along each source-target
    pair, so ``t`` and ``xt`` are accepted for API symmetry but do not affect
    the returned target.
    """
    del t, xt
    return x1 - x0

sample_noise_like 

sample_noise_like(tensor)

Sample standard Gaussian noise with the same shape as tensor.

Source code in src/versatil/models/layers/denoising/conditional_flow_matching.py 
96
97
98
def sample_noise_like(self, tensor: torch.Tensor) -> torch.Tensor:
    """Sample standard Gaussian noise with the same shape as ``tensor``."""
    return torch.randn_like(tensor)

sample_location_and_conditional_flow 

sample_location_and_conditional_flow(x0, x1, t=None, return_noise=False)

Sample the objects needed for one CFM training target.

If t is omitted, times are sampled uniformly in [0, 1]. Returns (t, x_t, u_t) by default, or (t, x_t, u_t, epsilon) when return_noise is enabled.

Source code in src/versatil/models/layers/denoising/conditional_flow_matching.py 
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def sample_location_and_conditional_flow(
    self,
    x0: torch.Tensor,
    x1: torch.Tensor,
    t: torch.Tensor | None = None,
    return_noise: bool = False,
) -> (
    tuple[torch.Tensor, torch.Tensor, torch.Tensor]
    | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
):
    """Sample the objects needed for one CFM training target.

    If ``t`` is omitted, times are sampled uniformly in ``[0, 1]``. Returns
    ``(t, x_t, u_t)`` by default, or ``(t, x_t, u_t, epsilon)`` when
    ``return_noise`` is enabled.
    """
    if t is None:
        t = torch.rand(x0.shape[0], device=x0.device, dtype=x0.dtype)
    if len(t) != x0.shape[0]:
        raise ValueError(
            f"Time batch size {len(t)} must match sample batch size {x0.shape[0]}."
        )
    epsilon = self.sample_noise_like(x0)
    xt = self.sample_xt(x0=x0, x1=x1, t=t, epsilon=epsilon)
    conditional_flow = self.compute_conditional_flow(x0=x0, x1=x1, t=t, xt=xt)
    if return_noise:
        return t, xt, conditional_flow, epsilon
    return t, xt, conditional_flow