Finetuning Pretrained Transformers into Variational Autoencoders (2021)

Abstract: Text variational autoencoders (VAEs) are notorious for posterior collapse, a phenomenon where the model’s decoder learns to ignore signals from the encoder. Because posterior collapse is known to be exacerbated by expressive decoders, Transformers have seen limited adoption as components of text VAEs. Existing studies that incorporate Transformers into text VAEs (Li et al., 2020; Fang et al., 2021) mitigate posterior collapse using massive pretraining, a technique unavailable to most of the research community without extensive computing resources. We present a simple two-phase training scheme to convert a sequence-to-sequence Transformer into a VAE with just finetuning. The resulting language model is competitive with massively pretrained Transformer-based VAEs in some internal metrics while falling short on others. To facilitate training we comprehensively explore the impact of common posterior collapse alleviation techniques in the literature. We release our code for reproducability.

## introduction properly tamed latent models offer explainable and interpolatable representations of observed data. recent works have shown such models to be especially useful in unsupervised learning settings. @xcite adapt a generative latent text model for successful unsupervised text style transfer and machine translation. @xcite achieve superior language modeling performance against common conditional counterparts. a popular variant of deep latent models is the variational autoencoder (vae) @xcite . for each observed x, the model assumes the existence of a corresponding multidimensional latent vector z. since the log evidence log p(x) is intractable for most interesting problems, the training process for vaes opts instead to minimize the log evidence lower bound (elbo): @xcite )]-d kl (q(z|x)||p(z)) (1) q(z|x) is a tractable, assumed posterior commonly modeled with a parametrized encoder q φ (z|x), while p(x|z) is the likelihood parametrized with a decoder p θ (x|z) that optimizes against reconstruction loss. while effective in theory, a common empirical challenge vaes present during training is posterior collapse -a phenomenon where the decoder ignores the latent signal from z (and thus the originating input) during reconstruction. posterior collapse can be diagnosed by checking if d kl (q(z|x)||p(z)) tends to zero during training. after @xcite adopted vae for text, subsequent studies have been introduced with attempts to mitigate posterior collapse in vae language models (lms). however, the brittle training process of vae lms remains an unsolved problem. @xcite present a method to utilize deep transformer @xcite models as components of vae lms. transformer-based vaes tap into the state-of-the-art capabilities of transformers while retaining representational advantages of vae lms. the paper mitigates posterior collapse by massive pretraining and a cyclical annealing schedule @xcite . while the study presents a promising outlook for transformer vaes, the suggested method is not accessible to researchers who lack access to large, target-domain-specific corpora or the computing power for massive lm pretraining. therefore, a demand arises for a way to finetune an existing transformer model into a vae lm with limited resources. our research attempts to fill this gap in the literature, and makes the following contributions: the resulting model extends existing transformer architectures and can be initialized from pretrained non-latent model checkpoints. ## model architecture we extend the t5 architecture @xcite into a vae. we modify a popular pretrained t5 model @xcite ) that deviates minimally from the original transformer (figure 1 ). hidden states from all layers of t5's encoder q φ (z|x) are mean-or max-pooled into a vector h @xmath0, where h is the encoder's hidden dimension. assumed prior q(z)'s mean µ and log variance σ vectors of dimension l is obtained from h pooled : as in a standard vae, a stochastic latent vector z is sampled using the reparameterization trick to enable back-propagation through sampling: we pass z into the decoder p θ (x|z) as the only component of decoder's cross attention. our method injects z into every layer of the decoder as in previous literature @xcite , but deviates in two important ways: first, we pass z as the sole key and value of encoder-decoder cross attention, instead of self-attention; second, we project z into the correct dimension (l × a × s, where l is the decoder layer count, a is the number of attention heads, and s is the embedding dimension per head) with a feed-forward network, instead of taking a copy of z to inject to each decoder layer. ## experiments during preliminary experiments, posterior collapse was observed in all training schemes without encoder warmup training. the decoder learns to ignore the initially noisy input signal from the encoder. thus, we compose our finetuning method in two separate phases. phase 1 -encoder warmup: weight of kl loss is set to zero, making our model's objective function similar to that of an ae. different input denoising percentages, encoder pooling strategies, latent dimension sizes, and decoder freezing configurations are compared. phase 2 -full finetuning: kl loss is reinstated and full vae training are conducted. we compare different input denoising percentages, encoder pooling strategies, kl annealing schedules, and kl thresholds. we run our proposed two-phase finetuning training scheme on four standard vae lm benchmark datasets: ptb @xcite , snli (bowman et al., 2016), yahoo @xcite , and yelp @xcite . following @xcite and @xcite , we perform intrinsic evaluation of our proposed transformer vae architecture. we report perplexity (ppl), kl-divergence between model posterior and assumed posterior (kl), and negative elbo on the test set. to assess the quality of learned latent codes, we also report mutual information (mi) (hoffman and johnson, 2016) and the number of active units (au) @xcite . mi measures the dependence of latent codes to encoder input. au measures the covariance between encoder input and latent codes. experimental hyperparameters such as specific annealing schedules and training epochs per phase are detailed in the appendix. ## conclusions and future work this paper explores common methods in the literature for combatting posterior collapse, and the extent to which they help in teaching latent information to pretrained transformer models. comprehensive experiments show that commonly employed posterior collapse mitigation techniques provide meaningful benefits in transforming existing language models into latent-aware architectures. among the tested procedures, we find that @xcite 's two-step training, coupled with @xcite 's denoising through token deletion, was the most impactful in mitigating posterior collapse. however, language models obtained via only finetuning exhibit consistent trade-offs between their latent representation metrics (mi, au) and language model metrics (ppl). optimizing our model to be competitive with massively pretrained baselines in one of the two metrics results in the model falling behind in the other. we also find that increasing training epochs further improves the impact of tested techniques, a result consistent with previous literature on largescale text vae pretraining. from our experiments, we identify several questions to be answered by future research. the impact of homogenizing finetuning (as suggested in this paper) and original pretraining objectives on language model metrics has to be further explored. while the original t5 architecture was also pretrained with a self-supervised denoising scheme, the model employs mask tokens for denoising, contrary to simple token deletions suggested by this paper. our findings also highlight the need for an established heuristic to interpret the quality of latent representations learned by language models. the research community has yet to decide on the optimal value of kl-divergence between the assumed prior and the model posterior to target during text vae training. empirical guidelines to dictate even a vague threshold for the kl-divergence, below which we declare the occurrence of posterior collapse, will help both training and evaluation of latent-aware language models.