AE

An autoencoder is a DNN trained to replicate an input vector $x \in \R^F$ at the output. It usually has a diabolo shape, with an encoder providing a low-dim latent representation $z \in \R^L$ from $x$ ($L \ll F$) and a decoder tries to reconstruct $x$ from $z$, outputting $\hat{x}$.

• Encoder: $p(z \vert x)$
• Decoder: $p(x \vert z)$

VAE

VAE can be seen as a probabilistic version of an AE, where the decoder outputs the parameters of a probability distribution of $x$. Thus, an VAE is capable of

• generating new $\hat{x}$ from unseen values of $z$
• transforming existing data $x_1, \dots, x_n$ by modifying their encoded latents.
• providing a prior distribution of $x$ in more complex Bayesian models.

Decoder

The VAE decoder is defined as

$p_\theta (x, z) = p_{\theta_x} (x \vert z) p_{\theta_z} (z) \tag{vae:dec}$

• Prior distribution for $z$: $p_{\theta_z}(z)$
• usually modelled with $\mathcal{N}(0, I)$
• modelled with Gamma distributions that better fit natural statistics of speech/audio power spectra
• Conditional distribution $p_{\theta_x} (x \vert z)$, as a non-linear function of $z$

Generation with VAE

If we consider the marginal distribution of $x$ by integrating on $p(x \vert z)$:

$p_\theta(x) = \int p_{\theta_x} (x \vert z) p_{\theta_z} (z) \mathrm{d}z \tag{vae:x-marginal}$

Because any conditional distribution $p_{\theta_x} (x \vert z)$ can provide a mode, $p_\theta (x)$ can be highly multi-modal.

💡 This only provides an analysis on $p_\theta(x)$. Later we will find that $p_{\theta_x} (z \vert x)$ is actually intractable, for which we actually cannot compute this integral.

Gaussian Prior

When the prior distribution is modelled with $\mathcal{N}(0, I)$, we have

\begin{aligned} p_{\theta_x}(x \vert z) &= \mathcal{N} \left(x; \mu_{\theta_x}(z), \mathrm{diag}\{ \sigma_{\theta_x}^2 (z) \}\right)\\ &=\prod_{i=1}^F p_{\theta_x} (x_f \vert z) \\ &=\prod_{i=1}^F \mathcal{N} \left(x_f; \mu_{\theta_x, f}(z), \mathrm{diag}\{ \sigma_{\theta_x, f}^2 (z) \}\right)\\ \end{aligned}

Encoder

To train the generative model in $\mathrm{(vae:dec)}$, we should learn $\theta$ from the training dataset $X = \{ x_n \in \R^F \}_{i=1}^N$, where **************************************it is assumed that $x_i \overset{\mathrm{iid}}{\sim} p(x)$
. We estimate $p(x)$ with $p_\theta(x)$ and by training on $x \in X$ we find

$\theta^* = \argmax_\theta \log p_\theta(x)$

EM: An Intractable Way

Using the classical Expectation-Maximization is not possible.

Variational Inference

Variational Inference is used to train the VAE. It is based on two principles:

• $p_\theta (z \vert x)$ in intractable, but the approximation posterior $q_\phi (z \vert x) \approx p_\theta(z \vert x)$ is acceptable
• The encoder and the decoder are jointly trained

A commonly chosen approximation is also Gaussian:

\begin{aligned} q_\phi (z \vert x) &= \mathcal{N} \left(z; \mu_\phi (x), \mathrm{diag} \{ \sigma_\phi^2 (x) \} \right) \end{aligned}

where the encoder network non-linearly maps $x$ to distribution parameters with two encoder networks (also referred to as recognition networks):

• $\mu_\phi: \R^F \to \R^L$
• $\sigma_\phi: \R^F \to \R^L_+$

Here, we only get parameters for the distribution for $z$, from which we need sampling to get $z$. However, sampling is not differentiable so back propagation cannot be done. In coding, this is tackled by the reparameterize trick that:

1. sample $\hat{z}$ from the Standard Normal Distribution $\mathcal{N}(0, I)$
2. transform $\hat{z}$ to $z$ by $z = \mu_\phi(x) + \sigma_\phi^2(x) x$

💡 Here, logvar that is actually $\log\sigma_\phi^2(x) \in \R$.

Training

The parameters $\theta, \phi$ are jointly estimated from $X$. The paper (Kingma and Welling, 2014) has shown that even if $\log p_\theta (X)$ is intractable, we can calculate an EvLdence Lower Bound (ELBO) depending on $\theta, \phi$:

The ELBO of an VAE is geven by:

\begin{aligned} \mathcal{L}(\theta, \phi; X) &= \mathbb{E}_{q_\phi (z \vert x)} \log p_\theta (x, z) - \mathbb{E}_{q_\phi (z \vert x)} \log q_\phi (x \vert z) \\ &= \mathbb{E}_{q_\phi (z \vert x)} \log p_\theta (x \vert z) - D_\mathrm{DL} \left( q_\phi (z \vert x) \Vert p_{\theta_z} (z) \right) \end{aligned}

Dynamical VAE

Dynamical VAEs (DVAEs) consider a sequence of observed random vectores $x_{1:T} = \{ x_t \in \R^F \}_{i=1}^T$ and a sequence of latent random vectors $z_{1:T} = \{ z_t \in \R^F \}_{i=1}^T$ that are assumed to be temporally correlated. Designing a DVAE consists in specifying the joint distributions $p_\theta (x_{1:T}, z_{1:T})$ that determine their dependencies.

