nifty8.minimization.kl_energies module#

SampledKLEnergy(position, hamiltonian, n_samples, minimizer_sampling, mirror_samples=True, constants=[], point_estimates=[], napprox=0, comm=None, nanisinf=True)[source]#

Provides the sampled Kullback-Leibler used for Variational Inference, specifically for geometric Variational Inference (geoVI) and Metric Gaussian VI (MGVI).

In geoVI a probability distribution is approximated with a standard normal distribution in the canonical coordinate system of the Riemannian manifold associated with the metric of the other distribution. The coordinate transformation is approximated by expanding around a point. The MGVI simplification occurs in case this transformation can be approximated using a linear expansion. In order to infer the optimal expansion point, a stochastic estimate of the Kullback-Leibler divergence is minimized. This estimate is obtained by sampling from the approximation using the current expansion point. During minimization these samples are kept constant; only the expansion point is updated. Due to the typically nonlinear structure of the true distribution these samples have to be updated eventually by instantiating a SampledKLEnergy again. For the true probability distribution the standard parametrization is assumed. The samples of this class can be distributed among MPI tasks.

Parameters:

position (nifty8.field.Field) – Expansion point of the coordinate transformation.
hamiltonian (nifty8.operators.energy_operators.StandardHamiltonian) – Hamiltonian of the approximated probability distribution.
n_samples (integer) – Number of samples used to stochastically estimate the KL.
minimizer_samp (DescentMinimizer or None) – Minimizer used to perform the non-linear part of geoVI sampling. If it is None, only the linear (MGVI) approximation for sampling is used and no further non-linear steps are performed.
mirror_samples (boolean) – Whether the mirrored version of the drawn samples are also used. If true, the number of used samples doubles. Mirroring samples stabilizes the KL estimate as extreme sample variation is counterbalanced. Default is True.
constants (list) – List of parameter keys that are kept constant during optimization. Default is no constants.
point_estimates (list) – List of parameter keys for which no samples are drawn, but that are (possibly) optimized for, corresponding to point estimates of these. Default is to draw samples for the complete domain.
napprox (int) – Number of samples for computing preconditioner for linear sampling. No preconditioning is done by default.
comm (MPI communicator or None) – If not None, samples will be distributed as evenly as possible across this communicator. If mirror_samples is set, then a sample and its mirror image will preferably reside on the same task if necessary.
nanisinf (bool) – If true, nan energies which can happen due to overflows in the forward model are interpreted as inf. Thereby, the code does not crash on these occasions but rather the minimizer is told that the position it has tried is not sensible.

Note

The two lists constants and point_estimates are independent from each other. It is possible to sample along domains which are kept constant during minimization and vice versa. If a key is in both lists, it will be inserted into the Hamiltonian and removed from the KL.

Note

Mirroring samples can help to stabilize the latent mean as it reduces sampling noise. But a mirrored sample involves an additional solve of the non-linear part of the transformation. When using MPI, the samples get distributed as evenly as possible over all tasks. If the number of tasks is smaller then the total number of samples (including mirrored ones), the mirrored pairs try to reside on the same task as their non mirrored partners. This ensures that at least the linear part of the sampling is re-used.