SlowGuess/ACL_2020_OCR huggingface dataset:
Shuning Jin $^{1*}$  Sam Wiseman $^{2}$  Karl Stratos $^{1}$  Karen Livescu $^{2}$
$^{1}$ Rutgers University  $^{2}$ Toyota Technological Institute at Chicago
shuning.jin@rutgers.edu, {swiseman,klivescu}@ttic.edu,stratos@cs.rutgers.edu
While much work on deep latent variable models of text uses continuous latent variables, discrete latent variables are interesting because they are more interpretable and typically more space efficient. We consider several approaches to learning discrete latent variable models for text in the case where exact marginalization over these variables is intractable. We compare the performance of the learned representations as features for low-resource document and sentence classification. Our best models outperform the previous best reported results with continuous representations in these low-resource settings, while learning significantly more compressed representations. Interestingly, we find that an amortized variant of Hard EM performs particularly well in the lowest-resource regimes. $^{1}$
Deep generative models with latent variables have become a major focus of NLP research over the past several years. These models have been used both for generating text (Bowman et al., 2016) and as a way of learning latent representations of text for downstream tasks (Yang et al., 2017; Gururanan et al., 2019). Most of this work has modeled the latent variables as being continuous, that is, as vectors in  $\mathbb{R}^d$ , in part due to the simplicity of performing inference over (certain) continuous latents using variational autoencoders and the reparameterization trick (Kingma and Welling, 2014; Rezende et al., 2014).
At the same time, deep generative models with discrete latent variables are attractive because the latents are arguably more interpretable, and because they lead to significantly more compressed
representations: A representation consisting of  $M$  floating point values conventionally requires  $M \times 32$  bits, whereas  $M$  integers in  $\{1, \ldots, K\}$  requires only  $M \times \log_2 K$  bits.
Unfortunately, discrete latent variable models have a reputation for being more difficult to learn. We conduct a thorough comparison of several popular methods for learning such models, all within the framework of maximizing the evidence lower bound (ELBO) on the training data. In particular, we compare learning such models either with a Vector Quantized-VAE (van den Oord et al., 2017, VQ-VAE), a more conventional VAE with discrete latent variables (Jang et al., 2017; Maddison et al., 2017), or with an amortized version of "Hard" or "Viterbi" Expectation Maximization (Brown et al., 1993), which to our knowledge has not been explored to date. We consider both models where the latents are local (i.e., per token) and where they are global (i.e., per sentence); we assess the quality of these learned discrete representations as features for a low-resource text classifier, as suggested by Gururangan et al. (2019), and in a nearest neighbor-based retrieval task.
Our classification experiments distinguish between (1) the setting where the classifier must consume only the discrete representation associated with each sentence (i.e., the discrete assignment that maximizes the approximate posterior), and (2) the setting where the classifier may consume the embeddings of this discrete representation learned by the VAE encoder. Note that the former setting is more flexible, since we need only store a sentence's discrete representation, and are therefore free to use task-specific (and possibly much smaller) architectures for classification. In case (1), we are able to effectively match the performance of Gururan-gan et al. (2019) and other baselines; in case (2), we outperform them. Our experiments also suggest that Hard EM performs particularly well in
*Work done as an intern at Toyota Technological Institute at Chicago.
  1Code available on GitHub: https://github.com/shuningjin/discrete-text-rep
