# Generative Particle Variational Inference via Estimation of Functional
Gradients

Neale Ratzlaff  Qinxun Bai  Li Fuxin  Wei Xu

###### Abstract

Recently, particle-based variational inference (ParVI) methods have gained
interest because they can avoid arbitrary parametric assumptions that are
common in variational inference. However, many ParVI approaches do not allow
arbitrary sampling from the posterior, and the few that do allow such sampling
suffer from suboptimality. This work proposes a new method for learning to
approximately sample from the posterior distribution. We construct a neural
sampler that is trained with the functional gradient of the KL-divergence
between the empirical sampling distribution and the target distribution,
assuming the gradient resides within a reproducing kernel Hilbert space. Our
generative ParVI (GPVI) approach maintains the asymptotic performance of ParVI
methods while offering the flexibility of a generative sampler. Through
carefully constructed experiments, we show that GPVI outperforms previous
generative ParVI methods such as amortized SVGD, and is competitive with ParVI
as well as gold-standard approaches like Hamiltonian Monte Carlo for fitting
both exactly known and intractable target distributions.

Bayesian Deep Learning

##  1 Introduction

Bayesian inference provides a powerful framework for reasoning and prediction
under uncertainty. However, computing the posterior is tractable with only a
few parametric distributions, making wider applications of Bayesian inference
difficult. Traditionally, MCMC and variational inference methods are utilized
to provide tractable approximate inference, but these approaches face their
own difficulties if the dimensionality of the space is extremely high. For
example, a recent case of interest is Bayesian neural networks (BNNs), which
applies Bayesian inference to deep neural network training in order to provide
a principled way to assess model uncertainty. The goal in this regime is to
model the posterior of every parameter in all the weight tensors from every
layer of a deep network. However, developing efficient computational
techniques for approximating this intractable posterior with extremely high
dimensionality remains challenging.

Recently, particle-based variational inference (ParVI) methods (Liu & Wang,
2016; Liu et al., 2019; Liu, 2017) have been proposed to represent the
variational distribution by a set of particles and update them through a
deterministic optimization process to approximate the posterior. While
achieving both asymptotic accuracy and computational efficiency, ParVI methods
are restricted by the fixed number of particles and lack the ability of
drawing new samples beyond the initial set of particles. To address this
issue, amortized ParVI methods (Wang & Liu, 2016) have been proposed to
amortize the ParVI gradients in training a neural sampler. While being
flexible in drawing samples, in Sec. 4 we show that amortized ParVI methods
cannot match the convergence behavior of ParVI methods.

In this work, we propose a generative particle variational inference (GPVI)
approach that addresses those issues. GPVI trains a neural sampler network by
directly estimating the functional gradient w.r.t. the KL-divergence between
the distribution of generated particles and the target distribution, and pulls
it back to update the neural sampler. As such it allows the neural sampler to
directly generate particles that match the posterior distribution, hence
achieving the asymptotic accuracy and computational efficiency of ParVI
methods. In figure 1, we show that the predictive distribution of 1D
regression functions sampled from GPVI nearly matches that from the ParVI
solution, while amortized ParVI fails.

The main computational challenge lies in a reliable estimate of the functional
gradient that involves the inverse of the input-output Jacobian of the neural
sampler. Instead of directly computing this term and paying a high
computational cost, we introduce a helper network to estimate the inverse
Jacobian vector product and train the helper network via gradient descent. By
alternating between this gradient step and the gradient update of the sampler
network, the computational cost is distributed over the whole training
procedure.

In experiments, our proposed approach achieves comparable convergence
performance as ParVI methods. It is considerably superior than that of
amortized ParVI methods, while still allowing efficient sampling from the
posterior. By directly applying our approach as a hypernetwork to generate
BNNs, we achieve competitive performance regarding uncertainty estimation.

In summary, our contributions are three-fold,

  * •

We propose GPVI, a new variational inference approach that trains a neural
sampler to generate particles from any posterior distribution. GPVI estimates
the functional gradient and uses it to update the neural sampler. It enjoys
the asymptotic accuracy and computational efficiency of ParVI methods.
Comparing with existing amortized ParVI methods, our approach enjoys the same
efficiency and flexibility while showing considerable advantage in convergence
behavior.

  * •

We design careful techniques for efficient gradient estimates that address the
challenges in approximating the product between the inverse of the Jacobian
and a vector.

  * •

We apply our approach to BNNs and achieve competitive uncertainty estimation
quality for deep neural networks.

(a) HMC

(b) SVGD

(c) GPVI

(d) Amortized SVGD

Figure 1: Predictive uncertainty of methods for a 1-D regression task. (a) HMC
predictive posterior matches the uncertainty in the data; (b) SVGD performs
comparably to HMC; (c) Our proposed GPVI performs similarly to SVGD, with the
additional capability of sampling new particles during inference; (d)
Amortized SVGD overestimates the uncertainty when the data is sparse.

##  2 Related Work

ParVI is a recent class of non-parametric variational inference methods. (Liu
& Wang, 2016; Liu, 2017) proposed stein variational gradient descent (SVGD),
which deterministically updates an empirical distribution of particles toward
the target distribution via a series of sequentially constructed smooth maps.
Each map defines a perturbation of the particles in the direction of steepest
descent towards the target distribution under the KL-divergence metric. Liu et
al. (2019) cast ParVI as a gradient flow on the space of probability measures
$\mathcal{P}_{2}$ equipped with the $W_{2}$ Wasserstein metric, and proposed
ParVI methods GFSD and GFSF by smoothing the density and the function
respectively. A potential limitation of SVGD and GFSF is the restriction of
the transport maps to an RKHS. The fisher neural sampler (Hu et al., 2018;
Grathwohl et al., 2020) lifts the function space of transportation maps to the
space of $\mathcal{L}_{2}$ functions, and performs a minimax optimization
procedure to find the optimal perturbation. As a non-parametric approach,
ParVI does not depend on the specific parameterization of the approximate
posterior. ParVI is also more efficient than MCMC methods due to its
deterministic updates and gradient-based optimization.

The limitation of ParVI is the inability to draw additional samples beyond the
initial set. To address this limitation, our approach allows for arbitrary
draws of new samples once the generator is trained. The most notable
generative ParVI approach prior to our work was amortized SVGD (Wang & Liu,
2016), which proposed using a generative model to approximately sample from
the SVGD particle distribution. Amortized SVGD performs a one-step update that
backpropagates the SVGD gradient back to the generator parameters, which is
often not optimal. A more detailed discussion of the differences between our
method and GPVI is presented in Sec. 3.5.

Performing direct Bayesian inference in the weight-space of neural networks
(Neal, 2012; Buntine & Weigend, 1991; MacKay, 1992) is known to be intractable
in the context of deep learning. Thus, approximations have been formulated
using Langevin Dynamics (Welling & Teh, 2011; Korattikara Balan et al., 2015),
Monte Carlo dropout (Gal & Ghahramani, 2016; Kingma et al., 2015), and
variational inference with mean-field Gaussian priors (Graves, 2011). The
variational approximation has been iterated upon extensively. Bayes by
Backprop (Blundell et al., 2015) proposed an unbiased gradient estimator that
allowed for training of deep Bayesian neural networks using the local
reparameterization trick from Kingma et al. (2015). However, fully factorized
Gaussian approximations are insufficient to capture the complex structure of a
high dimensional posterior (Yao et al., 2019). As a result, structured
approximate posteriors have been proposed using matrix Gaussians (Louizos &
Welling, 2016), normalizing flows (Louizos & Welling, 2017; Krueger et al.,
2017), and hypernetworks (Pawlowski et al., 2017).

Evaluating the performance of BNNs is difficult, as there is often no access
to the ground truth posterior. Therefore, prior work often considers a “gold-
standard” method such as HMC (Neal et al., 2011), or exact Gaussian processes
(Rasmussen, 2003) as a proxy for the true posterior. While these methods set a
high bar for BNNs, they are in general not scalable to large-scale problems.
Nonetheless there were attempts to bring HMC (Strathmann et al., 2015; Chen et
al., 2014) and GPs (Wilson et al., 2016; Cheng & Boots, 2017; Wang et al.,
2019; Hensman et al., 2013, 2015) in the realm of tractability for high
dimensional problems.

Variational inference via smooth transport maps is also used in normalizing
flows (Rezende & Mohamed, 2015). But the transport maps used in normalizing
flows and in Marzouk et al. (2016) are bijections which makes them unsuitable
for learning distributions over neural network functions.

##  3 Generative Particle Variational Inference

We want to learn a parametric generator:
$\boldsymbol{f}_{\boldsymbol{\theta}}:\mathcal{Z}\to\mathbb{R}^{d}$,
parameterized by $\boldsymbol{\theta}$, where
$\mathcal{Z}\subset\mathbb{R}^{d}$ is the convex space of input noise and
$\boldsymbol{x}=\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z})$
generates a sample $\boldsymbol{x}$ from an input noise $\boldsymbol{z}$. Let
$q(\boldsymbol{x})$ represents the implicit distribution of samples generated
by $\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z})$, where
$\boldsymbol{z}\sim N(\boldsymbol{0},I_{d})$. Let $p(\boldsymbol{x})$ be the
target distribution, we want to solve for
$\boldsymbol{f}_{\boldsymbol{\theta}}$ that minimizes the objective
$\text{KL}\left(q(\boldsymbol{x})\|p(\boldsymbol{x}\right))$.

Our approach treats this problem from a functional optimization perspective,
by first computing the functional gradient of the objective, and then pulling
it back to parameter space through the function parameterization. In the rest
of this section we introduce our algorithm in detail and compare it with
amortized ParVI approaches.

###  3.1 Functional gradient and its pullback

Let
$\mathcal{J}(\boldsymbol{f})=\text{KL}\left(q(\boldsymbol{x})\|p(\boldsymbol{x})\right)$
be the objective, where $\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})$. If
$\boldsymbol{f}$ is injective, by change of variables for probability measure,

$q(\boldsymbol{x})=\frac{p_{\boldsymbol{z}}(\boldsymbol{z})}{\left|\det\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)\right|},$
(1)

where $p_{\boldsymbol{z}}(\boldsymbol{z})$ is the distribution from which
$\boldsymbol{z}$ is sampled. The minimization objective becomes,

$\mathcal{J}(\boldsymbol{f})=\mathbf{E}_{\boldsymbol{z}}\left[-\log
p(\boldsymbol{f}(\boldsymbol{z}))+\log\frac{p_{\boldsymbol{z}}(\boldsymbol{z})}{\left|\det\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)\right|}\right].$
(2)

Consider some function approximation of $\boldsymbol{f}$, say
$\boldsymbol{f}=\boldsymbol{f}_{\boldsymbol{\theta}}$, then the minimization
objective becomes,

$\mathcal{J}(\boldsymbol{\theta})=\mathbf{E}_{\boldsymbol{z}}\left[-\log
p(\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}))+\log\frac{p_{\boldsymbol{z}}(\boldsymbol{z})}{\left|\det\left(\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}}\right)\right|}\right].$
(3)

Directly computing the gradient of $\mathcal{J}(\boldsymbol{\theta})$ w.r.t.
$\boldsymbol{\theta}$ involves not only an inverse of the Jacobian
$\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)^{-1}$, but
also second derivatives of $\boldsymbol{f}_{\boldsymbol{\theta}}$ (see
Appendix A for details), which is overly expensive to compute in practice.

We propose, instead, to first compute the functional gradient of (2) w.r.t.
$\boldsymbol{f}$, i.e., $\nabla_{\boldsymbol{f}}\mathcal{J}(\boldsymbol{f})$,
and then back-propagate it through the generator to get the gradient w.r.t.
$\boldsymbol{\theta}$, i.e.,

$\nabla_{\boldsymbol{\theta}}\mathcal{J}=\mathbf{E}_{\boldsymbol{z}}\left[\frac{\partial\boldsymbol{f}(\boldsymbol{z})}{\partial\boldsymbol{\theta}}\nabla_{\boldsymbol{f}}\mathcal{J}(\boldsymbol{f})(\boldsymbol{z})\right].$
(4)

The following theorem gives an explicit formula for computing the functional
gradient $\nabla_{\boldsymbol{f}}\mathcal{J}(\boldsymbol{f})$ when
$\boldsymbol{f}$ is chosen from a Reproducing kernel Hilbert space (RKHS). The
proof is provided in Appendix A.

######  Theorem 3.1.

