
\section{Details on Entropy Regularized Variational Autoencoders}
Adding an entropy regularization term to the E-step (equation \eqref{equ:evae}) 
equals to multiplying the repulsive force by a large weight coefficient to further strength its repulsive effect. 
The proof is following:
\begin{align*}
&\E_{\hat p(x)} \big[\mathrm{KL}[q_\eta(z ~| ~x) || p_\theta(z ~| ~x)]\big] - \alpha\E_{\hat p(x)q_\eta(z|x)}\big[- \log (q_\eta(z~|~x))\big] \\
& = \E_{\hat p(x)q_\eta(z|x)}\big[(1 + \alpha)\log q_\eta(z ~| ~x) - \log p _\theta(z ~| ~x)\big] \\
& \propto \E_{\hat p(x)q_\eta(z | x)}\big[\log q_\eta(z ~| ~x) - \frac{1}{1 + \alpha}\log p_\theta(z ~| ~x)\big] \\
& \propto \E_{\hat p(x)}\big[\mathrm{KL}[q_\eta(z~| ~x) || p^{\frac{1}{1 + \alpha}}_\theta(z~|~x)]\big].
\end{align*}
Together with equation \eqref{equ:update11}, the Stein variational gradients related to $\nabla_z \log p^{\frac{1}{1 + \alpha}}_\theta(z^{(i)}_k |x^{(i)})$ is 
\begin{align*}
\ff^*(z_k^{(i)}) = \frac{1}{n} \sum_{j=1}^n &[ \frac{1}{1 + \alpha} \nabla_{z} \log p(z_j^{(i)} ~|~ x^{(i)}) k(z_j^{(i)}, z_k^{(i)}) \\
	&+ \nabla_{z_j^{(i)}} k(z_j^{(i)}, z_k^{(i)}) ] \\
\propto \frac{1}{n} \sum_{j=1}^n & [\nabla_{z} \log p(z_j^{(i)} ~|~ x^{(i)}) k(z_j^{(i)}, z_k^{(i)}) \\
	&+ (1 + \alpha )\nabla_{z_j^{(i)}} k(z_j^{(i)}, z_k^{(i)})].
\end{align*}

\section{Derivation from Equation \eqref{equ:follow2} to \eqref{equ:follow3}}
Suppose $\delta$ in equation \eqref{equ:follow2} is very small, we can perform one step gradient descent of equation \eqref{equ:follow2}, which gives,
\begin{align*}
\Delta \eta^t \approx \delta - \epsilon\sum_{i=1}^m \big(\partial_\eta f(\xi_i;~\eta^t)  \delta   -  \ff^*( z_i ) \big)\partial_\eta f(\xi_i;~\eta^t) .
\end{align*}
Since $\delta$ is small enough, all items that invole $\delta$ in the above equation can be ignored. Thus we have
\begin{align*}
\eta^{t+1} &\gets \eta^t - \epsilon \Delta \eta^t \\
&\approx \eta^t + \epsilon  \sum_{i=1}^m   \partial_\eta f(\xi_i;~\eta^t) \ff^*(z_i).
\end{align*}
\end{comment}


\begin{figure*}[htbp]
\centering
\begin{tabular}{cccc}
\raisebox{2.em}{\rotatebox{90}{\small Log10 MSE}}\includegraphics[width=0.255\textwidth]{figures/steps/cost_ex.pdf} &
\includegraphics[width=0.255\textwidth]{figures/steps/cost_ex2.pdf} &
\includegraphics[width=0.255\textwidth]{figures/steps/cost_cosx.pdf} &
\raisebox{2.5em}{\includegraphics[width=0.11\textwidth]{figures/steps/cost_legend.pdf}} \\
{\small time (seconds)} & {\small time (seconds)} & {\small time (seconds)} &\\
(a) $\E(x)$ & (b) $\E(x^2)$ & (c)  $\E(\cos(wx+b))$  & \\
\end{tabular}
\vspace{-8bp}\caption{\small
Results when using different numbers of gradient decent steps for solving \eqref{equ:follow2}. 
The setting is the same as that in Figure~\ref{fig:gmm}, but 
we conduct experiments using 1, 5, 10, 20 gradient steps when solving \eqref{equ:follow2}, 
and show their corresponding training time in the $x$-axis, 
and their mean square error for estimating $\E_p h$ (for the ``testing'' distributions) in the $y$-axis. 
The Langevin samplers we used have $T=10$ layers (Langevin update steps). 
The results are evaluated by drawing 1,000 samples from the trained samplers at different iterations of SVGD. The dimension of the Gaussian Mixtures is $d=50$. 
}
\label{fig:computation}
\end{figure*}

\begin{figure*}[htbp]
\centering
\begin{tabular}{ccc}
\includegraphics[width=0.25\textwidth]{figures/vae/vae.png} \hspace{20bp} &
\includegraphics[width=0.25\textwidth]{figures/vae/entropy_vae.png} \hspace{20bp} &
\includegraphics[width=0.25\textwidth]{figures/vae/ours.png} \\
\small VAE-CNN & \small EVAE-CNN & \small ESteinVAE-CNN \\
\end{tabular}
\vspace{-8bp}\caption{\small Images generated by VAE-CNN, EVAE-CNN, and ESteinVAE-CNN}
\label{fig:vae_samples}
\end{figure*}
 
\section{Solving the Projection Step Using Different Numbers Gradient Steps}\label{sec:diffsteps}
Amortized SVGD requires us to solve the projection step 
using either \eqref{equ:follow1} or \eqref{equ:follow2} at each iteration. 
In practice, we approximately solve it using only one step of gradient descent 
starting from the old values of $\eta$ for the sake of computational efficiency. 

In order to study the trade-off of accuracy and computational cost here, 
we plot in Figure~\ref{fig:computation}
the results when we solve Eq~\eqref{equ:follow2} using different numbers of gradient descent steps (the result is almost identical when we solve Eq~\eqref{equ:follow1} instead).  
We can see that when using more gradient steps,
although the training time per iteration increases, 
the overall convergence speed may still improve, because it may take less iterations to converge. 
Figure~\ref{fig:computation} seems to suggest that using 5, 10,  20 steps gives better convergence than using a single step, but this may vary in different cases. We suggest to search for the best gradient step if the convergence speed is a primary concern. 
On the other hand, the number of gradient steps seems to have minor influence on the final result at the convergence as shown in Figure~\ref{fig:computation}. 

\section{Images Generated by Different VAEs}
Figure \ref{fig:vae_samples} shows  the images generated by the standard VAE-CNN, the entropy regularized VAE-CNN and ESteinVAE-CNN. We can see that both EVAE-CNN and ESteinVAE-CNN can generate images of good quality.


\section{APPLICATIONS OF AMORTIZED SVGD}\label{sec:apps}

With amortized SVGD, we can
use expressive inference networks to obtain better approximation and 
 explore new applications where traditional VI methods cannot be applied.  
In this section, we introduce two different applications of amortized SVGD. 
One is  training variational autoencoders \citep{kingma2013auto} 
with complex, non-Gaussian encoders, and the other is  
training ``smart'' MCMC samplers that adaptively improve their own hyper-parameters from past experience. 

\subsection{Amortized SVGD For Variational Autoencoders}\label{sec:vae}
Variational autoencoders (VAEs) \citep{kingma2013auto} are latent variable models of form
 $p_\theta(x) = \int_z p_\theta(x | \z) p_\theta(z) dz$ where 
 $\x$ is an observed variable and $\z$ is an un-observed latent variable. 
 Assume the empirical distribution of the observed variable is $\hat p(x)$, 
VAE learns the parameter $\theta$ using a variational EM algorithm 
which approximates the posterior distribution $p_\theta (z | x)$ with a simple encoder $q_\eta(z | x)$, and updates $\theta$ and $\eta$ alternatively by 
\begin{align}
& \theta \gets \argmax_{\theta}~ \E_{\hat p(x)q_\eta(z \mid x)}\big[\log p_\theta(x, ~ z)\big], \label{equ:m}\\
& \eta \gets \argmin_{\eta}~ 
\E_{\hat p(x)}\big[\mathrm{KL}(q_\eta(z \mid x) ~|| ~p_\theta(z \mid x))\big],  \label{equ:evae}
\end{align}
which alternates between updating $\theta$ by maximizing the joint likelihood (M-step \eqref{equ:m}) 
and approximating the posterior distribution $p_\theta (z|\x)$ given fixed $\theta$ with  variational inference (VI) (E-step \eqref{equ:evae}). 
In standard VAE,
 \eqref{equ:m} is performed using standard VI with the reparameterization trick \eqref{equ:rep},
 which requires $q_\eta$ to be tractable. 
Therefore, $q_\eta(z \mid x)$ is often defined as a Gaussian distribution with mean and diagonal covariance parameterized by neural networks with $x$ as input. 
This Gaussian assumption potentially limits the quality of the resulting generative models,
and more expressive encoders can improve the 
performance as shown in recent works \citep[e.g,][to name a few]{
 kingma2016improving, mescheder2017adversarial}.



By applying amortized SVGD to solve the posterior inference problem in \eqref{equ:evae},
we obtain simple algorithms that work with more complex inference networks. 
Specifically, we assume that $z \sim q_\eta(z|x)$ 
is generated by $z = f(\xi, x; ~ \eta)$, and optimize $\eta$ using update 
\eqref{equ:follow1}-\eqref{equ:follow3}. 
See Algorithm~\ref{alg:alg2}. 
In our experiment, we take $z = f(\xi, x; ~\eta)$
to be a deep neural network with binary Bernoulli dropout noise at the input layer of the network
which is more effective in approximating multi-modal posteriors than the simple Gaussian encoders. 


