lampe.inference.fmpe¶
Flow matching posterior estimation (FMPE) components.
The principle of FMPE is to train a regression network \(v_\phi(\theta, x, t)\) to approximate a vector field inducing a time-continuous normalizing flow between the posterior distribution \(p(\theta | x)\) and a standard Gaussian distribution \(\mathcal{N}(0, I)\).
After training, the normalizing flow \(p_\phi(\theta | x)\) induced by \(v_\phi(\theta, x, t)\) is used to evaluate the posterior density or generate samples.
References
Classes¶
Descriptions¶
- class lampe.inference.fmpe.FMPE(theta_dim, x_dim, freqs=3, build=<class 'zuko.nn.MLP'>, **kwargs)¶
Creates a flow matching posterior estimation (FMPE) network.
- Parameters:
theta_dim (int) – The dimensionality \(D\) of the parameter space.
x_dim (int) – The dimensionality \(L\) of the observation space.
freqs (int) – The number of time embedding frequencies.
build (Callable[[int, int], Module]) – The network constructor (e.g.
lampe.nn.ResMLP). It takes the number of input and output features as positional arguments.kwargs – Keyword arguments passed to the constructor.
- forward(theta, x, t)¶
- class lampe.inference.fmpe.FMPELoss(estimator, eta=0.001)¶
Creates a module that calculates the flow matching loss for a FMPE regressor.
Given a batch of \(N\) pairs \((\theta_i, x_i)\), the module returns
\[l = \frac{1}{N} \sum_{i = 1}^N \| v_\phi((1 - t_i) \theta_i + (t_i + \eta) \epsilon_i, x_i, t_i) - (\epsilon_i - \theta_i) \|_2^2\]where \(t_i \sim \mathcal{U}(0, 1)\) and \(\epsilon_i \sim \mathcal{N}(0, I)\).
- Parameters:
estimator (Module) – A regression network \(v_\phi(\theta, x, t)\).