[Paper Review] Cyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing

Hao Fu, Chunyuan Li+4 more|arXiv (Cornell University)|Mar 25, 2019

Algorithms and Data Compression36 references168 citations

TL;DR

Introduces a cyclical annealing schedule for the KL term in text VAEs, re-opening Path A multiple times to progressively learn meaningful latent codes and reduce KL vanishing with negligible extra cost.

ABSTRACT

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.

Motivation & Objective

• Explain KL vanishing in VAEs with autoregressive decoders and its causes.
• Propose cyclical annealing as a low-cost, effective training schedule for VAEs.
• Demonstrate the approach across language modeling, dialog generation, and unsupervised pre-training.

Proposed method

• Model VAE for text with an auto-regressive decoder and two latent learning paths (Path A via z and Path B via x_{<t}).
• Propose cyclical annealing: multiple cycles where beta increases from 0 to 1 and then resets, using warm restarts from prior cycles.
• Formalize beta_t with cycle-based scheduling: beta_t follows f(tau) then set to 1 for the remainder of the cycle; introduce M cycles and R proportion for annealing.
• Provide theoretical insight into beta’s impact via a lower bound analysis linking beta, mutual information I(z;n), and KL terms.
• Compare cyclical scheduling to constant and monotonic schedules, showing it encourages higher KL and richer latent representations.
• Validate on language modeling (PTB), dialog response generation (Switchboard), and unsupervised language pre-training (Yelp).

Experimental results

Research questions

• RQ1 What causes KL vanishing in VAEs with autoregressive decoders for text?
• RQ2 Can a cyclical beta schedule improve latent code learning and decoder utilization without extra computational cost?
• RQ3 How does cyclical annealing compare to constant and monotonic schedules across NLP tasks such as language modeling, dialog generation, and unsupervised pre-training?
• RQ4 Does re-opening Path A via cyclical scheduling yield more structured latent spaces and better downstream performance?

Key findings

• Cyclical annealing leads to lower reconstruction error and higher KL during training compared to monotonic schedules in experiments.
• The approach yields progressively better performance across cycles, indicating warm-start benefits from previous latent representations.
• In dialog generation, cyclical scheduling increases KL and BLEU scores, and produces more diverse responses than monotonic scheduling.
• For unsupervised pre-training, cyclical scheduling yields higher downstream accuracy and clearer latent clustering (visualized via t-SNE).
• The method achieves improvements with negligible additional computational cost relative to standard VAE training.
• Ablation suggests the gains mainly come from cyclical beta scheduling rather than cycling the learning rate.