Similar idea has also been explored in \citet{punips17steinvae}. Besides, we propose an \emph{entropy regularized VAE} 
to improve the standard VAE and get more diverse images by adding an entropy regularization term on the encoder networks. 
Here, we replace the $\eta$-update in \eqref{equ:evae} with 
\begin{align}
\eta 
& \gets \argmin_{\eta} \E_{\hat p(x)} [ \mathrm{KL}(q_\eta(z \mid x) ~|| ~p_\theta(z \mid x))\big] 
- \alpha H(q_\eta) \notag \\
& = \argmin_{\eta} \E_{\hat p(x)}  \big[ \mathrm{KL}(q_\eta(z \mid x) ~|| ~ p_\theta(z \mid x)^{\frac{1}{1+\alpha}})\big],   
\label{equ:entropy_vae}
\end{align}
where $H(q_\eta) =  - \E_{\hat p(x)q_\eta(z |x)}[\log q_\eta(z |x)]$ is the (conditional) 
entropy of $q_\eta(z | x)$, and $\alpha$ is a regularization coefficient. 
\eqref{equ:entropy_vae} can be solved by 
applying SVGD on the tempered distribution $p_\theta(z \mid x)^{\frac{1}{1+\alpha}}$, 
which can be done by $z_i ~ \gets ~ z_i  ~  + ~ \epsilon \ff^*(z_i)$ with 
\begin{align}\label{equ:update11alpha}
\begin{split}
\!\!\!\!\!\!\!
\ff^*(z_i) = \frac{1}{n} \sum_{j=1}^n [  &  \nabla_{z_j} \log p_\theta(z_j | x) k(z_j, z_i) ~~+  \\
& ~~~~~~~~~~~~~~~~~+ ~~ (1+\alpha) \nabla_{z_j} k(z_j, z_i) ], 
\end{split}
\end{align}
where the temperature parameter $(1+\alpha)$ 
becomes a weight coefficient of the repulsive force; 
a high temperature (or equivalent a large entropy regularization) yields 
a strong repulsive force and push the particles to be further away from each other. 


\begin{comment}
By applying amortized SVGD to solve the posterior inference problem in \eqref{equ:evae}, 
we obtain a simple algorithm that work with more complex inference networks. 
Specifically, we assume that $z \sim q_\eta(z|x)$ 
is generated by $z = f(\xi, x; ~ \eta)$, and optimize $\eta$ using update 
 \eqref{equ:follow1}-\eqref{equ:fnollow3}. See Algorithm~\ref{alg:alg2}.
\end{comment}

