Amortized inference for hierarchical models

Hierarchical models contain a growing collection of local latent variables: each experiment, specimen, or trajectory contributes its own unknown parameters. Assigning and optimizing a separate set of variational parameters for every group becomes expensive, and it provides no inference rule for a new group. Amortized variational inference replaces those free local parameters with the output of a shared inference network (Kingma and Welling, 2014).

Local and global latent variables

Consider \(M\) groups indexed by \(i=1,\ldots,M\). Let \(\boldsymbol{\alpha}\) denote global parameters shared across groups, let \(\mathbf{z}_i\) denote the local parameters for group \(i\), and let group \(i\) contain \(n_i\) observed input–response pairs,

\[ \mathcal{D}_i = \{(\mathbf{x}_{ij},\mathbf{y}_{ij})\}_{j=1}^{n_i}. \]

We treat the inputs \(\mathbf{X}_i=(\mathbf{x}_{i1},\ldots,\mathbf{x}_{in_i})\) as fixed and write \(\mathbf{Y}_i=(\mathbf{y}_{i1},\ldots,\mathbf{y}_{in_i})\) for the random responses. Define \(\mathcal{D}_{1:M}=(\mathcal{D}_1,\ldots,\mathcal{D}_M)\) and \(\mathbf{z}_{1:M}=(\mathbf{z}_1,\ldots,\mathbf{z}_M)\), with analogous notation for collections of inputs and responses. A common conditional factorization is

\[ p(\boldsymbol{\alpha},\mathbf{z}_{1:M},\mathbf{Y}_{1:M} \mid \mathbf{X}_{1:M}) = p(\boldsymbol{\alpha}) \prod_{i=1}^M \left[ p(\mathbf{z}_i\mid\boldsymbol{\alpha}) \prod_{j=1}^{n_i} p(\mathbf{y}_{ij}\mid\mathbf{x}_{ij},\mathbf{z}_i,\boldsymbol{\alpha}) \right]. \]

The global variable describes structure shared by all groups. Each local variable describes what remains specific to one group after conditioning on that shared structure.

Let an encoder with parameters \(\boldsymbol{\psi}_e\) map each data set to \(\mathbf{e}_i=E_{\boldsymbol{\psi}_e}(\mathcal{D}_i)\), and let a second encoder with parameters \(\boldsymbol{\psi}_s\) combine the collection \(\{\mathbf{e}_i\}_{i=1}^M\) into a global summary \(\mathbf{s}\). Let \(\boldsymbol{\psi}_g\) and \(\boldsymbol{\psi}_\ell\) parameterize the global and local guides, respectively, and collect all learned parameters in \(\boldsymbol{\psi}=(\boldsymbol{\psi}_e,\boldsymbol{\psi}_s,\boldsymbol{\psi}_g,\boldsymbol{\psi}_\ell)\). A structured variational guide can then retain the local–global distinction:

\[ q_{\boldsymbol{\psi}}( \boldsymbol{\alpha},\mathbf{z}_{1:M} \mid \mathcal{D}_{1:M}) = q_{\boldsymbol{\psi}_g}( \boldsymbol{\alpha}\mid\mathbf{s}) \prod_{i=1}^M q_{\boldsymbol{\psi}_\ell}( \mathbf{z}_i\mid\boldsymbol{\alpha},\mathbf{e}_i). \]

The local-guide parameters \(\boldsymbol{\psi}_\ell\) are shared across groups. The number of latent variables still grows with \(M\), but the number of learned parameters used to produce their approximate posteriors does not.

For example, a local Gaussian guide may use a network to produce

\[ (\boldsymbol{\mu}_i,\log\boldsymbol{\sigma}_i) = g_{\boldsymbol{\psi}_\ell}( \boldsymbol{\alpha},\mathbf{e}_i), \]

and then set

\[ q_{\boldsymbol{\psi}_\ell}( \mathbf{z}_i\mid\boldsymbol{\alpha},\mathbf{e}_i) = \mathcal{N}\!\left( \mathbf{z}_i\mid \boldsymbol{\mu}_i, \operatorname{diag}(\boldsymbol{\sigma}_i^2) \right). \]

