Automatic structured variational preconditioners for HMC

[fritzo]

This issue proposes a workflow speed up HMC in high dimensions: first train a structured Gaussian guide via SVI (which is cheap), then use the structured sqrt(covariance) as a preconditioner in HMC.

Components

• Add an AutoStructured guide and StructuredReparam #2812 An AutoStructured guide that allows users to create a structured Gaussian/Delta variational distribution that can be learned via SVI. Users specify the (1) learnable conditional distribution at each latent site, and (2) a set of learnable link-linear dependencies of each latent variable on a subset of other latent variables (this structure can be sparse and is not limited to structure present in the prior).
• Add an AutoStructured guide and StructuredReparam #2812 A StructuredReparam analogous to NeutraReparam but for the simpler family of Gaussian/Delta distributions. We can basically copy the pattern from Neutra (Hoffman et al. 2019). This is intended to be used by HMC and NUTS in lieu of full-rank mass matrix adaptation (which may be infeasible in high dimensions).
• Add an AutoStructured guide and StructuredReparam #2812 Generalizations to allow not only Gaussian/Delta but also normalizing flows as conditional distributions, and not only link-linear dependencies but also more general neural dependencies. This would allow (Ambrogioni et al. 2021).
• Add automatic dependency tracking to AutoStructured guide #2824 An "auto" option to AutoStructured to automatically determine posterior dependency structure from dependencies in the model (as @eb8680 has long suggested). These can be computed following (Webb et al. 2018).
• Use plate information in AutoStructured guide #2926 Inspect plates to determine when conditional independence can lead to sparsity in precision matrices.
• Support amortization by (1) including auxiliary variables for each non-constant observed site, and (2) reversing the order of posterior dependencies.