Let $\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})$, where $\boldsymbol{z}\sim
p_{\boldsymbol{z}}(\boldsymbol{z})$, vector function
$\boldsymbol{f}=(f^{1},\ldots,f^{d})\in\mathcal{H}^{d}$ with
$f^{i}\in\mathcal{H}$, where $\mathcal{H}$ is the RKHS with kernel
$k(\cdot,\cdot)$, $\mathcal{H}^{d}$ is equipped with inner product
$\langle\boldsymbol{f},\boldsymbol{g}\rangle_{\mathcal{H}^{d}}=\sum^{d}_{i=1}\langle
f^{i},g^{i}\rangle_{\mathcal{H}}$. For $\mathcal{J}(\boldsymbol{f})$ well-
defined by (2), we have

$\displaystyle\nabla_{\boldsymbol{f}}\mathcal{J}(\boldsymbol{f})(\boldsymbol{z})=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}^{\prime}}\bigg{[}-\nabla_{\boldsymbol{x}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z}^{\prime})}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$
(5)
$\displaystyle-\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}^{\prime}}\right)^{-1}\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})\bigg{]}.$

###  3.2 Reparameterization of $\boldsymbol{f}$

There are two considerations for reparameterizing $\boldsymbol{f}$. Firstly,
in order for the inverse
$\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}^{\prime}}\right)^{-1}$
in (5) to be well-defined, the Jacobian
$\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}^{\prime}}$ should be a
square matrix, i.e., input noise $\boldsymbol{z}$ should have the same
dimension as the output $\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})$. In
practice, especially those applications involving BNNs, each $\boldsymbol{x}$
represents parameters of a sampled neural network and therefore can be
extremely high dimensional. As a result, a high dimensional $\boldsymbol{f}$
can be computationally prohibitive. Secondly, in order for the change of
variables density formula (1) to hold, $\boldsymbol{f}$ needs to be injective,
which is in general not guaranteed for an arbitrary neural network function.

To overcome the above two concerns, we consider the following
parameterization,

$\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z})=\boldsymbol{g}_{\boldsymbol{\theta}}\left(\boldsymbol{z}^{(:k)}\right)+\lambda\boldsymbol{z},\quad\forall\boldsymbol{z}\in\mathbb{R}^{d},$
(6)

where $\boldsymbol{z}^{(:k)}\in\mathbb{R}^{k}$ denotes the vector consisting
of the first $k$ components of $\boldsymbol{z}$, and
$\boldsymbol{g}_{\boldsymbol{\theta}}:\mathbb{R}^{k}\to\mathbb{R}^{d}$ with
parameters $\boldsymbol{\theta}$ is a much slimmer neural network. In our
experiments, $\boldsymbol{g}_{\boldsymbol{\theta}}$ is designed with an input
dimension $k$ less than $30\%$ the size of $d$. For high dimensional open-
category experiments where $d>60,000$, we use a $k$ of less than $2\%$ of $d$.
For $\boldsymbol{f}_{\boldsymbol{\theta}}$ defined by (21)), the Jacobian is,

$\left[\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}}\right]_{d\times
d}=\left[\left[\frac{\partial\boldsymbol{g}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}^{(:k)}}\right]_{d\times
k}\bigg{|}\boldsymbol{0}_{d\times(d-k)}\right]_{d\times
d}+\lambda\boldsymbol{I}_{d},$ (7)

where $\lambda$ is a hyper-parameter. Note that for sufficiently large
$\lambda$, the Jacobian defined by (22) is positive definite. Since the domain
$\mathcal{Z}$ of $\boldsymbol{z}$ is convex, it is straightforward to show
that $\boldsymbol{f}_{\boldsymbol{\theta}}$ defined by (21) is injective. We
include the proof in Appendix A.

In practice we set $\lambda$ to be $1.0$ and find it sufficient throughout our
experiments.

###  3.3 Estimating the Jacobian inverse

The main computational challenge of (5) lies in computing the term

$\left(J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})\right)^{-1}\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z}),$
(8)

where
$J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})=\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}^{\prime}},$
especially considering that we need an efficient implementation for batched
$\boldsymbol{z}$ and $\boldsymbol{z}^{\prime}$.

Directly evaluating and storing the full Jacobian
$J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})$ for each
$\boldsymbol{z}^{\prime}$ of the sampled batch is not acceptable from the
standpoint of either time or memory consumption. There exist iterative methods
for solving the linear equation system
$\boldsymbol{y}=\left(J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})\right)^{-1}\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$
(Young, 1954; Fletcher, 1976), which involves computing the vector-Jacobian
product
$J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$
at each iteration. By alternating between this iterative solver and the
gradient update (4), it is possible to get a computationally amenable
algorithm. However, as shown in the Appendix C, such an algorithm does not
converge to the target distribution $p$ even for a simple Bayesian linear
regression task. This is due to the fact that batches of both $\boldsymbol{z}$
and $\boldsymbol{z}^{\prime}$ for evaluating (5) need to be re-sampled for
each gradient update (4) to avoid the cumulative sampling error. Therefore,
the above alternating procedure of iterative solver for
$\boldsymbol{y}=\left(J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})\right)^{-1}\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$
ends up shooting a moving target for different batches of $\boldsymbol{z}$ and
$\boldsymbol{z}^{\prime}$ at each iterate, which is difficult for convergence.

To overcome this computational challenge, we propose a helper network, denoted
by
$\boldsymbol{h}_{\boldsymbol{\eta}}(\boldsymbol{z}^{\prime},\nabla_{\boldsymbol{z}^{\prime}}k)$,
and parameterized by $\boldsymbol{\eta}$, that consumes both
$\boldsymbol{z}^{\prime}$ and
$\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$
and predicts
$\left(J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})\right)^{-1}\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$.
With the helper network, the functional gradient (5) can be computed by,

$\displaystyle\nabla_{\boldsymbol{f}}\mathcal{J}(\boldsymbol{f})(\boldsymbol{z})=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}^{\prime}}\bigg{[}-\nabla_{\boldsymbol{x}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z}^{\prime})}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$
(9)
$\displaystyle-\boldsymbol{h}_{\boldsymbol{\eta}}\left(\boldsymbol{z}^{\prime},\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})\right)\bigg{]}.$

We use the following loss to train the helper network,

$\mathcal{L}(\boldsymbol{\eta})=\|J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})\boldsymbol{h}_{\boldsymbol{\eta}}(\boldsymbol{z}^{\prime},\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z}))-\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})\|^{2},$
(10)

where
$J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})\boldsymbol{h}_{\boldsymbol{\eta}}$
can be computed by the following formula given the reparameterization (21) of
$\boldsymbol{f}_{\boldsymbol{\theta}}$,

$\boldsymbol{h}^{T}_{\boldsymbol{\eta}}\left(\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}^{\prime}}\right)=\boldsymbol{h}^{T}_{\boldsymbol{\eta}}\left(\frac{\partial\boldsymbol{g}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}^{\prime(:k)}}\right)+\lambda\boldsymbol{h}^{T}_{\boldsymbol{\eta}},$
(11)

where the Vector-Jacobian Product (VJP)
$\boldsymbol{h}^{T}_{\boldsymbol{\eta}}\left(\frac{\partial\boldsymbol{g}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}^{\prime(:k)}}\right)$
can be computed by one backward pass of the function
$\boldsymbol{g}_{\boldsymbol{\theta}}$.

###  3.4 Summary of the algorithm

Both $\boldsymbol{g}_{\boldsymbol{\theta}}$ and
$\boldsymbol{h}_{\boldsymbol{\eta}}$ can be trained with stochastic gradient
descent (SGD), and our algorithm alternates between the SGD updates of
$\boldsymbol{g}_{\boldsymbol{\theta}}$ and
$\boldsymbol{h}_{\boldsymbol{\eta}}$. Note that the helper network
$\boldsymbol{h}_{\boldsymbol{\eta}}$ only has to chase the update of
$\boldsymbol{g}_{\boldsymbol{\theta}}$, but no more extra moving targets due
to re-sampling of $\boldsymbol{z}$ and $\boldsymbol{z}^{\prime}$. Our
experiments in Sec. 4 and the appendix show that the helper network is able to
efficiently approximate (8) which enables the convergence of the gradient
update (4). Our overall Generative Particle VI (GPVI) algorithm is summarized
in Algorithm 1.

Initialize generator $\boldsymbol{g}_{\boldsymbol{\theta}}$, helper
$\boldsymbol{h}_{\eta}$, and learning rate $\epsilon$

while  _Not converge_ do

1\. sample two batches
$\\{\boldsymbol{z}_{i}\\},\\{\boldsymbol{z}^{\prime}_{i}\\}\sim
N(\boldsymbol{0},I_{d})$

2\. compute $k(\boldsymbol{z}^{\prime},\boldsymbol{z})$ and
$\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$

3\. forward $\boldsymbol{g}_{\boldsymbol{\theta}}$ to compute
$\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z})$ and
$\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}^{\prime})$ by (21)

4\. forward $\boldsymbol{h}_{\boldsymbol{\eta}}$ to compute
$\boldsymbol{h}_{\boldsymbol{\eta}}(\boldsymbol{z}^{\prime},\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z}))$

5\. backward $\boldsymbol{g}_{\boldsymbol{\theta}}$ to compute the VJP
$\boldsymbol{h}^{T}_{\boldsymbol{\eta}}\left(\frac{\partial\boldsymbol{g}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}^{\prime(:k)}}\right)$

lor and then construct
$J_{\boldsymbol{f}}(\boldsymbol{z}^{\prime})\boldsymbol{h}_{\boldsymbol{\eta}}$
by (11),

6\. update $\boldsymbol{h}_{\boldsymbol{\eta}}$ by
$\boldsymbol{\eta}\leftarrow\boldsymbol{\eta}-\epsilon\nabla_{\boldsymbol{\eta}}\mathcal{L}$,

lor where $\nabla_{\boldsymbol{\eta}}\mathcal{L}$ is computed by back-
propagating (10)

7\. compute the functional gradient by (9)

8\. update $\boldsymbol{\theta}$ by
$\boldsymbol{\theta}\leftarrow\boldsymbol{\theta}-\epsilon\nabla_{\boldsymbol{\theta}}\mathcal{J}$,

lor where $\nabla_{\boldsymbol{\theta}}\mathcal{J}$ is computed by (4)

end while

Algorithm 1 Generative Particle VI (GPVI)

###  3.5 Comparison with Amortized SVGD

Stein variational gradient descent (SVGD) represents $q(\boldsymbol{x})$ by a
set of particles $\\{\boldsymbol{x}_{i}\\}^{n}_{i=1}$, which are updated
iteratively by,

$\boldsymbol{x}_{i}\leftarrow\boldsymbol{x}_{i}+\epsilon\boldsymbol{\phi}^{*}(\boldsymbol{x_{i}}),$
(12)

where $\epsilon$ is a step size and
$\boldsymbol{\phi}^{*}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is a vector field
(perturbation) on the space of particles that corresponds to the optimal
direction to perturb particles, i.e.,

$\boldsymbol{\phi}^{*}=\operatorname*{arg\,min}_{\boldsymbol{\phi}\in\mathcal{F}}\left\\{\frac{d}{d\epsilon}\text{KL}(q_{[\epsilon\boldsymbol{\phi}]}(\boldsymbol{x})\|p(\boldsymbol{x}))\bigg{|}_{\epsilon=0}\right\\},$

where $q_{[\epsilon\boldsymbol{\phi}]}(\boldsymbol{x})$ denotes the density of
particles updated by (12) using the perturbation $\boldsymbol{\phi}$, where
the density of original particles is $q(x)$. When $\mathcal{F}$ is chosen to
be the unit ball of some RKHS $\mathcal{H}$ with kernel function
$k(\cdot,\cdot)$, SVGD gives the following closed form solution for
$\boldsymbol{\phi}^{*}$,

$\boldsymbol{\phi}^{*}(\boldsymbol{x})=\mathbf{E}_{\boldsymbol{x}^{\prime}\sim
q}\left[\nabla_{\boldsymbol{x}^{\prime}}\log
p(\boldsymbol{x}^{\prime})k(\boldsymbol{x}^{\prime},\boldsymbol{x})+\nabla_{\boldsymbol{x}^{\prime}}k(\boldsymbol{x}^{\prime},\boldsymbol{x})\right].$
(13)

