Discrete Auto-regressive Variational Attention Models for Text Modeling

Information underrepresentation is a common issue in applying VAEs for text modeling. Conventional VAEs build latent variables based on the last hidden state of LSTM encoder, i.e. $\mathbf{z}=z_{T}$. During the decoding process, we first sample $z_{T}$, from which new sentences $\hat{\mathbf{x}}=\hat{x}_{1:\hat{T}}$ can be generated:

where $\hat{T}$ is the length of reconstructed sentence $\hat{\mathbf{x}}$. However, the representation of $z_{T}$ is generally insufficient to summarize the semantic dependencies in $\mathbf{x}$, and thus deteriorates the reconstruction.

Posterior collapse usually arises as $D_{KL}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p(\mathbf{z}))$ diminishes to zero, where the local optimal gives $q_{\phi}(\mathbf{z}|\mathbf{x})=p(\mathbf{z})$. Posterior collapse happens inevitably as the ELBO contains both the reconstruction loss $\mathbb{E}_{\mathbf{z}\sim q_{\phi}(\mathbf{z}|\mathbf{x})}[\log p_{\theta}(\mathbf{x}|\mathbf{z})]$ and the KL-divergence $D_{KL}(q_{\phi}(\mathbf{z}|\mathbf{x})\|p(\mathbf{z}))$, as shown in Equation (1). When posterior collapse happens, $\mathbf{x}$ becomes independent of $\mathbf{z}$ as $p(\mathbf{x})p(\mathbf{z})=p(\mathbf{x})q_{\phi}(\mathbf{z}|\mathbf{x})=p(\mathbf{x})\frac{p(\mathbf{x},\mathbf{z})}{p(\mathbf{x})}=p(\mathbf{x},\mathbf{z})$. Therefore, the encoder learns a data-agnostic posterior without any information from $\mathbf{x}$, while the decoder fails to perform valid generation but purely based on random noise.

A key difference between variational auto-regressive attention models and conventional VAEs is the choice of a prior distribution. During the generation, the latent sequence $z_{1:T}$ are sampled from the prior unconditionally, and are then fed to the attention module together with $h_{1:\hat{T}}^{d}$. The most adopted prior $\mathcal{N}(0,I)$, however, is non-informative to generate well-correlated latent sequence for the attention as that during training. Therefore the decoder receives no informative supervision that gives reasonable generation.

To solve that, we deploy an auto-regressive prior $p_{\psi}(z_{1:T})=p_{\psi}(z_{1})\prod_{t=2}^{T}p_{\psi}(z_{t}|z_{1:t-1})$ parameterized by $\psi$ to capture the underlying semantic dependencies. Specifically, we take $p_{\psi}(z_{t}|z_{1:t-1})=\mathcal{N}(\hat{\mu}_{t},\hat{\sigma}_{t}I)$, where $(\hat{\mu}_{t},\hat{\sigma}_{t}I)$ is produced by a PixelCNN, a superior model in learning sequential data [23].

The training of GAVAM can be easily troubled by posterior collapse due to two aspects. To see this, similar to Equation 1, the minimization of the ELBO now can be written as:

On the one hand, the KL divergence scales linearly to the sequence length of $T$, which makes the training unstable across different input lengths. On the other hand, and more seriously, both $\phi$ and $\psi$ are used to minimize the KL divergence, which can easily trap the learned posteriors. To demonstrate this, for example, the KL divergence between two Gaussian distributions can be written as:

where $D$ is the latent dimension of $z_{t}$. Whenever $\sigma^{2}_{td}\rightarrow\hat{\sigma}^{2}_{td}$ and $\mu_{td}\rightarrow\hat{\mu}_{td}$ before $q_{\phi}(z_{1:T}|\mathbf{x})$ encodes anything from $\mathbf{x}$, both $q_{\phi}(z_{1:T}|\mathbf{x})$ and $p_{\psi}(z_{1:T})$ get stuck in local optimal and learn no semantic dependency for reconstruction.

Thanks to the nice properties of discreteness, the optimization of DAVAM does not suffer from posterior collapse. Specifically, the KL divergence of DAVAM can be written as:

where the third line is obtained with the entropy $H(q_{\phi}(z_{t}))=-1\log 1-0\log 0\rightarrow 0$. It can be found that $D_{KL}(q_{\phi}(z_{1:T}|\mathbf{x})\|p_{\psi}(z_{1:T}))$ is no longer relevant to posterior parameters $\phi$. Consequently, the update of variational posterior $q_{\phi}(z_{1:T}|\mathbf{x})$ does not rely on the prior but is determined purely by the reconstruction term. Therefore minimization of KL divergence will not lead to posterior collapse.

