AdaVAE: Exploring Adaptive GPT-2s in Variational Auto-Encoders for Language Modeling

VAE is famous for its continuous latent space, its extensions in the language domain have inspired new applications by exploiting many interesting properties of the model’s latent space. As a preliminary, the evidence lower bound (ELBO) of a VAE is:

where $\boldsymbol{X}$ is the texts to be modeled, and $\boldsymbol{z}$ is the latent variable sampled from latent space $\mathcal{Z}$.

In real world scenario, some defects limit the empirical performance of VAEs for language modeling including KL collapse issue Bowman et al. (2015), latent vacancy issue Xu et al. (2020) and token-latent inconsistency problem Shen et al. (2020). Many related theories and solutions to these drawbacks were proposed including optimizing decoder architectures Semeniuta et al. (2017); Li et al. (2020), inventing auxiliary objectives Xiao et al. (2018); Fang et al. (2019); Dai et al. (2020), novel encoder training schedule Bowman et al. (2015); Fu et al. (2019), incorporating flexible latent code posterior Wang et al. (2019), etc. These methods generally share the same goal: to impair the ability of a powerful decoder and strengthen the expression of latent space by reinforcing encoder ability.

With more powerful transformer-based language models arising, researchers start to march on combining language VAEs with transformers Vaswani et al. (2017). For example, transformers are recently considered in VAEs for classification Gururangan et al. (2019) and storytelling Wang and Wan (2019). Pre-training VAEs has been recently considered in conditional text generation to amortize the training of decoders and to allow easy adaptation in new generation tasks Duan et al. (2019).

Nevertheless, all aforementioned efforts utilize simple RNN Hopfield (1982) and shallow Transformer architectures thus they equip with limited model capacities. Pre-trained language models (PLMs) are recently introduced in VAEs to further boost both the generative and understanding abilities of VAEs. Li et al. (2020) proposed the first “big VAE” model at the same scale of BERT and GPT-2, their model connects two PLMs with different embedding spaces in a latent space, which demonstrate the efficacy for reducing related issue effectiveness in multiple NLP tasks. Park and Lee (2021) proposed to incorporate T5 Raffel et al. (2019) into VAEs. Fang et al. (2021) utilized GPT-2 Radford et al. (2019) in the VAE paradigm for controllable long story generation task. Fang et al. (2022) employed a discrete latent prior and additional noises to boost the control ability of a text VAE model. However, these methods essentially fine-tune two PLMs during model training, requiring a great amount of resources compared to RNN-based models or a single PLM.

For large pre-trained language models, massive training samples and huge parameter volumes help them gain unparalleled modeling ability. The conventional method to transfer a PLM to a specific data domain is fine-tuning, which mobilizes all parameters from the model. This can be intolerant in both computing and storage resources as the task load grows. Taming PLMs with high-efficiency w.r.t. distinct missions becomes one of the top trends in NLP. Various lightweight alternatives in PLMs were proposed. Houlsby et al. (2019) first came up with the idea to additionally add trainable adapter components with down sample and up sample layers in transformer blocks, which proved to achieve comparable results in NLP tasks with less than $10\%$ trainable parameters. Following this line, Pfeiffer et al. (2020) improved the method by changing the adapter components to different positions of transformer blocks. Li and Liang (2021) proposed to add trainable prefix to attention head in the model, while Hu et al. (2021) created a shortcut in the attention domain of transformers which consists of trainable down and up sampling layers. Recently, He et al. (2021) concluded all these methods into a unified paradigm, which can be formalized as: $h\leftarrow\lambda_{1}h+\lambda_{2}\Delta a$, where $h$ is the output of the attention layer or feedforward layer in one transformer block. The parameter $\lambda$ is varied according to different types of components, e.g., for Prefix tuning Li and Liang (2021), $\lambda_{1}=1-\lambda_{2}$ with $\lambda_{2}$ to be a pre-assigned scalar. As for Adapter tuning Houlsby et al. (2019), it does not employ a scalar, which means $\lambda_{1}=\lambda_{2}=1$. Finally, the $\Delta a$ decides the receiving information for the current component from the previous layers. When $\Delta a$ is transformed from states ahead of the current attention block, the information is said to be “parallelly” shared, otherwise, it is “sequantially” shared with the model. These methods generally share a similar approach of introducing additional trainable parameters to PLMs instead of activating the original pre-trained transformers.