To turn SVGD into a neural sampler, Amortized SVGD (Wang & Liu, 2016) first
samples particles from a generator and then back-propagates the particle
gradients (13) through the generator to update the generator parameters. Let
$\boldsymbol{x}=\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}),\
\boldsymbol{z}\sim N(\boldsymbol{0},I_{d})$ be the particle generating
process, where $\boldsymbol{f}_{\boldsymbol{\theta}}$ is the generator
parameterized by $\boldsymbol{\theta}$, amortized SVGD updates
$\boldsymbol{\theta}$ by,

$\boldsymbol{\theta}\leftarrow\boldsymbol{\theta}+\epsilon\sum_{i=1}^{m}\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}_{i})}{\partial\boldsymbol{\theta}}\boldsymbol{\phi}^{*}(\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}_{i})),$
(14)

where
$\boldsymbol{\phi}^{*}(\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}_{i}))$
is computed by (13).

The following Lemma gives an explicit view regarding what functional gradient
amortized SVGD back-propagates through the generator. The proof is given in
the Appendix A.

######  Lemma 3.2.

If particles are generated by $\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z}),\
\boldsymbol{z}\sim p_{\boldsymbol{z}}(\boldsymbol{z})$, Eqn. (13) is the
functional gradient of the KL objective w.r.t. the perturbation function
$\boldsymbol{\phi}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ applied on the output
space of $\boldsymbol{f}$, i.e.,

$\nabla_{\boldsymbol{\phi}}\bigg{|}_{\boldsymbol{\phi}=\text{id}}\text{KL}\left(q(\boldsymbol{x})\|p(\boldsymbol{x})\right)=-\boldsymbol{\phi}^{*},$
(15)

where $\boldsymbol{x}=\boldsymbol{\phi}(\boldsymbol{f}(\boldsymbol{z}))$ and
$\boldsymbol{\phi}=(\phi^{1},\ldots,\phi^{d})\in\mathcal{H}^{d}$ with
$\phi^{i}\in\mathcal{H}$, where $\mathcal{H}$ is the RKHS with kernel
$k(\cdot,\cdot)$, $\mathcal{H}^{d}$ is equipped with inner product
$\langle\boldsymbol{\phi},\boldsymbol{\xi}\rangle_{\mathcal{H}^{d}}=\sum^{d}_{i=1}\langle\phi^{i},\xi^{i}\rangle_{\mathcal{H}}$.

To see the difference between our GPVI and amortized SVGD, both methods learn
a generator $\boldsymbol{f}$ that generates particles $\boldsymbol{x}$ from
input noise $\boldsymbol{z}$. To update $\boldsymbol{f}$, gradients of the KL
objective $\mathcal{J}(\boldsymbol{f})$ have to be backpropagated to
$\boldsymbol{f}$ in both cases. Here, the best is to use the steepest descent
direction for $\boldsymbol{f}$, which GPVI does (proved in Thm 3.1). As Lemma
3.2 shows, the SVGD gradient $\boldsymbol{\phi}^{*}$ given by (13) is optimal
for a transformation $\boldsymbol{\phi}$ applied to the set of current
particles. In in both amortized SVGD and GPVI, such a set is not kept and
$\boldsymbol{\phi}$ does not exist since they only maintain and update
$\boldsymbol{f}$. But amortized SVGD directly uses
$\boldsymbol{\phi}^{*}(\boldsymbol{f})$ in (14) as the direction to update
$\boldsymbol{f}$, which is an unproven use – it only coincides with the
steepest descent direction for $\boldsymbol{f}$ used by GPVI in special cases
such as $\boldsymbol{f}=\text{id}$, where (5) and (13) are equivalent because
$\boldsymbol{z}=\boldsymbol{x}$ and the Jacobian is identity. Regarding the
definition of RKHS, GPVI’s RKHS approximates the tangent space of
$\boldsymbol{f}$, it is therefore naturally defined on $\boldsymbol{z}$.
(Amortized) SVGD uses RKHS to approximate the tangent space of
$\boldsymbol{\phi}$, which is naturally defined on $\boldsymbol{x}$.

The amortizing step applied in both methods, i.e., back-propagating some
functional gradient to update the generator parameters, is in general not
guaranteed to keep the original descent direction of the functional gradient.
However, GPVI amortizes the steepest descent direction which is optimal in
first order sense, while amortized SVGD amortizes a non-steepest descent
direction, which is more likely to result in a non-descending direction after
amortizing. In practice, as shown in the next section, our approach
consistently outperforms amortized SVGD in approximating the target
distribution and capturing model uncertainty.

##  4 Experiments

To demonstrate the effectiveness of our approach for approximate Bayesian
inference, we evaluated GPVI in two different settings: density estimation and
BNNs. In our density estimation experiment, we trained our generator to draw
samples from a target distribution, showing that GPVI can accurately fit the
posterior from the data. For our BNN experiments, we evaluated on regression,
classification, and high dimensional open-category tasks. Our experiments show
that among all methods compared, GPVI is the only method to excel in both
sampling efficiency and asymptotic performance. We compare GPVI with ParVI
methods: SVGD (Liu & Wang, 2016), GFSF (Liu et al., 2019), and KSD (Hu et al.,
2018; Grathwohl et al., 2020), as well as their corresponding amortized
versions. We also compared with a mean-field VI approach Bayes by Backprop
(MF-VI) (Blundell et al., 2015), as well as deep ensembles (Lakshminarayanan
et al., 2017), and HMC (Neal et al., 2011). We emphasize that our aim is not
to only maximize the likelihood of generated samples, rather we want to
closely approximate the posterior of parameterized functions given the data.
Thus, predictions w.r.t data unseen during training should reflect the
epistemic uncertainty of the model.

The details of our experimental setup are as follows. In our BNN experiments,
we parameterized samples from the target distribution as neural networks with
a fixed architecture. In the low dimensional regression and classification
settings, we drew 100 samples from the approximate posterior for both training
and evaluation, allowing us to evaluate the predictive mean and variance. For
methods without a sampler e.g. ParVI and deep ensembles, we initialized a 100
member ensemble. In the high dimensional open-category tasks, we instead used
10 samples due to the larger computation cost. For methods that utilize a
hypernetwork such as GPVI and amortized ParVI, we used Gaussian input noise
$z\sim\mathcal{N}(0,I)$ and varied the hypernetwork architecture depending on
the task. For ParVI and deep ensembles we randomly initialized each member of
the ensemble. In all methods except HMC we used the Adam optimizer (Kingma &
Ba, 2014). A detailed description of network architectures, chosen
hyperparameters, and experimental settings is given in the appendix.

###  4.1 Density Estimation

We first evaluate the ability of generative approaches to fit a target
distribution from data. We used 2D and 5D zero mean Gaussian distributions
with non-diagonal covariances as our target distributions. We consider sampler
networks with one hidden layer of width 2 and 5, respectively, and input noise
with the same dimensionality as the output. The error is defined as the
difference between the true variance and the variance of the sampler output
distribution $WW^{T}-\Sigma$, where $W$ is the weights of the sampler. As
shown in Table 1, while the estimation problem becomes harder with increased
dimension, our approach performs consistently well while amortized ParVI
methods suffer more or less from a performance drop and all end up inferior to
our approach. We have included more density estimation results in the
Appendix.

Method |  $\Sigma$ error (2d) $\downarrow$ |  $\Sigma$ error (5d) $\downarrow$
---|---|---
Amortized SVGD | 0.10 $\pm$ .09 | 0.37 $\pm$ .32
Amortized GFSF | 0.18 $\pm$ .04 | 0.21 $\pm$ .13
Amortized KSD | 0.28 $\pm$ .32 | 1.68 $\pm$ .52
GPVI | 0.14 $\pm$ .08 | 0.14 $\pm$ .04
Table 1: Comparison of generative Particle VI approaches for density
estimation of 2d and 5d Gaussian distributions.

###  4.2 Bayesian Linear Regression

We evaluated all methods on Bayesian linear regression to investigate how well
each method can fit a unimodal normal distribution over linear function
weights. The target function is a linear regressor with parameter vector
$\boldsymbol{\beta}\in\mathbb{R}^{d}$, i.e.,
$\boldsymbol{y}=\boldsymbol{X}\boldsymbol{\beta}+\epsilon$, where
$\boldsymbol{X}\sim\mathcal{N}(0,I_{d}),\ \epsilon\sim\mathcal{N}(0,I)$, and
$\beta^{i}\sim U(0,1)+5$. We chose such linear Gaussian settings where we can
explicitly compute the target posterior
$p(\boldsymbol{\beta}|\boldsymbol{X},\boldsymbol{y})\sim\mathcal{N}((\boldsymbol{X}^{T}\boldsymbol{X})^{-1}\boldsymbol{X}^{T}\boldsymbol{y},\boldsymbol{X}^{T}\boldsymbol{X})$
given observations $(\boldsymbol{X},\boldsymbol{y})$, allowing us to
numerically evaluate how well each method fits the target distribution. Each
method was trained to regress $\boldsymbol{y}$ from $\boldsymbol{X}$. For GPVI
and amortized ParVI methods that make use of a hypernetwork, we used a linear
generator with input noise $\boldsymbol{z}\in\mathbb{R}^{d}$. In this setting,
the output of a generator with bias $\boldsymbol{b}$ and weights $W$ follows
the distribution $\mathcal{N}(\boldsymbol{b},WW^{T})$, which should match the
posterior $p(\boldsymbol{\beta}|\boldsymbol{X},\boldsymbol{y})$. For ParVI
methods and deep ensembles, we initialized an ensemble of linear regressors
with $d$ parameters. We computed the mean and covariance of the learned
parameters to measure the quality of posterior fit. For MF-VI we chose a
standard normal prior on the weights, and likewise computed the mean and
covariance of weight samples to measure posterior fit. We also test a variant
of GPVI where we exactly compute the Jacobian inverse instead of using our
helper network. We set $d=3$ in our experiments. While simple, this
quantitative sanity check is crucial before evaluating on more complicated
domains where the true posterior is not available. It will be evident in later
experiments that the ability to closely approximate the true posterior under
this simple setup is indicative of performance on higher dimensional tasks.

Method |  $\mu$ error $\downarrow$ |  $\Sigma$ error $\downarrow$
---|---|---
SVGD | 0.006 $\pm$ .0024 |  0.125 $\pm$ .03
GFSF | 0.003 $\pm$ .0014 | 0.139 $\pm$ .07
KSD | 0.009 $\pm$ .0006 | 0.373 $\pm$ .11
Amortized SVGD |  0.002 $\pm$ .0009 | 0.158 $\pm$ .04
Amortized GFSF |  0.002 $\pm$ .0003 | 0.209 $\pm$ .05
Amortized KSD | 0.004 $\pm$ .0007 | 0.430 $\pm$ .10
MF-VI | 0.004 $\pm$ .0035 | 0.303 $\pm$ .04
Deep Ensemble | 0.004 $\pm$ .0002 | 1.0 $\pm$ 0
HMC | 0.009 $\pm$ .0003 | 0.181 $\pm$ .05
GPVI |  0.002 $\pm$ .0007 |  0.128 $\pm$ .04
GPVI Exact Jac |  0.002 $\pm$ .0004 |  0.106 $\pm$ .07
Table 2: Bayesian linear regression. Reported error is the $L_{2}$ norm of the
difference between the learned mean and covariance parameters, and the ground
truth after 50000 iterations.

Results are shown in table 2. Our GPVI outperforms amortized ParVI, MF-VI, and
surprisingly even HMC. GPVI is also very competitive with the best ParVI
methods, SVGD and GFSF. We note that GPVI is also competitive with the “Exact-
Jac” variant, where we invert the Jacobian explicitly, instead of using the
helper network. This “Exact-Jac” version of GPVI is not scalable, but it
serves as a benchmark to show that our helper network is able to closely
approximate the exact solution. Finally, as deep ensembles lacks any mechanism
for correctly estimating the covariance, it performs poorly as expected.

###  4.3 Multimodal Classification

The previous experiment possesses an analytically known target distribution.
However, most problems of interest involve distributions without closed-form
expressions. Therefore, we further tested on a 2-dimensional, 4-class
classification problem, where each class consists of samples from one
component of a mixture distribution. The mixture distribution is defined as
$p(x)=\sum_{i=1}^{4}\mathcal{N}(\mu_{i},0.3)$, with means
$\mu_{i}\in\\{(-2,-2),(-2,2),$ $(2,-2),(2,2)\\}$. We assigned labels
$y_{i}\in\\{1,2,3,4\\}$ according to the index of the mixture component the
samples were drawn from. For this task, samples are weight parameters
$\boldsymbol{\theta}$ of two-layer neural networks with 10 hidden units in
each layer and ReLU activations, denoted by
$f_{\boldsymbol{\theta}}:\mathbb{R}^{2}\rightarrow\mathbb{R}^{4}$.

