Discrete Variational Attention Models for Language Generation

Variational Autoencoders (VAEs) Kingma and Welling (2013) are well known class of generative models. Given observations $x$, we seek to infer latent variables $z$ from which new observations $\hat{x}$ can be generated. To achieve this, we need to maximize the marginal log likelihood $\log p_{\theta}(x)$, which is usually intractable due to the complex posterior $p(z|x)$. Consequently an approximate posterior $q_{\phi}(z|x)$ (i.e. the encoder) is introduced, and the evidence lower bound (ELBO) of the marginal likelihood is maximized as follows:

where $p_{\theta}(x|z)$ represents likelihood function conditioned on latent codes $z$, also known as the decoder. $\theta$ and $\phi$ are the corresponding parameters.

In language generation, VAEs are used to learn the mapping from the latent space of $z$ to the observations $x$. Based on this mapping, new sentences can be effectively generated. VAEs are usually armed with an RNN encoder and decoder, where the input sentences $x$ are given to the encoder, the latent variables $z$ are derived from the last hidden states $h^{e}_{T}$ via the reparameterization trick Kingma and Welling (2013) $z=\mu(h^{e}_{T})+\sigma(h^{e}_{T})\epsilon$ with $\epsilon\sim\mathcal{N}(0,I)$. To generate new sentences, we sample latent variables $z$ from the prior that are forwarded to the decoder to generate sequences $\hat{x}$ by

Modeling RNN-based language models with VAEs is prone to suffer from the diminishing effect of latent variables, i.e., the generated sequences $\hat{x}$ are loosely dependent on $z$ and thereon the learned latent space plays no role in generation. The diminishing effect largely comes from two aspects:

As shown in Figure 1, the proposed DVAM adopts two different RNNs as the encoder network and the decoder network, respectively. In order to build the connection between the encoder hidden states (denoted by $h^{e}_{1:T}$) and the decoder hidden states (denoted by$h^{d}_{1:T}$), we further involve an attention mechanism similar to sequence to sequence (seq2seq) models Bahdanau et al. (2014).

A proper variational posterior $q_{\phi}(z_{1:T}|x)$ plays an important role in effectively capturing the semantic information from the observations $x$. A simple idea would be choosing the widely used Gaussian distribution, which however easily leads to posterior collapse (as will be discussed in Section 3.3). As representations of languages are discrete in nature, we seek to quantize the latent space with a code book $\{e_{k}\}_{k=1}^{K}$, where $K$ is the code book size. The combination of code book vectors therefore represents the sequential dependency from observed sentences x. The advantages of quantized representation have also been verified in the recently proposed vector quantized variational autoencoder van den Oord et al. (2017); Razavi et al. (2019) for image generation. We set the variational posterior $q_{\phi}(z_{1:T}|x)$ as categorical distribution to index $\{e_{k}\}_{k=1}^{K}$ as follows:

Given the code book index $z_{t}$, the encoder hidden state $h_{t}^{e}$ is therefore quantized to $e_{z_{t}}$, and the attention scores can be computed as usual by

where $\tilde{\alpha}_{t}=v^{\top}\tanh(W_{e}e_{z_{1:T}}+W_{d}h^{d}_{t-1}+b)$ is the score before the softmax normalization.

We then compute the context vectors $c_{t}=\sum_{i=1}^{T}\alpha_{ti}e_{z_{i}}$ as an extra input to the decoder, and reformulate the generation process as

Therefore at each time step, the decoder receives the supervision from the context vector $c_{t}$ , which is a weighted sum from the code books $\{e\}_{k=1}^{K}$ of the whole sequence. Consequently, the variational posterior $q_{\phi}(z_{1:T}|x)$ encodes the sequential dependency from the observations that addresses the issue of information underrepresentation.

To allow new sentence generation, our model can also draw sequentially dependent samples from the prior $p(z_{1:T})$, which will be discussed in the next section.

Thanks to the nice property of discreteness, the optimization of DVAM does not suffer from posterior collapse. To see this, note that the ELBO in Equation (1) includes both the reconstruction loss and the KL divergence, whereas the KL divergence of DVAM can be written as

where the third line is obtained by $q_{\phi}(z_{t}=j_{t})=1$, $q_{\phi}(z_{t}\neq j_{t})=0$ and the fact that $H(q_{\phi}(z_{t}))=-1\log 1-0\log 0=0$. Consequently, $D_{KL}(q_{\phi}(z|x)\|p(z))$ is not differentiable w.r.t. the variational parameters $\phi$. The optimization of KL divergence does not play a role in the reconstruction loss optimization, and even we do not need to know the form of the prior $p(z_{1:T})$ in advance.