In a commonly-used mode called driven mode, by observing $u_{1:T} = \{ u_t \in \R^U \}_{t=1}^T$, the model generates $x_{1:T}$ as the output. This requires specifying $p_\theta(x_{1:T}, z_{1:T}, u_{1:T})$, but we usually assume $u_{1:T}$ is deterministic and $x_{1:T}, z_{1:T}$ are stochastic, turning the problem to specifying $p_\theta (x_{1:T}, z_{1:T} \vert u_{1:T})$

Ordering Dependencies

We naturally order the dependencies by a casual model:

• Causal: Generate the variable from the past - most popular for tasks like motion generation
• Non-causal: Generate the variable from both the past and the future
• Anti-causal: Generate the variable from the future

E.g. a very simple example of a causal model (the first line) and a non-causal model (the second line):

Causal DVAEs Modelling

For the causal DVAE, the joint distribution can be modelled with the chain rule:

\begin{aligned} q(x_{1:T}, z_{1:T} \vert u_{1:T}) &= \prod_{t=1}^T p(x_t, z_t \vert x_{1:t-1}, z_{1:t-1}, u_{1:t}) \\ &= \prod_{t=1}^T p(x_t \vert x_{1:t-1}, z_{1:t}, u_{1:t}) p(z_t \vert x_{1:t-1}, z_{1:t-1}, u_{1:t}) \end{aligned}

💡 Of note, $x_{1:0} = z_{1:0} = \empty$, therefore the first term is $p(x_1 \vert z_1, u_1) p(z_1 \vert u_1)$. The initialization requires specifying $p(z_1)$ and $p(z_1 \vert u_1)$.

Still, this depencency is complex that involves all the past and present $u$s, as well as all the past $x$s, $z$s.

Simplifications

There are possible simplifications, for example, the SSM family simplifies the dependencies as:

\begin{aligned} p(x_t \vert x_{1:t-1}, z_{1:t}, u_{1:t}) &= p(x_t \vert z_t) \\ p(z_t \vert x_{1:t-1}, z_{1:t-1}, u_{1:t}) &= p(z_t \vert z_{t-1}, u_t) \end{aligned}

But a more generally used may is:

\begin{aligned} p(x_t \vert x_{1:t-1}, z_{1:t}, u_{1:t}) &= p(x_t \vert x_{1:t-1}, u_t) \\ p(z_t \vert x_{1:t-1}, z_{1:t-1}, u_{1:t}) &= p(z_t \vert x_{1:t}, u_t) \end{aligned}

The probabilistic graph for it is like:

To further facterize the terms, RNNs with hidden states are introduced. Two possible implementations are sharing the same set of hidden states or use different sets:

Inference Model

Also, the posterior distribution of $z_{1:T}$: $p_\theta (z_{1:T} \vert x_{1:T}, u_{1:T})$ is intractable. We also need an inference model $q_\phi(z_{1:T} \vert x_{1:T}, u_{1:T})$ to approximate it. Exploiting D-Spearation, the inference can be formulated as:

\begin{aligned} p_\theta(z_{1:T} \vert x_{1:T}, u_{1:T}) &= \prod_{t=1}^T p_{\theta_z} (z_t \vert z_{1:t-1}, x_{1:T}, u_{1:T}) \\ q_\phi(z_{1:T} \vert x_{1:T}, u_{1:T}) &= \prod_{t=1}^T q_\phi (z_t \vert z_{1:t-1}, x_{1:T}, u_{1:T}) \end{aligned}

Training DVAE

The ELBO is extended to data sequences as:

\begin{aligned} \mathcal{L}(\theta, \phi; x_{1:T}, u_{1:T}) &= \mathbb{E}_{q_\phi(z_{1:T} \vert x_{1:T} \vert u_{1:T})} \ln p_\theta (x_{1:T}, z_{t:T} \vert u_{1:T}) \\ &-\mathbb{E}_{q_\phi(z_{1:T} \vert x_{1:T} \vert u_{1:T})} \ln q_\phi (z_{t:T} \vert x_{1:T}, u_{1:T}) \\ &= \sum_{t=1}^T \mathbb{E}_{q_\phi(z_{1:t} \vert x_{1:T} \vert u_{1:T})} \ln p_{\theta_x} (x_t \vert x_{1:t-1}, z_{1:t}, u_{1:t}) \\ &- \sum_{t=1}^T \mathbb{E}_{q_\phi(z_{1:t-1} \vert x_{1:T} \vert u_{1:T})} D_\mathrm{KL} \left(q_\phi(z_t \vert z_{1:t-1}, x_{1:T}, u_{1:T}) \Vert p_{\theta_z} (z_t \vert x_{1:t-1}, z_{1:t-1}, u_{1:t}) \right) \end{aligned}

💡 In the standard VAE, the regularization term has an analytical form for usual distributions like Gaussian. However, for DVAE, the reconstruction accuracy and regularization term require computing Monte Carlo estimates with samples drawn from $q_\phi (z_{1:\tau} \vert x_{1:T}, u_{1:T})$ where $\tau \in \{ 1, \dots, T \}$.

References

now publishers - Dynamical Variational Autoencoders: A Comprehensive Review