nifty.re.evidence_lower_bound module#

estimate_evidence_lower_bound(likelihood, samples, n_eigenvalues, *, compute_all=False, min_lh_eval=0.001, n_batches=10, tol=0.0, verbose=True, metric_jit=True, output_directory='', save_eigensystem_prefix='metric', resume_eigenvectors=None, resume_eigenvalues=None, orthonormalize_eigenvectors=True, orthonormalize_every_n_batches=8, orthonormalize_threshold=1e-06, orthonormalize_n_probes=2, trace_log_method='eigsh', trace_log_space='signal', slq_order=64, slq_num_samples=8, slq_key=None, slq_kwargs=None, use_radau_as_bound=False, slq_jit=False, analytic_prior_term=False)[source]#

Provides an estimate for the Evidence Lower Bound (ELBO).

Statistical inference deals with the problem of hypothesis testing, given some data and models that can describe it. In general, it is hard to find a good metric to discern between different models. In Bayesian Inference, the Bayes factor can serve this purpose. To compute the Bayes factor it is necessary to calculate the evidence, given the specific model \(p( d|\text{model})\) at hand. Then, the ratio between the evidence of a model A and the one of a model B represents how much more likely it is that model A represents the data better than model B.

The evidence for an approximated-inference problem can in principle be calculated. However, this is only practically feasible in a low-dimensional setting. What often can be computed is

\[\log(p(d)) - D_\text{KL} \left[ Q(\theta(\xi)|d) || p(\theta(\xi) | d) \right] = -\langle H(\theta( \xi), d)\rangle + \frac1 2 \left( N + \text{Tr } \log\Lambda\right),\]

where \(D_\text{KL} \left[ Q || p \right]\) is the Kullback-Leibler (KL) divergence between the approximating posterior distribution \(Q\) and the actual posterior \(p\). Since the Kullback-Leibler divergence \(D_\text{KL} [\cdot, \cdot] \geq 0\) is positive definite, it is convenient to consider the lower bound

\[\log(p(d)) \geq -\langle H(\theta(\xi), d)\rangle + \frac1 2 \left(N + \text{Tr } \log\Lambda \right),\]

which takes the name of Evidence Lower Bound (ELBO).

If the KL divergence is well minimized (which should always be the case when a Variational Inference approach is followed), then it is possible to utilize the ELBO (as a proxy for the actual evidences) and calculate the Bayes factors for model comparison.

Parameters:

