SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Anshuk Uppal
ansup@dtu.dk
Kristoffer Stensbo-Smidt
krss@dtu.dk
Wouter Boomsma
wb@di.ku.dk
Jes Frellsen
jefr@dtu.dk
In variational inference, the benefits of Bayesian models rely on accurately capturing the true posterior distribution. We propose using neural samplers that specify implicit distributions, which are well-suited for approximating complex multimodal and correlated posteriors in high-dimensional spaces. Our approach advances inference using implicit distributions by introducing novel approximate bounds by locally linearising the neural sampler. This is distinct from existing methods that rely on additional discriminator networks and unstable adversarial objectives. Furthermore, we present a new sampler architecture that, for the first time, enables implicit distributions over tens of millions of latent variables, addressing computational concerns by using differentiable numerical approximations. Our empirical analysis indicates our method is capable of recovering correlations across layers in large Bayesian neural networks, a property that is crucial for a network's performance but notoriously challenging to achieve. To the best of our knowledge, no other method has been shown to accomplish this task for such large models. Through experiments on downstream tasks, we demonstrate that our expressive posteriors outperform state-of-the-art uncertainty quantification methods, validating the effectiveness of our training algorithm and the quality of the learned implicit distribution.
In the Bayesian approach to statistical inference, information about the data is captured by the posterior distribution over a model's latent variables. By marginalising (averaging) over the variables weighted by the posterior, Bayesian inference can provide excellent generalisation and model calibration. This is particularly compelling for complex models, such as modern deep models, where extreme overparametrisation makes overfitting inevitable unless precaution is taken (Wilson et al., 2020).
Exact Bayesian inference is, however, intractable for all but the simplest models. Practitioners, therefore, rely on approximate inference frameworks, which can be broadly categorised as either sampling-based approaches or variational inference (VI, Saul et al., 1996; Peterson, 1987), the latter being the preferred approach for high-dimensional problems. VI seeks a tractable approximation to the intractable true posterior, limiting us to simple approximations like the isotropic Gaussian, which are often far too crude to satisfactorily cover the posterior, either only narrowly covering a single posterior mode or placing excessive mass in low-probability regions of the posterior (Foong et al., 2019, 2020).
Figure 1: A Bayesian neural network trained on a small regression dataset using four different approximation methods: mean-field variational inference (VI), Adversarial Variational Bayes (AVB, Mescheder et al., 2017a), our proposed model (LIVI), and Hamiltonian Monte Carlo (HMC). The dashed line represents the true underlying function from which the training data (black points) were sampled. For each model, we show the mean prediction together with the epistemic uncertainty (dark shaded region) and the combined epistemic and aleatoric uncertainties (light shaded region).
To achieve more expressive approximations, implicit variational inference (IVI) defines distributions implicitly by transforming samples from a simple distribution. This provides much greater flexibility at the cost of not having a tractable likelihood for the posterior approximation (Diggle et al., 1984; Mohamed et al., 2016) and thus often requires unstable adversarial objectives based on density ratio estimation, which are not easy to scale to high dimensions (Sugiyama et al., 2012; Huszar, 2017).
The different approaches to posterior approximation can be imagined lying on a spectrum illustrated in Fig. 1. On the left, the mean-field VI approximation is much too confident in regions without data. Sampling-based approaches, here represented by Hamiltonian Monte Carlo (HMC) on the right, converge to the true posterior given enough compute, which, however, is typically prohibitively large. Implicit VI approaches, here represented by our proposed approximation, LIVI (Section 4), and Adversarial Variational Bayes (AVB, Mescheder et al., 2017a), which uses a discriminator network for density ratio estimation, hit a tractable middle-ground between the two extremes.
LIVI can produce highly flexible implicit densities approximating the posteriors of complex models through general (non-invertible) stochastic transformations. We present novel, stable training objectives through approximating a lower bound to the model evidence as well as a novel generator architecture, which results in a vastly improved capacity to model expressive distributions, even in millions of dimensions. Specifically, our contributions are:
- We derive a novel approximate lower bound for variational inference using an implicit variational approximation to avoid unstable adversarial objectives.
- To keep computational overhead at large scale in check, we further lower bound the approximate lower bound reducing memory requirements and run times significantly.
- We devise a novel generator architecture capable of specifying highly correlated distributions in high dimensions whilst using a very manageable number of weights and biases, a necessity for modelling high-dimensional distribution.
The paper is organised as follows: Section 2 introduces the background on VI for Bayesian models in general, and VI with implicit distributions (IVI) in particular. Section 3 introduces our proposals to tackle well-known challenges in IVI. Integrating these proposals in Section 4 yields our novel bounds for implicit VI, and Section 5 places our paper in the context of the published literature. In Section 6, we experimentally evaluate our proposed method before summarising and concluding in Section 7.
2 Variational inference for Bayesian models
Consider the supervised learning setting, where we have a training set  $\mathcal{D} = \{(\pmb{x}_i,\pmb{y}_i)\}_{i=1}^n$ ,  $\pmb{X} = \{\pmb{x}_i\}_{i=1}^n$  are the covariates (inputs) and  $\pmb{Y} = \{\pmb{y}_i\}_{i=1}^n$  are the labels (outputs), which we wish to fit using a Bayesian model parametrised by  $\pmb{\theta} \in \mathbb{R}^m$ . That is, assuming a prior  $p(\pmb{\theta})$  over the parameters and a likelihood  $p(\mathcal{D}|\pmb{\theta})$ , we wish to compute the posterior  $p(\pmb{\theta}|\mathcal{D}) = \frac{p(\mathcal{D}|\pmb{\theta})p(\pmb{\theta})}{p(\mathcal{D})}$ . As
computing this posterior is intractable for all but the simplest models, we turn to the framework of variational inference (VI) which finds the best matching tractable variational distribution,  $q_{\gamma}(\pmb{\theta})$  parametrised by  $\gamma$  to the true posterior,  $p(\pmb{\theta} \mid \mathcal{D})$ . A sensible measure of their closeness is the (intractable) Kullback-Leibler (KL) divergence between the two distributions defined as
$$
\mathrm {D} _ {\mathrm {K L}} \left[ q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta}) \mid \mid p (\boldsymbol {\theta} \mid \mathcal {D}) \right] = \log p (\mathcal {D}) - \underbrace {\mathbb {E} _ {q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta})} \left[ \log p (\boldsymbol {\theta} , \mathcal {D}) - \log q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta}) \right]} _ {\mathcal {L} (\boldsymbol {\gamma})}. \tag {1}
$$
Since the evidence,  $p(\mathcal{D})$ , does not depend on  $\gamma$ , minimising the above divergence is equivalent to maximising  $\mathcal{L}(\gamma)$ , referred to as the evidence lower bound (ELBO), which then becomes the objective function in VI. It can be further expanded into a likelihood term and a regularisation term as
$$
\mathcal {L} (\boldsymbol {\gamma}) = \underbrace {\mathbb {E} _ {q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta})} \left[ \log p (\mathcal {D} \mid \boldsymbol {\theta}) \right]} _ {\text {l i k e l i h o o d t e r m}} - \underbrace {\mathrm {D} _ {\mathrm {K L}} \left[ q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta}) \right\rVert p (\boldsymbol {\theta})} _ {\text {r e g u l a r i s a t i o n t e r m}}, \tag {2}
$$
where the likelihood term encourages the variational approximation to model the data well, and the regularisation term tries to keep the posterior close to the prior.
2.1 Implicit variational inference
In implicit VI (IVI), the variational distribution is only implicitly defined through its generative process which allows sampling from this distribution, but neither its density nor the gradients of its density can be evaluated. One way to draw from an implicit distribution is to push a low-dimensional noise vector through a neural network (Mohamed et al., 2016; Tran et al., 2017; Shi et al., 2018). These models have also been called amortised neural samplers (Kumar et al., 2019), with the generative process  $z \sim q(z)$  and  $\theta = g_{\gamma}(z)$  and the corresponding density
$$
q _ {\gamma} (\boldsymbol {\theta}) = \frac {\partial}{\partial \theta_ {1}} \frac {\partial}{\partial \theta_ {2}} \dots \frac {\partial}{\partial \theta_ {m}} \int_ {\boldsymbol {z} | g _ {\gamma} (\boldsymbol {z}) \leq \boldsymbol {\theta}} q (\boldsymbol {z}) \mathrm {d} \boldsymbol {z}, \tag {3}
$$
where  $q(z)$  is a fixed base distribution and  $g_{\gamma}:\mathbb{R}^d\to \mathbb{R}^m$  is a non-linear, typically non-invertible mapping due to which this integral is non-trivial to solve. The resulting intractability of the density precludes the usage of popular objectives in probabilistic modelling. The likelihood term from Eq. (2) and its gradients can be estimated using Monte Carlo and the reparameterization trick (Kingma et al., 2014). However, the regularisation term involves the entropy of  $q_{\gamma}$ ,  $H_{q}[q_{\gamma}]$ :
$$
\mathrm {D} _ {\mathrm {K L}} \left[ \right. q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta}) \left\| \right. p (\boldsymbol {\theta}) \left. \right] = \mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta})} \left[ \log \frac {q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta})}{p (\boldsymbol {\theta})} \right] = \underbrace {\mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta})} \left[ \log q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta}) \right]} _ {- H _ {q} [ q _ {\boldsymbol {\gamma}} ]} - \mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta})} \left[ \log p (\boldsymbol {\theta}) \right], (4)
$$
which following the intractability of the density Eq. (3) is not explicitly available. Moreover, critical to consider,  $q_{\gamma}$  is not a well-defined density in the parameter space, as its support lies on a low dimensional manifold and has measure zero. Consequently, the KL divergence in Eq. (1) is not always well-defined (Arjovsky et al., 2017) giving rise to another noteworthy issue with IVI. To tackle this issue, standard IVI approaches (Sugiyama et al., 2012; Huszár, 2017) rely on density ratio estimators based on a discriminator (in GANs) to estimate the regularisation term. Geng et al. (2021) give a tractable and differentiable lower bound on this entropy.
3 A deep latent variable model and its entropy
In this section, we address the issues noted above, I: The implicit density lies on a low dimensional manifold, leading to ill-defined KL (Eq. (1)) and II: The entropy of the implicit density and its gradients are intractable.
3.1 Defining a proper KL divergence
To make the KL in Eq. (1) well-defined, we add Gaussian noise to the output of our neural sampler  $g_{\gamma}$  which has the effect of making the implicit density continuous in the ambient space of  $\theta$ , see Arjovsky et al. (Lemma 1, 2017). Accordingly, we switch to the corresponding variational distribution – a Gaussian deep latent variable model (DLVM). This is a special case of the semi-implicit distribution
(Yin et al., 2018; Titsias et al., 2019) and equivalent to the generative model in a Gaussian variational autoencoder (VAE, Kingma et al., 2014; Rezende et al., 2014):
$$
q _ {\gamma} (\boldsymbol {\theta}) = \int q _ {\gamma} (\boldsymbol {\theta} \mid \boldsymbol {z}) q (\boldsymbol {z}) \mathrm {d} \boldsymbol {z} = \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} [ q _ {\gamma} (\boldsymbol {\theta} \mid \boldsymbol {z}) ], \quad \text {w h e r e}, \tag {5}
$$
$$
q _ {\boldsymbol {\gamma}} (\boldsymbol {\theta} \mid \boldsymbol {z}) = \mathcal {N} (\boldsymbol {\theta} \mid g _ {\boldsymbol {\gamma}} (\boldsymbol {z}), \sigma^ {2} \boldsymbol {I} _ {m}), \quad g _ {\boldsymbol {\gamma}}: \mathbb {R} ^ {d} \rightarrow \mathbb {R} ^ {m}, \tag {6}
$$
$\pmb{\theta} \in \mathbb{R}^{m}$  and the latent variable  $\pmb{z} \in \mathbb{R}^{d}$ . We assume a Gaussian base density  $q(z) = \mathcal{N}(z \mid \mathbf{0}, \mathbf{I}_{d})$ .  $g_{\gamma}$  is the neural sampler (or generator) and  $\sigma^2 \in \mathbb{R}^+$  is the fixed, homoscedastic variance of the output density, which we take to be small. Thus, each sample from the base density,  $\pmb{z}$ , results in a Gaussian output density,  $q_{\gamma}(\pmb{\theta} \mid \pmb{z})$ . And the marginalisation leads to a very flexible  $q_{\gamma}(\pmb{\theta})$  which informally, is an infinite mixture of Gaussians (example 2, Im et al., 2017). Generally, we do not have a closed form for  $q_{\gamma}(\pmb{\theta})$  and consequently  $H[q_{\gamma}]$ , due to the non-linear function  $g_{\gamma}$  within  $q_{\gamma}(\pmb{\theta} \mid \pmb{z})$  in Eq. (5), so we are still left with challenge II, which we address next.
3.2 Approximating the intractable entropy
With the DLVM, we can expand the entropy of the variational approximation in Eq. (4) as,
$$
H [ q _ {\gamma} (\boldsymbol {\theta}) ] = - \mathbb {E} _ {\boldsymbol {z} \sim q _ {(z)}} \mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\gamma} (\boldsymbol {\theta} | z)} [ \log q _ {\gamma} (\boldsymbol {\theta}) ]. \tag {7}
$$
In principle, a Monte Carlo estimator of Eq. (5) could be used to approximate the  $q_{\gamma}(\theta)$  term and by extension, this entropy, but this estimator is biased and has high variance. This variance can be reduced with directed sampling via an encoder, but to avoid training an additional network, we derive an encoder-free approximation in the following. To do this, we need to work around the non-linearity of  $g_{\gamma}$  that causes the intractability.
Linearisation of the generator First, to estimate  $q_{\gamma}(\pmb{\theta})$  during training with lower variance than a direct Monte Carlo estimate, we consider an analytical approximation of Eq. (7) obtained via a local linearisation of the generator/neural sampler around  $\pmb{z}'$ :
$$
T _ {z ^ {\prime}} ^ {1} (\boldsymbol {z}) = g _ {\gamma} \left(\boldsymbol {z} ^ {\prime}\right) + \boldsymbol {J} _ {g} \left(\boldsymbol {z} ^ {\prime}\right) \left(\boldsymbol {z} - \boldsymbol {z} ^ {\prime}\right), \tag {8}
$$
where  $J_{g}(z^{\prime}) \in \mathbb{R}^{m \times d}$  is the Jacobian of  $g_{\gamma}$  evaluated in  $z^{\prime}$ , which we assume exists. We can approximate  $g_{\gamma}(z)$  by the linear function  $T_{z^{\prime}}^{1}(z)$  when  $z$  is close to  $z^{\prime}$  and obtain a Gaussian approximation of the output density,  $q_{\gamma}(\theta | z) \approx \tilde{q}_{z^{\prime}}(\theta | z) = \mathcal{N}(\theta | T_{z^{\prime}}^{1}(z), \sigma^{2}\pmb{I}_{m})$  and substitute this into Eq. (5) to approximate  $q_{\gamma}(\pmb{\theta})$ :
$$
q _ {\gamma} (\boldsymbol {\theta}) = \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \left[ q _ {\gamma} (\boldsymbol {\theta} \mid \boldsymbol {z}) \right] \approx \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \left[ \tilde {q} _ {\boldsymbol {z} ^ {\prime}} (\boldsymbol {\theta} \mid \boldsymbol {z}) \right] \tag {9}
$$
$$
= \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \left[ \mathcal {N} \left(\boldsymbol {\theta} \mid g _ {\gamma} \left(\boldsymbol {z} ^ {\prime}\right) + \boldsymbol {J} _ {g} \left(\boldsymbol {z} ^ {\prime}\right) \left(\boldsymbol {z} - \boldsymbol {z} ^ {\prime}\right), \sigma^ {2} \boldsymbol {I} _ {m}\right) \right] \tag {10}
$$
Now the expectation in Eq. (10) can be analytically solved by integrating over the latent variable (Tipping et al., 1999) and is Gaussian, which gives the closed form of the approximation
$$
q _ {\gamma} (\boldsymbol {\theta}) \approx \mathcal {N} \left(\boldsymbol {\theta} \mid \mu_ {\gamma} \left(\boldsymbol {z} ^ {\prime}\right), C _ {\gamma} \left(\boldsymbol {z} ^ {\prime}\right)\right) =: \tilde {q} _ {\boldsymbol {z} ^ {\prime}} (\boldsymbol {\theta}), \tag {11}
$$
$$
\mu_ {\gamma} \left(\boldsymbol {z} ^ {\prime}\right) = g _ {\gamma} \left(\boldsymbol {z} ^ {\prime}\right) - \boldsymbol {J} _ {g} \left(\boldsymbol {z} ^ {\prime}\right) \boldsymbol {z} ^ {\prime}, \quad C \left(\boldsymbol {z} ^ {\prime}\right) = \boldsymbol {J} _ {g} \left(\boldsymbol {z} ^ {\prime}\right) \boldsymbol {J} _ {g} \left(\boldsymbol {z} ^ {\prime}\right) ^ {\intercal} + \sigma^ {2} \boldsymbol {I} _ {m}. \tag {12}
$$
Approximation of the differential entropy We use the above result to approximate the entropy of the DLVM,
$$
H \left[ q _ {\gamma} (\boldsymbol {\theta}) \right] = - \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\gamma} (\boldsymbol {\theta} \mid \boldsymbol {z})} [ \log q _ {\gamma} (\boldsymbol {\theta}) ] \approx - \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\gamma} (\boldsymbol {\theta} \mid \boldsymbol {z})} [ \log \tilde {q} _ {\boldsymbol {z} ^ {\prime} = \boldsymbol {z}} (\boldsymbol {\theta}) ], \tag {13}
$$
where  $\log \tilde{q}_{z' = z}(\theta)$  is Eq. (11) evaluated at the samples  $z$  from the outer expectation. Importantly, we do the linearisation of  $q_{\gamma}(\theta)$  around the latent value  $z$  that is used to sample each  $\theta$  in the expectation, which means that for small  $\sigma^2$ -values,  $z$  is close to  $z'$ , and the linear approximation will be good. This amounts to using a point-wise Gaussian to approximate the entropy at  $z$  in the training objective, whilst sampling from the DLVM noted in Eq. (5). Hence, important to consider here, that the linearisation does not restrict the expressivity of the posterior that we finally obtain.
We could, in principle, evaluate the approximation in Eq. (13) numerically, but it involves the matrix inverse of  $C(z')$ , which is computationally demanding. In Appendix A, we show that for small values of the output variance  $\sigma^2$ , we can further approximate the differential entropy by
$$
H [ q _ {\gamma} (\boldsymbol {\theta}) ] \approx \frac {1}{2} \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \left[ \log \det  \left(\boldsymbol {J} _ {g} (\boldsymbol {z}) \boldsymbol {J} _ {g} (\boldsymbol {z}) ^ {\top} + \sigma^ {2} \boldsymbol {I} _ {m}\right) \right] + \frac {m}{2} + \frac {m}{2} \log 2 \pi \tag {14}
$$
4 Linearised Implicit Variational Inference
Finally, we arrive at our novel approximation to the ELBO for IVI using Eq. (5) as the variational distribution. Using the entropy approximation from Eq. (14), we obtain the approximate ELBO denoted the Linearised Implicit Variational Inference (LIVI) bound
$$
\mathcal {L} ^ {\prime} (\gamma) = \mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\gamma} (\boldsymbol {\theta})} \left[ \log p (\mathcal {D} | \boldsymbol {\theta}) + \log p (\boldsymbol {\theta}) \right] + \frac {1}{2} \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \left[ \log \det  \left(\boldsymbol {J} _ {g} (\boldsymbol {z}) \boldsymbol {J} _ {g} (\boldsymbol {z}) ^ {\intercal} + \sigma^ {2} \boldsymbol {I} _ {m}\right) \right] + c, \tag {15}
$$
where  $c = \frac{m}{2} + \frac{m}{2} \log 2\pi$ . As the determinant of a matrix is the product of its singular values, the log determinant is the sum of the log of these singular values. If  $s_d(z) \geq \ldots \geq s_1(z)$  are the non-zero singular values of the Jacobian  $J_g(z)$ , we can write the log determinant term as
$$
\frac {1}{2} \log \det  (\boldsymbol {J} _ {g} (\boldsymbol {z}) \boldsymbol {J} _ {g} (\boldsymbol {z}) ^ {\intercal} + \sigma^ {2} \boldsymbol {I} _ {m}) = \frac {1}{2} \sum_ {i = 1} ^ {d} \log \left(s _ {i} ^ {2} (\boldsymbol {z}) + \sigma^ {2}\right) + \frac {m - d}{2} \log \sigma^ {2} \tag {16}
$$
To avoid calculating all the singular values of a large Jacobian matrix, we can follow Geng et al. (2021, Eq. 10) and lower-bound Eq. (16) as
$$
\frac {1}{2} \sum_ {i = 1} ^ {d} \log \left(s _ {i} ^ {2} (z) + \sigma^ {2}\right) + \frac {m - d}{2} \log \sigma^ {2} \geq \frac {d}{2} \log \left(s _ {1} ^ {2} (z) + \sigma^ {2}\right) + \frac {m - d}{2} \log \sigma^ {2} \tag {17}
$$
We defer the details of Eq. (16) and the algorithm (LOBPCG, Knyazev, 2001) for obtaining  $s_1$  to Appendix C. This allows us to extend our variational approximation to high-dimensional parameter spaces, and using Eq. (17) we obtain a lower bound on  $\mathcal{L}'(\gamma)$  given by
$$
\mathcal {L} ^ {\prime \prime} (\gamma) = \mathbb {E} _ {\boldsymbol {\theta} \sim q _ {\gamma} (\boldsymbol {\theta})} \left[ \log p (\mathcal {D} \mid \boldsymbol {\theta}) + \log p (\boldsymbol {\theta}) \right] + \mathbb {E} _ {\boldsymbol {z} \sim q (\boldsymbol {z})} \left[ \frac {d}{2} \log \left(s _ {1} ^ {2} (\boldsymbol {z}) + \sigma^ {2}\right) \right] + c. \tag {18}
$$
We denote  $\mathcal{L}'(\gamma)$  the approximate LIVI bound with whole Jacobians and  $\mathcal{L}''(\gamma)$  the LIVI bound with a differentiable lower bound on the determinant. We draw comparisons between how well these bounds estimate the entropy of a DLVM model in Appendix F.5, and write the reparameterised version of both bounds in Appendix B. The relationships between the bounds and the evidence are
$$
\log p (\mathcal {D}) \geq \mathcal {L} (\gamma) \approx \mathcal {L} ^ {\prime} (\gamma) \geq \mathcal {L} ^ {\prime \prime} (\gamma) \tag {19}
$$
Based on the compute budget, the two LIVI bounds offer an accuracy vs compute trade-off. In both cases, entropy maximisation promotes the exploration of the latent space. It is also worthwhile to note that Eq. (15), which uses the entropy approximation in Eq. (13), is fundamentally different from other linearisation approaches in the literature (e.g., Immer et al., 2021), as we do not linearise the Bayesian model. Instead, our approach linearises the generator parameterising  $q_{\gamma}(\boldsymbol{\theta} \mid \boldsymbol{z})$  only when needed to approximate the differential entropy, thus leaving the variational approximation unchanged.
The usage of a secondary model to generate parameters of a primary model first appeared in the form of hypernetworks (Ha et al., 2017). Our approach is probabilistic and is hence closer to Bayesian hypernetworks (Krueger et al., 2017) that mandate invertibility of the generator. The formulation is exactly like normalising flows and thereby sidesteps the complexities of estimating the entropy term. Flow-based approximations require particular focus on the design of the variational approximation to curb the dimensionality of the flow and the Jacobian matrices. Louizos et al. (2017) use an expressive flow on the multiplicative factors of the weights in each layer and not on all weights jointly. Our method does not necessitate invertibility, making it more general. We also share parallels with works in manifold learning (Brehmer et al., 2020; Caterini et al., 2021) in that they also use noninvertible mappings and have to employ a general form of the standard change of variable formula  $\log \det(J_g(z)^{\mathsf{T}}J_g(z))$  in the approximating density. Kristiadi et al. (2022) infers an expressive distribution over a subset of weights by using a post hoc Laplace approximation as the base distribution for a normalising flow. This approach, however, inherits the usual scalability limitations of flows.