PaperlayerView exact revision on arXiv

Source-linked AI summary

Cyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing

Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, Lawrence Carin

arXiv:1903.10145v3cs.LGcs.AIcs.CLcs.CVstat.ML

TL;DR

Text VAEs with auto-regressive decoders can suffer KL vanishing because decoder training begins without good latent codes. The paper analyzes this mechanism and introduces cyclical β annealing, which repeatedly increases β and uses previous cycles as warm restarts. The approach improves performance across text modeling, dialog response generation, and unsupervised language pre-training with negligible extra computational cost.

  • Problem

    Auto-regressive text VAE decoders can ignore latent variables, causing KL vanishing, while existing scheduling methods lack a unified theoretical interpretation.

  • Method

    The paper analyzes the decoder's two learning paths and proposes cyclical annealing, repeatedly increasing β to reuse informative latent codes from previous cycles as warm restarts.

  • Results

    Cyclical annealing improves performance across text modeling, dialog response generation, and unsupervised language pre-training with negligible extra computational cost.

  • Takeaways & Limitations

    Progressively learned latent codes can be leveraged across β cycles to train the decoder more effectively.

  • Takeaways & Limitations

    The best way to schedule β remains largely unexplored, and monotonic annealing lacks a proper justification.

Abstract

from arXiv · show

Variational autoencoders (VAEs) with an auto-regressive decoder have been applied for many natural language processing (NLP) tasks. The VAE objective consists of two terms, (i) reconstruction and (ii) KL regularization, balanced by a weighting hyper-parameter β. One notorious training difficulty is that the KL term tends to vanish. In this paper we study scheduling schemes for β, and show that KL vanishing is caused by the lack of good latent codes in training the decoder at the beginning of optimization. To remedy this, we propose a cyclical annealing schedule, which repeats the process of increasing βmultiple times. This new procedure allows the progressive learning of more meaningful latent codes, by leveraging the informative representations of previous cycles as warm re-starts. The effectiveness of cyclical annealing is validated on a broad range of NLP tasks, including language modeling, dialog response generation and unsupervised language pre-training.

1 Introduction

VAEs support rich latent representations for NLP, but auto-regressive decoders can ignore those representations, causing KL vanishing. The paper explains this failure and proposes cyclical annealing to improve VAE training across several NLP tasks.

  • VAEs provide distribution-based latent representations that can model holistic sentence properties such as style, topic, and high-level linguistic and semantic features.
  • Auto-regressive decoders introduce KL vanishing by ignoring the latent variable during VAE training.The issue is associated with the sequential nature of text and the use of recurrent decoders.
  • Existing methods, including monotonic KL annealing, alleviate KL vanishing, but lack a unified theoretical interpretation.
  • KL vanishing occurs because the latent-conditioned path lacks good latent codes initially, while an expressive decoder can instead rely only on previously generated words.
  • Cyclical annealing repeats β annealing so informative latent codes from previous cycles serve as warm restarts for decoder training.The paper reports improved performance with negligible extra computational cost across text modeling, dialog response generation, and unsupervised language pre-training.

2 Preliminaries

Text VAEs combine an inference model and an auto-regressive decoder, optimizing reconstruction and KL regularization through a β-weighted objective. Fixed β can produce KL vanishing, while monotonic annealing offers a common but insufficiently justified remedy.

  • An auto-regressive language model generates each token conditioned on previously generated tokens.The notation x<t denotes all tokens before position t.
  • A text VAE uses an encoder qφ(z|x) to approximate the posterior and a decoder pθ(x|z) to generate text from latent vector z.
  • The VAE objective weights reconstruction error and KL regularization through Lβ = LE + βLR.LE is reconstruction error or negative log-likelihood, while LR is the KL regularizer.
  • With β = 1, the objective recovers the VAE; with β = 0 and delta-distributed qφ(z|x), it recovers an autoencoder.The autoencoder lacks prior regularization and can be prone to overfitting or plain dialog responses.
  • Under a constant β = 1 schedule, the KL term can vanish, making the encoder match the Gaussian prior and the decoder ignore z.
  • Monotonic annealing increases β from 0 to 1 and often improves the converged lower bound, but its justification and the best β schedule remain unresolved.

3 Cyclical Annealing Schedule

