Cyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing

Citation

• Authors: Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, and Lawrence Carin
• Year: 2019
• Venue: NAACL 2019
• arXiv: arXiv:1903.10145
• URL: https://aclanthology.org/N19-1021/

One-Line Takeaway

This paper proposes a simple but influential training trick for sequence VAEs: periodically reset the KL weight to reopen the latent path, so an autoregressive decoder repeatedly relearns to use informative latent codes instead of collapsing onto a latent-free solution.

Technical Core

• The underlying model remains a standard VAE with a $\beta$-weighted ELBO: $$\mathcal{L}_\beta(x) = \mathbb{E}_{q_\phi(z \mid x)}[\log p_\theta(x \mid z)] - \beta\, \mathrm{KL}(q_\phi(z \mid x)\,\|\,p(z)).$$
• The paper's diagnosis is that in autoregressive decoders, generation has two effective paths:
  • a latent path through $z$;
  • an autoregressive path through previously generated tokens.
• If $z$ is poor early in training, the decoder learns to rely on the autoregressive path, so the KL term vanishes and the latent becomes uninformative.
• The paper also emphasizes the standard decomposition $$\mathbb{E}_{p_{\mathrm{data}}(x)}\!\left[\mathrm{KL}(q_\phi(z \mid x)\,\|\,p(z))\right] = I_q(x;z) + \mathrm{KL}(q_\phi(z)\,\|\,p(z)),$$ which helps explain why large $\beta$ suppresses mutual information between data and latent code.

Method

• Divide training into $M$ cycles.
• In each cycle, let $\beta_t$ increase linearly from $0$ to $1$ over the first $R$ fraction of the cycle, then keep $\beta_t=1$ for the remaining fraction.
• Algorithmically, this means:
  1. reset $\beta$ to $0$ at the start of a cycle so the decoder is forced to use the current latent codes;
  2. gradually increase $\beta$ so the latent distribution is regularized back toward the prior;
  3. repeat, using the improved latent codes from the previous cycle as a warm start.
• No architecture change is required; the contribution is purely in the training schedule for the KL term.

Empirical Evaluation

• The paper tests the schedule on language modeling, dialogue generation, and unsupervised language pretraining settings.
• Relative to monotonic KL annealing, cyclical annealing tends to produce higher KL usage, better reconstruction, and more informative latent spaces.
• The strongest qualitative message is that repeated reopening of the latent path yields more diverse and less collapsed text generation.

Assumptions And Limits

This is an optimization heuristic, not a guarantee of semantic or identifiable latent structure. It is especially targeted at autoregressive sequence decoders and does not by itself solve the problem of what the latent variable should represent. It helps keep the latent active, but not necessarily interpretable.

Relevance To This Project

• Similarities:
  • The paper addresses exactly the regime the repo worries about: strong autoregressive decoders that can ignore a latent variable.
  • Its framing of a latent path versus a latent-free decoding path is closely related to the repo's concern that an entangled baseline already explains most of $p(y \mid x)$ without consulting $z$.
  • The method is directly relevant anywhere the repo uses KL schedules in continuous or discrete sequence-CVAE experiments.
• Differences:
  • Fu et al. study standard text VAEs, not a conditional router-generator factorization with $p(z \mid x)$ and $q(z \mid x,y)$.
  • Their method changes only the KL schedule; the repo also studies objective redesigns that explicitly reward strategy-sensitive latent usage.
  • The paper's goal is to keep $z$ informative, whereas the repo additionally wants $z$ to align with interpretable reasoning modes.

Comparison Notes

Bowman et al. (2016) provides the original text-VAE setting where KL collapse becomes obvious. He et al. (2019) gives a different optimization-dynamics diagnosis and proposes aggressive encoder updates instead of cyclic $\beta$ scheduling. Zhao et al. (2017) is a conditional text-generation paper that uses bag-of-words losses and KL annealing to fight the same failure mode.

Related Notes

• Generating Sentences from a Continuous Space
• Lagging Inference Networks and Posterior Collapse in Variational Autoencoders
• CVAE objective and loss terms