To train each method, we drew a total of 100 training points, and 200 testing
points from the target distribution. To evaluate the posterior predictive
distribution, we drew points from a grid spaced from $\\{-10,10\\}$, then plot
the predictive distribution as measured by the standard deviation in
predictions among model samples in figure 2.

In this setting, the true posterior $p(\boldsymbol{\theta}|\boldsymbol{X},y)$
is unknown and past work often relies on “gold-standard” approaches like HMC
to serve as the ground truth. In the previous Bayesian linear regression
tests, however, we saw that HMC was outperformed by GPVI as well as ParVI
methods SVGD and GFSF. In the current test, intuitively, a “gold-standard”
approach should yield a predictive distribution with high variance (high
uncertainty) in regions far from the training data, and low variance (low
uncertainty) in regions near each mixture component. As shown in figure 2,
GPVI and SVGD both have higher uncertainty in no-data regions than HMC, while
remaining confident on the training data. On the other hand, MF-VI and
amortized SVGD both underestimate the uncertainty. The performance of MF-VI is
as expected, as mean-field VI approaches are known to characteristically
underestimate uncertainty. Results of all methods as well as more tests on an
easier 2-class variant are included in the Appendix.

(a) GPVI

(b) SVGD

(c) HMC

(d) Amortized SVGD

(e) MF-VI

(f) Deep Ensemble

Figure 2: Predictive uncertainty of each method on the 4-class classification
task, as measured by the standard deviation between predictions of sampled
functions. Regions of high uncertainty are shown as darker, while lighter
regions correspond to lower uncertainty. The training data is shown as samples
from four unimodal normal distributions. It can be seen that amortized SVGD,
MF-VI and deep ensembles significantly underestimate the uncertainty in
regions with no training data

###  4.4 Open Category

To evaluate the scalability of our approach we turn to large-scale image
classification experiments. In this setting, it is not possible to exactly
measure the accuracy of posterior fits. We instead utilize the open-category
task to test if our uncertainty estimations can help detect outlier examples.
The open category task defines a set of inlier classes that are seen during
training, and a set of outlier classes only used for evaluation. While in
principle, the content of the outlier classes can be arbitrary, we use
semantically similar outlier classes by splitting the training dataset into
inlier and outlier classes. The open-category task is in general more
difficult than out-of-distribution experiments where different datasets are
used as outlier classes, since distributions of categories from the same
dataset may be harder to discriminate.

We evaluated on the MNIST and CIFAR-10 image datasets, following Neal et al.
(2018) to split each dataset into 6 inlier classes and 4 outlier classes. We
performed standard fully supervised training on the 6 inlier classes, and
measured uncertainty in the 4 unseen outlier classes. We evaluated the
uncertainty using two widely-used statistics: area under the ROC curve (AUC),
and the expected calibration error (ECE). The AUC score measures how well a
binary classifier discriminates between predictions made on inlier inputs, vs
predictions made on outlier inputs. A perfect AUC score of 1.0 indicates a
perfect discrimination, while a score of 0.5 indicates that the two sets of
predictions are indistinguishable. ECE partitions predictions into equally
sized bins, and computes the $L_{1}$ difference in expected accuracy and
confidence between bins, which represents the calibration error of the bin.
ECE computes a weighted average of the calibration error of each bin. Together
AUC and ECE tell us how well a model can detect outlier inputs, as well as how
well the model fits the training distribution.

MNIST consists of 70,000 grayscale images of handwritten digits at 28x28
resolution, divided into 60,000 training images and 10,000 testing images. We
further split the dataset by only using the first six classes for training and
testing. The remaining four classes are only used to compute the AUC and ECE
statistics. We chose the LeNet-5 classifier architecture for all models, and
trained for 100 epochs. Due to the larger computational burden, we only
consider 10 samples from each method’s approximate posterior for both training
and evaluation. Note that our approach is capable of generating many more, but
it would be computationally costly to train much more for ensembles and
particle VI approaches. For GPVI and amortized ParVI methods we used a 3 layer
MLP hypernetwork with layer widths $[256,512,1024]$, ReLU activations, and
input noise $\boldsymbol{z}\in\mathbb{R}^{256}$. We did not test HMC in this
setting, as the computational demand is too high.

Table 3 shows the results. All RKHS-based methods (GPVI, ParVI, amortized
ParVI) as well as deep ensembles achieve high supervised (“clean“) accuracy
with the LeNet architecture. MF-VI underfits slightly, while KSD struggled to
achieve competitive accuracy even after 100 epochs. All methods achieved an
AUC over 0.95, but SVGD, GFSF and GPVI have the highest AUC values,
respectively. In terms of calibration, GPVI and RKHS-based ParVI methods are
the best calibrated. KSD is the worst calibrated model, and deep ensembles/MF-
VI have middling performance.

Method | Clean | AUC$\uparrow$ | ECE $\downarrow$
---|---|---|---
SVGD | 99.3 |  .989 $\pm$ .001 |  .001 $\pm$ .0002
GFSF | 99.2 |  .988 $\pm$ .003 | .002 $\pm$ .0003
KSD | 97.7 | .964 $\pm$ .005 | .014 $\pm$ .0007
Amortized SVGD | 99.1 | .958 $\pm$ .015 |  .002 $\pm$ .0007
Amortized GFSF | 99.2 | .978 $\pm$ .005 | .004 $\pm$ .0013
Amortized KSD | 97.7 | .951 $\pm$ .008 | .017 $\pm$ .0010
MF-VI | 98.6 | .951 $\pm$ .008 | .014 $\pm$ .0027
Deep Ensemble | 99.3 | .972 $\pm$ .002 | .008 $\pm$ .0060
GPVI | 99.3 |  .988 $\pm$ .001 |  .001 $\pm$ .0005
Table 3: Results for open-category classification on MNIST. We show the result
of standard supervised training (Clean), as well as AUC and ECE statistics
computed from training on a subset of classes and testing on the rest of the
classes as outliers.

CIFAR-10 consists of 60,000 RGB images depicting 10 object classes at 32x32
resolution, divided into 50,000 training images and 10,000 testing images. We
adopted the same 6 inlier / 4 outlier split used in (Neal et al., 2018) for
the open-category setting. For this task we used a CNN with 3 convolutional
layers and two linear layers, which is much smaller than SOTA classifiers for
CIFAR-10. Though the classification accuracy would suffer, it allows us to
clearly evaluate our method without considering interactions with
architectural components such as BatchNorm or residual connections. We also
used 10 samples from each method’s approximate posterior for training and
evaluation. For GPVI and amortized ParVI methods we used the same hypernetwork
architecture as in the MNIST setup. Table 4 shows the results where it can be
seen that our GPVI almost match the performance of SVGD and GFSF in terms of
AUC while doing a bit better on ECE.

Method | Clean | AUC $\uparrow$ | ECE $\downarrow$
---|---|---|---
SVGD | 80.3 |  .683 $\pm$ .008 | .055 $\pm$ .004
GFSF | 80.6 |  .681 $\pm$ .004 | .068 $\pm$ .012
Amortized SVGD | 71.12 | .636 $\pm$ .018 | .073 $\pm$ .029
Amortized GFSF | 71.09 | .583 $\pm$ .007 | .042 $\pm$ .029
MF-VI | 70.0 | .649 $\pm$ .006 |  .016 $\pm$ .002
Deep Ensemble | 73.54 | .652 $\pm$ .018 | .033 $\pm$ .011
GPVI | 76.2 | .677 $\pm$ .008 |  .018 $\pm$ .015
Table 4: Open-category classification on CIFAR-10. We show results of standard
supervised training (Clean), as well as AUC and ECE of each method trained in
the open-category setting.

###  4.5 Discussion

Much of the focus regarding recent work on Bayesian neural networks concerns
their performance on open-category and out-of-distribution tasks with high
dimensional image datasets. Instead, we show that our method closely
approximates the target posterior, both in tasks where the posterior is
explicitly known, as well as when it is intractable. The Bayesian linear
regression and density estimation tasks served as sanity checks. Because the
posterior was known explicitly, we could quantitatively test how well each
method fit the posterior. For density estimation, GPVI outperforms amortized
approaches at matching the target covariance (Table 1). In Bayesian linear
regression, while the approximation was close for all methods, there was a
clear hierarchy in terms of which types of methods produced the tightest
approximation (Table 2). GPVI and RKHS-based ParVI achieved the best posterior
fit overall, and we see in further experiments that quality of fit here is
indicative of performance in more difficult tasks.

The four-class classification problem (Fig. 2), while seemingly simple, is
particularly difficult for most methods we evaluated since many methods tend
to overgeneralize to the corners. MF-VI underestimates the uncertainty as
expected (Yao et al., 2019; Minka, 2001; Bishop, 2006). Notably, amortized
SVGD also underestimates the uncertainty with a predictive distribution
resembling that of a standard ensemble. We believe this is due to the
compounding approximation error explained in Section 3.5, i.e., naively back-
propagating the Stein variational gradient to update the generator is more
likely to end up with a non-descent direction compared with back-propagating
the exact functional gradient as GPVI. As shown in Fig. 2, the posterior
approximation of GPVI is tighter, with uncertainty that better matches the
data distribution. SVGD performs the best in this task, with high uncertainty
everywhere except in regions near observed data. Surprisingly, GPVI and SVGD
outperform HMC here, with clearly higher uncertainty near the corners of the
sample space. Note that with the functional approximation by neural networks,
it is hard to determine if the true posterior exactly matches the intuitively
“ideal” uncertainty plot where low variance only shows up around each mixture
component. On the other hand, HMC is guaranteed to converge to the true
posterior over time (Durmus et al., 2017). Given a fixed computational budget,
it is possible that GPVI or SVGD could achieve better performance. GPVI also
has the extra benefit of the ability to draw additional samples, which is not
possible with ParVI or HMC.

Our final experimental setting, open-category, reveals that GPVI consistently
has a higher AUC than all other scalable sampling approaches. On MNIST, GPVI
is among the overall top performers, together with two ParVI methods SVGD and
GFSF. On CIFAR-10, GPVI is among the top two performers in ECE and performs
only slightly behind the two best ParVI methods in AUC. Most importantly, GPVI
outperforms with clear margins all amortized ParVI methods on both datasets
under all metrics. This is consistent with all other qualitative and
quantitative experiments we have conducted, which again shows the advantage of
GPVI over existing generative ParVI methods.

We found that KSD completely failed in CIFAR-10, achieving slightly better
than random accuracy. We believe this is due to the way that KSD performs the
particle update. Where SVGD has a closed form expression for the particle
transportation map, KSD parameterizes it with a critic network that has the
input-output dimensionality of the particle parameters. Training this critic
naturally becomes difficult when applying to neural network functions.

In the open-category setting, there is added time complexity for GPVI relative
to amortized methods due to training our helper network. Compared to amortized
SVGD, the extra time complexity of training GPVI is a constant factor less
than 1. In our MNIST experiments, this additive factor is 0.28; in CIFAR-10
experiments, this factor is 0.78. The difference in MNIST and CIFAR-10 is due
to that the output dimension of the helper network scales with the output
dimension of the generator. We believe this is acceptable given the
performance gain and the possible real (offline training) applications of
GPVI.

##  5 Conclusion

We have presented a new method that fuses the best aspects of parametric VI
with non-parametric ParVI. GPVI has asymptotic convergence on par with ParVI.
Additionally, GPVI can efficiently draw samples from the posterior. We also
presented a method for efficiently estimating the product between the inverse
of the Jacobian of a deep network, and a gradient vector. Our experiments
showed that GPVI performs on par with ParVI, and outperforms amortized ParVI
and other competing methods in Bayesian linear regression, a classification
task, as well as open-category tasks on MNIST and CIFAR-10. In the future we
want to explore the efficacy of our method applied to large scale tasks like
image generation.

### Acknowledgements

Neale Ratzlaff and Li Fuxin were partially supported by the Defense Advanced
Research Projects Agency (DARPA) under Contract No. HR001120C0011 and
HR001120C0022. Any opinions, findings and conclusions or recommendations
expressed in this material are those of the author(s) and do not necessarily
reflect the views of DARPA.