We follow the standard paradigm to minimize the ELBO of DAVAM. As shown in Equation 7, since KL divergence is neither relevant to $\phi$ nor $\theta$, only the reconstruction term should be concerned. In the meanwhile, as the latent variables $z_{1:T}$ are determined based on Euclidean distances between $h^{e}_{1:T}$ and code book $\{e_{k}\}_{k=1}^{K}$, we further regularize them to stay close via a Frobenius norm. The training objective for stage one is

where $\beta$ is the regularizer, and $\mathrm{{sg(\cdot)}}$ stands for stop-gradient operation. Note that the quantization in Equation 6 is non-differentiable. To allow the back-propagation algorithm to proceed, we adopt the widely employed straight through estimator (STE) [25] to copy gradients from $e_{z_{t}}$ to $h_{t}^{e}$, as is shown in Figure 2.

For the update of code book $\{e_{k}\}_{k=1}^{K}$, we first apply K-means algorithm to calculate the average over all latent variables $h_{1:T}^{e}$ that are closest to $\{e_{k}\}_{k=1}^{K}$, and then take exponential moving average over the mini-batch update of the code book so as to stabilize the optimization.

After the convergence of DAVAM, we resort to update the auto-regressive prior $p_{\psi}(z_{t}|z_{1:t-1})$. To mimic the semantic dependency in the learned posterior $q_{\phi}(z_{1:T}|\mathbf{x})$, the prior is supposed to fit the latent sequence $z_{1:T}\sim q_{\phi}(z_{t}|\mathbf{x})$. This can be realized by the minimizing their KL-divergence w.r.t. $\psi$ as

which can be simplified to the cross-entropy loss between $z_{1:T}$ and $\gamma_{1:T}$ according to Equation (7).

We compare the proposed DAVAM against a number of baselines, including the classical LSTM-based Language Modeling-(LSTM-LM), vanilla VAE [1], and its advanced variants: annealing VAE [16], cyclic annealing VAE¹¹1https://github.com/haofuml/cyclical_annealing [11], lagging VAE²²2https://github.com/jxhe/vae-lagging-encoder [12], Free Bits (FB) [17] and pretraining+FBP VAE³³3https://github.com/bohanli/vae-pretraining-encoder [6]. All these baselines do not use the attention module in their architectures.

For ablation studies, we further compare to 1) GAVAM, which takes Gaussian distribution instead of the one-hot categorical distribution over $z_{1:T}$ to verify the advantages of discreteness; 2) We also remove the attention mechanism (denoted as DAVAM-q) to test the effect of discreteness on the last latent variable $z_{T}$. 3) Finally, to check the choice of prior, we replace the auto-regressive prior with uninformative Gaussian priors, which leads to variational attention models (VAE+Attn) first proposed by [13, 13].

We evaluate language modeling using three metrics: 1) Reconstruction loss (Rec) $\mathbb{E}_{\boldsymbol{z}\sim q_{\phi}(\boldsymbol{z}|\boldsymbol{x})}[\log p_{\theta}(\boldsymbol{x}|\boldsymbol{z})]$ that measures the ability to recover data from latent space; 2) Perplexity (PPL) measuring the capacity of language modeling; Both lower Rec and PPL give better models in general; and 3) KL divergence (KL) indicating whether posterior collapse occurs.

For baselines, we keep the same hyper-parameter settings to pretraining+FBP VAE [6], e.g., the dimension of latent space, word embeddings and hidden states of LSTM. Since our latent variables are discrete, we cannot use importance weighted samples to approximate the reconstruction loss in Lagging VAE and pretraining+FBP VAE.

For DAVAM and its ablation counterparts, we keep the same set of hyper-parameters. By default, we set the code-book size $K$ as $512$. We first warm up the training for 10 epochs, and then gradually increase $\beta$ in Equation (8) from 0.1 to $\beta_{max}=5.0$, in a similar spirit to annealing VAE [16]. For all experiments, we use the SGD optimizer with initial learning rate $1.0$, and decay it until five counts if the validation loss does not decrease for $2$ epochs. For the auto-regressive prior, we use a 16-layer PixelCNN with one-dimensional convolution followed by residual connections.