Here \(\boldsymbol{\mu}_i\) and the componentwise positive \(\boldsymbol{\sigma}_i\) have the same dimension as \(\mathbf{z}_i\), and \(\mathcal{N}(\mathbf{z}\mid\boldsymbol{\mu},\boldsymbol{\Sigma})\) denotes the normal density with mean \(\boldsymbol{\mu}\) and covariance \(\boldsymbol{\Sigma}\).

The inference network performs the repeated local optimization implicitly. For a new group generated by the same process, one forward pass produces an approximate posterior. The network output can also initialize separate group-specific variational parameters when local refinement is desired.

Permutation-invariant set encoders

The order of independent observations within a group carries no information. An encoder for \(\mathcal{D}_i\) should therefore produce the same representation after any permutation of its elements. A simple construction is

\[ \mathbf{e}_i = \rho_{\boldsymbol{\psi}_\rho} \!\left( \sum_{j=1}^{n_i} h_{\boldsymbol{\psi}_h} (\mathbf{x}_{ij},\mathbf{y}_{ij}) \right), \]

where the same map \(h_{\boldsymbol{\psi}_h}\) processes each observation and \(\rho_{\boldsymbol{\psi}_\rho}\) processes their pooled representation; \(\boldsymbol{\psi}_e\) collects these parameters. Sum pooling is permutation invariant. It can also preserve group size if \(h_{\boldsymbol{\psi}_h}\) emits a constant coordinate; otherwise \(n_i\) or \(\log n_i\) should be supplied separately. Mean pooling is also permutation invariant, but it discards group size unless that information is added explicitly. This matters because posterior uncertainty generally depends on the number of observations. A second pooling operation over \(\mathbf{e}_1,\ldots,\mathbf{e}_M\), together with \(M\) when needed, can provide the global summary \(\mathbf{s}\).

This shared-map-and-pooling construction is the basic Deep Sets architecture (Zaheer et al., 2017).

Using a permutation-invariant encoder expresses an exchangeability assumption about the jointly permuted input–response pairs. It is appropriate for unordered replicate measurements. Time series and spatial sequences carry order or geometry, so their encoders must preserve that structure instead of treating the observations as an unordered set.

Training objective and limitations

The guide is trained by maximizing the hierarchical evidence lower bound. Under the factorization above, it can be written as

\[ \mathbb{E}_{q_{\boldsymbol{\psi}}} \left[ \log p(\boldsymbol{\alpha}) - \log q_{\boldsymbol{\psi}_g}(\boldsymbol{\alpha}\mid\mathbf{s}) + \sum_{i=1}^M \left\{ \log p(\mathbf{z}_i\mid\boldsymbol{\alpha}) - \log q_{\boldsymbol{\psi}_\ell} (\mathbf{z}_i\mid\boldsymbol{\alpha},\mathbf{e}_i) + \sum_{j=1}^{n_i} \log p(\mathbf{y}_{ij}\mid \mathbf{x}_{ij},\mathbf{z}_i,\boldsymbol{\alpha}) \right\} \right]. \]

If the full summary \(\mathbf{s}\) is available, a uniformly sampled minibatch \(\mathcal{B}\) of \(B\) groups gives an unbiased estimator by replacing the group sum with \(\frac{M}{B}\sum_{i\in\mathcal{B}}\) while leaving the global prior and global-guide terms unscaled. The scaling therefore applies to each complete group contribution: its local conditional-prior term, local-guide term, and likelihood terms. Otherwise, computing the global guide requires a separately justified stochastic or streaming approximation. Naively scaling the entire hierarchical ELBO is not unbiased.

Amortization creates a new source of error. For a fixed guide family, the total inference gap can be decomposed into an approximation gap, caused by restrictions of that family, and an amortization gap, caused by the inference network failing to return the best member of the family for a particular data set (Cremer et al., 2018). Comparing the amortized result with carefully optimized non-amortized inference in the same guide family estimates the amortization component. Posterior predictive checks and simulation-based calibration assess overall adequacy, but do not by themselves identify which component is responsible.

The following inverse-kinematics notebook illustrates amortized inference for a fixed observation vector (Karumuri and Bilionis, 2024). The hierarchical construction here extends the same shared-inference-network idea to a variable number of related groups and to unordered sets of observations within each group.