## References

  * Bishop (2006) Bishop, C. M.  _Pattern recognition and machine learning_.  springer, 2006.
  * Blundell et al. (2015) Blundell, C., Cornebise, J., Kavukcuoglu, K., and Wierstra, D.  Weight uncertainty in neural networks.  _arXiv preprint arXiv:1505.05424_ , 2015.
  * Buntine & Weigend (1991) Buntine, W. L. and Weigend, A. S.  Bayesian back-propagation.  _Complex systems_ , 5(6):603–643, 1991.
  * Chen et al. (2014) Chen, T., Fox, E., and Guestrin, C.  Stochastic gradient hamiltonian monte carlo.  In _International conference on machine learning_ , pp. 1683–1691, 2014.
  * Cheng & Boots (2017) Cheng, C.-A. and Boots, B.  Variational inference for gaussian process models with linear complexity.  In _Advances in Neural Information Processing Systems_ , pp. 5184–5194, 2017.
  * Durmus et al. (2017) Durmus, A., Moulines, E., and Saksman, E.  On the convergence of hamiltonian monte carlo.  _arXiv preprint arXiv:1705.00166_ , 2017.
  * Fletcher (1976) Fletcher, R.  Conjugate gradient methods for indefinite systems.  In _Numerical analysis_ , pp. 73–89. Springer, 1976.
  * Gal & Ghahramani (2016) Gal, Y. and Ghahramani, Z.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In _international conference on machine learning_ , pp. 1050–1059, 2016.
  * Grathwohl et al. (2020) Grathwohl, W., Wang, K.-C., Jacobsen, J.-H., Duvenaud, D., and Zemel, R.  Learning the stein discrepancy for training and evaluating energy-based models without sampling.  In _International Conference on Machine Learning_ , pp. 3732–3747. PMLR, 2020.
  * Graves (2011) Graves, A.  Practical variational inference for neural networks.  _Advances in neural information processing systems_ , 24:2348–2356, 2011.
  * Hensman et al. (2013) Hensman, J., Fusi, N., and Lawrence, N. D.  Gaussian processes for big data.  _arXiv preprint arXiv:1309.6835_ , 2013.
  * Hensman et al. (2015) Hensman, J., Matthews, A., and Ghahramani, Z.  Scalable variational gaussian process classification.  2015\.
  * Hu et al. (2018) Hu, T., Chen, Z., Sun, H., Bai, J., Ye, M., and Cheng, G.  Stein neural sampler.  _arXiv preprint arXiv:1810.03545_ , 2018.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kingma et al. (2015) Kingma, D. P., Salimans, T., and Welling, M.  Variational dropout and the local reparameterization trick.  _Advances in neural information processing systems_ , 28:2575–2583, 2015.
  * Korattikara Balan et al. (2015) Korattikara Balan, A., Rathod, V., Murphy, K. P., and Welling, M.  Bayesian dark knowledge.  _Advances in Neural Information Processing Systems_ , 28:3438–3446, 2015.
  * Krueger et al. (2017) Krueger, D., Huang, C.-W., Islam, R., Turner, R., Lacoste, A., and Courville, A.  Bayesian hypernetworks.  _arXiv preprint arXiv:1710.04759_ , 2017.
  * Lakshminarayanan et al. (2017) Lakshminarayanan, B., Pritzel, A., and Blundell, C.  Simple and scalable predictive uncertainty estimation using deep ensembles.  In _Advances in neural information processing systems_ , pp. 6402–6413, 2017.
  * Liu et al. (2019) Liu, C., Zhuo, J., Cheng, P., Zhang, R., and Zhu, J.  Understanding and accelerating particle-based variational inference.  In _International Conference on Machine Learning_ , pp. 4082–4092. PMLR, 2019.
  * Liu (2017) Liu, Q.  Stein variational gradient descent as gradient flow.  In _Advances in neural information processing systems_ , pp. 3115–3123, 2017.
  * Liu & Wang (2016) Liu, Q. and Wang, D.  Stein variational gradient descent: A general purpose bayesian inference algorithm.  _Advances in neural information processing systems_ , 29:2378–2386, 2016.
  * Louizos & Welling (2016) Louizos, C. and Welling, M.  Structured and efficient variational deep learning with matrix gaussian posteriors.  In _International Conference on Machine Learning_ , pp. 1708–1716, 2016.
  * Louizos & Welling (2017) Louizos, C. and Welling, M.  Multiplicative normalizing flows for variational bayesian neural networks.  _arXiv preprint arXiv:1703.01961_ , 2017.
  * MacKay (1992) MacKay, D. J.  A practical bayesian framework for backpropagation networks.  _Neural computation_ , 4(3):448–472, 1992.
  * Marzouk et al. (2016) Marzouk, Y., Moselhy, T., Parno, M., and Spantini, A.  An introduction to sampling via measure transport.  _arXiv preprint arXiv:1602.05023_ , 2016.
  * Minka (2001) Minka, T. P.  _A family of algorithms for approximate Bayesian inference_.  PhD thesis, Massachusetts Institute of Technology, 2001.
  * Neal et al. (2018) Neal, L., Olson, M., Fern, X., Wong, W.-K., and Li, F.  Open set learning with counterfactual images.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , pp. 613–628, 2018.
  * Neal (2012) Neal, R. M.  _Bayesian learning for neural networks_ , volume 118.  Springer Science & Business Media, 2012.
  * Neal et al. (2011) Neal, R. M. et al.  Mcmc using hamiltonian dynamics.  _Handbook of markov chain monte carlo_ , 2(11):2, 2011.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al.  Pytorch: An imperative style, high-performance deep learning library.  _arXiv preprint arXiv:1912.01703_ , 2019.
  * Pawlowski et al. (2017) Pawlowski, N., Brock, A., Lee, M. C., Rajchl, M., and Glocker, B.  Implicit weight uncertainty in neural networks.  _arXiv preprint arXiv:1711.01297_ , 2017.
  * Rasmussen (2003) Rasmussen, C. E.  Gaussian processes in machine learning.  In _Summer School on Machine Learning_ , pp. 63–71. Springer, 2003\.
  * Rezende & Mohamed (2015) Rezende, D. and Mohamed, S.  Variational inference with normalizing flows.  In _International Conference on Machine Learning_ , pp. 1530–1538. PMLR, 2015.
  * Saad (2003) Saad, Y.  _Iterative methods for sparse linear systems_.  SIAM, 2003.
  * Strathmann et al. (2015) Strathmann, H., Sejdinovic, D., Livingstone, S., Szabo, Z., and Gretton, A.  Gradient-free hamiltonian monte carlo with efficient kernel exponential families.  _Advances in Neural Information Processing Systems_ , 28:955–963, 2015.
  * Wang & Liu (2016) Wang, D. and Liu, Q.  Learning to draw samples: With application to amortized mle for generative adversarial learning.  _arXiv preprint arXiv:1611.01722_ , 2016.
  * Wang et al. (2019) Wang, K., Pleiss, G., Gardner, J., Tyree, S., Weinberger, K. Q., and Wilson, A. G.  Exact gaussian processes on a million data points.  In _Advances in Neural Information Processing Systems_ , pp. 14648–14659, 2019.
  * Welling & Teh (2011) Welling, M. and Teh, Y. W.  Bayesian learning via stochastic gradient langevin dynamics.  In _Proceedings of the 28th international conference on machine learning (ICML-11)_ , pp. 681–688, 2011.
  * Wilson et al. (2016) Wilson, A. G., Hu, Z., Salakhutdinov, R. R., and Xing, E. P.  Stochastic variational deep kernel learning.  _Advances in Neural Information Processing Systems_ , 29:2586–2594, 2016.
  * Yao et al. (2019) Yao, J., Pan, W., Ghosh, S., and Doshi-Velez, F.  Quality of uncertainty quantification for bayesian neural network inference.  _arXiv preprint arXiv:1906.09686_ , 2019.
  * Young (1954) Young, D.  Iterative methods for solving partial difference equations of elliptic type.  _Transactions of the American Mathematical Society_ , 76(1):92–111, 1954.

## Appendix for Generative Particle Variational Inference via Estimation of
Functional Gradients

##  Appendix A Proofs

The summation convention is used on all repeated indices (e.g.
$a^{i}b^{i}:=\sum_{i}a^{i}b^{i}$) in the following proofs.

###  A.1 Proof of Theorem 3.1

To compute the gradient of
$\mathcal{J}(\boldsymbol{f})=\mathbf{E}_{\boldsymbol{z}}\left[-\log
p(\boldsymbol{f}(\boldsymbol{z}))+\log\frac{p_{\boldsymbol{z}}(\boldsymbol{z})}{\left|\det\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)\right|}\right],$
for any $\boldsymbol{\phi}\in T_{\boldsymbol{f}}\mathcal{H}$,

$\displaystyle d\mathcal{J}_{\boldsymbol{f}}(\boldsymbol{\phi})$
$\displaystyle=$
$\displaystyle\frac{d}{dt}\bigg{|}_{t=0}\mathcal{J}(\boldsymbol{f}+t\boldsymbol{\phi})$
(16) $\displaystyle\stackrel{{\scriptstyle\leavevmode\hbox to6.95pt{\vbox
to6.95pt{\pgfpicture\makeatletter\hbox{\hskip 3.4727pt\lower-3.4727pt\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{3.2727pt}{0.0pt}\pgfsys@curveto{3.2727pt}{1.80748pt}{1.80748pt}{3.2727pt}{0.0pt}{3.2727pt}\pgfsys@curveto{-1.80748pt}{3.2727pt}{-3.2727pt}{1.80748pt}{-3.2727pt}{0.0pt}\pgfsys@curveto{-3.2727pt}{-1.80748pt}{-1.80748pt}{-3.2727pt}{0.0pt}{-3.2727pt}\pgfsys@curveto{1.80748pt}{-3.2727pt}{3.2727pt}{-1.80748pt}{3.2727pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-1.75pt}{-2.25555pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{1}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}}{{=}}$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[\frac{d}{dt}\bigg{|}_{t=0}\left(-\log
p((\boldsymbol{f}+t\boldsymbol{\phi})(\boldsymbol{z}))\right)\right]-\mathbf{E}_{\boldsymbol{z}}\left[\frac{d}{dt}\bigg{|}_{t=0}\log\left|\det\left(\frac{\partial(\boldsymbol{f}+t\boldsymbol{\phi})}{\partial\boldsymbol{z}}\right)\right|\right]$
$\displaystyle=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[-\nabla_{x^{i}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})}\frac{d}{dt}\bigg{|}_{t=0}(f^{i}+t\phi^{i})(\boldsymbol{z})\right]-\mathbf{E}_{\boldsymbol{z}}\left[\text{Tr}\left(\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)^{-1}\frac{d}{dt}\bigg{|}_{t=0}\frac{\partial(\boldsymbol{f}+t\boldsymbol{\phi})}{\partial\boldsymbol{z}}\right)\right]$
$\displaystyle=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[-\nabla_{x^{i}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})}\phi^{i}(\boldsymbol{z})\right]-\mathbf{E}_{\boldsymbol{z}}\left[\left(\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)^{-1}\right)^{j}_{i}\frac{\partial\phi^{i}}{\partial
z^{j}}\right]$ $\displaystyle\stackrel{{\scriptstyle\leavevmode\hbox
to6.95pt{\vbox to6.95pt{\pgfpicture\makeatletter\hbox{\hskip
3.4727pt\lower-3.4727pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{
}{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{3.2727pt}{0.0pt}\pgfsys@curveto{3.2727pt}{1.80748pt}{1.80748pt}{3.2727pt}{0.0pt}{3.2727pt}\pgfsys@curveto{-1.80748pt}{3.2727pt}{-3.2727pt}{1.80748pt}{-3.2727pt}{0.0pt}\pgfsys@curveto{-3.2727pt}{-1.80748pt}{-1.80748pt}{-3.2727pt}{0.0pt}{-3.2727pt}\pgfsys@curveto{1.80748pt}{-3.2727pt}{3.2727pt}{-1.80748pt}{3.2727pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{
} }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-1.75pt}{-2.25555pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{2}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} }
\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}}{{=}}$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[-\nabla_{x^{i}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})}\langle
k(\boldsymbol{z},\cdot),\
\phi^{i}(\boldsymbol{z})\rangle_{\mathcal{H}}\right]-\mathbf{E}_{\boldsymbol{z}}\left[\left\langle\left(\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)^{-1}\right)^{j}_{i}\nabla_{z^{j}}k(\boldsymbol{z},\cdot),\
\phi^{i}(\boldsymbol{z})\right\rangle_{\mathcal{H}}\right]$ $\displaystyle=$
$\displaystyle\Bigg{\langle}\mathbf{E}_{\boldsymbol{z}}\Bigg{[}-\nabla_{x^{i}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})}k(\boldsymbol{z},\cdot)-\left(\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)^{-1}\right)^{j}_{i}\nabla_{z^{j}}k(\boldsymbol{z},\cdot)\Bigg{]},\phi^{i}\Bigg{\rangle}_{\mathcal{H}},$

