Improving Variational Autoencoder for Text Modelling with Timestep-Wise Regularisation (2020)

Abstract: The Variational Autoencoder (VAE) is a popular and powerful model applied to text modelling to generate diverse sentences. However, an issue known as posterior collapse (or KL loss vanishing) happens when the VAE is used in text modelling, where the approximate posterior collapses to the prior, and the model will totally ignore the latent variables and be degraded to a plain language model during text generation. Such an issue is particularly prevalent when RNN-based VAE models are employed for text modelling. In this paper, we propose a simple, generic architecture called Timestep-Wise Regularisation VAE (TWR-VAE), which can effectively avoid posterior collapse and can be applied to any RNN-based VAE models. The effectiveness and versatility of our model are demonstrated in different tasks, including language modelling and dialogue response generation.

## introduction variational autoencoders (vae) @xcite , together with other deep generative models, including generative adversarial networks @xcite and autoregressive models @xcite , have attracted a mass of attention in the research community as they have shown their ability to learn compact representations from complex, high-dimensional unlabelled data. vaes have been widely used in many nlp tasks, such as text modelling @xcite @xcite @xcite , style transfer @xcite , and response generation @xcite . in addition, vaes are also useful to several downstream tasks, e.g., classification @xcite @xcite , transfer learning @xcite , etc. however, there is a challenging optimisation issue of vaes known as posterior collapse (a.k.a. kl loss vanishing), where the variational posterior collapses to the prior and the latent variable is ignored by the model during generation @xcite . this is particularly prevalent when employing vae-rnn architectures for text modelling. when posterior collapse happens, the decoder will be downgraded to a simpler language model and the vae cannot learn good latent representations of data @xcite . different strategies have been proposed to address this issue, such as annealing the kl term in the vae loss function @xcite @xcite , replacing the recurrent decoder with convolutional neural networks (cnns) @xcite , using a sophisticated prior distribution such as the von mises-fisher (vmf) distribution @xcite ; and adding mutual information into the vae objectives @xcite . while the aforementioned strategies have shown effectiveness in tackling the posterior collapse issue to some extent, they either require careful engineering between the reconstruction loss and the kl loss @xcite @xcite , or designing more sophisticated model structures @xcite @xcite . in this paper, we propose a simple and robust architecture called timestep-wise regularisation vae (twr-vae), which can effectively alleviate the vae posterior collapse issue in text modelling. existing vae-rnn models for text modelling only impose kl regularisation on the latent variable of the rnn encoder at the final timestep, forcing the latent variable to be close to a gaussian prior. in contrast, our twr-vae imposes kl regularisation on the latent variables of every timestep of the rnn encoder, which we dub timestep-wise regularisation. we hypothesise that timestep-wise regularisation is crucial to avoid posterior collapse and to learn good representations of data, and allows a more robust model learning process. in addition, the proposed timestep-wise regularisation strategy is generic and in theory can be applied to any existing vae-rnn models, e.g., vanilla rnn and gru-based vae models. twr-vae shares some similarity with existing vae-rnn models, where the input to the decoder is the latent variable sample from the variational posterior at the final timestep of the encoder. while this is a reasonable design choice, we also explore two model variants of twr-vae, namely, twr-vae mean and twr-vae sum . at each time step, both model variants sample a latent variable from the timestep dependent variational posterior of the encoder. twr-vae mean averages the sampled latent variables which is then used as input to the decoder, whereas twr-vae sum performs vector addition on the sampled latent variables instead. to demonstrate the effectiveness of our method, we select a number of strong baseline models and conduct comprehensive evaluations in two benchmark tasks involving five public datasets. for the language modelling task, experimental results show that our twr-vae model can effectively alleviate the posterior collapse issue and consistently give better predictive performance than all the baselines as evidenced by both quantitative (e.g., negative log likelihood and perplexity) and qualitative evaluation. for the dialogue response generation task, our model yields better or comparable performance to the state-ofthe-art baselines based on three evaluation metrics (i.e. bleu @xcite , bow embedding @xcite and dist @xcite ). manual examination also shows that the dialogue responses generated by our model are more diverse and contentful than the baselines, as well as being simpler in model design. our two model variants also show comparable performance to the best baseline, although not as strong as twr-vae. in summary, the contribution of our paper are three-fold: (1) we propose a simple and robust method, which can effectively alleviate the posterior collapse issue of vae via timestep-wise regularisation; (2) our approach is generic which can be applied to any rnn-based vae models; (3) our approach outperforms the state-of-art on language modelling and yields better or comparable performance on dialogue response generation. the code of twr-vae is available at: https://github.com/ruizheliuoa/twr-vae . ## related work variational autoencoder @xcite yields great performance when it was applied to image generation @xcite , facial attribute style transfer @xcite @xcite , etc. it also has been applied to many natural language processing tasks, including text generation @xcite @xcite , dialogue response generation @xcite @xcite @xcite , and style transfer @xcite @xcite for all these applications, there is a common issue called posterior collapse (or kl loss vanishing) @xcite . several different types of methods were proposed to address this issue. kl annealing is the most common and basic solution used in almost all works @xcite @xcite @xcite . another type of approaches attempt to weaken the decoder of vae to avoid posterior collapse, such as introducing word dropout and historyless decoding into the decoder @xcite , replacing the decoder with different cnns @xcite , and adding skip connections in the decoder @xcite . others tried to solve this issue by introducing new regularisers @xcite @xcite , using more sophisticated prior distributions @xcite , etc. more recently, @xcite used a cyclical annealing schedule to alleviate the kl loss vanishing issue. @xcite proposed a lagging inference network to update the encoder multiple times before a single decoder update to address the issue from the perspective of training dynamics. @xcite applied the batch normalisation to the parameters of the approximate posterior and ensured that the lower bound of the expectation of the kl is positive to avoid posterior collapse. in contrast, our approach only imposes the kl regularisation on timestep-wise latent variables in the encoder, which is simpler without changing the vae training mode, introducing more complicated posterior distributions or adding a kl annealing as a warm-up setup. ## methodology in this section, we introduce the proposed timestep-wise regularisation vae (twr-vae) model as well as its two model variants. we briefly introduce the background of vae before describing the technical details of the proposed models. ## conclusion in this paper, in order to solve posterior collapse issue of vae in text modelling, we propose a simple and generic model called timestep-wise regularisation vae, which imposes the kl regularisation on the latent variables of every timestep of the encoder. empirical results in language modelling show that our model can give better performance than all baselines while avoiding posterior collapse. ablation studies show that the timestep-wise regularisation can be easily applied into different rnn-based vae models and improve their performance. in addition, we evaluate the timestep-wise regularisation in dialogue response generation task, and the results suggest that our model yields better or comparable performance to the state-of-the-art and can generate relevant, contentful and diverse responses. a the derivation of the elbo (eq. 1) the elbo can be directly derivated from the marginal log likelihood log p θ (x i ):