Understand VAE with VamPrior I am currently reading this paper.
The authors propose to use as an prior this expression:
$$
p_\lambda(z) = \frac{1}{K} \sum^K_{k=1} q_\phi (z\mid u_k)
$$
where $q_\phi$ is the encoder, and $u_k$ is a learned "pseudo-input". What I dont understand is:


*

*What is $q_\phi(z\mid u_k)$? How do we sample from it? Do we predict like for a normal VAE just the $\mu(u_k)$ and $\sigma(u_k)$ of a normal distribution and then sample from this one?

*How do we train that? For a single Gaussian we can calculate the KL Divergence analytically, in terms of $\mu$ and $\sigma$ and then back propagate through the reparametrization trick. How can we do it here?   

 A: 
What is $q_\phi(z|u_k)$?

It's the variational approximation to the posterior, conditioned on the learned pseudo-input $u_k$. It's the same as if $u_k$ were a real data point, and passed to the stochastic encoder (inference model) $q_\phi$.

How do we sample from it? 

Since $u_k$ is a fixed "data point", just treat it as a regular data point, and sample from the posterior as you normally would.

Do we predict like for a normal VAE just the $μ(u_k)$ and $σ(u_k)$ of a normal distribution and then sample from this one?

Yes, assuming $q_\phi$ is parametrized as a normal distribution.

How do we train that? 

Same way you normally would, by a Monte Carlo estimate of the loss:
$$ \mathcal{L}(\phi,\theta,\lambda) = \mathbb{E}_{x\sim q(x)}\left[  
\frac{1}{L}\sum_\ell \left(
\ln p_\theta(x|z^{(\ell)}_\phi) + \ln p_\lambda(z^{(\ell)}_\phi) - \ln q_\phi( z^{(\ell)}_\phi | x)
\right)
\right] $$
where $\lambda$ are the parameters of the prior (normally, there are no parameters since $p_\lambda(z)=\mathcal{N}(0,I)$).
Assuming $q_\phi$ is reparametrizable, then so is $$ p_\lambda(z_\phi) = \frac{1}{K}\sum_k q_\phi(z_\phi|u_k), $$
and thus we can back-propagate through this loss (well you also need differentiability wrt to $u_k$). See the appendix of the paper for details.

For a single Gaussian we can calculate the KL Divergence analytically, in terms of μ and σ and then back propagate through the reparametrization trick. How can we do it here?

See above, except in this case it is not analytic of course. You can see this via:
\begin{align*}
\mathbb{E}_{x\sim q(x)}  
\left[
\frac{1}{L}\sum_\ell 
\ln p_\lambda(z^{(\ell)}_\phi) - \ln q_\phi( z^{(\ell)}_\phi | x)
\right]
&\approx 
\mathbb{E}_{x\sim q(x)}  
\mathbb{E}_{z_\phi\sim q_\phi(z|x)}  
\left[
\ln 
\frac{
p_\lambda(z^{}_\phi) }{ q_\phi( z_\phi | x)}
\right] \\
&= -\mathbb{E}_{x\sim q(x)}  
\mathfrak{D}_\text{KL}[ q_\phi( z_\phi | x) || p_\lambda(z^{}_\phi)]
\end{align*}
(Recall we want to maximize $\mathcal{L}$; the literature should probably use a different symbol)
