{"id":578602,"date":"2019-04-15T08:51:05","date_gmt":"2019-04-15T15:51:05","guid":{"rendered":"https:\/\/cm-edgetun.pages.dev\/en-us\/research\/?p=578602"},"modified":"2019-06-26T14:19:52","modified_gmt":"2019-06-26T21:19:52","slug":"less-pain-more-gain-a-simple-method-for-vae-training-with-less-of-that-kl-vanishing-agony","status":"publish","type":"post","link":"https:\/\/cm-edgetun.pages.dev\/en-us\/research\/blog\/less-pain-more-gain-a-simple-method-for-vae-training-with-less-of-that-kl-vanishing-agony\/","title":{"rendered":"Less pain, more gain: A simple method for VAE training with less of that KL-vanishing agony"},"content":{"rendered":"

\"\"<\/a><\/p>\n

There is a growing interest in exploring the use of variational auto-encoders (VAE), a deep latent variable model (opens in new tab)<\/span><\/a>, for text generation. Compared to the standard RNN-based language model (opens in new tab)<\/span><\/a> that generates sentences one word at a time without the explicit guidance of a global sentence representation, VAE is designed to learn a probabilistic representation of global language features such as topic, sentiment or language style, and makes the text generation more controllable. For example, VAE can generate sentences with a specific tense, sentiment or topic.<\/p>\n

However, training VAE on languages is notoriously difficult due to something called KL vanishing<\/em>. While VAE is designed to learn to generate text using both local context and global features, it tends to depend solely on local context and ignore global features when generating text. When this happens, VAE is essentially behaving like a standard RNN language model.<\/p>\n

In \u201cCyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing (opens in new tab)<\/span><\/a>,\u201d to be presented at 2019 Annual Conference of the North American Chapter of the Association for \u00a0 Computational Linguistics (NAACL) (opens in new tab)<\/span><\/a>, researchers at Microsoft Research AI<\/a> and Duke University (opens in new tab)<\/span><\/a> propose an extremely simple remedy to KL vanishing as well as their proposal to make the code publicly available on Github (opens in new tab)<\/span><\/a>. The remedy is based on a new scheduling scheme called Cyclical Annealing Schedule<\/em>. Intuitively, during the course of VAE training, we periodically adjust the weight of the KL term in the objective function, providing the model opportunities to learn to leverage the global latent variables in text generation, thus encoding as much global information in the latent variables as possible. The paper briefly describes KL vanishing and why it happens, introduces the proposed remedy, and illustrates the VAE learning process using a synthetic dataset.<\/p>\n

What is KL vanishing and why does it happen?<\/h3>\n

VAEs aim to learn probabilistic representations z<\/strong> of natural languages x<\/strong>, with an objective consisting of two terms: (1) reconstruction<\/em> to guarantee the inferred latent feature z<\/strong> can represent its corresponding observed sentence; and (2) KL regularization<\/em> to leverage the prior knowledge to modulate language understanding. The two terms are balanced by a weighting hyper-parameter \u03b2:<\/p>\n

\"\"<\/a><\/p>\n

When applied on text corpora, VAEs typically employ an auto-regressive decoder, which sequentially generates the word tokens based on ground-truth words in the previous steps, in conjunction with latent z<\/strong>. Recent work (opens in new tab)<\/span><\/a> has found that na\u00efve training of VAEs (keeping constant \u03b2=1) leads to model degeneration\u2014the KL term becomes vanishingly small. This issue causes two undesirable outcomes: (1) the learned features are almost identical to the uninformative Gaussian prior, for all observed languages; and (2) the decoder completely ignores the latent feature, and the learned model reduces to a simpler neural language model (opens in new tab)<\/span><\/a>. Hence, the KL vanishing<\/em> issue.<\/p>\n

This negative result is so far poorly understood. We developed a two-path<\/em> competition interpretation to shed light on the issue. Let\u2019s first look at the standard VAE in Figure 1 (a), below. The reconstruction of sequence x<\/strong>=[x1<\/sub> ,…,xT<\/sub>] depends only on one path passing through the encoder \u03d5, latent representation z<\/strong> and decoder \u0398. However, when an auto-regressive decoder is used in a VAE, there are two paths from observed x<\/strong> to its reconstruction, as shown in Figure 1(b). Path A is the same as that in the standard VAE, where z<\/strong> serves as the global representation that controls the generation of x<\/strong>; Path B leaks the partial ground-truth information of x<\/strong> at every time step of the sequential decoding. It generates xt<\/sub> conditioned on x<t<\/i><\/sub>=[x1<\/sub>,…,xt-1<\/sub>]. Therefore, Path B can potentially bypass Path A to generate xt<\/sub>, leading to KL vanishing. From this perspective, we hypothesize that the KL vanishing problem is related to the low quality of z<\/strong> at the beginning phase of decoder training. This is highly possible when the naive constant schedule of \u03b2=1 is used, as the KL term pushes z<\/strong> close to an uninformative prior, less representative of the corresponding observations. This lower quality z<\/strong> introduces more difficulties in reconstructing x<\/strong>, and eventually blocks the information flow via Path A. As a result, the model is forced to learn an easier solution to decoding\u2014generating x<\/strong> via Path B only.<\/p>\n

 <\/p>\n

\"\"<\/a>

Figure 1: Illustration of information flows on (a) one path in a standard VAE, and (b) two paths in a VAE with an auto-regressive decoder.<\/p><\/div>\n

Cyclical Annealing Schedule<\/h3>\n

A simple remedy via scheduling \u03b2 during VAE training was proposed by Bowman, et al, as shown in Figure 2(a). It starts with \u03b2=0 at the beginning of training, and gradually increases \u03b2 until \u03b2=1 is reached. This monotonic schedule (opens in new tab)<\/span><\/a> of \u03b2 has become the de facto<\/em> standard in training text VAEs, and has been widely adopted in many NLP tasks. Why does it improve the performance empirically? When \u03b2<1, z<\/strong> is trained to focus more on capturing useful information for reconstruction of x<\/strong>. When the full VAE objective is considered (\u03b2=1), z<\/span><\/strong> learned earlier can be viewed as VAE initialization; such latent features are much more informative than the random start in constant schedule and thus are ready for the decoder to use.<\/p>\n

\"Figure<\/a>

Figure 2: Annealing \u03b2 with (a) the monotonic schedule and (b) the cyclical schedule.<\/p><\/div>\n

Is there a better schedule? It is key to have meaningful latent z<\/strong> at the beginning of training the decoder, so that Path A is utilized. The monotonic schedule under-weights the prior regularization when \u03b2<1; the learned z<\/strong> tends to collapse into a point estimate. This underestimation can result in sub-optimal decoder learning. A natural question concerns how one can get a better distribution estimate for z<\/strong> as initialization, while decoder has the opportunity to leverage such z<\/strong> in learning.<\/p>\n

Our proposal is to use the latent z<\/strong> trained under the full VAE objective as initialization. To learn to progressively improve z<\/strong> we propose a cyclical schedule<\/em> for \u03b2 that simply repeats the monotonic schedule multiple times as shown in Figure 2(b). We start with \u03b2=0, increase \u03b2 at a fast rate, and then stay at \u03b2=1 for subsequent learning iterations. This completes one period of monotonic schedule. It encourages the model to converge towards the VAE objective, and infers its first raw full latent distribution. Unfortunately, \u03b2=1 gradually blocks Path A, forbidding more information from passing through z<\/strong>. Crucially, we then start the second period of \u03b2 annealing and training is continued at \u03b2=0 again. This perturbs the VAE objective, dislodges it from the convergence, and reopens Path A. Importantly, the decoder now (1) has the opportunity to directly leverage z<\/strong>, without obstruction from KL; and (2) is trained with the better latent z<\/strong> than point estimates, as the full distribution learned in the previous period is fed in. We repeat this \u03b2 annealing process several times to achieve better convergences.<\/p>\n

Visualization of learning dynamics in the latent space<\/h3>\n

To visualize the learning processes on an illustrative problem, let\u2019s consider a synthetic dataset (opens in new tab)<\/span><\/a> consisting of 10 different sequences, as well as a VAE model with a 2-dimensional latent space, and an LSTM encoder and decoder.<\/p>\n

We visualize the resulting division of the latent space for different training steps in Figure 3, where each color corresponds to the latent probabilistic representation of a sequence. We observe that:<\/p>\n