Skip to content

FlowSampler ODE 101

intermediate · Integrate a closed-form velocity field from noise to a Gaussian target; study steps vs fidelity without training.

Run it:

1
python examples/10-sampling/03-flow/01-flow-sampler-ode/main.py

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
"""FlowSampler: integrate a velocity field from noise to data, no training needed.

For a Gaussian target N(mu, s^2 I) under the linear path x_t = t x1 + (1-t) x0,
the marginal velocity is closed-form:

    u(x, t) = mu + (t s^2 - (1 - t)) / (t^2 s^2 + (1 - t)^2) * (x - t mu)

so the sampler's mechanics can be studied in isolation from learning.
"""

import torch
from torch import nn

from torchebm.samplers import FlowSampler

MU = torch.tensor([2.0, -1.0])
S = 0.5


class AnalyticVelocity(nn.Module):  # the exact marginal velocity field u(x, t)
    def forward(self, x, t, **kwargs):
        t = torch.as_tensor(t, dtype=x.dtype, device=x.device)
        if t.ndim == 0:
            t = t.expand(x.size(0))
        t = t.view(-1, 1)
        var = t**2 * S**2 + (1 - t) ** 2
        return MU + (t * S**2 - (1 - t)) / var * (x - t * MU)


flow = FlowSampler(
    model=AnalyticVelocity(),
    interpolant="linear",
    prediction="velocity",
    integrator="euler",
)

# More steps, smaller discretization error: the samples converge to N(mu, s^2 I).
torch.manual_seed(0)
z = torch.randn(8192, 2)
for num_steps in (2, 8, 64):
    samples = flow.sample(x=z, n_steps=num_steps)
    mean_err = (samples.mean(0) - MU).norm()
    std_err = (samples.std(0) - S).abs().max()
    print(f"num_steps={num_steps:3d}  mean error {mean_err:.3f}   std error {std_err:.3f}")

# Swap integrator="euler" for "heun" or "dopri5" (adaptive) and the error at few
# steps collapses; trained velocity models (see 20-training) plug in identically.