the following identities are used in step 1 and 2,

$d\log|\det A|=\text{Tr}(A^{-1}dA),$ $\phi^{i}(\boldsymbol{z})=\langle
k(\boldsymbol{z},\cdot),\phi^{i}(\boldsymbol{z})\rangle_{H},$
$\nabla_{z^{j}}\phi^{i}(\boldsymbol{z})=\langle\nabla_{z^{j}}k(\boldsymbol{z},\cdot),\phi^{i}(\boldsymbol{z})\rangle_{H}.$

By definition of the gradient,

$\nabla_{\boldsymbol{f}}\mathcal{J}(\boldsymbol{f})(\boldsymbol{z})=\mathbf{E}_{\boldsymbol{z}^{\prime}}\left[-\nabla_{\boldsymbol{x}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z}^{\prime})}k(\boldsymbol{z}^{\prime},\boldsymbol{z})-\left(\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}^{\prime}}\right)^{-1}\right)\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})\right].$
(17)

###  A.2 Proof of Lemma 3.2

To compute the gradient of
$\mathcal{J}(\boldsymbol{\phi})=\mathbf{E}_{\boldsymbol{z}}\left[-\log
p(\boldsymbol{\phi}\left(\boldsymbol{f}(\boldsymbol{z}))\right)+\log\frac{p_{\boldsymbol{z}}(\boldsymbol{z})}{\left|\det\left(\frac{\partial(\boldsymbol{\phi}\circ\boldsymbol{f})}{\partial\boldsymbol{z}}\right)\right|}\right]$
at $\boldsymbol{\phi}=\text{id}$, for any $\boldsymbol{v}\in
T_{\boldsymbol{\phi}}\mathcal{H}$,

$\displaystyle d\mathcal{J}_{\boldsymbol{\phi}}(\boldsymbol{v})$
$\displaystyle=$
$\displaystyle\frac{d}{dt}\bigg{|}_{t=0}\mathcal{J}(\boldsymbol{\phi}+t\boldsymbol{v})\bigg{|}_{\boldsymbol{\phi}=\text{id}}$
(18) $\displaystyle=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[\frac{d}{dt}\bigg{|}_{t=0}\left(-\log
p((\boldsymbol{f}+t\boldsymbol{v}\circ\boldsymbol{f})(\boldsymbol{z}))\right)\right]-\mathbf{E}_{\boldsymbol{z}}\left[\frac{d}{dt}\bigg{|}_{t=0}\log\left|\det\left(\frac{\partial(\boldsymbol{f}+t\boldsymbol{v}\circ\boldsymbol{f})}{\partial\boldsymbol{z}}\right)\right|\right]$
$\displaystyle=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[-\nabla_{x^{i}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})}\frac{d}{dt}\bigg{|}_{t=0}(f^{i}+t(\boldsymbol{v}\circ\boldsymbol{f})^{i}(\boldsymbol{z})\right]-\mathbf{E}_{\boldsymbol{z}}\left[\text{Tr}\left(\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)^{-1}\frac{d}{dt}\bigg{|}_{t=0}\frac{\partial(\boldsymbol{f}+t\boldsymbol{v}\circ\boldsymbol{f})}{\partial\boldsymbol{z}}\right)\right]$
$\displaystyle=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[-\nabla_{x^{i}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})}v^{i}(\boldsymbol{f}(\boldsymbol{z}))\right]-\mathbf{E}_{\boldsymbol{z}}\left[\left(\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right)^{-1}\right)^{j}_{i}\frac{\partial
f^{k}}{\partial z^{j}}\frac{\partial v^{i}}{\partial x^{k}}\right]$
$\displaystyle=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[-\nabla_{x^{i}}\log
p(\boldsymbol{x})\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})}\langle
k(\boldsymbol{f}(\boldsymbol{z}),\cdot),\
v^{i}(\boldsymbol{f}(\boldsymbol{z}))\rangle_{\mathcal{H}}\right]-\mathbf{E}_{\boldsymbol{z}}\left[\left\langle\nabla_{x^{i}}k(\boldsymbol{x},\cdot)\bigg{|}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})},\
v^{i}(\boldsymbol{f}(\boldsymbol{z}))\right\rangle_{\mathcal{H}}\right]$
$\displaystyle=$
$\displaystyle\bigg{\langle}\mathbf{E}_{\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})}\left[-\nabla_{x^{i}}\log
p(\boldsymbol{x})k(\boldsymbol{x},\cdot)-\nabla_{x^{i}}k(\boldsymbol{x},\cdot)\right],v^{i}\bigg{\rangle}_{\mathcal{H}}.$

By definition of the gradient,

$\nabla_{\boldsymbol{\phi}}\mathcal{J}(\boldsymbol{\phi})(\boldsymbol{x})\Big{|}_{\boldsymbol{\phi}=\text{id}}=\mathbf{E}_{\boldsymbol{x}^{\prime}}\left[-\nabla_{\boldsymbol{x}^{\prime}}\log
p(\boldsymbol{x}^{\prime})k(\boldsymbol{x}^{\prime},\boldsymbol{x})-\nabla_{\boldsymbol{x}^{\prime}}k(\boldsymbol{x}^{\prime},\boldsymbol{x})\right],$
(19)

where $\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{z})$ and
$\boldsymbol{x}^{\prime}=\boldsymbol{f}(\boldsymbol{z}^{\prime})$.

###  A.3 Directly computing the gradient $\mathcal{J}(\boldsymbol{\theta})$

As a completion of the discussion, we also derive the gradient of
$\mathcal{J}(\boldsymbol{\theta})=\mathbf{E}_{\boldsymbol{z}}\left[-\log
p(\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}))+\log\frac{p_{\boldsymbol{z}}(\boldsymbol{z})}{\left|\det\left(\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}}\right)\right|}\right]$
w.r.t. $\boldsymbol{\theta}$, for any $\boldsymbol{v}\in
T_{\boldsymbol{\theta}}W$,

$\displaystyle d\mathcal{J}_{\boldsymbol{\theta}}(\boldsymbol{v})$
$\displaystyle=$
$\displaystyle\frac{d}{dt}\bigg{|}_{t=0}\mathcal{J}(\boldsymbol{\theta}+t\boldsymbol{v})$
(20) $\displaystyle=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[\frac{d}{dt}\bigg{|}_{t=0}\left(-\log
p(\boldsymbol{f}_{\boldsymbol{\theta}+t\boldsymbol{v}}(\boldsymbol{z}))\right)\right]-\mathbf{E}_{\boldsymbol{z}}\left[\frac{d}{dt}\bigg{|}_{t=0}\log\det\left(\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}+t\boldsymbol{v}}}{\partial\boldsymbol{z}}\right)\right]$
$\displaystyle=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[-\nabla_{x^{j}}\log
p(\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}))\frac{d}{dt}\bigg{|}_{t=0}f^{j}_{\boldsymbol{\theta}+t\boldsymbol{v}}(\boldsymbol{z})\right]-\mathbf{E}_{\boldsymbol{z}}\left[\text{Tr}\left(\left(\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}}\right)^{-1}\frac{d}{dt}\bigg{|}_{t=0}\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}+t\boldsymbol{v}}}{\partial\boldsymbol{z}}\right)\right]$
$\displaystyle=$
$\displaystyle\mathbf{E}_{\boldsymbol{z}}\left[-\nabla_{x^{j}}\log
p(\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}))\frac{\partial
f^{j}}{\partial\theta^{i}}v^{i}i\right]-\mathbf{E}_{\boldsymbol{z}}\left[\text{Tr}\left(\left(\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}}\right)^{-1}\frac{\partial^{2}\boldsymbol{f}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}\partial\boldsymbol{\theta}}\boldsymbol{v}\right)\right]$
$\displaystyle=$
$\displaystyle\Bigg{\langle}\mathbf{E}_{\boldsymbol{z}}\Bigg{[}-\nabla_{x^{j}}\log
p(\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}))\frac{\partial
f^{j}}{\partial\theta^{i}}-\left(\left(\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}}\right)^{-1}\right)^{j}_{k}\left(\frac{\partial^{2}f^{j}_{\boldsymbol{\theta}}}{\partial
z^{k}\partial\theta^{i}}\right)\Bigg{]},v^{i}\Bigg{\rangle}_{E},$

where $\langle\cdot,\cdot\rangle_{E}$ denotes the Euclidean inner product.
Therefore,

$\nabla_{\theta^{i}}(\mathcal{J}(\boldsymbol{\theta}))=\mathbf{E}_{\boldsymbol{z}}\left[-\frac{\partial
f^{j}}{\partial\theta^{i}}\nabla_{x^{j}}\log
p(\boldsymbol{f}_{\boldsymbol{\theta}}(\boldsymbol{z}))-\left(\left(\frac{\partial\boldsymbol{f}_{\boldsymbol{\theta}}}{\partial\boldsymbol{z}}\right)^{-1}\right)^{j}_{k}\left(\frac{\partial^{2}f^{j}_{\boldsymbol{\theta}}}{\partial
z^{k}\partial\theta^{i}}\right)\right].$

Computation of the second term in the expectation involves second derivatives
of $\boldsymbol{f}_{\boldsymbol{\theta}}$, which is not amenable in large-
scale problems. As a result, we avoid this direct approach and choose to
pullback the functional gradient (17) instead.

###  A.4 Proof of Injectivity of $\boldsymbol{f}$

We parameterize $\boldsymbol{f}$ as follows,

$\boldsymbol{f}(\boldsymbol{z})=\boldsymbol{g}\left(\boldsymbol{z}^{(:k)}\right)+\lambda\boldsymbol{z},\quad\forall\boldsymbol{z}\in\mathbb{R}^{d},$
(21)

where $\boldsymbol{z}^{(:k)}\in\mathbb{R}^{k}$ denotes the vector consisting
of the first $k(\ll d)$ components of $\boldsymbol{z}$, and
$\boldsymbol{g}:\mathbb{R}^{k}\to\mathbb{R}^{d}$ is a much slimmer neural
network. For any experimental setting, our parameterization of
$\boldsymbol{g}$ uses a dimension $k$ less than $30\%$ the size of $d$. For
our high dimensional open-category experiments where $d>60000$, we use a $k$
of less than $2\%$ of $d$. The Jacobian is,

$\left[\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}}\right]_{d\times
d}=\left[\left[\frac{\partial\boldsymbol{g}}{\partial\boldsymbol{z}^{(:k)}}\right]_{d\times
k}\bigg{|}\boldsymbol{0}_{d\times(d-k)}\right]_{d\times
d}+\lambda\boldsymbol{I}_{d}.$ (22)

We now show that our practical choice of $\boldsymbol{f}$ is indeed injective.

From (21), for sufficiently large $\lambda$, the Jacobian $J_{\boldsymbol{f}}$
given by (22) is positive definite. Since the domain of $\boldsymbol{z}$ is
convex, it is straightforward to show $\boldsymbol{f}$ is injective,

For any two different points
$\boldsymbol{z}_{1},\boldsymbol{z}_{2}\in\mathcal{Z}$, consider the line
segment
$\boldsymbol{z}_{1}+t(\boldsymbol{z}_{2}-\boldsymbol{z}_{1}),t\in[0,1]$, which
lies in $\mathcal{Z}$ given the convexity of $\mathcal{Z}$. From the
Fundamental Theorem of calculus,

$\displaystyle(\boldsymbol{z}_{2}-\boldsymbol{z}_{1})^{T}(\boldsymbol{f}\left(\boldsymbol{z}_{2})-\boldsymbol{f}(\boldsymbol{z}_{1})\right)$
$\displaystyle=$
$\displaystyle(\boldsymbol{z}_{2}-\boldsymbol{z}_{1})^{T}\left(\int_{0}^{1}J_{\boldsymbol{f}}(\boldsymbol{z}_{1}+t(\boldsymbol{z}_{2}-\boldsymbol{z}_{1}))dt\right)\cdot(\boldsymbol{z}_{2}-\boldsymbol{z}_{1})$
$\displaystyle=$
$\displaystyle\int_{0}^{1}(\boldsymbol{z}_{2}-\boldsymbol{z}_{1})^{T}J_{\boldsymbol{f}}(\boldsymbol{z}_{1}+t(\boldsymbol{z}_{2}-\boldsymbol{z}_{1}))(\boldsymbol{z}_{2}-\boldsymbol{z}_{1})dt>0$