On the other hand, since the latent variables $z_{1:T}$ are obtained based on the Euclidean distance to the code books, we encourage the hidden states to stay close to $\{e_{k}\}_{k=1}^{K}$. Therefore, the training objective for DVAM can be formulated as

where $\beta$ is the regularizer, and $\mathrm{{sg(\cdot)}}$ stands for stop-gradient operation. Note that since the quantization is non-differentiable, we adopt the widely used straight through estimator (STE) Bengio et al. (2013)] to copy gradients from $e_{z_{t}}$ to $h_{t}^{e}$, as is shown in Figure 1, In terms of the code books $\{e_{k}\}_{k=1}^{K}$, since Equation (7) is non-differentiable with them, we first apply the K-means algorithm to calculate the average over all latent variables $h_{1:T}^{e}$ that are closest to $\{e_{k}\}_{k=1}^{K}$, then we take exponential moving average over the code books so as to stabilize the mini-batch update.

As mentioned earlier in Section 3.1, a simple idea is to directly assign Gaussian distributions over the attention, i.e. $z_{t}=\mu_{t}(h^{e}_{t})+\sigma_{t}(h^{e}_{t})\epsilon$ for $\epsilon\sim\mathcal{N}(0,I)$, leading to the Gaussian Variational Attention Model (GVAM).

To analyze the defects of GVAM, we first recall the KL divergence between two Gaussian distributions, which can be written as follows:

where $D$ is the latent dimension of $z_{t}$. We denote $\{\mu_{td},\sigma_{td}\}$ and $\{\hat{\mu}_{td},\hat{\sigma}_{td}\}$ as the parameters of the posterior and prior distributions, respectively. Unlike the objective of DVAM in Equation (3.2), we find that Equation (3.3) is differentiable w.r.t to variational parameters $\phi$, and therefore should be included for the minmization of ELBO. As we also need an auto-regressive prior during generation, such differentiation brings the challenge for GVAM. Unlike DVAM that learns the posterior $q_{\phi}(z_{1:T}|x)$ in advance to teach the prior, GVAM needs to involve the posterior and prior jointly during the optimization, i.e.,

Such optimization can be easily troubled by posterior collapse due to two aspects. Firstly, the scale of the KL divergence increases linearly to the length of the sequence, and a large scale usually overlooks the minimization of the reconstruction loss. The training thereon is unstable across various observations $x$ with different lengths. Secondly and more seriously, both $\phi$ and $\psi$ are used to minimize the KL divergence. Whenever $q_{\phi}(z_{1:T}|x)$ collapses to $p_{\psi}(z_{1:T})$ before the $q_{\phi}(z_{1:T}|x)$ learns any sequential dependency from the observations, both $q_{\phi}(z_{1:T}|x)$ and $p_{\psi}(z_{1:T})$ are trapped to the local optimal that cannot provide structural supervision to the decoder during both training and generation.

We take three benchmark datasets of language modelling for verification: Yahoo Answers  Xu and Durrett (2018), Penn Tree Marcus et al. (1993), and a down-sampled version of SNLI Bowman et al. (2015a). A summary of dataset statistics is shown in Table 1.

We compare the proposed DVAM against a number of baselines, including the classical LSTM-LM, vanilla VAE Kingma and Welling (2013), as well as its advanced variants, e.g. annealing VAE Bowman et al. (2015b), cyclic annealing VAE²²2https://github.com/haofuml/cyclical_annealing Fu et al. (2019), lagging VAE³³3https://github.com/jxhe/vae-lagging-encoder He et al. (2019), Free Bits (FB) Kingma et al. (2017) and pretraining+FBP VAE⁴⁴4https://github.com/bohanli/vae-pretraining-encoder Li et al. (2019). We also compare to our closest counterpart, i.e., Gaussian Variational Attention Model (GVAM) so as to directly verify the advantages of discreteness in the variational attention model.

We evaluate the performance of language generation models using three metrics, i.e., the reconstruction loss (Rec) calculated as $\mathbb{E}_{\boldsymbol{z}\sim q_{\phi}(\boldsymbol{z}|\boldsymbol{x})}[\log p_{\theta}(\boldsymbol{x}|\boldsymbol{z})]$ for measuring the ability to recover data from latent space (the lower the better), the Perplexity (PPL) measuring the capacity of language modelling (the lower the better), and the KL divergence (KL) between the posterior $q(z|x)$ and the prior $p(z)$ indicating whether posterior collapse occurs.

By default, all the ablation studies are conducted on Yahoo dataset with the default parameter settings.