Finetuning Pretrained Transformers into Variational Autoencoders

Citation

• Authors: Seongmin Park and Jihwa Lee
• Year: 2021
• Venue: Insights from Negative Results in NLP 2021
• arXiv: arXiv:2108.02446
• URL: https://aclanthology.org/2021.insights-1.5/
• Citation count: 22 (Google Scholar, retrieved 2026-04-20)

One-Line Takeaway

This paper shows that a pretrained sequence-to-sequence Transformer can be turned into a text VAE by finetuning alone, using a careful two-phase schedule rather than Optimus-scale pretraining.

Technical Core

• Start from a pretrained T5 model and insert a latent bottleneck between encoder and decoder.
• Pool encoder hidden states into $h_{\mathrm{pooled}}$, then predict Gaussian posterior parameters $$\mu = h_{\mathrm{pooled}} W_\mu, \qquad \log \sigma = h_{\mathrm{pooled}} W_\sigma,$$ and sample $$z = \mu + \sigma \odot \epsilon, \qquad \epsilon \sim \mathcal{N}(0, I).$$
• The decoder cross-attends only to the latent code rather than to ordinary encoder states: the latent is projected into decoder-layer key/value memories, $$(K_{\mathrm{ca}}, V_{\mathrm{ca}}) = (z W_{\mathrm{proj}}, z W_{\mathrm{proj}}).$$
• The training objective is the usual sequence ELBO: $$\mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \mathrm{KL}(q_\phi(z \mid x)\,\|\,p(z)).$$

Method

• Phase 1: encoder warmup. Set the KL weight to zero and train the model like an autoencoder so the latent path starts carrying useful information before the decoder learns to ignore it.
• Phase 2: full VAE finetuning. Reintroduce the KL term and continue training with collapse-mitigation heuristics.
• The paper studies several anti-collapse interventions: input denoising, KL annealing, KL thresholding, pooling choices, latent dimensionality, and decoder freezing.
• A main practical finding is that encoder warmup plus denoising is more reliable than relying on KL schedules alone.

Empirical Evaluation

• The model is evaluated on PTB, SNLI, Yahoo, and Yelp, using perplexity, KL, mutual information, and active units.
• The two-phase recipe can maintain non-collapsed latents with only finetuning, which is the main result.
• The tradeoff remains sharp: settings that improve latent usage often worsen language-model metrics, and vice versa.
• The paper explicitly reports that cyclical annealing alone was not sufficient in their low-resource setting to prevent collapse.

Assumptions And Limits

This is a practical recipe paper, not a claim that finetuning solves text-Transformer VAEs cleanly. It still treats posterior collapse as the central difficulty, and the best models do not simultaneously dominate pretrained heavy baselines on both generation quality and latent-information metrics. The setting also remains unconditional reconstruction rather than a disentangled conditional factorization.

Relevance To This Project

This paper is unusually relevant because it is about the exact engineering problem this repo faces in one form or another: how to introduce a meaningful latent variable into an already pretrained autoregressive Transformer without retraining the whole model from scratch.

Its main lesson for this project is negative but useful: even when the architecture is straightforward, the adaptation landscape strongly prefers collapse, so training schedule and latent-interface choices matter as much as the nominal ELBO.

Comparison Notes

Optimus makes large-scale pretrained language VAEs work by spending much more pretraining compute. Park and Lee ask the cheaper question: what can be salvaged by careful finetuning alone? Relative to the repo's methodology, this paper is still solving a simpler problem because it learns a VAE language model rather than refactoring an already-competent conditional generator into a strategy-aware factorization.

Related Notes

• OPTIMUS: Organizing Sentences via Pre-trained Modeling of a Latent Space
• Lagging Inference Networks and Posterior Collapse in Variational Autoencoders
• CVAE objective and loss terms
• Challenges of Learning a Latent Variable Factorization of Pre-trained Language Models