the last inequality is due to the positive definiteness of
$J_{\boldsymbol{f}}$ on $\mathcal{Z}$. Therefore,
$\boldsymbol{f}(\boldsymbol{z}_{1})\neq\boldsymbol{f}(\boldsymbol{z}_{2})$,
$\boldsymbol{f}$ is injective.

##  Appendix B Method Details

We provide architectures and computational details of our generator and helper
networks.

###  B.1 Helper Network

The helper network $\boldsymbol{h}_{\eta}$ takes $\boldsymbol{z}^{\prime}$ and
$\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$ as
inputs and outputs
$\left(\frac{\partial\boldsymbol{f}(\boldsymbol{z}^{\prime})}{\partial\boldsymbol{z}^{\prime}}\right)^{-1}\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$.
Given the specific parameterization of $\boldsymbol{f}$ defined by (21), the
Jacobian (22) only depends on $\boldsymbol{z}^{{}^{\prime}(:k)}$, the first
$k$ components of $\boldsymbol{z}^{\prime}$. The helper network applies a
fully connected layer with ReLU activation to the input
$\boldsymbol{z}^{{}^{\prime}(:k)}$ and
$\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$
respectively, then the outputs from these two branches are concatenated and
sent to a three-layer fully connected network. We optimize the helper network
with the Adam optimizer Kingma & Ba (2014) and a learning rate of $1e-4$. We
found that regularization techniques like batchnorm and weight decay hurt
performance, presumably because we are optimizing for an exact solution and
thus do not want a smoothed estimate. While our helper network makes
computation of the Jacobian inverse tractable, it nonetheless adds the
complexity of an extra optimization problem per training step of
$\boldsymbol{f}$. To mitigate the increased training time we update
$\boldsymbol{h}_{\eta}$ just once per training step of $\boldsymbol{f}$;
taking one step also keeps the network from over-committing to a single input.

###  B.2 Kernel Selection

We follow Liu & Wang (2016) and use an RBF kernel
$k(x,x^{\prime})=\exp{(\frac{1}{h}||x-x^{\prime}||^{2}_{2})}$ for all kernel-
based methods including GPVI, SVGD, GFSF, and their amortized variants. The
bandwidth $h$ is computed using the median method also proposed by Liu & Wang
(2016): $h=med^{2}/\log n$, where $med$ is the median of pairwise distances
between samples $\\{x_{i}\\}_{i=1}^{n}$.

##  Appendix C Comparing with an Alternative Solver for
$\left(\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}^{\prime}}\right)^{-1}\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$

Here we justify the use of our helper network $\boldsymbol{h}_{\eta}$, by
comparing its performance to a more traditional linear solver: the stabilized
biconjugate gradient method (BiCGSTAB) (Saad, 2003). BiCGSTAB is a well-known
iterative algorithm for solving systems of the form $Ax=b$. It is similar to
the conjugate gradient method, but does not require $A$ to be self-adjoint,
giving BiCGSTAB wider applicability. In figure 3, we compare our helper
network (”Network”) against BiCGSTAB in the Bayesian linear regression setting
from section 4.2 in the main text. For BiCGSTAB, we solve $B*B$ systems of the
form
$JJ^{-1}\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})=\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$
for each iteration of training our generator, where
$J=\frac{\partial\boldsymbol{f}}{\partial\boldsymbol{z}^{\prime}}$ and $B*B$
is the effective batch size of
$\nabla_{\boldsymbol{z}^{\prime}}k(\boldsymbol{z}^{\prime},\boldsymbol{z})$.
Due to the greatly increased training time from solving $B*B$ independent
problems per training iteration, we only run BiCGSTAB for one step, and warm
start from the previous solution at each new generator training iteration. In
addition to increased training time, BiCGSTAB suffers from instability due to
the constantly changing
$\frac{\partial\boldsymbol{f}(\boldsymbol{z}^{\prime})}{\partial\boldsymbol{z}^{\prime}}$
as well as
$\nabla_{\boldsymbol{z^{\prime}}}k(\boldsymbol{z^{\prime}},\boldsymbol{z})$,
due to the resampling of $\boldsymbol{z}$ and $\boldsymbol{z^{\prime}}$ at
each step. As seen in figure 3, when using BiCGSTAB, GPVI is unable to fit the
target distribution, while using our helper network (”Network”) we are able to
efficiently minimize the mean and covariance error.

We also considered using a normalizing flow as a replacement for our generator
network, as normalizing flows are invertible with lower triangular Jacobians.
Unfortunately, the efficiency of normalizing flows comes from setting the
input-output dimensionality to be equal to force the Jacobian to be square.
When using our method for BNNs, as explained in Section B, we cannot afford to
store a generator with equal input-output dimensionality, making normalizing
flows an inefficient choice at best for our own method. Our helper network
stands as an efficient, novel solution for explicitly estimating the Jacobian
inverse vector product when needed. Nevertheless, we compare GPVI with
normalizing flow method for density estimation in the next section.

(a) Mean Error

(b) Cov Error

Figure 3: Comparing our helper network with BiCGSTAB in the Bayesian linear
regression setting.

##  Appendix D Additional Results

We show additional results on density estimation and classification.

###  D.1 Density Estimation

Here we provide additional results in the density estimation setting.

#### Comparison with Explicit Jacobian

Here we perform the experiments from section 4.1 of the main paper, and
compare GPVI with a variant where the Jacobian and its inverse are explicitly
computed. The inverse is computed with the PyTorch (Paszke et al., 2019)
function torch.inverse, which uses LAPACK routines getrf and getri. We can see
in table 5 that GPVI is competitive with the “Exact-Jac” variant in the 2D
setting, and even performs better in the 5D setting. This is unexpected, as
explicitly computing the Jacobian and its inverse should be the correct way to
compute the functional gradient. We hypothesize that the Jacobian of our
generator network is ill-conditioned. In this case any inversion algorithm is
more prone to large numerical error. Because our helper network is updated
once per training iteration, we may avoid the large gradients that come from
numerical error. The increased performance of GPVI over the “Exact-Jac”
variant in the 5D setting may be due to this smoother training.

Method |  $\Sigma$ error (2d) $\downarrow$ |  $\Sigma$ error (5d) $\downarrow$
---|---|---
Amortized SVGD | 0.10 $\pm$ .09 | 0.37 $\pm$ .32
GPVI | 0.14 $\pm$ .08 | 0.14 $\pm$ .04
GPVI Exact Jac | 0.15 $\pm$ .09 | 0.49 $\pm$ .17
Table 5: Comparison of generative Particle VI approaches for density
estimation of 2d and 5d Gaussian distributions.

#### Energy Potentials

We use the four non-Gaussian energy potentials defined in Rezende & Mohamed
(2015), and train GPVI as a sampler. In figures 4-7, we see (from left to
right) the target density, GPVI, amortized SVGD, and normalizing flows
(Rezende & Mohamed, 2015). We found that while our method has more parameters,
we don’t need nearly as deep a model as with normalizing flows. We used just 2
hidden layers for GPVI, while we needed a flow of depth 32 to get comparable
performance. We trained each method for 200000 steps with a batch size of 100.
To generate plots, we sampled 20000 points from each model. We detail the rest
of the hyperparameters in section F. In figures 5, 6, 7 we see that GPVI is
able to capture the variance of the target distribution, while amortized SVGD
samples mostly from the mean. For the Sin Bisect setting, GPVI is the only
method that captures some of both halves of the middle section. But in Sin
Split, only normalizing flows capture both halves of the split section.

(a) Twin Moons Target

(b) GPVI

(c) Amortized SVGD

(d) Normalizing Flow

Figure 4: Twin Moons

(a) Sinusoid Target

(b) GPVI

(c) Amortized SVGD

(d) Normalizing Flow

Figure 5: Sinusoid

(a) Sin Bisect Target

(b) GPVI

(c) Amortized SVGD

(d) Normalizing Flow

Figure 6: Sin Bisect

(a) Sin Split Target

(b) GPVI

(c) Amortized SVGD

(d) Normalizing Flow

Figure 7: Sin Split

###  D.2 Classification

In figure 8 we show the results of all evaluated methods on the four-class
classification problem from section 4.3 of the main text.

(a) GPVI

(b) SVGD

(c) GFSF

(d) KSD

(e) HMC

(f) Amortized SVGD

(g) Amortized GFSF

(h) Amortized KSD

(i) MF-VI

(j) Deep Ensembles

Figure 8: Predictive uncertainty of each method on the 4-class classification
task, as measured by the standard deviation between predictions of sampled
functions.

In addition to the four-class classification problem we presented in the main
paper, we show results on a simpler two class variant. In this setting, we
generate data in the same way as in the four-class setting, but we use a
mixture distribution with two components. Specifically, the mixture
distribution is defined as $p(x)=\sum_{i=1}^{2}\mathcal{N}(\mu_{i},0.3)$, with
means $\mu_{i}\in\\{(-2,-2),(2,2)\\}$. We assigned labels $y_{i}\in\\{1,2\\}$
according to the index of the mixture component the samples were drawn from.
We show the results of each method in figure 9. We can see that GPVI again
gives better uncertainty estimates than other sampling based approaches.
Again, we do not know what the true posterior over classifications looks like,
but GPVI and RKHS-based ParVI approaches give uncertainty estimates that
closely match our intuition for this problem.

(a) GPVI

(b) SVGD

(c) GFSF

(d) KSD

(e) HMC

(f) Amortized SVGD

(g) Amortized GFSF

(h) Amortized KSD

(i) MF-VI

(j) Deep Ensembles

Figure 9: Predictive uncertainty of each method on the 2-class classification
task, as measured by the standard deviation between predictions of sampled
functions.

##  Appendix E Experimental Details

Here we provide some additional details regarding the setup and reporting of
our experiments and their results.

HMC Details We use the same settings for HMC for the Bayesian linear
regression and 2/4 class classification tasks. We sampled the momentum from a
standard normal distribution, and used 25 leap-frog steps per sample. We ran
each experiment for 25K steps in total, with a burn-in of 20K steps, and a
step size of $0.0005$. We thinned each chain after burn-in, and tuned the
number of leapfrog steps and step-size for each experiment. To check for
convergence, we checked that the means of each chain were similar, indicating
mixing.

1D Regression For our 1D regression task (figure 1 of the main paper), we
generate a dataset of 80 samples $X$ with targets $Y$. We draw 76 samples of
$X$ uniformly from $[-6,-2]\cup[2,6]$, and draw the 4 remaining samples from
$[-2,2]$. The targets $Y$ are computed as $Y=-(1+X)\sin(1.2X)+\epsilon$, where
$\epsilon\sim\mathcal{N}(0,0.04)$. For all methods we use 100 posterior
samples, and train for 50K iterations. For GPVI and amortized SVGD we use a
generator with two linear layers $[32,32]$ and Gaussian input noise of the
same dimension.

Density Estimation
In the density estimation setting (section 4.1 of the main text), we randomly
generate $2d$ and $5d$ target covariance matrices. Specifically, the target
covariance we wish to fit is computed as
$\Sigma^{*}=\Sigma\Sigma^{T},\Sigma\sim\mathcal{N}(0,I_{d}),d\in\\{2,5\\}$. We
use an MLP with 1 hidden layer of width 2 to approximately sample from a
unimodal Gaussian distribution with covariance $\Sigma^{*}$. Given Gaussian
input noise, the variance of the output distribution of the linear generator
is computed as $W^{T}W$, which should match $\Sigma^{*}$ after training.

Open Category Tasks
For the open-category tasks, we computed two metrics to evaluate each method:
AUROC (AUC) and ECE. In this context, AUC measures how well a binary
classifier can discriminate between predictions made on inlier vs outlier test
inputs. An AUC score of $1.0$ indicates that the predictions made by a model
are perfectly separable, meaning that we could set a threshold on predicted
probabilities to detect every outlier test input. In practice we first make
predictions on the inlier and outlier test sets, then compute the variance of
the predictions over the sampled predictors. The AUC score is then computed
against the predictive variance using the scikit-learn library.

