Requesting feature: "Laplace family autoguides"

The existing "AutoLaplaceApproximation" uses a delta guide to find the MLE of x, then creates a Laplace approximation around that MLE. But the resulting Laplace approximation is optimal in no way except its MLE. Either term of its ELBO can be arbitrarily bad.

The feature request I'm making here is for a "Laplace family autoguide". Say we have a model p(θ|x), where θ is the vector of model parameters and x is the data. The guide q_{φ,ψ} would be a multivariate normal with mean φ and precision set to be the observed information of p(θ|x) at θ = φ, with an adjustment parametrized by ψ to ensure that it is positive definite. That is:

q_{φ,ψ} = 𝒩[ φ, (ℐ_θ(φ)+M_ψ)⁻¹],

where

ℐ_θ(φ) := -H[log p(θ|x)] | evaluated at φ,

and M_ψ is a positive definite matrix such that (ℐ+M_ψ) is also positive definite. (The optimal choice of M_ψ is not obvious and would need to be left up to the user; I will discuss one possible choice briefly at the end of this feature request.)

In practice, this would mean that the guide has to make direct calls to the model. Something like:

    hessCenter = pyro.condition(model, phi_dict)
    trace1 = poutine.trace(hessCenter)
    trace2 = trace1.get_trace(*args, **kwargs)
    loss = -trace2.log_prob_sum()
    thetaParts = [v for k,v in phi_dict]
    H = hessian.hessian(loss, thetaParts)
    M_matrix = calculateM(H,psi_scale_params)   # calculateM should be chosen by the user, as explained in appendix

    thetaMean = catToOneVector(thetaparts)
    theta = pyro.sample('theta',
                    dist.MultivariateNormal(thetaMean, precision_matrix=H+M_matrix),
                    infer={'is_auxiliary': True})]

(Note that this pseudocode is drawn from a case where data is treated in an unusual way, so may need to be adjusted to condition on data normally.)

This is the basic idea. To make it more useful, there are several possible extensions: sequential sampling, amortizing, and subsampling:

• Sequential sampling refers to the idea that the guide should sample the high-level parameters first, then sample lower-level parameters in conditionally-independent blocks by conditioning on the higher-level parameters. This can substantially reduce the dimension of the matrices involved and thus speed the linear algebra needed for sampling and/or calculating density.
• Amortizing refers to the idea that some nuisance parameters in θ could be set deterministically to their MLE conditional on the data and the other parameters. The MLE of these nuisance parameters would need to be calculated twice: once with the parameters above them set to φ to calculate the Hessian; and once with all other parameters in θ set to their sampled values, in order to calculate the density p(θ|x) for the ELBO.
• Subsampling refers to the idea that the Hessian for a large amount of iid data could be estimated using a subsample of that data. This would mean that the above pseudocode where the guide calls the model would have to be able to separate the model's prior component (which doesn't need rescaling from subsample N to overall N) from its likelihood component (which does need such rescaling).

Appendix: making the precision matrix positive definite

In order to ensure that the covariance matrix of q is positive definite, we're adding the positive definite matrix M to the observed information ℐ. We construct M as a function of both ℐ and some set of parameters ψ.

One possible construction is to have one positive parameter ψ_i per dimension of ℐ, then define:

M_ii:=ψ_i*logsumexp[0, -ℐ_ii/ψ_i, -[abs(ℐ_ii)-Σ_{j≠i}abs(ℐ_ij)] /ψ_i]

This construction ensures that ℐ+M_ψ has positive diagonal entries and is strongly diagonally dominant; these two conditions in turn ensure that it is positive definite. By including each ψ_i into the set of parameters over which the ELBO is maximized, we are also allowing the fitted variational posterior estimate q to have higher curvature than the observed information, even in cases where the observed information is in fact positive definite. This would deal with cases where the true posterior had high kurtosis and thus a Laplace approximation badly overestimates the tail probability.

It may be possible to improve this construction by using a single ψ parameter for multiple "similar" dimensions of ℐ, but discussing this possibility is beyond the scope of this feature request. Constructions using other criteria for positive definiteness are also possible and we hope will be a subject of further research.