[Paper Review] AdaVAE: Exploring Adaptive GPT-2s in Variational Auto-Encoders for Language Modeling

Haoqin Tu, Zhongliang Yang+2 more|arXiv (Cornell University)|May 12, 2022

Topic Modeling4 citations

TL;DR

AdaVAE proposes a parameter-efficient VAE framework that unifies GPT-2 as both encoder and decoder using adaptive adapters, enabling effective latent space learning via Latent Attention. It achieves state-of-the-art language modeling and competitive performance in generation and classification tasks with only 14.66% of parameters activated during training.

ABSTRACT

Variational Auto-Encoder (VAE) has become the de-facto learning paradigm in achieving representation learning and generation for natural language at the same time. Nevertheless, existing VAE-based language models either employ elementary RNNs, which is not powerful to handle complex works in the multi-task situation, or fine-tunes two pre-trained language models (PLMs) for any downstream task, which is a huge drain on resources. In this paper, we propose the first VAE framework empowered with adaptive GPT-2s (AdaVAE). Different from existing systems, we unify both the encoder\&decoder of the VAE model using GPT-2s with adaptive parameter-efficient components, and further introduce Latent Attention operation to better construct latent space from transformer models. Experiments from multiple dimensions validate that AdaVAE is competent to effectively organize language in three related tasks (language modeling, representation modeling and guided text generation) even with less than $15\%$ activated parameters in training. Our code is available at \url{https://github.com/ImKeTT/AdaVAE}.

Motivation & Objective

• To address the high computational cost and inefficient parameter usage in existing large-scale VAEs that fine-tune full pre-trained language models (PLMs).
• To improve latent space construction in transformer-based VAEs by designing a method compatible with PLM representations, overcoming limitations of RNN-based latent fusion techniques.
• To unify both encoder and decoder components using the same GPT-2 architecture with shared word embeddings for model consistency and efficiency.
• To enable effective representation learning and generation in low-resource and multi-task settings with minimal trainable parameters.

Proposed method

• Introduces adaptive adapter modules inserted between feedforward layers and attention block outputs in GPT-2, enabling parameter-efficient fine-tuning.
• Proposes Latent Attention, a mechanism that incorporates encoder representations directly into the attention computation to better construct meaningful latent codes.
• Employs two parameter-efficient GPT-2 models as unified encoder and decoder, sharing the same word embedding space for consistency.
• Explores two latent knowledge infusion strategies: conditioning the decoder with latent vectors and injecting them into hidden states during generation.
• Uses a variational inference objective with reparameterization to optimize the evidence lower bound (ELBO), enabling end-to-end training.
• Applies adapter-based fine-tuning while freezing all original GPT-2 parameters, drastically reducing trainable parameters to 14.66% of total.

Experimental results

Research questions

• RQ1 Can a unified, parameter-efficient GPT-2 architecture serve as both encoder and decoder in a VAE without sacrificing performance?
• RQ2 How effective is Latent Attention in constructing a meaningful and disentangled latent space compared to standard latent vector injection?
• RQ3 What is the optimal balance between encoder and decoder depth in a VAE when using parameter-efficient PLMs?
• RQ4 Can AdaVAE achieve state-of-the-art performance in language modeling while using significantly fewer trainable parameters than existing 'big VAE' models?
• RQ5 How well does AdaVAE generalize across diverse NLP tasks, including low-resource classification and controllable text generation?

Key findings

• AdaVAE achieves a test perplexity (PPL) of 15.49 on the WikiText-2 dataset, outperforming previous SOTA models with significantly fewer trainable parameters.
• The model achieves a -ELBO of 125.56 with 14.66% of parameters activated, demonstrating strong optimization efficiency and representation learning.
• Latent space interpolation and analogy tasks show smooth semantic and syntactic transitions between generated sentences, confirming disentangled and meaningful latent representations.
• Ablation studies show that 8 encoder layers and 12 decoder layers yield the best performance, with diminishing returns beyond 8 encoder layers.
• The model maintains strong performance in low-resource text classification and controllable generation, confirming its robustness across diverse NLP tasks.
• The use of adapters enables effective fine-tuning with minimal computational overhead, making multi-task training feasible and efficient.