Expected Calibration Error (ECE) measures how well calibrated a model’s
predictions are. Given predictions on inlier test inputs, ECE partitions
predictions into $M=15$ equally sized bins according to their confidence. For
each bin $b_{m}$ we compute the difference between their accuracy
$1/|b_{m}|\sum_{x_{i}\in b_{m}}(\hat{y}-y)$ and confidence
$1/b_{m}\sum_{x_{i}\in m}\hat{p}$. Where $\hat{p},\hat{y}$ are predicted
probabilities and associated label with respect to a data-point $x_{i}$ and
true label $y$. This difference in expected accuracy and confidence between
each bin represents the bin’s calibration error.

For all experiments we report an average of three runs using the
hyperparameter settings given in section F.

##  Appendix F Hyperparameter Settings

In tables LABEL:tab:density-energy-params \- LABEL:tab:cifar10-params we
detail the hyperparameters chosen for each method in each experimental
setting. We refer to the sampler network as $\boldsymbol{f}$, and the
predicting classifier as $\boldsymbol{g}$. We generally use the same structure
of $\boldsymbol{f}$ and input noise for GPVI and amortized ParVI methods.

Table 6: Hyperparameters for density estimation of energy potentials. Density Estimation (energy potentials) |
---|---
Hyperparameter | Value
Common |
| Posterior Samples | $20000$
| Learning Rate | $1e-4$
|  $\boldsymbol{f}$ Hidden Layers | $2$
|  $\boldsymbol{f}$ Hidden Width | $[500,500]$
|  $\boldsymbol{f}$ Input Noise Stdev | $\sigma\in\\{1.0,2.0,6.0\\}$
| Training Steps | $200e3$
| Minibatch Size | $100$
GPVI |
| Optimizer | Adam
|  $\boldsymbol{h}_{\eta}$ Hidden Width | $500$
|  $\boldsymbol{h}_{\eta}$ Hidden Layers | $3$
|  $\boldsymbol{h}_{\eta}$ Learning Rate | $1e-4$
Normalizing Flow |
| Optimizer | RMSProp
| Flow Length | $32$
| Weight Decay | $1e-3$
| Base Dist Stdev | $\sigma\in\\{1.0,2.0,6.0\\}$
| Flow Architecture | Planar Flow
. |  |
Table 7: Hyperparameters for Bayesian linear regression task Bayesian Linear Regression |
---|---
Hyperparameter | Value
Common |
| Optimizer (all) | Adam
| Posterior Samples | $100$
| Learning Rate | $1e-3$
|  $\boldsymbol{f}$ Hidden Layers | None
|  $\boldsymbol{f}$ Input Noise $\boldsymbol{z}$ | $\mathcal{N}(0,I_{3})$
|  $\boldsymbol{g}$ Hidden Layers | None
| Training Steps | $50e3$
| Minibatch Size | $10$
GPVI |
|  $\boldsymbol{h}_{\eta}$ Hidden Width | $10$
|  $\boldsymbol{h}_{\eta}$ Hidden Layers | $3$
|  $\boldsymbol{h}_{\eta}$ Learning Rate | $1e-4$
KSD (Amortized and ParVI) |
| Critic Hidden Layers | $1$
| Critic Hidden Width | $100$
| Critic Learning Rate | $1e-3$
| Critic $L_{2}$ Weight | $10$
MF-VI (Bayes by Backprop) |
| Weight Prior | Scale Mixture
| Mixture Weight $\pi$ | $0.5$
| $\sigma_{1}$ | $1.0$
| $\sigma_{2}$ | $0$
. |  |
Table 8: Hyperparameters for 2/4 class classification task 4/2-class Classification |
---|---
Hyperparameter | Value
Common |
| Optimizer (all) | Adam
| Posterior Samples | $100$
| Learning Rate | $1e-3$
| Training Steps | 500e3
| Minibatch Size | 100
|  $\boldsymbol{f}$ Hidden Layers | 2
|  $\boldsymbol{f}$ Hidden Width | 64
|  $\boldsymbol{f}$ Nonlinearity | ReLU
|  $\boldsymbol{f}$ Input Noise $\boldsymbol{z}$ | $\mathcal{N}(0,I_{64})$
|  $\boldsymbol{g}$ Hidden Layers | 2
|  $\boldsymbol{g}$ Hidden Width | 10
|  $\boldsymbol{g}$ Nonlinearity | ReLU
GPVI |
|  $\boldsymbol{h}_{\eta}$ Hidden Width | $544$
|  $\boldsymbol{h}_{\eta}$ Hidden Layers | $3$
|  $\boldsymbol{h}_{\eta}$ Learning Rate | $1e-4$
KSD (Amortized and ParVI) |
| Critic Hidden Layers | $2$
| Critic Hidden Width | $100$
| Critic Learning Rate | $1e-3$
| Critic $L_{2}$ Weight | $10$
MF-VI (Bayes by Backprop) |
| Weight Prior | Scale Mixture
| Mixture Weight $\pi$ | $0.5$
| $\sigma_{1}$ | $1.0$
| $\sigma_{2}$ | $\exp{(-6)}$
. |  |
Table 9: Hyperparameters for MNIST open category task _Open-Category (MNIST)_ |
---|---
Common |
| Optimizer (all) | Adam
| Posterior Samples | $10$
| Training Epochs | $100$
| Minibatch Size | $50$
|  $\boldsymbol{f}$ Learning Rate | $1e-5$
|  $\boldsymbol{f}$ Hidden Layers | $3$
|  $\boldsymbol{f}$ Hidden Width | $[256,512,1024]$
|  $\boldsymbol{f}$ Nonlinearity | ReLU
|  $\boldsymbol{f}$ Input Noise $\boldsymbol{z}$ | $\mathcal{N}(0,I_{256})$
|  $\boldsymbol{g}$ Learning Rate | $1e-5$
|  $\boldsymbol{g}$ Architecture | LeNet-5
|  $\boldsymbol{g}$ Nonlinearity | ReLU
GPVI |
|  $\boldsymbol{h}_{\eta}$ Hidden Width | $512$
|  $\boldsymbol{h}_{\eta}$ Hidden Layers | $3$
|  $\boldsymbol{h}_{\eta}$ Learning Rate | $1e-4$
KSD (Amortized and ParVI) |
| Critic Hidden Layers | $2$
| Critic Hidden Width | $512$
| Critic Learning Rate | $1e-4$
| Critic $L_{2}$ Weight | $1.0$
MF-VI (Bayes by Backprop) |
| Weight Prior | Scale Mixture
| Mixture Weight $\pi$ | $0.5$
| $\sigma_{1}$ | $1.0$
| $\sigma_{2}$ | $\exp{(-6)}$
. |  |
Table 10: Hyperparameters for CIFAR-10 open category task _Open-Category (CIFAR-10)_ |
---|---
Common |
| Optimizer (all) | Adam
| Posterior Samples | $10$
| Training Epochs | $200$
| Minibatch Size | $50$
|  $\boldsymbol{f}$ Learning Rate | $1e-5$
|  $\boldsymbol{f}$ Hidden Layers | $3$
|  $\boldsymbol{f}$ Hidden Width | $[400,600,1000]$
|  $\boldsymbol{f}$ Nonlinearity | ReLU
|  $\boldsymbol{f}$ Input Noise $\boldsymbol{z}$ | $\mathcal{N}(0,I_{400})$
|  $\boldsymbol{g}$ Hidden Layers |  $3$ conv, $2$ linear
|  $\boldsymbol{g}$ Hidden Width | $[32,64,64,128,10]$
|  $\boldsymbol{g}$ Nonlinearity | ReLU
GPVI |
|  $\boldsymbol{h}_{\eta}$ Hidden Width | $512$
|  $\boldsymbol{h}_{\eta}$ Hidden Layers | $3$
|  $\boldsymbol{h}_{\eta}$ Learning Rate | $1e-4$
KSD (Amortized and ParVI) |
| Critic Hidden Layers | $2$
| Critic Hidden Width | $512$
| Critic Learning Rate | $1e-4$
| Critic $L_{2}$ Weight | $1.0$
MF-VI (Bayes by Backprop) |
| Weight Prior | Scale Mixture
| Mixture Weight $\pi$ | $0.5$
| $\sigma_{1}$ | $1.0$
| $\sigma_{2}$ | $\exp{(-6)}$
. |  |


Varun Narasimhachar, Jayne Thompson, Jiajun Ma, Gilad Gour, and Mile Gu.
Quantifying memory capacity as a quantum thermodynamic resource.
Physical Review Letters, 122(6), Feb 2019.

[17]
Elias M Stein and Rami Shakarchi.
Real analysis: measure theory, integration, and Hilbert spaces.
Princeton University Press, 2009.

[18]
Joel G. Smith.
The information capacity of amplitude- and variance-constrained sclar
gaussian channels.
Information and Control, 18(3):203–219, Apr 1971.

[19]
S.(Shitz) Shamai.
Capacity of a pulse amplitude modulated direct detection photon
IEE Proceedings I Communications, Speech and Vision,
137(6):424, 1990.

[20]
S. Shamai and I. Bar-David.
The capacity of average and peak-power-limited quadrature gaussian
IEEE Transactions on Information Theory, 41(4):1060–1071, Jul

[21]
J. Huang and S.P. Meyn.
Characterization and computation of optimal distributions for channel
IEEE Transactions on Information Theory, 51(7):2336–2351, Jul

[22]
Joel Gorham Smith.
On the information capacity of peak and average power
constrained Gaussian channels.
University of California, Berkeley, 1969.

[23]
Alex Dytso, Mert Al, H. Vincent Poor, and Shlomo Shamai Shitz.
On the capacity of the peak power constrained vector gaussian
channel: An estimation theoretic perspective.
IEEE Transactions on Information Theory, 65(6):3907–3921, Jun

[24]
Semih Yagli, Alex Dytso, H. Vincent Poor, and Shlomo Shamai Shitz.
An upper bound on the number of mass points in the capacity achieving
distribution for the amplitude constrained additive gaussian channel.
2019 IEEE International Symposium on Information Theory (ISIT),
Jul 2019.

[25]
Alexander S Holevo and Reinhard F Werner.
Evaluating capacities of bosonic gaussian channels.
Physical Review A, 63(3):032312, 2001.

[26]
Filippo Caruso, Vittorio Giovannetti, and Alexander S Holevo.
One-mode bosonic gaussian channels: a full weak-degradability
New Journal of Physics, 8(12):310, 2006.

[27]
Maksim E Shirokov.
On the energy-constrained diamond norm and its application in quantum
information theory.
Problems of Information Transmission, 54(1):20–33, 2018.

[28]
Simon Becker and Nilanjana Datta.
Convergence rates for quantum evolution and entropic continuity
bounds in infinite dimensions.
Communications in Mathematical Physics, 374(2):823–871, 2020.

[29]
Ranjith Nair.
Quantum-limited loss sensing: Multiparameter estimation and bures
distance between loss channels.
Physical Review Letters, 121(23):230801, 2018.

[30]
Ranjith Nair, Guo Yao Tham, and Mile Gu.
Optimal gain sensing of quantum-limited phase-insensitive amplifiers.
Physical Review Letters, 128(18):180506, 2022.

[31]
Ronald J Evans and J Boersma.
The entropy of a poisson distribution (c. robert appledorn).
SIAM Review, 30(2):314–317, 1988.

[32]
Jerrold E Marsden, Michael J Hoffman, Terry Marsden, et al.
Basic complex analysis.
Macmillan, 1999.

[33]
Borzoo Rassouli and Bruno Clerckx.
An upper bound for the capacity of amplitude-constrained scalar awgn
IEEE Communications Letters, 20(10):1924–1926, Oct 2016.

[34]
A.L. McKellips.
Simple tight bounds on capacity for the peak-limited discrete-time
International Symposium onInformation Theory, 2004. ISIT 2004.
Proceedings., 2004.

[35]
Andrew Thangaraj, Gerhard Kramer, and Georg Bocherer.
Capacity bounds for discrete-time, amplitude-constrained, additive
white gaussian noise channels.
IEEE Transactions on Information Theory, 63(7):4172–4182, Jul

[36]
Amos Lapidoth and Stefan M. Moser.
On the capacity of the discrete-time poisson channel.
IEEE Transactions on Information Theory, 55(1):303–322, Jan

[37]
Amos Lapidoth, Stefan M. Moser, and MichÈle A. Wigger.
On the capacity of free-space optical intensity channels.
IEEE Transactions on Information Theory, 55(10):4449–4461, Oct