The encoder of a VAE should extract features from given contents to produce meaningful latent space, while the decoder ought to generate fluent sentences with given latent representations. In order to obtain a unified word embedding space in AdaVAE, we construct both encoder and decoder using GPT-2, which leaves us no worry about connecting two word embedding spaces from different models as in Li et al. (2020). To make GPT-2 a qualified encoder, we take advice from mighty extractors such as BERT, one of its architectural advantages lies in the unmasked/bi-directional transformer layer. Thus we remove the causal mask in GPT-2 transformer layers to make it an encoder of AdaVAE with full vision of input contexts, this mask-free operation is widely used in the encoders of existing PLMs and VAEs Raffel et al. (2019); Lewis et al. (2019); Fang et al. (2021). As for the decoder, we employ GPT-2, a powerful generative transformer model by design.

The paradigm of fine-tuning two separate PLMs in large-scale VAEs requires a lot more computing resources than a single PLM, and the storage requirements will become too heavy to tolerate as the task loads increase. To avoid such dilemma, we propose and explore different parameter-efficient components including different types or insertion methods into encoder and decoder layers, which means only additional minimum parameters need to be activated for every task. Specifically, we propose a parallel adapter placed after the feedforward layers (see Section 2.2) of each attention block in both encoder&decoder for AdaVAE. And we further compare our approach with different adapters and Prefix tuning method Li and Liang (2021) for ablation study in Section 5.1. Overall, these two settings make AdaVAE more elegant to be constructed and more efficient to be trained compared with existing “big VAEs”.

How to form the latent space from the encoder and utilize it in the decoder to narrow the gap between discrete sentences to the continuous latent embedding is a key problem. Li et al. (2020) and Park and Lee (2021) employed a pooled feature from the last encoder layer and pass it to a linear transformation to obtain latent space, which may be not sufficient to leverage the knowledge learned from transformer layer. Fang et al. (2021) used the last state from the encoder as both the key and value vectors to conduct averaged attention by matrix multiplication. Their model learns both prior and posterior of the latent space from the same type of input (i.e., the same textual content), and shares most of the learning parameters in this process. We contend that 1. vectors from different domains in the attention operation (i.e., key, value, query) should be distinct to carry specific knowledge. 2. To avoid potential KL collapse issue, one ought to produce latent posterior and prior from different types of input sources Lucas et al. (2019).

We thus propose the improved Latent Attention operation to generate meaningful latent space in AdaVAE: to get latent vector $\boldsymbol{v_{z}}$, we adopt the last hidden state $\boldsymbol{v_{x}}$ from the encoder and assign:

where $\mathbf{E}$ is a identity matrix with the same size of $\boldsymbol{v_{x}}$, $f(\cdot)$ is a linear transformation for mapping $\boldsymbol{v_{x}}$ to the key vector space, and finally the $\boldsymbol{v_{z}}$ is from the attention operation between derived $\mathbf{Q_{z},K_{z},V_{z}}$. Then the latent vector $\boldsymbol{v_{z}}$ is taken to reparameterize the mean ($\boldsymbol{\mu}$) and variance ($\boldsymbol{\sigma}$) of $\mathcal{Z}$:

where $\boldsymbol{z}$ is a latent vector sampled from space $\mathcal{Z}$, $f_{\mu}$ and $f_{\sigma}$ are two linear transformations, $\odot$ is the element-wise multiplication. Note that, we only use it to model the posterior of latent space and leave its prior to be a normalized Gaussian Bowman et al. (2015). This setting takes full advantage of learned information from the encoder and reduces the possibility of KL vanishing problem that may occur in the previous work Fang et al. (2021).