likelihood (nifty.re.likelihood.Likelihood) – Log-likelihood of the model.
samples (nifty.re.evi.Samples) – Collection of samples from the posterior distribution.
n_eigenvalues (int) – Maximum number of eigenvalues to compute exactly. If trace_log_method=”slq”, the remaining trace-log term is estimated via stochastic Lanczos quadrature. If trace_log_method=”eigsh”, the remainder is approximated using the smallest computed eigenvalue and at least one exact eigenvalue is therefore required. Zero exact eigenvalues are supported in SLQ mode. Note that if n_eigenvalues equals the total number of relevant degrees of freedom of the problem, all relevant eigenvalues are always computed irrespective of other stopping criteria.
compute_all (bool) – If True, compute all eigenvalues and eigenvectors of the relevant metric subspace. Overrides n_eigenvalues.
min_lh_eval (float) – Smallest eigenvalue of the likelihood to be considered. If the estimated eigenvalues become smaller than 1 + min_lh_eval, the eigenvalue estimation terminates and uses the smallest eigenvalue as a proxy for all remaining eigenvalues in the trace-log estimation. Default is 1e-3.
n_batches (int) – Number of batches into which the eigenvalue estimation gets subdivided. The batch schedule is defined for the total n_eigenvalues; when resuming, the remaining work continues with the tail of this schedule. Only after completing one batch the early stopping criterion based on min_lh_eval is checked for.
tol (Optional[float]) – Tolerance on the eigenvalue calculation. Zero indicates machine precision. Default is 0.
verbose (Optional[bool]) – Print list of eigenvalues and summary of evidence calculation. Default is True.
metric_jit (bool or callable) – Whether to jit the metric. Default is True.
output_directory (Optional[str]) – If set, saves the cumulative metric eigenvalues and eigenvectors after each batch to {output_directory}/{prefix}_eigenvalues.npy and {output_directory}/{prefix}_eigenvectors.npy.
save_eigensystem_prefix (str) – Prefix for eigensystem filenames. A suffix _signal or _data is appended automatically depending on the selected trace-log space. Default is “metric”.
resume_eigenvectors (Optional[np.ndarray]) – Precomputed eigenvectors to resume the eigenvalue calculation. The array is expected to be shaped (metric_size, n_vectors).
resume_eigenvalues (Optional[np.ndarray]) – Eigenvalues corresponding to resume_eigenvectors. If not provided, they are estimated via the metric. Required when orthonormalize_eigenvectors is True and resume_eigenvectors is set. Supplied eigenvalues are validated against the selected signal- or data-space operator.
orthonormalize_eigenvectors (bool) – If True, re-orthonormalize the cumulative eigenvector basis before projecting out the corresponding subspace. Default is True.
orthonormalize_every_n_batches (int) – Re-orthonormalize every N batches when orthonormalize_eigenvectors is True. Default is 8.
orthonormalize_threshold (Optional[float]) – Re-orthonormalize whenever a randomized probe of V.T @ V deviates from the identity by more than this threshold. Set to None to disable this check. Default is 1e-6.
orthonormalize_n_probes (int) – Number of randomized probe vectors used to estimate the orthonormality error. Default is 2.
trace_log_method ({"eigsh", "slq"}) – Strategy for estimating the trace-log term. “eigsh” uses exact eigenvalues plus a minimum-eigenvalue approximation for the remainder. “slq” uses exact eigenvalues plus stochastic Lanczos Gauss quadrature for the remainder. Radau diagnostics are computed only when use_radau_as_bound=True.
trace_log_space ({"auto", "signal", "data"}) – Space in which to estimate the trace-log term. “signal” uses the full metric in signal space with f=log. “data” uses the implicit data-space operator LSM^T LSM with f=log1p. “auto” selects the smaller of data- and signal-space dimensions (falls back to “signal” if data-space shapes are unavailable). Default is “signal”.
slq_order (int) – Lanczos order for SLQ (only used when trace_log_method=”slq”).
slq_num_samples (int) – Number of Hutchinson probes for SLQ (only used when trace_log_method=”slq”). At least two probes are required whenever an SLQ remainder is estimated so its stochastic uncertainty is defined.
slq_key (int or jax.Array) – PRNG seed or key for SLQ. If None, a deterministic default is used.
slq_kwargs (Optional[dict]) – Additional kwargs forwarded to the internal SLQ implementation (excluding A, f, order, num_samples, key, n, deflate_eigvecs, lam_min, lam_max, extra_fns, and compute_radau).
use_radau_as_bound (bool) – If True, compute two-endpoint Radau diagnostics and use the upper Radau estimate on the remaining trace-log to build a more conservative lower_error. At least one exact eigenvalue is required to provide an upper spectral endpoint. The result is not guaranteed to be a rigorous bound for arbitrary f (including log/log1p) unless the corresponding conditions for Gauss–Radau bounds are met. Default is False, in which case the one-sigma SLQ standard error is used.
slq_jit (bool) – If True, JIT-compile the SLQ computation. This closes over A, f, order, and num_samples so they are treated as static. Compilation is local to this estimator invocation and is not cached across calls; for one-off runs the compile cost may outweigh the speedup. Default is False.
analytic_prior_term (bool) – If True, compute the quadratic prior term \(\frac{1}{2}\langle \xi^\dagger \xi \rangle_q\) analytically as \(\frac{1}{2}(\mathrm{Tr}\,\Sigma + \bar\xi^\dagger\bar\xi)\) with \(\Sigma = \Lambda^{-1}\). The trace is estimated from the exact eigenvalues plus (if needed) a centered SLQ remainder using \(f(x)=1/x-1\) (signal space) or \(f(x)=1/(1+x)-1\) (data space). The removed identity contribution is added back explicitly. Requires trace_log_method=”slq” or compute_all=True when not all eigenvalues are computed. Default is False.

Returns:

`elbo_samples` (np.array) – List of elbo samples from the posterior distribution. The samples are returned to allow for more accurate elbo statistics.
stats (dict) – Dictionary with a summary of the statistics of the estimated ELBO. The keys of this dictionary are:
- elbo_mean: mean of the ELBO samples.
- elbo_std: standard deviation of the ELBO samples (sampling spread).
- elbo_se: standard error of the ELBO mean (elbo_std / sqrt(n_samples)).
- elbo_up: elbo_mean + elbo_std (one-sigma upper envelope).
- elbo_lw: elbo_mean - elbo_std - lower_error (one-sigma lower estimate).
- lower_error: estimate of the residual trace-log uncertainty. For trace_log_method=”eigsh” this is the legacy lower-bound term based on the smallest computed eigenvalue. For trace_log_method=”slq” it is a one-sigma stochastic uncertainty based on the SLQ standard error of the remaining trace-log, not a rigorous bound. If analytic_prior_term=True, an additional term from the SLQ standard error of the trace-inverse estimate is included.
- trace_inv_exact, trace_inv_slq, trace_inv_const, trace_inv_se, trace_inv_total, prior_mean_sq, prior_term: returned when analytic_prior_term=True.

Warning

To perform Variational Inference there is no need to take into account quantities that are not explicitly dependent on the inferred parameters. Explicitly calculating these terms can be expensive, therefore they are neglected in NIFTy. Since in most cases they are also not required for model comparison, the provided estimate may not include terms which are constant in these parameters. Only when comparing models for which the likelihood includes (possibly data-dependent) constants (or when the ELBO is needed to approximate the true evidence) these contributions have to be considered. For example, for a Gaussian distributed signal and a linear problem (Wiener Filter problem) the only term missing is \(-\frac1 2 \log \det |2 \pi N|\), where \(N\) is the noise covariance matrix.