\begin{algorithm}[t]                     
\caption{Amortized SVGD for training VAE
\label{alg:alg2}                          
\begin{algorithmic}                   
\STATE Set batch size $m$ and kernel $k(\z, \z')$. 
\STATE Initialize encoder parameter $\eta$ and decoder parameter $\theta$.
\FOR {iteration $t$}
\STATE Pick input $x^{(i)}$ from the training data.
\STATE Draw $\{\xi_k\}_{k=1}^m\sim q_0$, calculate $z_k^{(i)}=f(x^{(i)}, \xi_k;~\eta)\vspace{.1cm}$, and their Stein variational gradients $\ff^*(z_k^{(i)})$ in \eqref{equ:update11} or \eqref{equ:update11alpha} 
related to gradient $\nabla_z \log p_\theta(z_k^{(i)}|x^{(i)})$.
\STATE Update $\eta$ using either \eqref{equ:follow1}, \eqref{equ:follow2} or \eqref{equ:follow3}. 
\STATE Update $\theta$ with \\ $~~~~~~~~~ \theta \gets \theta +  \frac{1}{m} \sum_{k=1}^{m}\nabla_\theta \log p_\theta(x^{(i)}, ~ z_k^{(i)})$.
\ENDFOR
\end{algorithmic}
\end{algorithm}





\subsection{Training Langevin Samplers}
\label{sec:trainlangevin}
By viewing typical MCMC procedures as a simulator $f(\cdot)$, 
we can apply amortized SVGD to adaptively improve hyperparameters in MCMC inference. 
This is useful when we need to 
perform Bayesian inference on a large number of different, but similar datasets or posteriors,
where we can adaptively improve the MCMC sampler for future tasks by leveraging the information of the previous tasks.  
An example of this, which we consider in this work, is adaptively learning optimal step sizes for Langevin dynamics. 

To specify the general framework,
we assume that we are interested in drawing samples from a set of distributions 
$$\Q = \{ p_\vartheta(\z)  \colon \vartheta \in \Theta\},$$ 
indexed by parameter $\vartheta$. 
We are interested in learning 
a network $f(\xi, p_\vartheta; ~\eta)$ 
which maps the distribution $p_\vartheta$ to stochastic posterior samples. 
Note that this is a generalization of Problem 1 which 
focuses on approximating an individual distribution. 
In practice, $p_\vartheta$ could be the posterior distributions of unknown parameters of interest conditioning on different observed data, 
or models of different individuals in hierarchical models.  

In order to learn the network $f(\xi, p_\vartheta; ~\eta)$, we modify Algorithm~\ref{alg:alg1}, 
to perform amortized SVGD on a randomly selected $p_\vartheta$ from $\Q$ at each iteration.  
See Algorithm~\ref{alg:alg3}. 
In this way, we expect that the trained network $f(\xi, p_\vartheta; ~\eta)$ can perform well on similar $p_\vartheta$ drawn from the same distribution, 
but never seen by the training algorithm.

A useful perspective is that typical MCMC methods can be viewed 
as neural networks $f(\xi, p_\vartheta; ~\eta)$ 
handcrafted by researchers, with nice theoretical properties. 
We can leverage the structure of existing MCMC to 
design the architecture of $f(\cdot)$, 
and use amortized SVGD to adaptively improve its parameters across different tasks. 

 As an example, 
Langevin dynamics draws samples from $p_\vartheta$ 
by starting with an initial sample $z^0$ and performing iterative random updates of form  
$z^{t+1} \gets f_t(z^t)$ with 
\begin{align}\label{equ:langevin} 
& f_t(z^t) = \z^{t} + \eta^t \odot  \nabla_\z \log p_\vartheta(\z^t)   + \sqrt{2\eta^t} \odot \xi^t, 
\end{align}
which performs gradient ascent with a Gaussian perturbation.  
Here $\eta^t$ denotes a vector-valued step size at the $t$-th iteration and ``$\odot$'' denotes element-wise product, 
and $\xi^t$ is a standard Gaussian random vector of the same size as $\z^{t}$. 



\begin{algorithm}[tbp]                     
\caption{Amortized SVGD for Learning to Sample a distribution family $\Q = \{p_\vartheta \colon \vartheta \in \Theta \}$
\label{alg:alg3}                          
\begin{algorithmic}                   
\STATE Goal: Learn parameter $\eta$ to train $z = f(\xi, p_\vartheta; ~\eta)$ drawing samples from $p_\vartheta$. 
\STATE Set batch size $m$, and kernel $k(\z,\z')$. Initialize $\eta^0$.  
\FOR {iteration $t$}
\STATE Randomly select $p_\vartheta$ from $\Q$.
\STATE Draw random seed $\{\xi\}_{i=1}^{m}\sim q_0.$
\STATE Calculate $z_i = f(\xi_i, ~ p_\vartheta; ~\eta)$, 
and the Stein variational gradient $\ff^*(z_i)$ in \eqref{equ:update11}. 
\STATE  Update parameter $\eta$ using either \eqref{equ:follow1}, \eqref{equ:follow2} or \eqref{equ:follow3}. 
\ENDFOR
\end{algorithmic}
\end{algorithm}

We can view $T$ iterations of Langevin dynamics as a $T$-layer neural network: 
\begin{align}\label{equ:ls} 
z^T = f(\xi, ~ p_\vartheta; ~ \eta) = f_{T-1} \circ \cdots \circ f_0(\xi, ~ p_\vartheta; ~ \eta), 
\end{align}
in which 
the initial samples and the Gaussian noise form the random seeds of the network, that is, $\xi =\{\xi^t\}_{t = 0}^{T-1}\cup \{z^0\}$, and 
the step sizes $\eta=\{\eta^t\}_{t = 0}^{T-1}\vspace{0.1 cm}$ form the parameters which we can estimate using Algorithm~\ref{alg:alg3}. 
In cases when 
it is difficult to calculate $\nabla_z \log p_\vartheta$ exactly, such as the case of Bayesian inference with 
big datasets, 
we can use stochastic gradient Langevin dynamics \citep{welling2011bayesian}
to approximate $\nabla_z \log p_\vartheta$ with subsampling, which introduces another source of randomness. 
This, however, does not influence the application of Algorithm~\ref{alg:alg3}, 
since our method does not need to know the structure of the random seed distribution. 


In practice, a large value of $T$ would result in 
a deep network and cause a vanishing gradient problem. 
We address this problem by partitioning the $T$ layers into small blocks of size $5$ or $10$, 
and evaluate the gradient of the parameters in each block by back-propagating the Stein variational gradient from the output of its own block. 


















\section{AMORTIZED SVGD: TOWARDS AN AUTOMATIC NEURAL SAMPLER}
\label{sec:amortizedsvgd}

SVGD and other particle-based methods become inefficient when we need to apply them repeatedly on 
a large number of different, but similar target distributions for multiple tasks,
because they can not leverage the similarity between the different distributions 
and may require a large memory to restore a large number of particles. 
This problem can be addressed by training a neural network $f(\xi;~\eta)$ to output particles that would 
have been produced by SVGD;
this amounts to  ``amortizing'' or compressing the nonparametric SVGD 
into a parametric network, 
yielding a solution for wild variational inference in Problem~\ref{pro:prob1} of Section~1.

One straightforward way to achieve this is to run SVGD until convergence and 
train $f(\xi;~\eta)$ to fit the resulting SVGD particles (e.g., by using generative adversarial networks (GAN) \citep{goodfellow2014generative}). 
This, however, requires running many epochs of fully converged SVGD and can be slow in practice. 
We instead propose an \emph{incremental approach} in which $\eta$ is iteratively adjusted 
so that the network outputs $z = f(\xi;~\eta)$ improves by moving along the Stein variational gradient direction in \eqref{equ:update11}, 
in order to move towards the target distribution. 

Specifically, denote by $\eta^t$ the parameter estimated at the $t$-th iteration of our method; 
each iteration of our method 
draws a batch of random inputs $\{\xi_i\}_{i=1}^m$ 
and calculate their corresponding output $z_i = f(\xi_i;~\eta^t)$ based on $\eta^t$, where $m$ is a mini-batch size (e.g., $m=100$). 
The Stein variational gradient $\ff^*(z_i)$ in 
\eqref{equ:update11} would then ensure that $z'_i = z_i + \epsilon \ff^*( z_i)$ forms a better approximation of the target distribution $p$. 
Therefore, we should adjust $\eta$ to make it output $\{z'_i\}$ instead of $\{z_i\}$, 
that is, we want to update $\eta$ by
\begin{align}\label{equ:follow1}
\eta^{t+1} \gets & \argmin_\eta  \sum_{i=1}^m || f(\xi_i;~\eta)  - z_i' ||_2^2,
\end{align}
where $z_i'  = z_i  + \epsilon \ff^*(z_i)$. 
This process is repeated until convergence, 
in which case the outputs of network $f$
can no longer be improved by SVGD and hence should form a good approximation of the target $p$. 
See Algorithm~\ref{alg:alg1}.

If we assume $\epsilon$ is very small, then 
\eqref{equ:follow1} can be approximated by a least square optimization. To see this, note that 
$f(\xi_i;~\eta) \approx f(\xi_i;~\eta^t) + \partial_\eta f(\xi_i;~\eta^t) (\eta - \eta^t)$ by Taylor expansion. 
Since $z_i = f(\xi_i;~\eta^t)$, we have
$$
 ||  f(\xi_i;~\eta)  - z_i' ||_2^2 \approx || \partial_\eta f(\xi_i;~\eta^t) (\eta  - \eta^t)  - \epsilon \ff^*(z_i)  ||_2^2. 
$$
As a result, \eqref{equ:follow1} reduces to a least square optimization: 
\begin{align}\label{equ:follow2}
\begin{split}
&  \eta^{t+1} \gets  \eta^t - \epsilon \Delta \eta^t, \\ 
& \Delta \eta^t = \argmin_{\delta}  \sum_{i=1}^m   || \partial_\eta f(\xi_i;~\eta^t)  \delta   -  \ff^*( z_i ) ||_2^2.
\end{split}
\end{align}
It may be still slow to solve a least square problem at each iteration. 
We can derive a more computationally efficient approximation by performing only one step of gradient descent of \eqref{equ:follow1} starting at $\eta^t$ (or equivalently \eqref{equ:follow2} starting at $\delta = 0$), which gives 
\begin{align}\label{equ:follow3}
\eta^{t+1} \gets \eta^t + \epsilon  \sum_{i=1}^m   \partial_\eta f(\xi_i;~\eta^t) \ff^*(z_i). 
\end{align}
Although update \eqref{equ:follow3} is derived as an approximation of \eqref{equ:follow1} or \eqref{equ:follow2}, 
it is computationally faster and it works effectively in practice; 
this is because when $\epsilon$ is small, one step of gradient update can be sufficiently close to the optimum.  

Update \eqref{equ:follow3} has a simple and intuitive interpretation: it can be thought as \emph{a ``chain rule'' that back-propagates the Stein variational gradient to the network parameter $\eta$}. 
This can be justified by considering the special case when we use only a single particle $(n=1)$ 
in which case $\ff^*(z_i)$ in \eqref{equ:update11} reduces to the typical gradient $\nabla_z \log p(z_i)$,
and update \eqref{equ:follow3} reduces to the typical gradient ascent for maximizing $\E_{\xi}[\log p(f(\xi;~\eta))],$ 
 in which case $f(\xi;~\eta)$ is trained to maximize $\log p(z)$ (that is, \emph{learning to optimize}), 
instead of \emph{learning to draw samples from $p$} for which it is crucial to use the Stein variational gradient $\ff^*( z_i)$ to diversify the outputs to capture the uncertainties in $p$. 

Update \eqref{equ:follow3} also has a close connection with the typical variational inference with the reparameterization trick \citep{kingma2013auto}. Let $q_\eta(z)$ be the density function of $z = f(\xi;~\eta)$, $\xi\sim q_0$. Using the reparameterization trick, the gradient of $\KL(q_\eta~||~p)$ w.r.t. $\eta$ equals
\begin{align*}
\nabla_\eta\KL(q_{\eta}~||~p) 
= -\E_{\xi \sim q_0}[ \partial_\eta f(\eta;~\xi)
	 \nabla_z \log (p(z)/q_\eta(z))]. 
\end{align*}
With $\{\xi_i\}$ i.i.d. drawn from $q_0$ and $z_i = f(\xi_i; ~\eta), ~\forall i$, we can obtain a standard stochastic gradient descent for minimizing the KL divergence:  
\begin{align}\label{equ:rep}
\begin{split}
&\eta^{t+1} \gets \eta^t +  \sum_{i=1}^m \partial_\eta f(\xi_i;~\eta^t) \tilde \ff^*(z_i), \\
\text{where} ~~~~ 
&\tilde \ff^*(z_i) = \nabla_z \log p(z_i) - \nabla_z \log q_{\eta^t}(z_i). 
\end{split}
\end{align}
This is similar to \eqref{equ:follow3}, but replaces the Stein gradient $\ff^*(z_i)$ defined in \eqref{equ:update11} with 
$\tilde \ff^*(z_i)$. 
However, because $\tilde \ff^*(z_i)$ depends on the density $q_{\eta^t}$, 
which is assumed to be intractable in Problem 1, \eqref{equ:rep} is not directly applicable in our setting. 
Further insights can be obtained by noting that 
\begin{align}
\label{equ:tmp}
\ff^*(z_i) 
& \approx \E_{z\sim q_{\eta^t}}[\nabla_z \log p(z)k(z,z_i) + \nabla_zk(z,z_i)] \notag \\
& =  \E_{z\sim q_{\eta^t}}[(\nabla_z \log p(z) - \nabla_z \log q_{\eta^t}(z))k(z,z_i)]   \notag \\
& = \E_{z\sim q_{\eta^t}} [ \tilde \ff^*( z) k(z, z_i)], 
\end{align}
where \eqref{equ:tmp} is obtained by using Stein's identity \eqref{equ:steinid}. 
Therefore, $\ff^*(z_i)$ can be treated as a smoothed version of $\tilde \ff^*(z_i) $ obtained by convolving it with kernel $k(z,z')$. 





\section{EXPERIMENTS}

We apply amortized SVGD to the two different applications mentioned in section \ref{sec:apps}, 
 and demonstrate that our method can train expressive inference networks to draw samples from intractable posterior distributions. For all our experiments, we use the standard RBF kernel $k(z, z')=\exp(- \frac{1}{h}|| z- z'||^2_2)$, and take the bandwidth to be $h = \mathrm{med}^2 / \log n$, where $\mathrm{med}$ is the median of the pairwise distance between the current points $\{z_i\}^n_{i=1}$. 
 We use update \eqref{equ:follow3} in Algorithm 1, which solves Eq \eqref{equ:follow1}-Eq \eqref{equ:follow2} using a single gradient step. 
 We find that using more gradient steps does not change the final results  significantly, 
 but may potentially increase the convergence speed (see Appendix \ref{sec:diffsteps}).


\subsection{Training Langevin Samplers}\label{sec:sampler}
In this section, 
we use amortized SVGD (Algorithm~\ref{alg:alg3}) to learn the step size parameters in the Langevin sampler in \eqref{equ:ls}.  
We test a number of distribution families, 
including Gaussian Mixture, Gaussian Bernoulli RBM, Bayesian logistic regression and Bayesian neural networks. 
In all the cases, we train the sampler 
with a set of ``training distributions'' and evaluate the sampler 
on ``test distributions'' that are not seen by the algorithm during training. 
We compare our method to the 
typical Langevin sampler 
with power decay step size, selected to be the best from $\eta^t =  10^a / (t + b)^{\gamma}$
where $\gamma = 0.55$, $a \in \{-6,...,2\}$, $b \in \{0,...,9\}$. 

\paragraph{Gaussian Mixture}\label{para:gmm}
We first train the Langevin samplers 
to learn to sample from simple Gaussian mixtures. 
We consider a family of Gaussian Mixtures  
$q_\vartheta(z) = \frac{1}{10} \sum_{i=1}^{10} \N(z; \vartheta_i, 0.1^2)$, 
where $\vartheta$ is the mean parameter. 

We train the sampler by drawing random elements of $\vartheta$ from $\mathrm{Uniform}(-1,1)$ at each iteration of Algorithm~\ref{alg:alg3},
and evaluate the quality of the samplers on new values of 
$\vartheta$ drawn from the same distribution, but not seen during training. 
In Figure~\ref{fig:gmm},
we find that the sampler 
trained by amortized SVGD obtains good approximation 
with $T=5$ or $15$ steps of Langevin updates (Figure~\ref{fig:gmm}(d)-(e)). 
In comparison, the typical Langevin sampler with the best power decay step size performs much worse (Figure~\ref{fig:gmm}(b)).
\begin{comment} 
To gain a further understanding, 
we tested a particle version of KSD variational inference that directly 
fit a set of sample $\{z_i\}$ with a target distribution $p$ by $\min_{\{z_i\}}\KL(
\{z_i\}~||~ p)$ using gradient descent, and compare the particle version of SVGD \eqref{equ:update11}; 
in this way, we obtain a comparison between the KSD and KL objectives 
without the influence of the network structure of $f(\cdot)$. 
Figure \ref{fig:ksd} shows result 
on a toy example where we find that the particles obtained by 
the gradient descent of KSD tends to concentrate on a lower dimensional manifold,
which the SVGD particles tend to perform much better. 
Note that both the particles given by SVGD and minimizing KSD 
reaches steady state when KSD equals zero (since KSD is the gradient manguitude of KL as shown in Section~\ref{sec:ksd}). 
We think the degeneracy of the KSD particles is due to the non-convexity when directly minimizing KSD. 
In a related work \citep{liu2016black}, 
it is shown that KSD provides a good discriminative objective for minimizing the importance weights when the optimization is a convex optimization. 
\end{comment}
\begin{comment}
\begin{SCfigure}
\centering
\caption{\small 
Particles obtained by SVGD (red) and 
particles obtained by gradient descent for minimizing the KSD 
$\min_{\{x_i\}} \S(\{x_i\} ~||~ p)$} (black).
\includegraphics[width=0.2\textwidth]{../figures/ksd_svgd.pdf} 
\label{fig:ksd}
\end{SCfigure}
\end{comment}



\begin{figure*}[ht]
\centering
\hspace{-.1\textwidth}\begin{tabular}{ccc}
\raisebox{3em}{\rotatebox{90}{\small Test Accuracy}}  \includegraphics[width=0.20\textwidth]{figures/lr/lr_acc}&
\raisebox{3em}{\rotatebox{90}{\small Test Accuracy}} \includegraphics[width=0.20\textwidth]{figures/lr/nn_acc} &
\raisebox{1.5em}{\includegraphics[width=0.016\textwidth]{figures/lr/nn_legend.pdf} }
\put(2, 60){\small Amortized SVGD (10 steps)}
\put(2, 49){\small Power Decay ($10^4$ steps)}
\put(2, 39){\small Power Decay ($10^3$ steps)}
\put(2, 28){\small Power Decay ($10^2$ steps)}
\put(2, 17){\small Power Decay (10 steps)} \\
\vspace{-.5\baselineskip}
\small ~~~~~~Steps of Langevin updates & \small ~~~~~~Steps of Langevin updates & \\
(a) \small Bayesian logistic regression & (b) \small Bayesian neural networks &
\end{tabular}
\vspace{-10bp}\caption{\small
By training the step size using amortized SVGD on dataset a9a, 
we obtain good performance ``on'' datasets a1a-a8a using Langevin sampler with only $T=10$ steps of Langevin updates. 
}
\label{fig:classification}
\end{figure*}




\paragraph{Restricted Boltzmann Machine (RBM)}\label{sec:rbm}
We test our method on 
Gaussian-Bernoulli RBM which
is high dimensional  and multi-modal. 
Gaussian-Bernoulli RBM is a hidden variable model consisting of a continuous observable variable $z\in \R^d$ and a 
binary hidden variable $h \in \{\pm 1\}^{\ell}$ with joint probability
$p_\vartheta (z, h) \propto \exp( z^\top B h + b^\top z + c^\top h - \frac{1}{2} || z ||_2^2),$
where the parameters include $\vartheta = \{B, b, c\}.$ 
We obtain the marginal distribution of $p_\vartheta(z)$ by summing over $h$:
$p_{\vartheta}(z) \propto \exp\big [ b^\top z - \frac{1}{2}|| z ||_2^2 + \sigma(B^\top z + c) \big ],$
where $\sigma(h) = \sum_{i=1}^\ell\log (\exp(h_i) +\exp(- h_i))$. 
In our experiments, we take $\ell=10$ hidden variables 
and $d=100$ observable variables,
and randomly draw $b$ and $c$ from $\normal(0, I)$ and 
$B$ uniformly from $\{\pm 0.1\}$ in the training time. 
The evaluation is on a new set of parameters $\vartheta = \{B, b, c\}$ drawn from  the same distribution. 

Figure~\ref{fig:rbm} shows the result when we use the trained samplers to estimate integral quantities of form 
$\E_{p_\vartheta} [ h] $ with different testing functions $h$.  
The plots show the MSE for estimating $\E_{p_\vartheta} [h]$
using $\sum_i h(z_i)/n$ with $z_i$ generated from the trained sampler $z_i = f(\xi_i, p_\vartheta; ~\eta)$. 
Here the sample size $n$ is the number of i.i.d. samples generated 
from the trained sampler in the testing time. 
In the case of typical, non-adaptive Langevin samplers, it is the number of 
Langevin Markov chains that we run in parallel. 
In Figure \ref{fig:rbm} (a)-(c), 
we generally find that amortized SVGD allows us to train high quality Langevin samplers with a small number $T$ of Langevin update steps. 
We observe that we can always further refine the 
result of our trained samplers using additional typical MCMC steps. 
For example, in Figure  \ref{fig:rbm}, 
we find that when using 20 steps of typical Langevin dynamics (with a power decay step size)
 to refine the output of the $T=30$ layer Langevin sampler trained by amortized SVGD achieves results close to the exact Monte Carlo. 



\paragraph{Bayesian Classification}\label{sec:bayes}
We test our method on Bayesian Logistic Regression and Bayesian neural networks for binary classification on real world datasets. 
In this case, the distribution of interest 
has a form of $p_\vartheta(z) = p(z | D)$, 
where $z$ is the network weights in logistic regression and neural networks, 
and $D$ is the dataset for binary classification, which we view as the parameter $\vartheta$, 
that is, different dataset $D$ yields different posterior $p(z |D)$, 
and we are interested in training the Langevin sampler
on a set of available datasets, 
and hope it performs well on future datasets that have similar structures. 
This setting can be useful, for example, 
in the streaming setting where we use 
existing datasets to adaptively improve the Langevin sampler. 
In our experiment, 
we take nine similar datasets (a1a-a9) from the libsvm repository\footnote{\url{https://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/}}; 
we train our Langevin sampler on a9a, and evaluate the sampler on the remaining 
8 datasets (a1a-a8a). Our training and evaluation steps are as follows: 

1. Estimate the step sizes $\{\eta^t\}_{t=0}^{T-1}$ of the Langevin sampler using amortized SVGD based on dataset a9a. 

2. Apply the Langevin sampler with the estimated step size to the training subsets of a$k$a, $k=1,\ldots,8$, to obtain posterior samples $\{z_i^k\}$ of the classification weights. 

3. Calculate the test likelihood of $\{z_i^k\}$ on the testing subsets of  a$k$a, $k=1,\ldots,8$. Report the averaged testing likelihood averaged on the 8 datasets in Figure \ref{fig:classification}. 

Because each dataset is relatively large, we use the stochastic gradient approximation as suggested by \citet{welling2011bayesian} (with a minibatch size of 100) in Langevin samplers.  
We find that the $T = 10$ Langevin samplers trained by amortized SVGD on a9a
is comparable with the $T = 10^3$ Langevin sampler with the best power decay step size. 






\subsection{Training VAE With Amortized SVGD}\label{steinvae}
We compare the entropy regularized VAE trained with amortized SVGD (denoted by \textit{ESteinVAE}), 
which is Algorithm~\ref{alg:alg2} with update \eqref{equ:update11alpha}, 
with the standard VAE and entropy regularized standard VAE (denoted by \textit{EVAE}) on the dynamically binarized MNIST dataset \citep{burda2015importance}. 
We tested the following settings:

1. A standard VAE (VAE-f) with a fully connected encoder consisting of one hidden layer with 400 hidden units, and a Gaussian output hidden variable with diagonal covariance. 

2. A standard  convolutional Gaussian VAE (VAE-CNN) with a convolutional encoder consisting of 2 convolution layers with $5\times 5$ filters, stride 2 and $[16, 32]$ features maps,
followed by a fully connected layer with 512 hidden units. 

3. A convolutional Gaussian entropy regularized (EVAE-CNN) with the same encoder and decoder structures as the standard convolutional VAE (VAE-CNN).

4. Entropy regularized VAEs trained by our amortized SVGD in Algorithm~\ref{alg:alg2}, 
 with the same encoder architecture as VAE-f and VAE-CNN, respectively, but removing the Gaussian noise on the top layer and adding a multiplicative Bernoulli noise with a dropout rate of 0.3 to the each layer of the encoder. These two cases are denoted by ESteinVAE-f, and ESteinVAE-CNN, respectively. 

In all these cases, we use $\mathrm{dim}(\z) = 32$ latent variables and the decoders have symmetric architectures as the encoders in all models. Adam
is used with a learning rate tuned on the training images. 
For each training iteration, a batch of 128 images is used, 
and for each image we draw a batch of $m=5$ samples to apply amortized SVGD.



\paragraph{Marginal likelihood}
Table~\ref{tbl:mnist} reports the test log-likelihood of all the methods estimated using Hamiltonian annealed importance sampling (HAIS) \citep{wu2016quantitative} 
with 100 independent AIS chains and 10,000 intermediate transitions, averaged on 5000 test images. 
We find that ESteinVAE-f significantly outperforms VAE-f, 
and ESteinVAE-CNN slightly outperforms VAE-CNN and EVAE-CNN. 
Table~\ref{tbl:mnist} also reports the effective sample size (ESS) of the HAIS estimates. 
The fact that the effective sample sizes of all the methods are close suggests that accuracy of the different NLL estimates are comparable. 



\begin{table}[t]
\caption{\small Negative log-likelihood on binarized mnist test dataset.}
\centering
\vspace{3bp}\begin{tabular}{ccc}
\hline
\textbf{Model} & \textbf{NLL/nats} & \textbf{ESS} \\\hline
\small VAE-f & 90.32 &  84.11 \\%  90.10 
\small ESteinVAE-f & 88.85 & 83.49 \\
\hline
\small VAE-CNN & 84.68 & 85.50 \\ 
\small EVAE-CNN & 84.43 & 84.91 \\
\small ESteinVAE-CNN & 84.31 & 86.57 \\% 83.85 \\
\hline
\end{tabular}
\label{tbl:mnist}
\end{table}

\begin{table}[t]
\caption{\small Quantitative imputation experiment based on 2,000 true images and 500 reconstructed images per true image.} 
\centering
\vspace{3bp}\begin{tabular}{ccc}
\hline
\textbf{Model} & \textbf{Accuracy} & \textbf{Entropy} \\
\hline
\small ESteinVAE-CNN   & 0.84			   & 0.501 \\
\small EVAE-CNN 		& 0.82			   & 0.382 \\
\small VAE-CNN        & 0.83              & 0.340   \\
\hline
\end{tabular}
\label{tb:imputation}
\end{table}

\begin{comment}
\begin{table}[t]
\caption{\small Quantitative imputation experiment based on 2,000 true images and 500 reconstructed images per true image.} 
\centering
\vspace{3bp}\begin{tabular}{ccc}
\hline
\textbf{Model} & \textbf{Accuracy} & \textbf{Entropy} \\
\hline
\small ESteinVAE-CNN   & 0.84			   & 0.570 \\
\small EVAE-CNN 	   & 0.83		   & 0.442 \\
\small VAE-CNN        & 0.82	    & 0.348   \\
\hline
\end{tabular}
\label{tb:imputation}
\end{table}

If I increase the imputation area, we can see that both the entropy term of EVAE and SteinVAE increases but the entropy of VAE stays, which makes sense because we now have larger uncertainty. For Accuracy, there may be two reasons that cause mismatching (generate more diverse images from the imputated ambiguous images or all the generated images are the same label but different with original image label.) The first one may cause the low accuracy in SteinVAE. For VAE, there still may exsits some images, even with imputation, it still can have a good reconstructed images (for example digit "1"), so we can not only see the image accuracy. What's mroe. the imputation process did not suggest that we should have a higher accuracy(each generated images are indepdent). 

In the new experiment, the entropy of the VAE stays while the accuracy drops a lot, which indicates it generates images with the same label but different with original one.
\end{comment}

\begin{figure}[t]
\centering
\begin{tabular}{c}
\includegraphics[width=0.45\textwidth]{figures/vae/steinvae_NMAR.png} \\
\small ESteinVAE-CNN \\
\includegraphics[width=0.45\textwidth]{figures/vae/entropy_vae_NMAR.png} \\
\small EVAE-CNN \\
\includegraphics[width=0.45\textwidth]{figures/vae/vae_NMAR.png} \\
\small VAE-CNN\\
\end{tabular}
\vspace{-10bp}\caption{\small Imputation results from benchmark VAE-CNN, EVAE and ESteinVAE-CNN. The first column shows the original image. Col. 2 shows the missing image.
The remaining columns show imputations for 20 reconstructed samples.}
\label{fig:denosing}
\end{figure}



\paragraph{Missing data imputation}
We demonstrate that ESteinVAE is able to learn multi-modal latent representation with the binary dropout noise.  
We consider missing data imputation with pixel missing in a square sub-regions of the image (the second column of Figure~\ref{fig:denosing}). 
We use a simple method to reconstruct image by applying one step of encode-decode operation starting from uniform noise $\mathrm{Uniform}(0, 1)$. 
Specifically, let $x = [x^{v}, x^m]$ where $x^v$ and $x^m$ denotes the visible and missing parts, respectively. Our imputation procedure consists of
the following  steps: 
\begin{enumerate}

\item Draw $x^c$ from $\mathrm{Uniform}(0, 1)$;

\item Draw $z \sim q_\phi(z | x^c,~x^{v})$; 

\item Draw $x \sim p(x | z)$.
\end{enumerate}

Each column (starting from the third column) of Figure~\ref{fig:denosing} shows an independent run of this procedure, 
where we can see that ESteinVAE-CNN is able to generate diverse construction when ambiguity exists, 
while the EVAE-CNN and VAE-CNN tend to be trapped in a local mode. 
This suggests that the diagonal variance of the latent variables in the Gaussian encoder of VAE-CNN tends to be small, 
underestimating the posterior uncertain, 
while ESteinVAE can capture the multi-modal posterior due to the dropout noise. 


Table \ref{tb:imputation} is the quantitative result of the imputation experiment, 
in which  the ``accuracy'' column denotes the number of original images whose digit is in its reconstructed images.
and the ``entropy'' column denotes  the entropy of the probability of the reconstructed images belonging to different digit classes. 
EsteinVAE obtains more diverse images and slightly more accurate reconstructed images.





\section{CONCLUSION}

We propose a new method to train neural samplers for given distributions, together with various applications to learning to draw samples using neural samplers. Future directions include exploring more efficient neural architectures and theoretical understanding of our method.



\section{INTRODUCTION}

Modern machine learning increasingly relies on highly complex probabilistic models to reason about uncertainty.  
A key computational challenge is to develop efficient inference techniques to approximate, or draw samples from complex distributions. 
Currently, most inference methods, including MCMC and variational inference, are \emph{hand-designed} by researchers or domain experts. 
This makes it difficult to fully optimize the choice of different methods and their parameters, and exploit the structures in the problems of interest in an automatic way. 
The hand-designed algorithm can also be inefficient when there is a need to perform fast inference repeatedly on a large number of different distributions with similar structures. 
This happens, for example, when we need to reason about a number of observed datasets in settings like online learning or personalized prediction,  
or need fast inference as inner loops for other algorithms such as  learning latent variable models (such as variational autoencoder \citep{kingma2013auto}) or unnormalized distributions. 
Therefore, it is highly desirable to develop intelligent probabilistic inference systems that can adaptively improve their own performance to fully optimize the computational efficiency, and generalize to new tasks with similar structures. 
Developing such systems requires solving 
the following \emph{learning-to-sample} problem: 
\begin{problem}\label{pro:prob1}
Given a distribution with density $p(z)$ on set $\Z$
and a simulator $z= f(\xi;~\eta)$, such as a neural network, 
which takes a parameter $\eta$ and a random seed $\xi$ drawn from $q_0$ and outputs a value $z$ in $\Z$, 
we want to find an optimal parameter $\eta$ so that the density of the random output  $z = f(\xi;~\eta)$ with $\xi\sim q_0$ closely matches the target $p$. 
\end{problem}
Here, we  assume that we do not know the analytical form 
of the simulator $f(\cdot)$ (which we call inference network),  
and we can only query it through the output value $f(\xi;~\eta)$ and derivative $\partial_\eta f(\xi;~\eta)$ for given $\eta$ and $\xi$.  
We also assume that the random seed distribution $q_0$ is unknown and we can  
only access it through the draws of the random input $\xi$; 
that is, $q_0$ can be arbitrarily complex, and can be discrete, continuous or hybrid. 

Because of the above assumption, 
we cannot directly calculate the density $q_{\eta}(z)$ of the output variable $z = f(\xi;~\eta)$; 
this makes it difficult to solve Problem~\ref{pro:prob1} using the typical variational inference (VI) methods. 
Recall that VI finds an optimal parameter $\eta$ to approximate the target $p$ with $q_{\eta}$, 
by minimizing the KL divergence:
\begin{align}\label{equ:kl}
\KL(q_\eta ~||~p) = \E_{q_\eta}[\log (q_\eta/p)].
\end{align}
However equation \eqref{equ:kl} requires calculating the density $q_\eta(z)$ or its derivative, which is intractable by our assumption (and is called implicit models in \citet{mohamed2016learning}). 
This holds true even when Monte Carlo gradient estimates \citep{hoffman2013stochastic} and the reparametrization trick \citep{kingma2013auto} are applied. 

\begin{comment}{}
\begin{figure}[t]
   \centering
   \scalebox{1}{
   \includegraphics[width=.6\textwidth]{figures/neuralsampler1
   }
   \begin{picture}(0,0)(0,0)
   \put(-250,-10){\figcapsize \emph{Given distribution}}
      \put(-170,-10){{\figcapsize \emph{Black-box neural sampler}}}
   \put(-50,-10){\figcapsize \emph{Samples}}
   \end{picture}\\[.5em]
   \caption{\figcapsize   
   Our methods ``learn to draw samples'',  
   constructing black-box neural samplers for given distributions. 
   It allows us to automatize the hyper-parameter tuning of Bayesian inference, 
 speed up the inference inner loops of learning algorithms, 
and eventually replace hand-designed inference algorithms with more efficiently one that is trained on past tasks and is improved adaptively over time.}
   \label{fig:example}
\end{figure}
\end{comment}

This requirement of calculating $q_\eta(z)$
makes it difficult for practitioners to use variational inference in an automatic way,  
especially in domains where it is critical to use 
expressive inference networks to achieve good  
approximation quality. 
Methods that do not require to explicitly calculate $q_\eta(z)$, 
referred to as \emph{wild variational inference}, or \emph{variational programming} \citep{operator},  
can significantly simplify the design and expand the applications of VI methods, 
allowing practioners to focus more on choosing proposals that work best with their specific tasks.
\begin{comment}
this is orthogonal to 
 \emph{black-box variational inference} \citep{ranganath2013black}
which refers to methods that work for generic target distributions $p(z)$ without gradient calculation, or significant model-by-model consideration (but still require to calculate the proposal density $q_\eta(z)$). 
\end{comment}


\begin{figure}[t]
   \centering
   \scalebox{1}{
   \includegraphics[width=.45\textwidth]{figures/neuralsampler1
   }
   \begin{picture}(0,0)(0,0)
   \put(-230,-10){\figcapsize \emph{Given distribution}}
      \put(-145,-10){{\figcapsize \emph{Inference network}}}
   \put(-45,-10){\figcapsize \emph{Samples}}
   \end{picture}\\[.5em]
   \caption{
   \figcapsize   
   Wild variational inference allows us 
   to train general stochastic neural inference networks to learn to draw (approximate) samples from the target distributions, 
without restriction on the computational tractability of the density function of the neural inference networks.   
   }
   \label{fig:example}
\end{figure}


A similar problem also appears in importance sampling (IS), 
where it requires calculating the IS proposal density $q(z)$ in order to calculate the importance weight $w(z) = p(z)/q(z)$. 
However,
there exist methods that use no explicit information of the proposal $q(z)$, and, seemingly counter-intuitively, give better asymptotic variance or converge rate than the typical IS that uses the proposal information \citep[e.g.,][]{liu2016black, briol2015probabilistic, henmi2007importance, delyon2014integral}.
Discussions on this phenomenon date back to \citet{o1987monte}, who 
argued that ``Monte Carlo (that uses the proposal information) is fundamentally unsound'',
and developed Bayesian Monte Carlo \citep{o1991bayes} as 
an instance that uses no information of proposal $q(z)$, yet gives better convergence rate than the typical Monte Carlo $\Od(n^{-1/2})$ 
convergence rate 
\citep{briol2015probabilistic}. 
Despite the substantial difference between IS and VI, 
these results intuitively suggest the possibility of developing efficient variational inference without using the information of density function $q(z)$ explicitly. 

\paragraph{Main Idea} 
In this work, we study a simple algorithm for Problem~\ref{pro:prob1}, 
motivated by a recent Stein variational gradient descent (SVGD) algorithm. 
Briefly speaking, 
SVGD is a nonparametric functional gradient descent algorithm which solves $\min_{q} \KL(q ~|| ~p)$ without parametric assumption on $q$,
and approximates the functional gradient, called the Stein variational gradient, using a set of samples (or particles) $\{z_i\}_{i=1}^n$ which iteratively evolves.  
See Section~\ref{sec:svgd}
for more introduction on SVGD. 
We can then view Problem~\ref{pro:prob1}
as a constrained optimization problem,
\begin{align}
\min_q \KL(q ~|| ~p), ~~~~ s.t. ~~~ q = q_\eta ~~\text{for some $\eta$},   
\end{align}
which motivates us to develop a project gradient like algorithm that 
iteratively calculates the Stein variational gradient, and projects it to the finite dimensional $\eta$-space to update parameter $\eta$. 
At the convergence, the samples drawn from $q_\eta$ 
reach the equilibrium state of SVGD, and hence form 
a good approximation of $p$. 
We can view this method as ``amortizing'' or distilling the SVGD dynamics 
using parametric family $q_\eta$ or its inference network $f(\xi; ~\eta)$
and call it amortized SVGD. 
See
Section \ref{sec:amortizedsvgd} for the detailed description of amortized SVGD.  

Our algorithm provides a simple approach for 
the wild inference problem in Problem 1, enabling wide application in approximate learning and inference. 
We explore two examples in this paper. 
In Section~\ref{sec:vae}, 
we apply amortized SVGD to learn complex encoder functions in variational autoencoder (VAE), 
allowing it to 
model complex latent variable space.  
In Section~\ref{sec:trainlangevin}
we use amortized SVGD to learn hyper-parameters of MCMC samplers, 
which allows us to adaptively improve the efficiency of Bayesian computation 
when performing a large number of similar tasks.   
\begin{comment}
\red{\small MOVE Critically, the SVGD gradient includes a repulsive term to ensure that the generated samples have the right amount of variability that matches $p(x).$
In this way, we ``amortize SVGD'' using a neural network, which makes it possible for our method to adaptively improve its own efficiency  
by leveraging fast experience, especially in cases when it needs to perform fast inference repeatedly on a large number of similar tasks. 
}
\end{comment}


\paragraph {Related Work} 

There has been a number of very recent work 
that studies advanced variational inference methods that do not require explicitly calculating $q_\eta(\z)$ (Problem 1). 
This includes 
adversarial variational Bayesian \citep{mescheder2017adversarial} which approximates 
the KL divergence with a density ratio estimator, operator variational inference \citep{operator} which replaces
the KL divergence with an alternative operator variational objective 
that is equivalent to Stein discrepancy (see 
Appendix~\ref{sec:ksd} for more discussion), 
and amortized MCMC \citep{li2017approximate} which
proposes to amortize arbitrary 
MCMC dynamics to make it applicable to discrete models and gradient-free settings. 


The key advantage of our method is its simplicity. 
both \citet{mescheder2017adversarial} and \citet{operator} 
require training some type of auxiliary networks, 
while our main algorithm (Algorithm~\ref{alg:alg1} with update \eqref{equ:update11})  
is very simple, and is essentially a generalization of the typical gradient descent rule that replaces the typical gradient with the stein variational gradient. 

It is also possible to adopt the auxiliary variational inference methods \citep[e.g.,][]{agakov2004auxiliary,salimans2015markov} 
to solve Problem 1 by treating $\xi$ as a hidden variable. 
however, in order to frame a tractable joint distribution $p(z, \xi)$, one would need to add additional noise on $z$, e.g., assume $z = f(\xi; ~\eta) +  \normal(0, \sigma^2)$ with a careful choice of noise variance $\sigma$, and also need to introduce an additional auxiliary network 
to approximate $q(\xi ~|~ z)$, which makes the algorithm more complex than ours.

The idea of amortized inference \citep{gershman2014amortized}
 has been recently applied in various domains of probabilistic reasoning, 
 including both amortized variational inference \citep[e.g.,][]{kingma2013auto, jimenez2015variational} 
and date-driven designs of Monte Carlo based methods \citep[e.g.,][]{paige2016inference}, 
 to name only a few. 

 \citet{balan2015bayesian} 
  also explored the idea of amortizing or distil MCMC samplers to obtain compact fast posterior representation. 
Most of these methods uses typical variational inference methods and hence need to use simple $q_\eta$ to ensure tractability.
  
There is a large literature on traditional adaptive MCMC methods \citep[e.g.,][]{andrieu2008tutorial, roberts2009examples} which can be used to adaptively adjust the proposal distribution of MCMC by exploiting the special theoretical properties of MCMC (e.g., by minimizing the autocorrelation). 
Our method is simpler, more generic, and works efficiently in practice thanks to the use of gradient-based back-propagation. 
{Finally, 
connections between stochastic gradient descent and variational inference have been discussed and exploited in 
\citet{mandt2016variational, maclaurin2015early}. }
 

\begin{comment}
\red{below from two methods}


 There is a number of very recent works that attacks the same problem of relaxing the proposal requirement for $q_\eta$. 

 \citep{operator} avoids calculating $q(z)$ by minimizing a neural-based Stein discrepancy; 
 we discuss a similar approach based on minimizing kernelized Stein discrepancy \citep[][]{gorham2015measuring, liu2016kernelized, oates2014control, chwialkowski2016kernel}. 
  in Section~\ref{sec:} and Section~\ref{sec:} 
  
  \red{[Combine with the related work in two methods in wild inference here].}

\red{1. add auxarlity variational inference}

\red{2. add adverarial Bayes, amortized MCMC and Blei's deep hierical paper}
 

 There is also a raising interest recently on a similar problem of ``learning to optimize'' \citep[e.g.,][]{andrychowicz2016learning, daniel2016learning, li2016learning}, which is technically easier than the more general problem of ``learning to sample''.
An attractive property of our method is that it reduces to ``learning to optimize'' when only one particle is used in SVGD. Therefore, 
with a small number of particles, we can at least achieve the performance of ``learning to optimize'' for maximum a posteriori (MAP). 

The idea of amortized inference \citep{gershman2014amortized}
 has been recently applied in various domains of probabilistic reasoning, 
 including both amortized variational inference \citep[e.g.,][]{kingma2013auto, jimenez2015variational}, 
and data-driven proposals for (sequential) Monte Carlo methods \citep[e.g.,][]{paige2016inference}, 

 to name only a few. Most of these methods, however, 
 require to explicitly calculate $q(z)$ (or its gradient). 
\end{comment}


\section{KSD Variational Inference}
\label{sec:ksd}

Kernelized Stein discrepancy (KSD) provides a discrepancy measure between distributions 
and can be in principle used as a variational objective function in replace of KL divergence. 
In fact, thanks to the special form of KSD (\eqref{equ:ksdexp}-\eqref{equ:uv}), 
one can derive a standard stochastic gradient descent for minimizing KSD without needing to estimate $q_\eta(z)$ explicitly, 
which provides a conceptually simple wild variational inference algorithm. 
Although
this work mainly focuses on 
amortized SVGD which we find to be easier to implement and tend to perform superior to KSD variational inference in practice (see Figure~\ref{fig:ksd_compare}),   
we think the KSD approach is of theoretical interest and hence give a brief discussion here.
 


Specifically, take $q_\eta$ to be the density of the random output $\z = f(\xi;~\eta)$ when $\xi\sim q_0$, 
and we want to find $\eta$ to minimize $\S(q_\eta~||~p)$. 
Assuming $\{\xi_i\}$ is i.i.d. drawn from $q_0$, we can approximate $\S^2(q_\eta~||~p)$ unbiasedly 
using the U-statistics in \eqref{equ:uv}, and derive a standard gradient descent
\begin{gather}\label{equ:aksd}
\eta \gets \eta -   \epsilon \frac{2}{n(n-1)}\sum_{i\neq j} \partial_\eta f(\xi_i;~\eta) ~ \nabla_{\z_i}\kappa_p(\z_i, \z_j),
\end{gather}
where $\z_i = f(\xi_i;~\eta)$.
This enables a wild variational inference method based on directly minimizing $\eta$ with standard (stochastic) gradient descent. 
We call this algorithm \emph{amortized KSD}.  
Note that \eqref{equ:aksd} is similar to \eqref{equ:follow3} in form, 
but replaces $\ff^*(\z_i)$ with 
$$\bar{\ff}^*(\z_i)  \overset{def}{=}  - 2\sum_{j\colon i\neq j} \nabla_{\z_i}\kappa_p(\z_i, \z_j)/(n(n-1)).$$   
Here $\bar \ff^*$ depends on the second order derivative of $\log p$ because $\kappa_p(z,z')$ depends on $\nabla \log p$, 
which makes it more difficult to implement amortized KSD than amortized SVGD. 



Intuitively, 
minimizing KSD can be viewed as 
seeking a stationary point of KL divergence under SVGD updates. 
To see this, recall that $q_{[\epsilon \ff]}$ denotes the density of $\z' = \z+\epsilon \ff(\z)$ when $\z  \sim q$. 
From \eqref{equ:ff00}, we have for small $\epsilon$, 

\begin{align*} 
\S^2(q ~||~ p) 
 \approx \frac{1}{\epsilon}\max_{\phi\in \F}\big\{{\KL(q ~|| ~p) - \KL(q_{[\epsilon\ff]} ~||~ p)}\big\}.  
\end{align*}
That is, KSD measures the maximum degree of decrease in the KL divergence when we update the particles along the optimal SVGD perturbation direction $\ff^*$.
If $q = p$, then the decrease of the KL divergence equals zero and $\S^2(q~||~p)$ equals zero. 
In fact, 
KSD can be explicitly represented as the magnitude of the functional gradient of the KL divergence w.r.t. $\ff$ in RKHS  \citep{liu2016stein}, 
 $$
\S(q~||~p) = \Big|\Big| 
\nabla_\ff  F(0)
\Big|\Big|_{\H^d}, ~~~~~ F(\ff) \overset{def}{=}  \KL(q_{[\epsilon\ff]} ~|| ~ p) , 
$$
where $ \nabla_{\ff}F(\ff)$ denotes the functional gradient of the function $F(\ff)$ w.r.t. $\ff$ defined in RKHS $\H^d$, and $ \nabla_{\ff}F(\ff)$ is also an element in $\H^d$. 
Therefore,   
in contrast to amortized SVGD which attends to minimize the KL objective $F(\ff)$, 
KSD variational inference minimizes the gradient magnitude $||\nabla_\ff  F(0)  ||_{\H^d}$ of KL divergence. 


This idea is closely related to 
the operator variational inference \citep{operator}, which directly minimizes the variational form of Stein discrepancy in \eqref{equ:ff00} and \eqref{equ:ksd} 
with $\F$ replaced by sets of parametric neural networks. 
Specifically, 
\citet{operator} assumes $\F$ consists of a neural network $\ff_\tau(\z)$ with parameter $\tau$, and find $\tau$ jointly with $\eta$ by solving 
a min-max game: 
$$\min_{\eta} \max_{\tau} \E_{\z\sim q_\eta} [\sumsteinp \ff_\tau(\z)].$$ 
This yields a more challenging computation problem, although it is possible that the neural networks provide stronger discrimination than RKHS in practice.  
The main advantage of the KSD based approach is that it leverages the closed form solution in RKHS, 
yields a simpler optimization formulation based on standard gradient descent.  



Figure \ref{fig:ksd_compare} shows results of Langevin samplers trained by amortized SVGD and amortized KSD, respectively,  
for learning simple Gaussian mixtures under the same setting as that in Section \ref{sec:sampler}. 
We find that amortized KSD tends to perform worse (Figure~\ref{fig:ksd_compare}(c)) than amortized SVGD (Figure~\ref{fig:ksd_compare}(b)); 
given that it is also less straightforward to implement amortized KSD (for requiring calculating $\nabla_{z}\kappa_p(z,z')$ in \eqref{equ:aksd}), we did not test it in our other experiments.

\begin{figure}[t]
\centering
\begin{tabular}{ccc}
\includegraphics[width=0.14\textwidth]{figures/gmm/posterior/truth.pdf} &
\includegraphics[width=0.14\textwidth]{figures/gmm/posterior/layer15.pdf} &
\includegraphics[width=0.14\textwidth]{figures/gmm/posterior/ksdu.pdf}  \\
\scriptsize \begin{tabular}{c} (a) Test Distribution \\(truth) \end{tabular} 
  & \scriptsize \begin{tabular}{c} (b)  Amortized SVGD\\ (15 steps) \end{tabular} 
 & \scriptsize \begin{tabular}{c} (c) Amortized KSD \\ (15 steps) \end{tabular}
\end{tabular}
\vspace{-10bp}\caption{\small 
Learning to sample from GMM.
The Langevin sampler 
with step size trained by amortized SVGD (b)
obtains close approximation with $T=15$ is close to the true test distribution (a) while
amortized KSD (c) which we use equation \eqref{equ:aksd} to perform does not work as well as amortized SVGD.
}

\label{fig:ksd_compare}
\end{figure}

\begin{comment}
To better understand the nature of KSD-based inference methods. We also consider a \emph{particle-based} KSD variational inference, in which we find a set of particles $\{z_i\}_{i=1}^n$ (instead of a parametric sampler) to approximate the target distribution $p$, by minimizing 
the empirical KSD: 
$$
\min_{\{z_i\}_{i=1}^n} \bigg\{ \hat \S_u^2( \{z_i\}_{i=1}^n ~||~ p)  \overset{def}{=} \frac{1}{n(n-1)} \sum_{i\neq j} \kappa_p (z_i, z_j) \bigg\}, 
$$
which can be solved by initializing $\{z_i\}_{i=1}^n$ randomly, and updating it iteratively using gradient descent. This idea is similar to the particle-based SVGD algorithm in that it also defines a deterministic particle dynamics for pproximating distribution $p$. It is related to the black-box Stein importance sampling \citep{liu2016kernelized}, which minimizes KSD to optimize the importance weights, with fixed particle locations. 

The performance of such particle-based KSD minimization method, however, is disappointing. 
Figure~{\ref{fig:ksd}} shows the result on a toy example where we find that the particles obtained by 
the gradient descent of KSD (black) tends to concentrate on a lower dimensional manifold,
which the SVGD particles (red) tend to perform much better. 
More theoretical understandings 
are needed to understand the such degeneration of 
the particle-based KSD minimization method. 
Note that in the case of amortized KSD, because the particles are constrained to be drawn from a given random sampler, it does not suffer from the degeneration problem shown in Figure~\ref{fig:ksd}, but the result in Figure \ref{fig:ksd} may still reveal some key problem of KSD minimization methods.  

\begin{figure}
\centering
\includegraphics[width=0.2\textwidth]{../figures/ksd_svgd.pdf} 
\caption{ 
Particles obtained by SVGD (red) and 
particles obtained by gradient descent for minimizing the KSD 
$\min_{\{x_i\}} \S(\{x_i\} ~||~ p)$}.
\label{fig:ksd}
\end{figure}
\end{comment}


\section{Background}
\section{STEIN VARIATIONAL GRADIENT DESCENT}
\label{sec:svgd}
Stein variational gradient descent (SVGD) \citep{liu2016stein} is a 
nonparametric variational inference algorithm that 
iteratively transports a set of particles $\{z_i\}_{i=1}^n$ 
to approximate the target distribution $p$ by performing a type of functional gradient descent on the KL divergence.  
We give a quick overview of the SVGD in this section.  

Let $p(z)$ be a positive density function on $\R^d$ which we want to approximate with a set of particles 
$\{ z_i\}_{i=1}^n$. 
SVGD starts with a set of initial particles $\{ z_i\}_{i=1}^n$, 
and 
updates the particles iteratively by 
\begin{align}\label{equ:xxii}
z_i  \gets z_i +  \epsilon \ff(z_i),  ~~~~ \forall i = 1, \ldots, n,  
\end{align}
where $\epsilon$ is a step size, and 
$\ff\colon \RR^d \to \RR^d$ is a velocity field which should be chosen to push the particle distribution closer to the target distribution. 
Assume the current particles are drawn from a distribution $q$, and 
let $q_{[\epsilon\ff]}$ be the distribution of the updated particles $z^\prime = z + \epsilon \ff(z)$ when $z \sim q$. 
The optimal choice of $\ff$ can be framed as the following optimization problem:
\begin{align}\label{equ:ff00}
\ff^* =   \argmax_{\ff \in \F}  \bigg\{  -   \frac{d}{d\epsilon} \KL(q_{[\epsilon\ff]} ~|| ~ p)~ \bigg |_{\epsilon = 0}  \bigg\}, 
\end{align}
that is, $\ff$ should yield a maximum decreasing rate on the KL divergence between the particle distribution and the target distribution. 
Here, $\F$ is 
a function set that includes 
the possible velocity fields and is chosen to be  
the unit ball of a vector-valued reproducing kernel Hilbert space (RKHS) $\H = \H_0 \times \cdots \times \H_0$,
where  $\H_0$ is a RKHS formed by scalar-valued functions associated with a positive definite kernel $k(z,z')$, that is, 
$\F = \{\ff \in \H \colon ||\ff||_\H\leq  1 \}$. 
This choice of $\F$ allows us to 
consider velocity fields in infinite dimensional function spaces while still obtaining a closed form solution. 
\citet{liu2016stein} showed that 
the objective function 
in \eqref{equ:ff00} equals a simple linear functional of $\ff$: 
\begin{align}\label{equ:klstein00}
&~ - \frac{d}{d\epsilon} \KL(q_{[\epsilon\ff]} ~|| ~ p) \big |_{\epsilon = 0}  = \E_{x\sim q}[\sumstein_p \ff(x)], \\[.5\baselineskip]
&~~\text{with}~~~ \sumstein_p \ff(z)  = \nabla_x \log p(z) ^\top \ff (z)+ \nabla_z^\top \ff(z),  
\end{align}
where $\sumstein_p$ is a linear operator acting on a velocity field $\ff$ and returns a scalar-valued function;
$\sumstein_p$ is called the Stein operator in connection with Stein's identity which shows that 
the RHS of \eqref{equ:klstein00} equals zero if $p = q$: 
\begin{align}\label{equ:steinid}
\E_{p}[\sumstein_p \ff] =\E_{p}[ \nabla_z \log p ^\top \ff + \nabla_z \cdot \ff] = 0. 
\end{align}
This is a result of integration by parts assuming the values of $p(z)\ff(z)$ vanish on the boundary of the integration domain.   
Therefore, the optimization in \eqref{equ:ff00} reduces to 
\begin{align}\label{equ:ksd}
\!\!\!\!\S(q || p) \overset{def}{=} \max_{\ff \in \H} \{ \E_{z\sim q} [\sumstein_p \ff(z)]  ~~~s.t.~~~~ ||\ff ||_{\H} \leq 1\}, 
\end{align}
where $\S(q ~||~ p)$ is the kernelized Stein discrepancy defined in \citep{liu2016kernelized, chwialkowski2016kernel}, which equals zero if and only if $p = q$ under proper conditions
that ensure the function space $\H$ is rich enough.  

Observe that \eqref{equ:ksd} 
is ``simple'' in that it is a linear functional optimization on a 
unit ball of a Hilbert space. 
Therefore, it is not surprise to derive a closed form solution:  

\begin{align}\label{equ:ffss} 
\ff^*(\cdot) \propto  \E_{z\sim q}[\nabla_z \log p(z)k(z,\cdot) + \nabla_x k(z,\cdot)], 
\end{align} 
where $k(z,z')$ is the positive definite kernel associated with RKHS $\H_0$. 
See \citet{liu2016kernelized} for the derivation.
We call $\ff^*$ the Stein variational gradient direction since it provides the  
optimal direction for pushing the particles towards the target distribution $p$. 


In the practical SVGD algorithm, 
we start with a set of initial particles, 
calculate its corresponding $\ff^*$ by replacing the expectation under $q$ with the empirical average of particles, 
and use it to update the particles: 
\begin{align}\label{equ:update11}
&&&z_i ~ \gets ~ z_i  ~  + ~ \epsilon \ff^*(z_i), 
~~~~~~~~\forall i = 1, \ldots, n,  
\\
\!\!\!\!\!\!\!
&&&\ff^*(z_i) = \frac{1}{n} \sum_{j=1}^n [  \nabla_{z_j} \log p(z_j) k(z_j, z_i) + \nabla_{z_j} k(z_j, z_i) ]. \notag
\end{align}
The two terms in $\ff^*(z_i)$ play two different roles: 
the term with the gradient $\nabla_z \log p(z)$ drives the particles toward the high probability regions of $p(z)$, 
while the term with $\nabla_z k(z,z_i)$ serves as a repulsive force to encourage different particles to be different from each other as shown in \citet{liu2016stein}. 
Overall, this procedure provides diverse points for approximating distribution $p$ when it converges. 

It is easy to see from \eqref{equ:update11} that $\ff^*(z_i)$ reduces to the typical gradient $\nabla_z \log p(z_i)$ when there is only a single particle ($n=1$) and $\nabla_z k(z,z_i) = 0$ when $z=z_i$,  
in which case SVGD reduces to the standard gradient ascent for maximizing $\log p(z)$ (i.e., maximum \emph{a posteriori} (MAP)). 


\paragraph{Computing the Kernelized Stein Discrepancy}

By substituting the $\ff^*$ in \eqref{equ:ffss} into \eqref{equ:ksd},  
one can show that \citep{liu2016kernelized, chwialkowski2016kernel, oates2014control}
\begin{align}
\label{equ:ksdexp} 
{\S(q  ~||~ p)}^2 =  \E_{\z, \z' \sim ~ q} [\kappa_p (\z,\z')] , 
\end{align}
where $\kappa_p(\z,\z')$ is a positive definite kernel obtained by applying Stein operator on $k(\z,\z')$ twice, as a function of $z$ and $z'$, respectively. 
It has the following computationally tractable form: 
\begin{align*}
\kappa_p(\z, \z') 
= &\score_p(\z)^\top k(\z,\z') \score_p(\z') + \score_p(\z)^\top \nabla_{\z'} k(\z, \z') \notag \\ 
&+\nabla_\z k(\z,\z')^\top \score_p(\z') + \nabla_\z^\top (\nabla_{\z'} k(\z,\z')),  
\end{align*}
where $\score_p(\z) = \nabla_\z\log p(\z)$.
The form of KSD in \eqref{equ:ksdexp} provides a computationally tractable way for estimating the Stein 
 discrepancy between a set of samples $\{\z_i\}$ (e.g., drawn from an unknown $q$) and a distribution $p$ specified by its score function $\nabla_\z\log p(\z)$ (which is independent of its normalization constant),
\begin{align}\label{equ:uv}
\hat \S_u^2(\{\z_i\}~||~p) =  \frac{1}{n(n-1)}\sum_{i\neq j} [\kappa_p(\z_i, \z_j)],
\end{align}
where $\hat \S_u^2(q~||~p)$ provides an unbiased estimator (hence called a $U$-statistic) for $\S^2(q~||~p)$.
As a side result of this work, 
we will discuss the possibility of using KSD as an objective function for wild variational inference in Appendix~\ref{sec:ksd}.


\begin{algorithm}[t]                     
\caption{Amortized SVGD for Problem~\ref{pro:prob1}
\label{alg:alg1}                          
\begin{algorithmic}                   
\STATE Set batch size $m$, step-size scheme $\{\epsilon_t\}$ and kernel $k(z,z')$. Initialize $\eta^0$. 
\FOR {iteration $t$}
\STATE  Draw random $\{\xi_i\}_{i=1}^m$, calculate $z_i = f(\eta^t;~\xi_i)$, 
and the Stein variational gradient $\ff^*(z_i)$ in \eqref{equ:update11}. 
\STATE  Update parameter $\eta$ using either \eqref{equ:follow1}, \eqref{equ:follow2} or \eqref{equ:follow3}. 
\ENDFOR
\end{algorithmic}
\end{algorithm}

\section{Stein Variational Gradient Descent (SVGD)}
Stein variational gradient descent (SVGD) \citep{liu2016stein} is a general purpose Bayesian inference algorithm motivated by 
Stein's method \citep{stein1972, barbour2005introduction} and kernelized Stein discrepancy \citep{liu2016kernelized, chwialkowski2016kernel, oates2014control}. 
It uses an efficient \emph{deterministic} gradient-based update 
to iteratively evolve a set of particles $\{x_i\}_{i=1}^n$ to minimize the KL divergence with the target distribution. 
SVGD has a simple form that reduces to the typical gradient descent for maximizing $\log p$ when using only one particle $(n=1)$, 
and hence can be easily combined with the successful tricks for gradient optimization, 
including stochastic gradient, adaptive learning rates (such as adagrad), and momentum.

To give a quick overview of the main idea of SVGD, let $p(x)$ be a positive density function on $\R^d$ which we want to 
approximate with a set of particles 
$\{ x_i\}_{i=1}^n$. 
SVGD initializes the particles by sampling from some simple distribution $q_0$, and 
updates the particles iteratively by 
\begin{align}\label{equ:xxii}
x_i  \gets x_i +  \epsilon \ff(x_i),  ~~~~ \forall i = 1, \ldots, n,  
\end{align}
where $\epsilon$ is a step size, and 
$\ff(x)$ is a ``particle gradient direction''
chosen to maximumly decrease the KL divergence between the distribution of particles and the target distribution, in the sense that  
\begin{align}\label{equ:ff00}
\ff =   \argmax_{\ff \in \F}  \bigg\{  -   \frac{d}{d\epsilon} \KL(q_{[\epsilon\ff]} ~|| ~ p) \big |_{\epsilon = 0}  \bigg\}, 
\end{align}
where $q_{[\epsilon \ff]}$ denotes the density of the updated particle $x^\prime = x + \epsilon \ff(x) $ when the density of the original particle $x$ is $q$, and $\F$ is the set of perturbation directions that we optimize over. 
We choose $\F$ to be the unit ball of a vector-valued reproducing kernel Hilbert space (RKHS) $\H^d = \H \times \cdots \times \H$ with each $\H$ associating with a positive definite kernel $k(x,x')$; 
note that $\H$ is dense in the space of continuous functions with universal kernels such as the Gaussian RBF kernel. 

Critically, the gradient of KL divergence in \eqref{equ:ff00} 
equals a simple linear functional of $\ff$, allowing us to obtain a closed form solution for the optimal $\ff$. 
\citet{liu2016stein} showed that 
\begin{align}\label{equ:klstein00}
&- \frac{d}{d\epsilon} \KL(q_{[\epsilon\ff]} ~|| ~ p) \big |_{\epsilon = 0}  = \E_{x\sim q}[\sumstein_p \ff(x)], \\[.5\baselineskip]
&~~~~\text{with}~~~~~ \sumstein_p \ff(x)  = \nabla_x \log p(x) ^\top \ff (x)+ \nabla_x \cdot \ff(x),  
\end{align}
where $\sumstein_p$ is considered as a linear operator acting on function $\ff$ and is called the Stein operator in connection with Stein's identity which shows that 
the RHS of \eqref{equ:klstein00} equals zero if $p = q$: 
\begin{align}\label{equ:steinid}
\E_{p}[\sumstein_p \ff] =\E_{p}[ \nabla_x \log p ^\top \ff + \nabla_x \cdot \ff] = 0. 
\end{align}
This is a result of integration by parts assuming the value of $p(x)\ff(x)$ vanishes on the boundary of the integration domain.   

Therefore, the optimization in \eqref{equ:ff00} reduces to 
\begin{align}\label{equ:ksd}
\S(q ~||~ p) \overset{def}{=} \max_{\ff \in \H^d} \{ \E_{x\sim q} [\sumstein_p \ff(x)]  ~~~s.t.~~~~ ||\ff ||_{\H^d} \leq 1\}, 
\end{align}
where $\S(q ~||~ p)$ is the kernelized Stein discrepancy defined in \citet{liu2016kernelized}, which equals zero if and only if $p = q$ under mild regularity conditions. 
Importantly, the optimal solution of \eqref{equ:ksd} yields a closed for

$$
\ff^*(x') \propto  \E_{x\sim q}[\nabla_x \log p(x)k(x,x') + \nabla_x k(x,x')]. 
$$
By approximating the expectation under $q$ with the empirical average of the current particles $
\{x_i\}_{i=1}^n$,  SVGD admits a simple form of update: 
\begin{align}\label{equ:update11}
&& &~~~~~~~~~~~ x_i ~ \gets ~ x_i  ~  + ~ \epsilon \Delta x_i, 
~~~~~~~~\forall i = 1, \ldots, n,  \notag
\\
~
&&& \text{where~~~~~}\Delta x_i = \hat \E_{x\in \{x_i\}_{i=1}^n} [  \nabla_{x} \log p(x) k(x, x_i) + \nabla_{x} k(x, x_i) ], 
\end{align}
and $\hat\E_{x\sim \{x_i\}_{i=1}^n}[f(x)] = \sum_i f(x_i)/n$. 
The two terms in $\Delta x_i$ play two different roles: 
the term with the gradient $\nabla_x \log p(x)$ drives the particles toward the high probability regions of $p(x)$, 
while the term with $\nabla_x k(x,x_i)$ serves as a repulsive force to encourage diversity;
to see this, consider a stationary kernel $k(x,x') = k(x-x')$, then the second term reduces to $\hat \E_x \nabla_{x} k(x,x_i) = - \hat \E_x \nabla_{x_i} k(x,x_i)$, 
which can be treated as the negative gradient for minimizing the average similarity $\hat \E_x k(x,x_i)$ in terms of $x_i$. 
Overall, this particle update produces diverse points for distributional approximation and uncertainty assessment, and 
also has an interesting ``momentum'' effect in which the particles move collaboratively to escape the local optima.

It is easy to see from \eqref{equ:update11} that $\Delta x_i$ reduces to the typical gradient $\nabla_x \log p(x_i)$ when there is only a single particle ($n=1$) and $\nabla_x k(x,x_i)$ when $x=x_i$,  
in which case SVGD reduces to the standard gradient ascent for maximizing $\log p(x)$ (i.e., maximum \emph{a posteriori} (MAP)). 


\begin{algorithm}[t]                     
\caption{Amortized SVGD for Problem~\ref{pro:prob1}
\label{alg:alg1}                          
\begin{algorithmic}                   
\STATE Set batch size $m$, step-size scheme $\{\epsilon_t\}$ and kernel $k(x,x')$. Initialize $\eta^0$. 
\FOR {iteration $t$}
\STATE  Draw random $\{\xi_i\}_{i=1}^m$, calculate $x_i = f(\eta^t;~\xi_i)$, 
and the Stein variational gradient $\Delta x_i$ in \eqref{equ:update11}. 
\STATE  Update parameter $\eta$ using \eqref{equ:follow1}, \eqref{equ:follow2} or \eqref{equ:follow3}. 
\ENDFOR
\end{algorithmic}
\end{algorithm}