The way to infuse learned latent knowledge from the latent space into the generative process determines the effectiveness of the decoder employing learned representations. Inspired by existing methods, we investigate two different frames to add latent variables into decoder layers. For a latent variable $\boldsymbol{z}\in\mathbb{R}^{d}$ drawn from $\mathcal{Z}$, we investigate two different fusion methods:

• •

  Add to Memory (AtM) Li et al. (2020) projects $\boldsymbol{z}$ to both attention key and value spaces by a unified linear layer $f(\cdot)$, and concatenate them with key and value vector in each attention layer:

  		$\displaystyle\boldsymbol{k_{z}}=\boldsymbol{v_{z}}=f(\boldsymbol{z})\in\mathbb{R}^{d\times l},$		(4)
  		$\displaystyle\boldsymbol{k=[k;k_{z}]}\in\mathbb{R}^{l},\boldsymbol{v=[v;v_{z}]}\in\mathbb{R}^{l},$

  where $l,\boldsymbol{k},\boldsymbol{v}$ are the size of both key and value spaces, Key vector and value vector severally.

• •

  Pseudo Self-Attention (PSA) Fang et al. (2021) shares a similar idea with AtM, but it uses separate convolutional transformations with $\boldsymbol{z}$ as input to make sure that $k_{z}\not=v_{z}$, then PSA concatenates them with respect vectors just like AtM to conduct the pseudo self-attention in decoder layers.

The training loss of our model is based on the plain VAE. To maximize the potential of our model, we further incorporate free bit threshold and KL annealing techniques during model training.

For model architecture and initialization, the encoder and decoder were separately 8 layers and 12 layers GPT-2 transformers initialized from pre-trained weights in Hugginface.¹¹1https://huggingface.co/gpt2 The parameter efficient components were chosen from Adapter Houlsby et al. (2019) and Prefix Li and Liang (2021), the hidden size of the Adapter was chosen to be 128 or 512 depending on different tasks, while the hidden size of the Prefix was 30 for the ablation study in language modeling task. The hidden size of latent space $\mathcal{Z}$ was set to 32 for language modeling and 768 for classification and generation.

For training details, we first activated the parameter-efficient components in the encoder and parameters in latent spaces for the first $1/6$ training steps and then added parameter-efficient components in the decoder for the rest of the training time, this setting is helpful for training a VAE model Li et al. (2019). Similarly, we also employed linear warming-up procedure for learning rate Popel and Bojar (2018) to make it increases linearly from 0 to $5\times 10^{-5}$ in the first $1/6$ training iterations. We train the model with the batch size of around 64 on one TITAN X GPU with 12G memory.

The detailed dataset statistics are in TABLE 1. We conduct three generation tasks and an understanding task span from six datasets. For language modeling task, we use YELP, YAHOO and PTB from Optimus Li et al. (2020) directly. For controllable generation, we take a shorter version of YELP dataset (denoted as $\texttt{YELP}_{S}$), which is originally designed for the style transfer learning with labels Shen et al. (2020). As for style transfer task, we use the same $\texttt{YELP}_{S}$ dataset as in the mentioned generation task. For language classification (understanding) task, we apply $\texttt{YELP}_{S}$ data as well as additionally introduced WNLI and SST-2 dataset from GLUE benchmark Wang et al. (2018).

For the evaluation of language modeling ability, we took Perplexity (PPL), the negative evidence lower bound (ELBO), mutual information between input sentence and $\mathcal{Z}$ (MI) and activated units in the latent space (AU) as measurements. All metrics were implemented exactly following the public codebase²²2https://github.com/ChunyuanLI/Optimus for fair comparisons. While PPL and ELBO measure the fluency of generated sentences, MI and AU indicate the representation learning capacity of a latent variable model.

We first explore the effects of different types of parameter-efficient components as well as latent space generation&infusion methods for the proposed model. We explored 7 types of VAE models in Table 2, their visual illustrations are in Figure 2:

1. 1.

   AdaVAE: Uses the proposed parallel adapter for feedforward layer in transformers, Latent Attention for latent space construction, PSA for representation infusion in the decoder.

2. 2.

   w/ Fine-tuning: Removes adapters and fine-tunes the original model.

3. 3.

   w/ LG: Uses the pooled feature from encoder and a linear layer to form $\mathcal{Z}$.

4. 4.

   w/ LG; +AtM: Uses LG and both the PSA&AtM methods together for latent infusion.

5. 5.

   +Prefix: Adds Prefix components to attention blocks.

6. 6.

   w/ parallel_attn: Replaces the original adapters with parallel adapters for attention outputs.

7. 7.

   w/ sequential_attn: Replaces the original adapters with sequential adapters for attention outputs.

Table 2 shows the proposed model with different adding components or training schemes, AdaVAE achieves a good trade-off between language modeling and representation learning ability among presented models. Focusing on specific model structures, (1) From the first row in Table 2, fine-tuning the model does not bring significant improvement with significantly larger training parameters ($100\%$ vs. $14.66\%$), and even perform worse on PPL (YELP and YAHOO) and one MI metric (YELP) compared with the proposed parameter-efficient training method. This demonstrates our design of adaptive GPT-2 encoder&decoder in AdaVAE is efficient, which leads the way to resolve the excessive resource consumption problem faced by existing “big VAEs”. (2) At the second row, we could find that replacing the proposed Latent Attention with LG from Optimus makes the overall model performance worse. This demonstrates the unfitness to employ simple linear transformation on learned transformer features for space $\mathcal{Z}$. Adding AtM to infuse latent representation with PSA generally boosts model’s representation learning with higher MI scores. This is ascribed to the cumulative use of learned latent knowledge with added trainable parameters ($5\%$ more model parameters) with AtM and PSA. (3) At the last row in Table 2, we focus on the gain of model performance brought by different parameter-efficient components. Adding the Prefix component will introduce more training parameters, but with little or even negative performance gain in the learning process (lower MI and higher PPL on YAHOO). Besides, replacing the original parallel adapters for the feedforward layers with either parallel or sequential adapters for attention output lags behind the proposed adapter in the model. These help us to understand the adaptive tuning method in our proposed model is practical to use. Overall, AdaVAE with proposed adaptive GPT-2s and the Latent Attention shows more consistent performance in the capacity trade-off among ablation components.

We further compare our model with different baselines from conventional VAEs with state-of-the-art PLM-based VAEs. Our baseline models are listed as follows:

• •

  iVAE Fang et al. (2019): a VAE model considers implicit posterior representation instead of the explicit form.

• •

  GPT-2 Radford et al. (2019): a large-scale LM pre-trained on large scale real-world dataset and fine-tuned on each dataset for 1 epoch.

• •

  T5 VAE Park and Lee (2021): a PLM-based VAE model that fine-tunes the encoder and decoder of T5 Raffel et al. (2019) model into VAE setting.

• •

  Optimus Li et al. (2020): the first PLM-based VAE model that connects pre-trained BERT and GPT-2 by organizing a continuous latent space.

• •

  DPrior Fang et al. (2022): a PLM-based VAE model uses discrete latent prior and additional noise in the latent space to strengthen control ability under Optimus.

Table 3 shows the comparison between the proposed model and some SOTA VAEs. We draw the following conclusions based on it: firstly, to find the best value for every metric on each dataset, the proposed model holds the lowest PPL and -ELBO values on two datasets among all baselines, demonstrating the advantages in language modeling of our proposed approach. Though the performance of proposed model surges ahead of most baselines on all metrics, there still remain a small margin between our model and Optimus on MI scores (less than 0.15 when $\lambda=0.5$ on PTB, YELP). We argue that it is essential to consider both LM and representation-related statistical results simultaneously. Optimus achieves the highest MI score with much worse PPL or -ELBO results compared with AdaVAE, demonstrating that Optimus actually sacrifices large amount of LM ability to gain some improvements on MI. While our model stays a steady trade-off in this circumstance (generally better PPL and -ELBO and slightly lower MI). Also note that only AdaVAE activates partial model parameters during training, which is another huge advantage compared to strong fine-tuned baselines (e.g., DPrior, Optimus).

