SlowGuess/NAACL_2021_OCR huggingface dataset:
Xiaoan Ding
Chicago, IL, 60637, USA
xiaoanding@uchicago.edu
Kevin Gimpel
kgimpel@ttic.edu
Variational autoencoders (VAEs) are widely used for latent variable modeling of text. We focus on variations that learn expressive prior distributions over the latent variable. We find that existing training strategies are not effective for learning rich priors, so we add the importance-sampled log marginal likelihood as a second term to the standard VAE objective to help when learning the prior. Doing so improves results for all priors evaluated, including a novel choice for sentence VAEs based on normalizing flows (NF). Priors parameterized with NF are no longer constrained to a specific distribution family, allowing a more flexible way to encode the data distribution. Our model, which we call FlowPrior, shows a substantial improvement in language modeling tasks compared to strong baselines. We demonstrate that FlowPrior learns an expressive prior with analysis and several forms of evaluation involving generation.
Variational autoencoders (VAEs; Kingma and Welling, 2014) have been widely applied to many natural language processing tasks (Bowman et al., 2016; Zhang et al., 2016; Shen et al., 2017; Kim et al., 2018; Fang et al., 2019; Chen et al., 2019). VAEs provide statistical transparency in describing observations in a latent space and flexibility when used in applications that require directly manipulating the learned representation (Hu et al., 2017). Recent work (Li et al., 2020) has combined VAEs with BERT/GPT in representation learning and guided generation. However, the representation capacity of VAEs is still limited for modeling sentences due to two main reasons.
One is known as the posterior collapse problem, in which the posterior "collapses" to the prior and the generator learns to ignore the latent variable (Bowman et al., 2016). Many methods have been developed to address it: annealing (Fu et al., 2019),
weakening the capacity of the generator (Semeniuta et al., 2017; Yang et al., 2017), manipulating training objectives (Burda et al., 2016; Higgins et al., 2017; Zhao et al., 2017), including the use of free bits (FB) (Kingma et al., 2016; Li et al., 2019), and changing training (He et al., 2019).
The other reason is the restrictive assumption of the parametric forms for the prior and approximate posterior. While these forms are computationally efficient, they limit the expressivity of the model. The main existing solutions (Kingma et al., 2016; Tomczak and Welling, 2018; Razavi et al., 2019) focus on enriching the variational posterior, while other work focuses on learning an expressive prior (Tomczak and Welling, 2018; Serban et al., 2017; Chen et al., 2017).
In this paper, we follow the latter line of research and draw upon methods in building and learning expressive priors. We first show empirically that the original VAE objective, the evidence lower bound (ELBO), is not effective when learning priors. The issue is not solely due to posterior collapse since it is not resolved by using modifications based on free bits. To address this issue, we propose using a combined objective, adding to the ELBO a second objective (denoted  $\mathbf{M}_{\mathrm{IS}}$ ) which is a different lower bound on the log marginal likelihood obtained using importance sampling (Burda et al., 2016).
Using the combination of the ELBO and  $\mathbf{M}_{\mathrm{IS}}$ , we compare multiple choices for the prior, including a mixture of Gaussians, a prior based on a variational mixture of posteriors (VampPrior; Tomczak and Welling, 2018), and a prior based on normalizing flows (NF), specifically real NVP transformations (Dinh et al., 2016). Using a real NVP prior entails creating an invertible mapping from a simple base distribution to the prior distribution of the latent variable in a VAE. This choice allows a flexible prior distribution that is not constrained to a specific parametric family. The hope is that it would be better at modeling the data distribution.
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 3242-3258 June 6-11, 2021. ©2021 Association for Computational Linguistics
We perform an empirical evaluation of priors and objective functions for training VAE sentence models on four standard datasets. We find the best performance overall when using the flow-based prior and the combined objective in the training objective. We refer to this setting as FlowPrior. The generation of prior samples with FlowPrior comports to the training distribution while maintaining a higher diversity than competing models in our quantitative and qualitative evaluation.
To summarize, this paper contributes: (1) a strategy for improved training of sentence VAEs based on combining multiple lower bounds on the log marginal likelihood; (2) the first results applying real NVP to model the prior in sentence VAEs; and (3) comprehensive evaluation and analysis with three expressive priors and training objective variations.
Variational autoencoders (VAEs; Kingma and Welling, 2014) are a popular framework for learning latent variable models with continuous latent variables. Let  $x$  be the observed variable and  $z$  the latent variable. The model factorizes the joint distribution over  $x$  and  $z$  into a prior  $p_{\psi}(z)$  and a generator  $p_{\theta}(x \mid z)$ . Maximizing the log marginal likelihood  $\log p(x)$  is intractable in general, so VAEs introduce an approximate posterior  $q_{\phi}(z \mid x)$  parameterized using a neural network (i.e., an "inference network"), and replace the log marginal likelihood with the evidence lower bound (ELBO):
$$
\begin{array}{l} \log p (x) \geq \mathbb {E} _ {q _ {\phi} (z | x)} [ \log p _ {\theta} (x | z) ] \\ - \operatorname {K L} \left(q _ {\phi} (z | x) \mid p _ {\psi} (z)\right) \tag {1} \\ \end{array}
$$
Maximizing the right-hand side of the equation above can be viewed as a regularized autoencoder in which the first term is the negative reconstruction error and the second is the negative KL divergence between the approximate posterior  $q_{\phi}(z|x)$  and the latent variable prior  $p_{\psi}(z)$ . It is common in practice to fix the prior  $p_{\psi}(z)$  to be a standard Gaussian distribution and only learn  $\theta$  and  $\phi$  (Bowman et al., 2016; Yang et al., 2017; Shen et al., 2017).
While constraining the prior to be a fixed standard Gaussian is common, it is not necessary for tractability. Researchers have found benefit from using richer priors and posteriors (Rezende and Mohamed, 2015; Kingma et al., 2016; Chen et al., 2017; Ziegler and Rush, 2019; Ma et al., 2019). In this paper, we consider investigating alternative
prriors while still using the standard Gaussian form for the approximate posterior.
3 Choices for Prior Families
We now describe the three kinds of priors we will compare in our experiments. The first two are based on Gaussian mixtures (Sec. 3.1) and the third is based on normalizing flows (Sec. 3.2). We take these three prior families into consideration because they represent the three main categories of work in learning priors: simple Gaussian mixtures (usually as baselines), defining the prior as a function of the approximate posterior (Tomczak and Welling, 2018; Chen et al., 2018), and flow-based priors (Chen et al., 2017; Ziegler and Rush, 2019; Ma et al., 2019; Lee et al., 2020). Note that we do not make any changes to the approximate posterior distribution. That is, the approximate posterior follows a Gaussian distribution with a diagonal covariance matrix as in standard VAEs.
3.1 Gaussian Mixture Priors
Our first choice is a uniform mixture of  $K$  Gaussians (MoG):
$$
p _ {\psi} (z) = \frac {1}{K} \sum_ {k = 1} ^ {K} f \left(z; \mu_ {k}, \operatorname {d i a g} \left(\sigma_ {k} ^ {2}\right)\right) \tag {2}
$$
where  $f(z;\mu ,\Sigma)$  is the density function of a  $d$ -dimensional Gaussian with mean  $\mu$  and covariance matrix  $\Sigma$ . The  $\mu_{k}$  and  $\sigma_{k}$  are learnable parameter vectors with dimensionality  $d$  (which is 32 in our experiments). This prior was used as a baseline by Tomczak and Welling (2018). We refer to a VAE that uses this prior as MoG-VAE.
Tomczak and Welling (2018) extend MoG-VAE to a "Variational Mixture of Posteriors" prior (VampPrior). This approach parameterizes the prior using a mixture of Gaussians with components given by a variational posterior conditioned on learnable "pseudo-inputs":
$$
p _ {\psi} (z) = \frac {1}{K} \sum_ {k = 1} ^ {K} q _ {\phi} (z \mid u _ {k}) \tag {3}
$$
where  $K$  is the number of pseudo-inputs, each of which is denoted  $u_{k}$ . Pelsmaeker and Aziz (2020) applied this idea to text modeling and we follow their strategy for defining pseudo-inputs. That is, each  $u_{k}$  consists of a sequence of embeddings that have the same dimensionality as word embeddings. For each component  $k$ , the lengths of pseudo-inputs
can vary; they are sampled based on the statistics of the lengths in the training set. We refer to a VAE with this prior as Vamp-VAE.
3.2 Flow-based Priors
Our third choice for a prior distribution is to leverage normalizing flows (NF). A normalizing flow is a sequence of invertible, deterministic transformations. By repeatedly applying the rule for change of variables (see the Appendices for details), the base density is transformed into a more complex one. Networks parameterized using NF can be trained through exact maximum log-likelihood computation. Exact sampling is performed by drawing a sample from the base distribution and performing the chain of transformations. This allows a flexible prior and is expected to have more expressive latent components compared to those based on Gaussian mixtures.
Computing the Jacobian of functions with high dimension and the determinants of large matrices (i.e., the two main computations in NF) are very expensive. Our flow-based prior uses real-valued non-volume preserving (real NVP) transformations (Dinh et al., 2016) which are efficient in both training and sampling. The transformations are based on scale and translation operations. It is worth noting that these two operations are not used in computing the Jacobian determinant and inverse. So one can design arbitrarily complex operations that allow a flexible transformation without incurring large computational cost.
More specifically, we apply real NVP as a prior by creating an invertible mapping between a base distribution  $p_0(z_0)$  (in our case,  $z_0 \sim \mathcal{N}(0, \mathrm{I})$ ) and the prior distribution  $p_\psi(z_L)$  in the VAE:
$$
z _ {L} = f _ {L} \circ f _ {L - 1} \circ \dots \circ f _ {1} (z _ {0}) \tag {4}
$$
where  $z_{L}$  is the sentence latent variable and  $f_{1},f_{2},\ldots ,f_{L}$  are all bijective functions.
Using the change-of-variables theorem, given a latent variable  $z_{L}$ , we can compute the exact density under the prior with the "image"  $z_{0}$  acquired by inverting the transformation:
$$
z _ {0} = f _ {1} ^ {- 1} \circ \dots \circ f _ {L - 1} ^ {- 1} \circ f _ {L} ^ {- 1} (z _ {L}) \tag {5}
$$
$$
\log p _ {\psi} \left(z _ {L}\right) = \log p _ {0} \left(z _ {0}\right) - \sum_ {l = 1} ^ {L} \log | \det (\frac {\partial f _ {l} \left(z _ {l - 1}\right)}{\partial z _ {l - 1}}) | \tag {6}
$$
We refer to a VAE with a real NVP prior as real NVP-VAE. We find our best setting to consist of a real NVP prior and the combined objective in Section 4.1 and we refer to this setting as FlowPrior.
4 Objectives for Learning Priors in VAEs
ELBO. Our preliminary experiments found that, when training with the standard ELBO, using more sophisticated priors does not improve perplexity compared to standard Gaussian priors (Table 3). Though these priors could potentially be highly multimodal, the learned prior parameters yield approximately unimodal forms (Figure 1, left).
Several approaches have been proposed to mitigate or avoid collapse in the approximate posterior. One method that we include in our experiments is a variation of KL divergence known as "free bits" (FB) KL (Li et al., 2019; Kingma et al., 2016). Posterior collapse is mitigated, but the VAE models still do not benefit much from expressive priors (Tables 1-2). Pelsmaeker and Aziz (2020) made similar observations with an improved FB objective. We speculate that these undesirable results are due to the lack of learning signal for the prior parameters.
Marginal Likelihood via Importance Sampling. In the ELBO, the prior distribution only appears in the KL term. As a consequence, the prior parameters receive a limited amount of learning signal. The posterior network, by contrast, receives gradient updates from both the reconstruction and KL terms. When minimizing the KL term the potentially expressive prior density can "collapse" to a unimodal form, as this may facilitate minimizing the KL divergence between the approximate posterior and prior.
We consider optimizing another objective, a different lower bound on the log marginal likelihood obtained using importance sampling (Burda et al., 2016):
$$
\log \frac {1}{N} \sum_ {i = 1} ^ {N} \frac {p _ {\theta} (x | z ^ {(i)}) p _ {\psi} (z ^ {(i)})}{q _ {\phi} (z ^ {(i)} | x)}, \mathrm {s . t .} z ^ {(i)} \sim q _ {\phi} (z | x)
$$
where  $x$  is an input in the training data and  $N$  is the number of samples in use. This objective was proposed as the training objective in the importance-weighted autoencoder (IWAE; Burda et al., 2016), and was shown to be a tighter lower bound on the log marginal likelihood than the ELBO. In this paper, we denote this objective by  $\mathbf{M}_{\mathrm{IS}}$ .
1More details about normalizing flows and real NVP are in the Appendices.
In addition to providing a tighter lower bound,  $\mathbf{M}_{\mathrm{IS}}$  also increases the flexibility of the approximate posterior, as shown by Cremer et al. (2017). By increasing  $N$ , the approximate posterior has an implicit complex distribution that approaches the true posterior, which may also be beneficial in learning an expressive prior.
Combination of the Two. However,  $\mathbf{M}_{\mathrm{IS}}$  is not necessarily optimal by itself for training VAEs. Rainforth et al. (2018) prove that using  $\mathbf{M}_{\mathrm{IS}}$  with a large value of  $N$  is detrimental in learning the posterior, which is also shown in our empirical evaluation in Table 3. If we only have  $\mathbf{M}_{\mathrm{IS}}$ , the approximate posterior  $q$  only appears in the denominator so learning seeks to make samples from the posterior  $q$  less likely under  $q$ , which could cause  $q$  to become a poor proposal distribution. The ELBO, with its reconstruction loss, appears helpful in learning a better posterior. Therefore, we optimize the sum of the ELBO and  $\mathbf{M}_{\mathrm{IS}}$ , which was proposed by Rainforth et al. (2018).
4.1 Combined Training Objective
Our combined training objective then contains three terms:  $\mathbf{M}_{\mathrm{IS}}$ , reconstruction, and sample-based KL. We draw  $N$  samples from  $q_{\phi}(z|x)$ , and compute the three terms using the same samples:
$$
\begin{array}{l} \mathcal {L} (\theta , \phi , \psi ; x) = \log \frac {1}{N} \sum_ {i = 1} ^ {N} \frac {p _ {\theta} (x | z ^ {(i)}) p _ {\psi} (z ^ {(i)})}{q _ {\phi} (z ^ {(i)} | x)} \\ + \frac {1}{N} \sum_ {i = 1} ^ {N} \log p _ {\theta} (x | z ^ {(i)}) - \mathrm {K L} _ {\phi , \psi} (x, \{z ^ {(i)} \} _ {i = 1} ^ {N}) \\ s. t. z ^ {(i)} \sim q _ {\phi} (z | x) \tag {7} \\ \end{array}
$$
When training with the ELBO alone, one typically uses a single sample from  $q_{\phi}(z|x)$ . However, since we draw multiple samples anyway in order to compute  $\mathbf{M}_{\mathrm{IS}}$ , we use those same samples for the reconstruction term, which can lead to more robust gradients of that term than the standard approach of using one sample.
The reason we use sample-based estimates for the KL divergence is because our choices for the prior preclude the possibility of a closed form for the KL. We consider two different approaches when computing sample-based KLs: standard KL and a modified one inspired by free bits (Li et al., 2019; Pelsmaeker and Aziz, 2020; Kingma et al., 2016), which we refer to as  $FBKL$ .
For standard KL, we use Monte Carlo estimation in computing the KL divergence with  $N$  samples:
$$
\begin{array}{l} \mathrm {K L} _ {\phi , \psi} (x, \{z ^ {(i)} \} _ {i = 1} ^ {N}) = \\ \frac {1}{N} \sum_ {i = 1} ^ {N} \left(\log q _ {\phi} \left(z ^ {(i)} \mid x\right) - \log p _ {\psi} \left(z ^ {(i)}\right)\right) \tag {8} \\ \end{array}
$$
For the FB KL, we follow prior work (Kingma et al., 2016) that replaces the KL with a hinge loss term in each latent dimension:
$$
\begin{array}{l} \mathrm {F B K L} _ {\phi , \psi} (x, \{z ^ {(i)} \} _ {i = 1} ^ {N}) = \\ \sum_ {j = 1} ^ {d} \max  \left(\lambda , \mathrm {K L} _ {\phi , \psi} ^ {j} (x, \{z ^ {(i)} \} _ {i = 1} ^ {N})\right) \tag {9} \\ \end{array}
$$
where  $\mathrm{KL}_{\phi, \psi}^{j}$  denotes the KL computed only for dimension  $j$  of the latent variable, and  $\lambda$  is the "target rate" hyperparameter.
4.2 Training Procedure
We describe our training procedure below for FlowPrior, which combines a real NVP prior with the objective in Eq. 7. For simplicity, our description only uses one input  $x$ . In practice, we use minibatches with a stochastic gradient based optimizer. All the parameters  $(\theta, \phi, \psi)$  are updated simultaneously during training.
1. Draw  $N$  samples  $z_{L}^{(1)}, z_{L}^{(2)}, \ldots, z_{L}^{(N)}$  from the inference network using the reparameterization trick.
2. Perform the inverse transformation to get the image of each point under the base distribution:  $z_0^{(1)}, z_0^{(2)}, \ldots, z_0^{(N)}$ .
3. Compute the exact log likelihood of the sample prior with change of variable theorem (Eq. 6).
4. Compute and backpropagate the loss (Eq. 7).
When using the other priors (standard Gaussian, MoG, and VampPrior), we do not need steps 2 and 3 above because those priors can be computed directly without the inverse transformation or change of variable theorem.
We consider four widely-used, publicly available English datasets: the Penn Treebank (PTB) (Marcus et al., 1993; Bowman et al., 2016), Yahoo (Yang et al., 2017; He et al., 2019), Yelp sentiment (Shen et al., 2017), and SNLI (Bowman et al., 2015).
Our baselines include standard VAE with linear KL annealing (Bowman et al., 2016); Cyc-VAE (Fu et al., 2019) in which the KL term is reweighted with a cyclical annealing schedule; Lag-VAE (He et al., 2019) which updates the encoder multiple times before each decoder update;  $VAE + FB$  (Kingma et al., 2016; Chen et al., 2017) which replaces the standard KL with FB (i.e., Eq. 9 with  $N = 1$ ); and Pre-VAE+FB (Li et al., 2019) which initializes the VAE with a pretrained autoencoder and replaces standard KL with FB. We evaluated these baselines using their open source implementations.
In addition, we include two prior-learning baselines: MoG-VAE (Eq. 2) and Vamp-VAE (Eq. 3). We follow Pelsmaeker and Aziz (2020) and set 100 components/pseudo-inputs. Unlike the earlier baselines, for which we used open source codebases, we implemented the MoG-VAE and Vamp-VAE models on top of our standard VAE implementation, which was also used for FlowPrior.
5.3 Implementation and Training Details
Across all the experiments for our implemented baselines (i.e., standard VAE, MoG-VAE, VampVAE) and our proposed model FlowPrior, we follow prior work (Kim et al., 2018; He et al., 2019; Li et al., 2019) and use a single-layer LSTM encoder and decoder with a 32-dimensional latent variable. We use a batch size of 32 and train using SGD.
Our evaluation measures language modeling performance, the use of the latent variable, and the quality and diversity of generations from the prior and posterior. The metrics are listed below:
PPL: We estimate log marginal likelihood using importance sampling (Burda et al., 2016) and calculate perplexity on the test set.
KL: We report the KL term in the ELBO on the test set. When training with FB KL, we still report standard KL. For standard VAE, we compute KL with its closed-form expression. Otherwise, we report the KL estimated with samples.
Model	PPL(↓)	Recon(↓)	KL	AU(↑)	MI(↑)
VAE	101.40	101.28	0.00	0	0.00
Cyc-VAE	107.73	101.17	2.01	5	1.24
Lag-VAE	100.25	100.41	1.04	3	0.79
VAE + FB	101.56	99.84	4.46	32	0.90
Pre-VAE + FB	96.35	94.52	8.15	32	6.30
MoG-VAE	98.22	100.54	0.00	0	0.00
MoG-VAE + FB	97.50	99.44	2.35	32	0.68
Vamp-VAE	98.27	100.56	0.00	0	0.00
Vamp-VAE + FB	97.83	99.53	2.31	32	0.72
FlowPrior	94.72	98.46	3.28	2	2.25
FlowPrior + FB	93.58	99.20	7.21	31	2.83
Table 1: Language modeling results on PTB dataset.
MI: We follow Hoffman and Johnson (2016) and report estimated mutual information between the observation and its latent variable.
AU: A dimension  $z$  in the latent variable is considered "active" if  $\mathrm{Cov}_x(\mathbb{E}_{z\sim q(z|x)}[z]) > 10^{-2}$ . AU is then the number of active latent dimensions (Burda et al., 2016).
F-PPL and R-PPL: These metrics measure the correspondence between generated sentences from the model and the training corpus. We evaluate both F-PPL and R-PPL by estimating 5-gram language models using the KenLM toolkit (Heafield, 2011) with its default smoothing method. For F-PPL, we estimate language models from the actual text and compute the perplexity of the generated samples. For R-PPL, we estimate language models from the generated samples and compute the perplexity of the actual text.
Self-BLEU: The self-BLEU metric is one measure of the diversity of a set of samples (Zhu et al., 2018). It is calculated by averaging the BLEU scores computed between all pairs of samples.
6.1 Language Modeling
We first perform language modeling tasks to characterize models' efficacy at modeling texts in terms of modeling the distribution of language data and making use of the latent variable. We refer to our model as FlowPrior, which uses the training objective in Eq. 7 which includes  $\mathbf{M}_{\mathrm{IS}}$  and the standard KL (Eq. 8). We use FlowPrior + FB to refer to our model with the FB KL (Eq. 9).
2The links to their implementations are in the Appendix.
3We use the open source implementations for other baselines. All models are trained with the simple linear annealing schedule, with same hyperparameter search space. We run each setting with 5 random seeds and report the medians.
<sup>4</sup>We use 1000 samples which appears to be more than sufficient for estimation; Ziegler and Rush (2019) found that using more than 50 samples did not even show much difference.
<sup>5</sup>Our R-PPL is slightly different from that in Fang et al. (2019). For R-PPL, we always concatenate the training set vocabulary (one word per line) to the set of samples from the models to ensure LMs have seen the entire vocabulary.
Model	PPL(↓)	Recon(↓)	KL	AU(↑)	MI(↑)
Yahoo
VAE	65.77	333.17	0.00	0	0.00
MoG-VAE	64.60	332.90	0.00	0	0.00
Vamp-VAE	74.81	344.61	0.01	0	0.00
FlowPrior	62.49	331.57	1.43	4	1.62
FlowPrior + FB	68.29	345.68	10.99	25	0.61
Yelp
VAE	35.10	35.18	0.00	0	0.00
MoG-VAE	35.18	35.20	0.01	0	0.00
Vamp-VAE	34.99	35.15	0.00	0	0.00
FlowPrior	31.82	30.25	4.15	2	2.46
FlowPrior + FB	39.03	36.87	10.13	32	2.57
SNLI
VAE	25.97	41.34	0.00	0	0.00
MoG-VAE	28.05	40.96	0.44	1	0.41
Vamp-VAE	25.98	41.35	0.00	0	0.00
FlowPrior	22.41	37.89	3.83	3	0.97
FlowPrior + FB	26.19	43.56	7.59	32	3.16
Comparison to baselines. Table 1 shows results on the PTB dataset for several VAEs from prior work and our implemented models. Since our contributions lie in learning the prior instead of changing the training procedure or manipulating the KL term, we set the baselines as standard VAE, MoG, and VampPrior for the rest of the paper. We report the performance of FlowPrior and those baselines on Yahoo, Yelp, and SNLI in Table 2.