The paper explains KL vanishing through competing decoder paths and proposes cyclical annealing to repeatedly rebuild useful latent representations before enforcing the full VAE objective.

  • Motivation: The auto-regressive decoder can bypass latent-variable conditioning through previously generated words, causing KL vanishing when latent codes are initially poor.Path A uses latent codes, while Path B conditions on previous words and can generate sequences without them.
  • Cyclical Annealing Schedule: Cyclical annealing repeatedly lowers β, increases it to 1, and uses latent codes learned under the full objective as warm restarts.Each cycle begins with β = 0, anneals toward β = 1, then fixes β = 1 for the remaining iterations.
  • Cyclical Annealing Schedule: The schedule divides training into M cycles, with M = 4 illustrated and R = 0.5 used for within-cycle increases.Existing constant and monotonic schedules are special cases of the cyclical schedule under particular M and R settings.
  • Cyclical Annealing Schedule: Within each cycle, annealing encourages representative latent codes, while fixing β = 1 regularizes their distribution toward the prior and optimizes the full VAE objective.The schedule uses R as the proportion of a cycle devoted to increasing β.
  • Impact of β: The objective analysis decomposes KL regularization into mutual information and marginal KL, clarifying how β balances informative latent codes against prior matching.Higher mutual information increases correlation between latent variables and data and reduces the degree of KL vanishing.

4 Visualization of Latent Space

On an illustrative sequence dataset, cyclical annealing progressively improves reconstruction and produces increasingly structured latent clusters across cycles.

  • Experimental Setup: The illustrative experiment uses 10 one-hot sequences, a 2-dimensional latent space, two-layer 64-unit LSTMs, and 40K training iterations.The schedules are compared by visualizing learning dynamics and latent-space structure.
  • Learning Curves: The three schedules have similar ELBO, reconstruction-error, and KL learning curves, but cyclical annealing yields lower reconstruction error and higher KL divergence.Its improvements become progressively larger than those of the preceding cycle, with periodic patterns across cycles.
  • Latent Space: Constant scheduling keeps latent codes heavily mixed, while monotonic scheduling forms 10 cluttered Gaussian clusters that remain cluttered.The cyclical schedule resembles monotonic scheduling during its first cycle.
  • Latent Space: From the second cycle onward, cyclical annealing forms more separated clusters, and later cycles capture clearer latent-space patterns than earlier ones.β < 1 can create holes between clusters, which are alleviated when β returns to 1.

5 Related Work

Prior work addresses KL vanishing by weakening the decoder’s word-history path or strengthening its latent-code path. This paper instead applies cyclical scheduling to β, using prior cyclical optimization ideas but reopening the latent path for a different purpose.

  • Techniques for KL vanishing: Existing methods either weaken Path B so the decoder relies more on latent codes, or improve Path A so latent codes are easier to use.Word dropping and dilated CNNs target Path B; skip connections, auxiliary losses, and semi-amortized training target Path A.
  • Techniques for KL vanishing: Semi-amortized training shares the goal of obtaining better latent codes but incurs high stochastic-inference cost, whereas scheduling adds no extra compute overhead.The paper describes scheduling methods as complementary and reports that cyclical scheduling can further improve these techniques.
  • β-VAE: β-VAE research primarily uses β > 1 to learn disentangled latent representations, unlike this paper’s focus on scheduling β during VAE training.The supplied passage characterizes β-VAE as a related but distinct line of work.
  • Cyclical schedules: Earlier cyclical schedules primarily vary learning rates; this paper is the first described here to apply cyclical scheduling to β in VAEs.Its motivation is to reopen Path A and train the decoder with high-quality z, rather than to use a learning-rate warm restart.

6 Experiments

Experiments evaluate cyclical annealing on language modeling, dialog response generation, and unsupervised language pre-training. Across these settings, the schedule improves latent-code use or task performance while adding almost no cost, and its gains are attributable to cyclical β rather than cyclical learning rate.

  • Language modeling: The PTB language-modeling experiments compare cyclical annealing with monotonic annealing and semi-amortized VAE training under matched experimental settings.The monotonic baseline linearly anneals β from 0 to 1 over 10 epochs; SA-VAE uses 10 SVI steps.
  • Language modeling: Under similar ELBO and PPL, cyclical training yields lower reconstruction errors and higher KL values than monotonic training.The learning curves also show progressively better performance after each cycle.
  • Language modeling: Cyclical scheduling outperforms monotonic scheduling for both standard VAE and SA-VAE training, while adding almost zero cost.SA-VAE reduces KL vanishing but takes 472s per epoch versus 30s for standard VAE training.
  • Language modeling: With a monotonic maximum β of 0.5, reconstruction and KL terms improve individually but ELBO becomes substantially worse than with β = 1; cyclical scheduling still improves performance.The experiment is used to distinguish cyclical scheduling from simply using a lower β.
  • Dialog response generation: For dialog generation, cyclical CVAE responses are highly diverse and cover multiple plausible dialog acts, whereas monotonic responses repeat plain answers.The paper evaluates CVAE and a VAE augmented with a bag-of-words loss.
  • Dialog response generation: Cyclical scheduling reduces KL vanishing in CVAE and produces the highest BLEU scores, while the auxiliary BoW loss is unnecessary for higher-fidelity responses at lower cost.BoW can nevertheless be further improved when integrated with cyclical scheduling.
  • Unsupervised language pre-training: On Yelp unsupervised language pre-training, cyclical scheduling consistently produces the highest fine-tuning accuracy and clearer latent-code clusters.The latent codes z are used as text features for downstream classification.
  • Ablation study: An ablation finds that cyclical β, rather than cyclical learning rate η, contributes to improved performance.The learning-rate cycle is introduced to equalize optimization-trajectory length across β cycles, but the isolated β cycle provides the improvement.