Then we focus on the effect of $\lambda$. As the free bits threshold $\lambda$ increases, MI value generally yields a better performance in both Optimus and AdaVAE. This is because a larger KL threshold brings up the amount of knowledge that should be learned in the latent space. While the trend of PPL value is monotonous with $\lambda$ in Optimus, there is a rebound in AdaVAE with the optimal PPL with $\lambda=0.50$ on two corpus. With a fairly high MI score, we believe it is the suitable $\lambda$ value for a good trade-off between model’s language modeling and representation learning ability.

We further conduct controllable generation on $\texttt{YELP}_{S}$ dataset collected by Shen et al. (2020), which contains 444K training sentences, and we use separated datasets of 10k sentences for validation/testing respectively as in Optimus. The goal is to generate text reviews given the positive/negative sentiment from the $\texttt{YELP}_{S}$ dataset. For evaluation metrics, we employed the following metrics for automatic evaluation: Accuracy (Acc.) is measured by a pre-trained classifier on the $\texttt{YELP}_{S}$ dataset, indicating the controllability of the model. BLEU (B.) for the quality evaluation of generated sentences. G-score (G-S.) computes the geometric mean of Accuracy and BLEU, measuring the comprehensive quality of both content and style. Self-BLEU (S-B.) measures the diversity of the generated sentences. And for BLEU-F1 (B-F1), we have: $\text{BLEU-F1}=\frac{2\times\text{BLEU}\times(1-\text{Self-BLEU})}{\text{BLEU}+(1-\text{Self-BLEU})}$, which evaluates the overall metric involving text quality and diversity simultaneously.

The baselines are described as follows:

• •

  Control-Gen Hu et al. (2017): employs discriminator on latent space to control the generation process.

• •

  ARAE Zhao et al. (2018): learns an auto-encoder first, and then train a GAN to produce the latent vectors.

• •

  NN-Outlines Subramanian et al. (2018): uses a general purpose encoder for text generation, and a second stage training with labeled examples for the controllable generation task.

• •

  Optimus Li et al. (2020): finetunes the model for text general generation, then conducts second training stage with a conditional GAN Mirza and Osindero (2014) on latent space based on given labels.

To verify our model’s capacity on the task of controllable generation, we took a similar experimental setup as ARAE and Optimus: we first trained our AdaVAE on $\texttt{YELP}_{S}$ dataset without label information, then we employed a conditional GAN on the learned latent representation of our model to produce latent vector $z_{y}$ based on label $y$. And finally, we explored two types of training settings for the decoder to produce controllable texts, i.e., AdaVAE only activates parameter-efficient adapters in the decoder during generation, which equips $12.03\%$ trainable parameters. $\textsc{AdaVAE}_{dec}$ activates all parameters in the decoder during generation, which equips $46.79\%$ trainable parameters compared with the original LM.

The results are shown in Table 4. We can draw the following conclusion: (1) $\textsc{AdaVAE}_{dec}$ with its decoder fully activated is much more competent than AdaVAE in all metrics, we believe this is because the guided text generation process depends more on decoder updating to produce sentences with different text labels. However, neither $\textsc{AdaVAE}_{dec}$ nor AdaVAE finetunes all parameters in the encoder, indicating our adaptive GPT-2 encoder generates persist and robust $\mathcal{Z}$ space for further generation. (2) Compared with baseline models, $\textsc{AdaVAE}_{dec}$ can achieve the best performance in both G-score and BLEU-F1, which demonstrates that the proposed model is capable of generating controllable contexts that preserve human-like text features (high BLEU score). (3) However, there is still a big margin between our $\textsc{AdaVAE}_{dec}$ and baselines on Self-BLEU for text diversity. This is because the motivation to fully activate the decoder in the first place is to produce higher quality contents, we can access the model performance via the overall metric BLEU-F1 that takes this trade-off into consideration.