4831
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 4831-4842 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
case (1) when there is little supervised data, and that VQ-VAE struggles in this setting.
Our work builds on recent advances in discrete representation learning and its applications. In particular, we are inspired by recent success with VQ-VAEs outside NLP (van den Oord et al., 2017; Razavi et al., 2019). These works show that we can generate realistic speech and image samples from discrete encodings, which better align with symbolic representations that humans seem to work with (e.g., we naturally encode continuous speech signals into discrete words). Despite its success in speech and vision, VQ-VAE has not been considered as much in NLP. One exception is the translation model of Kaiser et al. (2018) that encodes a source sequence into discrete codes using vector quantization. But their work focuses on making inference faster, by decoding the target sequence from the discrete codes non-autoregressively. To our knowledge, we are the first that explores general text representations induced by VQ-VAEs for semi-supervised and transfer learning in NLP.
In addition to exploring the viability of VQ-VAEs for text representation learning, an important part of this paper is a systematic comparison between different discretization techniques. Gumbel-Softmax (Jang et al., 2017; Maddison et al., 2017) is a popular choice that has been considered for supervised text classification (Chen and Gimpel, 2018) and dialog generation (Zhao et al., 2018). In the binary latent variable setting, straight-through estimators are often used (Dong et al., 2019). Another choice is "continuous decoding" which takes a convex combination of latent values to make the loss differentiable (Al-Shedivat and Parikh, 2019). Yet a less considered choice is Hard EM (Brown et al., 1993; De Marcken, 1995; Spitkovsky et al., 2010). A main contribution of this work is a thorough empirical comparison between such different choices in a controlled setting.
To demonstrate the usefulness of our models, we focus on improving low-resource classification performance by pretraining on unlabeled text. Previous best results are obtained with continuous latent-variable VAEs, e.g., VAMPIRE (Gururangan et al., 2019). We show that our discrete representations outperform these previous results while being significantly more lightweight.
We consider generative models of a sequence  $x = x_{1:T}$  of  $T$  word tokens. We assume our latents to be a sequence  $z = z_{1:L}$  of  $L$  discrete latent vectors, each taking a value in  $\{1,\dots ,K\} ^M$ ; that is,  $z\in \{1,\ldots ,K\}^{M\times L}$ . As is common in VAE-style models of text, we model the text autoregressively, and allow arbitrary interdependence between the text and the latents. That is, we have  $p(x,z;\pmb {\theta}) = p(z)\times \prod_{t = 1}^{T}p(x_{t}\mid x_{< t},z;\pmb {\theta})$ , where  $\pmb{\theta}$  are the generative model's parameters. We further assume  $p(z)$  to be a fully factorized, uniform prior:  $p(z) = \frac{1}{K^{ML}}$ .
Maximizing the marginal likelihood of such a model will be intractable for moderate values of  $K, M$ , and  $L$ . So we consider learning approaches that maximize the ELBO (Jordan et al., 1999) in an amortized way (Kingma and Welling, 2014; Rezende et al., 2014):
$$
\operatorname {E L B O} (\pmb {\theta}, \phi) = \mathbb {E} _ {q (z | x; \phi)} \left[ \log \frac {p (x , z ; \pmb {\theta})}{q (z | x ; \phi)} \right],
$$
where  $q(z\mid x;\phi)$  is the approximate posterior given by an inference or encoder network with parameters  $\phi$ . The approaches we consider differ in terms of how this approximate posterior  $q$  is defined.
Mean-Field Categorical VAE (CatVAE) A standard Categorical VAE parameterizes the approximate posterior as factorizing over categorical distributions that are independent given  $x$ . We therefore maximize:
$$
\begin{array}{l} \mathbb {E} _ {q (z \mid x; \phi)} [ \log p (x \mid z; \boldsymbol {\theta}) ] - \sum_ {m, l} \mathrm {K L} (q _ {m l} | | p _ {m l}) \\ = \mathbb {E} _ {q (z \mid x; \phi))} [ \log p (x \mid z; \boldsymbol {\theta}) ] \\ + \sum_ {m, l} H (q _ {m l}) - M L \log K, \\ \end{array}
$$
where  $q(z\mid x;\phi) = \prod_{m = 1}^{M}\prod_{l = 1}^{L}q_{ml}(z_{ml}\mid x;\phi),$ $p_{ml} = 1 / K$  , and  $H$  is the entropy.
We approximate the expectation above by sampling from the  $q_{ml}$ , and we use the straight-through gradient estimator (Bengio et al., 2013; Jang et al., 2017) to compute gradients with respect to  $\phi$ . We find this approach to be more stable than using the REINFORCE (Williams, 1992) gradient estimator, or a Concrete (Maddison et al., 2017; Jang et al., 2017) approximation to categorical distributions. Specifically, we sample from a categorical distribution using the Gumbel-Max trick (Maddison et al., 2014) in the forward pass, and approximate the
gradient using softmax with a small temperature. This approach is also referred to as straight-through Gumbel-Softmax (Jang et al., 2017).
VQ-VAE A VQ-VAE (van den Oord et al., 2017; Razavi et al., 2019) can also be seen as maximizing the ELBO, except the approximate posterior is assumed to be a point mass given by
$$
q _ {m l} (z _ {m l} | x) = \left\{ \begin{array}{l l} 1 & \text {i f} z _ {m l} = \hat {z} _ {m l} \\ 0 & \text {o t h e r w i s e} \end{array} \right.,
$$
$$
\hat {z} _ {m l} = \underset {j \in \{1, \dots , K \}} {\arg \min } | | \mathbf {e} _ {j} ^ {(m)} - \operatorname {e n c} (x) _ {m l} | | _ {2}, \tag {1}
$$
and  $\mathbf{e}_j^{(m)}\in \mathbb{R}^d$  is an embedding of the  $j^{\mathrm{th}}$  discrete value  $z_{ml}$  can take on, and  $\mathrm{enc}(x)_{ml}\in \mathbb{R}^d$  is an encoding corresponding to the  $ml^{\mathrm{th}}$  latent given by an encoder network. These  $\mathbf{e}_j^{(m)}$  embedding vectors are often referred to as a VQ-VAE's "code book". In our setting, a code book is shared across latent vectors.
VQ-VAEs are typically learned by maximizing the ELBO assuming degenerate approximate posteriors as above, plus two terms that encourage the encoder embeddings and the "code book" embeddings to become close. In particular, we attempt to maximize the objective:
$$
\begin{array}{l} \log p (x \mid \hat {z}) - \sum_ {m, l} | | \operatorname {s g} (\operatorname {e n c} (x) _ {m l}) - \mathbf {e} _ {\hat {z} _ {m, l}} ^ {(m)} | | _ {2} ^ {2} \tag {2} \\ - \beta \sum_ {m, l} | | \mathrm {e n c} (x) _ {m l} - \mathrm {s g} (\mathbf {e} _ {\hat {z} _ {m, l}} ^ {(m)}) | | _ {2} ^ {2}, \\ \end{array}
$$
where  $\mathrm{sg}$  is the stop-gradient operator, and  $\hat{z} = \hat{z}_{1:L}$  is the sequence of minimizing assignments  $\hat{z}_{m,l}$  for each  $\mathrm{enc}(x)_{ml}$ . The loss term following the  $\beta$  is known as the "commitment loss". Gradients of the likelihood term with respect to  $\mathrm{enc}(x)$  are again estimated with the straight-through gradient estimator.
Hard EM We train with an amortized form of Hard EM. First we define a relaxed version of  $z$ ,  $\tilde{z}$ , where each  $\tilde{z}_{ml}$  is a softmax over  $K$  outputs (rather than a hard assignment) and is produced by an inference network with parameters  $\phi$ . In the E-Step, we take a small, constant number of
Figure 1: Discrete VAE architectures with  $M = 2$ . The Local (middle) and Global (bottom) models are two different encoder-decoder setups. The top row shows the procedure of converting continuous output from encoder into discrete input to decoder by drawing discrete samples: VQ-VAE (top left) draws samples from point mass distributions using nearest neighbor lookup from the code books; CatVAE (top right) samples from categorical distributions directly.
gradient steps to maximize  $\log p(x|\tilde{z};\pmb{\theta})$  with respect to  $\phi$  (for a fixed  $\pmb{\theta}$ ). In the M-Step, we take a single gradient step to maximize  $\log p(x|\hat{z};\pmb{\theta})$  with respect to  $\pmb{\theta}$ , where  $\hat{z}$  contains the element-wise argmaxes of  $\tilde{z}$  as produced by the inference network (with its most recent parameters  $\phi$ ). Thus, Hard EM can also be interpreted as maximizing the (relaxed) ELBO. We also note that taking multiple steps in the hard E-step somewhat resembles the recently proposed aggressive training of VAEs (He et al., 2019).
4 Models and Architectures
Recall that the latent sequence is  $z = z_{1:L}$ , where  $z_l \in \{1, \dots, K\}^M$ . We consider two generative models  $p(x | z; \boldsymbol{\theta})$ , one where  $L = T$  and one where  $L = 1$ . Each latent in the former model corresponds to a word, and so we refer to this as a "local" model, whereas in the second model we view the latents as being "global", since there is one latent vector for the whole sentence. We use the following architectures for our encoders and decoder, as illustrated in Figure 1.
2Note this assumes our generative model can condition on such a relaxed latent variable.
4.1 Encoder
The encoder (parameterized by  $\phi$ ) maps an example  $x$  to the parameters of an approximate posterior distribution. Our encoder uses a single-layer Transformer (Vaswani et al., 2017) network to map  $x = x_{1:T}$  to a sequence of  $T$  vectors  $\mathbf{h}_1, \ldots, \mathbf{h}_T$ , each in  $\mathbb{R}^d$ .
Mean-Field Categorical VAE For the local model, we obtain the parameters of each categorical approximate posterior  $q_{mt}$  as  $\mathrm{softmax}(\mathbf{W}_m\mathbf{h}_t)$ , where each  $\mathbf{W}_m \in \mathbb{R}^{K\times d}$  is a learned projection. For the global model, we obtain the parameters of each categorical approximate posterior  $q_{m1}$  as  $\mathrm{softmax}\left(\frac{\sum_t\mathbf{W}_m\mathbf{h}_t}{T}\right)$ ; that is, we pass token-level  $\mathbf{h}_t$  vectors through learned projections  $\mathbf{W}_m$ , followed by mean-pooling.
VQ-VAE For the local model, let  $\tilde{d} = d / M$ . We obtain  $\mathrm{enc}(x)_{mt}$ , the encoding of the  $mt^{\mathrm{th}}$  latent variable, as  $\mathbf{h}_{t,(m - 1)\tilde{d}:m\tilde{d}}$ , following Kaiser et al. (2018). That is, we take the  $m^{\mathrm{th}}\tilde{d}$ -length subvector of  $\mathbf{h}_t$ . For the global model, let  $\tilde{d} = d$ . We first project  $\mathbf{h}_t$  to  $\mathbb{R}^{Md}$ , mean-pool, and obtain  $\mathrm{enc}(x)_{m1}$  by taking the  $m^{\mathrm{th}}\tilde{d}$ -length subvector of the resulting pooled vector. A VQ-VAE also requires learning a code book, and we define  $M$  code books  $\mathbf{E}^{(m)} = [\mathbf{e}_1^{(m)^{\top}};\dots ;\mathbf{e}_K^{(m)^{\top}}]\in \mathbb{R}^{K\times \tilde{d}}$ .
Hard EM We use the same encoder architecture as in the mean-field Categorical VAE case. Note, however, that we do not sample from the resulting categorical distributions. Rather, the softmax distributions are passed directly into the decoder.
4.2 Decoder
In the case of the mean-field Categorical VAE, we obtain a length-  $L$  sequence of vectors  $z_{l}\in$ $\{1,\dots ,K\} ^M$  after sampling from the approximate posteriors. For the VQ-VAE, on the other hand, we obtain the sequence of  $\hat{z}_l$  vectors by taking the indices of the closest code book embeddings, as in Equation (1).
In both cases, the resulting sequence of discrete vectors is embedded and consumed by the decoder. In particular, when learning with a VQ-VAE, the embedding of  $\hat{z}_{ml}$  is simply  $\mathbf{e}_{\hat{z}_{ml}}^{(m)}$ , whereas for the Categorical VAE each discrete latent is embedded using a trained embedding layer. In the local model, when  $M > 1$ , we concatenate the  $M$  embeddings to form a single real vector embedding for the  $l^{\mathrm{th}}$  latent variable. In the global model, we use the  $M$
embeddings directly. This resulting sequence of  $T$  or  $M$  real vectors is then viewed as the source side input for a standard 1-layer Transformer encoder-decoder model (Vaswani et al., 2017), which decodes  $x$  using causal masking.
As above, for Hard EM, we do not obtain a sequence of discrete vectors from the encoder, but rather a sequence of softmax distributions. These are multiplied into an embedding layer, as in the Categorical VAE case, and fed into the Transformer encoder-decoder model.
5 Evaluating Latent Representations
Similar to Gururangan et al. (2019), we evaluate the learned latent representations by using them as features in a text classification system. We are in particular interested in using latent representations learned on unlabeled text to help improve the performance of classifiers trained on a small amount of labeled text. Concretely, we compare different discrete latent variable models in following steps:
1. Pretraining an encoder-decoder model on indomain unlabeled text with an ELBO objective, with early stopping based on validation perplexity.
2. Fixing the encoder to get discrete latents for the downstream classification task, and training a small number of task-specific parameters on top, using varying amounts of labeled data. As noted in the introduction, we consider both reembedding these latents from scratch, or using the embeddings learned by the encoder.
The datasets we use for classification are AG News, DBPedia, and Yelp Review Full (Zhang et al., 2015), which correspond to predicting news labels, Wikipedia ontology labels, and the number of Yelp stars, respectively. The data details are summarized in Table 1. For all datasets, we randomly sample 5,000 examples as development data. To evaluate the efficiency of the latent representation in low-resource settings, we train the classifier with varying numbers of labeled instances: 200, 500, 2500, and the full training set size (varies by dataset). We use accuracy as the evaluation metric.
In preprocessing, we space tokenize, lowercase, and clean the text as in Kim (2014), and then truncate each sentence to a maximum sequence length of 400. For each dataset, we use a vocabulary of the 30,000 most common words.
Dataset	# Classes	Train	Dev	Test
AG News	4	115K	5K	7.6K
DBPedia	14	555K	5K	70K
Yelp Review Full	5	645K	5K	50K
Table 1: The number of classes and the numbers of examples in each data subset, for the classification tasks.
5.2 Transfer Paradigm
When transferring to a downstream classification task, we freeze the pretrained encoder and add a lightweight classifier on top, viewing each sentence as an  $L$ -length sequence of vectors in  $\{1,\ldots ,K\} ^M$  as described in Section 4. For instance, the sentence (from the DBPedia dataset) "backlash is a 1986 australian film directed by bill bennett" is encoded as [90, 114, 30, 111] under a global model with  $M = 4$  and as [[251, 38], [44, 123], [94, 58], [228, 53], [88, 55], [243, 43], [66, 236], [94, 72], [172, 61], [236, 150]] under a local model with  $M = 2$ .
As noted in the introduction, we consider two ways of embedding the integers for consumption by a classifier. We either (1) learn a new task-specific embedding space  $\mathbf{E}_{\mathrm{task}}^{(m)}$  (i.e., reembedding) or (2) use the fixed embedding space  $\mathbf{E}^{(m)}$  from pretraining. The first setting allows us to effectively replace sentences with their lower dimensional discrete representations, and learn a classifier on the discrete representations from scratch. In the local model, we obtain token-level embedding vectors by concatenating the  $M$  subvectors corresponding to each word. The resulting embeddings are either averaged, or fed to a Transformer and then averaged, and finally fed into a linear layer followed by a softmax.
6 Experimental Details
6.1 Baselines
We first experiment with three common text models: CBOW (Mikolov et al., 2013), bidirectional LSTM (Hochreiter and Schmidhuber, 1997), and a single-layer Transformer encoder. We find CBOW (with 64-dimensional embeddings) to be the most robust in settings with small numbers of labeled instances, and thus report results only with this baseline among the three. Further, we compare to VAMPIRE (Gururangan et al., 2019), a framework of pretraining VAEs for text classification using continuous latent variables. We pretrain VAMPIRE
models on in-domain text for each dataset with 60 random hyperparameter search (with same ranges as specified in their Appendix A.1), and select best models based on validation accuracy in each setting.
6.2 Hyperparameters
In our experiments, we use Transformer layers with  $d_{\mathrm{model}} = 64$ . For optimization, we use Adam (Kingma and Ba, 2015), either with a learning rate of 0.001 or with the inverse square-root schedule defined in Vaswani et al. (2017) in pretraining. We use a learning rate of 0.0003 in classification. We tune other hyperparameters with random search and select the best settings based on validation accuracy. For the latent space size, we choose  $M$  in  $\{1,2,4,8,16\}$  and  $K$  in  $\{128,256,512,1024,4096\}$ . Model specific hyperparameters are introduced below.
6.3 VQ-VAE
In VQ-VAE, an alternative to the objective in Equation (2) is to remove its second term, while using an auxiliary dictionary learning algorithm with exponential moving averages (EMA) to update the embedding vectors (van den Oord et al., 2017). We tune whether to use EMA updates or not. Also, we find small  $\beta$  for commitment loss to be beneficial, and search over  $\{0.001, 0.01, 0.1\}$ .
6.4 Mean-Field Categorical VAE
We find that using the discrete analytic KL divergence term directly in the ELBO objective leads to posterior collapse. The KL term vanishes to 0 and the  $q_{ml}$  distributions converge to the uniform priors. To circumvent this, we modify the KL term to be  $\max (\mathrm{KL},\lambda)$ . This is known as Free Bits (Kingma et al., 2016; Li et al., 2019), which ensures that the latent variables encode a certain amount of information by not penalizing the KL divergence when it is less than  $\lambda$ . We set  $\lambda = \gamma ML\log K$  where  $\gamma$  is a hyperparameter between 0 and 1. That is, we allocate a "KL budget" as a fraction of  $ML\log K$ , which is the upper bound of KL divergence between  $ML$  independent categorical distributions and uniform prior distributions. Since in this case  $\mathrm{KL}(q_{ml}(z_{ml}|x)||p_{ml}(z_{ml})) = \log K - H[q_{ml}(z_{ml}|x)]$ , this is equivalent to thresholding  $H[q_{ml}(z_{ml}|x)]$  by  $(1 - \gamma)\log K$ . We experiment with  $\gamma \in \{0.2,0.4,0.6,0.8,1\}$ .
Note that when  $\gamma \geq 1$  the VAE reduces to an autoencoder.
Figure 2: The accuracies obtained by Hard EM, Categorical VAE, and VQ-VAE representations, averaged over the AG News, DBPedia, and Yelp Full development datasets, for different numbers of labeled training examples. Triangular and circular markers correspond to global and local models, respectively. Unshaded and shaded markers correspond to reembedding from scratch and using encoder embeddings, respectively.
Figure 3: The averaged accuracies obtained from using Hard EM, Categorical VAE, and VQ-VAE representations and 200 labeled examples, for different  $M$  values.
6.5 Hard EM
We vary the number of gradient steps in the E-step in  $\{1,3\}$ . At evaluation time, we always take the argmax of  $\tilde{z}$  to get a hard assignment.
In Figure 2, we compare the accuracy obtained by the representations from our Hard EM, Categorical VAE, and VQ-VAE models, averaged over the development datasets of AG News, DBPedia, and Yelp Full. In particular, we plot the best accuracy obtained over all hyperparameters (including  $M$ ) for different numbers of labeled examples; we distinguish between local and global models,
and between when the discrete representations are reembedded from scratch and when the encoder embeddings are used.