7 Conclusions

The paper explains KL vanishing through a two-path decoder analysis and attributes it to poor latent codes early in decoder training. It proposes repeatedly reopening the latent path so prior-cycle codes serve as warm restarts, with effectiveness shown across three NLP tasks.

  • Contributions: A two-path interpretation identifies poor latent codes at the beginning of decoder training as the source of KL vanishing.The interpretation is used to explain existing β-scheduling schemes and motivate cyclical annealing.
  • Contributions: Cyclical scheduling progressively improves performance by reopening the path at β = 0 and using good latent codes from previous cycles as warm restarts.This mechanism is presented as the central explanation for the schedule’s progressive gains.
  • Contributions: The approach is demonstrated on three NLP tasks and is reported as superior to or complementary to other techniques.
  • Schedule design: The cyclical schedule can be instantiated from monotonic schedules such as linear, sigmoid, or cosine by repeating annealing and fixing stages.Figure 7 compares the monotonic schedules with their corresponding cyclical versions.

B Proofs on the β and MI

The appendix develops probability-based identities for the β-weighted VAE objective and related latent-variable models. It decomposes the KL term into mutual-information and marginal-KL components and gives supporting definitions for the proof setup.

  • Proof setup: The appendix introduces the β-weighted training objective and rewrites its terms separately for the subsequent proofs.
  • Proof setup: The proof uses encoder and decoder probability measures p and q and expresses reconstruction loss for z in negative-log-likelihood form.
  • Probability identities: A lemma and corollary establish probability identities involving conditional entropy and mutual information for random variables x and z.
  • KL decomposition: The KL term is decomposed into mutual information and marginal KL components, labeled F1 and F2.
  • Extensions: The appendix also formulates ELBO-related objectives for dialog variables and semi-supervised learning, with α controlling classification-loss regularization and ψ denoting classifier parameters.
  • Extensions: Good latent codes are especially important for classification when classifiers are simple or labeled data are limited.

D.1 CVAE for Dialog Response Generation

On the SW dialog-response task, cyclical scheduling outperforms monotonic scheduling while producing better reconstruction, KL, and BLEU behavior under similar ELBO values. The SW dataset contains 70 topics and uses 2,316/60/62 dialogs for training, validation, and testing.

  • Dataset and implementation: The SW dataset contains 70 available topics and is split into 2,316 training, 60 validation, and 62 test dialogs.The implementation uses code published by Zhao et al. (2017).
  • Results: The cyclical schedule outperforms the monotonic schedule in both tested CVAE settings.Full BLEU scores are reported in Table 5.
  • Results: Under similar ELBO results, cyclical scheduling yields lower reconstruction errors, higher KL values, and higher BLEU values.These metrics are reported for the SW dialog-response experiments.
  • Results: The monotonic schedule tends to overfit, whereas the cyclical schedule does not, particularly on reconstruction errors.The reported learning curves are shown in Figure 8.

D.3 Hyper-parameter tuning

The paper evaluates cyclical schedules through schedule comparisons, ablations, and hyper-parameter studies. More cycles and a larger annealing proportion generally improve performance, while cyclical scheduling produces more structured latent representations than the alternatives.

  • Number of cycles M: A larger number of cycles M provides higher performance across various proportion values R, with greater improvement when R is small.When R is large, the improvement from increasing M is small.
  • Schedule design: Cyclical schedules repeat annealing across multiple cycles, illustrated with M = 4 and R = 0.5.The top row shows traditional monotonic schedules, while the bottom row shows corresponding cyclical schedules.
  • Schedule comparison: Under similar ELBO results, cyclical scheduling provides lower reconstruction errors, higher KL values, and higher BLEU values than monotonic scheduling.The monotonic schedule tends to overfit, while the cyclical schedule does not.
  • Latent representations: Cyclical β produces much more separated latent structures than the other two methods on Yelp tSNE embeddings.More structured latent patterns usually lead to better classification performance.
  • Annealing proportion R: A larger annealing proportion R leads to better performance for various values of M.Small R performs worse because the schedule becomes more similar to a constant schedule; M = 1 recovers the monotonic schedule.
Loading 1903.10145v3…