We also present generated sentences with given sentiment in Table 5. For the negative label, there are words like “terrible” and “not impressed” representing negative sentiment. While for the positive label, words including “great”, “perfect” and “happy” indicate the positive sentiment of generated sentences.

To validate that the proposed model is qualified as a textual feature extractor even with minimum trainable parameters. We conducted full-sized as well as low resource classification task. The percentage of activated parameter is #params.. In Table 6, FB means only the linear layer of a classifier was activated in training. All results were averaged on 5 runs with different random seeds.

We view the model performances by the size of training corpus, (1) when the number of labeled training sample is very low (full WNLI / 10 or 100 training samples from $\texttt{YELP}_{S}$), AdaVAE can achieve better classification accuracy than fine-tuned AdaVAE, and is even superior than both fine-tuned and parameter-efficient BERT or Optimus. (2) For middle sized training data (full / 1,000$\sim$10,000 training samples from $\texttt{YELP}_{S}$), AdaVAE shows competitive performance compared with BERT and Optimus and generally better than fine-tuned AdaVAE. (3) For large-scale dataset (SST-2), the performance of AdaVAE is inferior than BERT and Optimus by around $6\%$.

Though the performance of AdaVAE with adaptive components falls behind baselines on the large-scale dataset, it only activates very few parameters in the encoder to fulfill this task. These statistics demonstrate that AdaVAE with few activated parameters is competent to extract textual features like specialized PLM such as BERT or Optimus. We ascribe it to the structural modification of unmasked GPT-2 transformers as the encoder of AdaVAE as well as Latent Attention’s powerful knowledge learning ability from transformers.

From the concluded results, as the increase of training data size, models with adaptive parameter-efficient components gradually loses their advantage on small-sized dataset compared to the fine-tuning models. This phenomenon is also reported by a concurrent work Chen et al. (2022), which verifies our observations. Since we did not change the adapter structure significantly, the training time of our model is $60\%$ to $100\%$ compared with fine-tuning Ding et al. (2022).

Further, we visualize the distribution of $\mathcal{Z}$ using T-SNE Van der Maaten and Hinton (2008) on $\texttt{YELP}_{S}$ in Figure 4. Firstly, the representations from AdaVAE with adaptive settings (Figure 4 (a), (b)) can be better separated compared with fine-tuned one. Secondly, latent distribution with a bigger adapter size (Figure 4 (b)) yields a more compact clustering in the figure but also means more activated training parameters. This demonstrates a trade-off between training parameters and representation learning ability of AdaVAE.

We also conducted latent analogy and interpolation task. For a given sentence triplet $S_{A},S_{B},S_{C}$, the analogy task generates a sentence from source $S_{A}$ to target $S_{B}$ and with a similar style of $S_{C}$ as examples are shown in Table 7. We can tell from the table, that generated sentences absorb all given sentence features. For example, three output texts in the analogy task talk about food, which are relevant to Target $S_{B}$. When Input $S_{C}$ turns to negative, the Output $S_{D}$ also steers to the negative emotion. As for interpolation task, given a sentence pair $S_{A},S_{B}$, latent interpolation generates texts with styles transfer from $S_{A}$ to $S_{B}$ by latent space traversal as examples are shown in Table 8. From the interpolated example, generated texts mix the semantics and syntax of the two initial sentences (when $\tau$ is 0.0 or 1.0) and smoothly change from one to the other.

We conducted experiments to select the best encoder&decoder settings w.r.t. the number of layers in the transformer models. As shown in Table 9, we varied both encoder layer number (denote as Enc.) and decoder layer number (denote as Dec.) from $[6,8,10,12]$. And we find that: (1) models with fewer decoder layers show worse language modeling capacities, e.g., the model with 8 encoder layers and 12 decoder layers reaches the lowest PPL value as well as -ELBO value among other decoder settings. (2) AdaVAE with more encoder layers is not necessary, as the model with the best performance equips 8 encoder layers.