Fuse It More Deeply! A Variational Transformer with Layer-Wise Latent Variable Inference for Text Generation (2022)

Abstract: The past several years have witnessed Variational Auto-Encoder’s superiority in various text generation tasks. However, due to the sequential nature of the text, auto-regressive decoders tend to ignore latent variables and then reduce to simple language models, known as the \textit{KL vanishing} problem, which would further deteriorate when VAE is combined with Transformer-based structures. To ameliorate this problem, we propose Della, a novel variational Transformer framework. Della learns a series of layer-wise latent variables with each inferred from those of lower layers and tightly coupled with the hidden states by low-rank tensor product. In this way, Della forces these posterior latent variables to be fused deeply with the whole computation path and hence incorporate more information. We theoretically demonstrate that our method can be regarded as entangling latent variables to avoid posterior information decrease through layers, enabling Della to get higher non-zero KL values even without any annealing or thresholding tricks. Experiments on four unconditional and three conditional generation tasks show that Della could better alleviate KL vanishing and improve both quality and diversity compared to several strong baselines.

## introduction variational autoencoder (vae) @xcite has proven to be successful in generating various kinds of text, such as stylistic text @xcite , dialogue @xcite , story @xcite and poetry @xcite . the sequential nature of the text leads to typically used autoregressive decoders in vae for language generation. however, such strong decoders tend to evade the difficulty of learning meaningful latent codes by heavily relying on previously generated words * corresponding author. email: sms@tsinghua.edu.cn and hence ignore latent variables @xcite , known as kl vanishing or posterior collapse. this problem causes two drawbacks: (a) the posterior distribution quickly turns into the prior one (usually standard gaussian), falling to build expressive latent representations; (b) the decoder reduces to a naive language model, resulting in monotonous generated text @xcite . to ameliorate this problem, researchers have designed various techniques. among them, three broadly used methods include weakening @xcite @xcite , kl annealing @xcite and kl threshold @xcite @xcite . nonetheless, the weakening of decoders restrains models' language modelling capability; annealing hyperparameters are hard to tune; kl threshold introduces a non-smooth objective with some optimization difficulties. in the era of rnn, vae can be easily incorporated by using the latent variable as the initial decoder state, while how to combine vae with recently prevalent transformer @xcite architectures, which have made a breakthrough in text generation, still remains an open challenge. as shown in fig. 1 , existing methods of integrating transformer into vae fall into three main paradigms: (a) directly adding latent variables to input token embeddings (abbr. embedding) @xcite ; (b) using latent variables as a separate memory token vector to be attended by selfattention in each layer (abbr. memory) @xcite ; (c) combining latent variables with the last-layer decoder states before output softmax (abbr. softmax) @xcite . however, paradigm (a) brings noise for self-attention. in paradigm (b), memory vectors tend to be ignored by attention, even exacerbating kl vanishing. in paradigm (c), latent variables couldn't deeply interfere with the whole computation path. sec.3.3 presents more detailed analyses. to better incorporate transformer into vae and theoretically ameliorate the kl vanishing problem, we propose della foot_0 , a novel variational transformer framework. della learns a series of layerwise latent variables in a transformer encoder, and each is inferred from those of lower layers and then tightly coupled with the hidden states in the corresponding decoder layer by low-rank tensor product. our method theoretically stimulates the entanglement of latent variables and hence allows propagation of undiminished latent information through layers. as a result, della forces posterior latent variables to be deeply fused with the entire computation path and encode richer information of input text, achieving higher kl values even without any annealing or threshold training tricks. in summary, our contributions are as follow: (i) we are the first to propose layer-wise inferred latent variables in transformer-based architecture to mitigate kl vanishing; we (ii) innovatively inject latent variables using low-rank tensor product, (iii) provide a theoretical validity of our method and (iv) demonstrate its effectiveness on four unconditional and three conditional generation tasks. our codes are available at https://github.com/openvlg/della.git . ## related work thanks to the representation capacity of latent space, vae has been widely adopted for both image generation @xcite ) and text generation @xcite . in the early stage, vae was combined with rnn decoders for gener-in spite of extensive applications, vae suffered from kl vanishing in the scenario of text generation @xcite . several lines of techniques have been proposed to alleviate this problem. the first line is to avoid a too fast decrease of the kl divergence by re-weighting. kl annealing @xcite linearly increased the weight of kl term from 0 to 1 during the warm-up period. @xcite further proposed cyclical annealing, which repeats the warm-up process multiple times. the second line guarantees a positive lower bound of the kl term. kl thresholding @xcite achieved a fixed minimum by combining a hinge loss, while bn-vae @xcite learned more flexible ones via batch normalization. δ-vae @xcite chose to restrain the family of posterior distributions. the third line aims to constraint decoders to force a more informative latent variable. @xcite introduced an auxiliary bow (bag-of-words) loss. he et al. (2019 in these years, the powerful transformer has been integrated with vae to benefit diverse tasks, including text classification @xcite , story generation @xcite and dialogue generation @xcite . optimus @xcite further bridged the pre-trained bert @xcite and gpt-2 @xcite with vae for pre-training. most existing works inject latent variables into the transformer decoder by the three paradigms, embedding @xcite , memory @xcite and softmax @xcite , as discussed in sec. 1, while these methods shallowly fuse the latent variables with hidden states. to achieve deeper fusion and ameliorate kl vanishing, we propose della. the most relevant architecture to our model is hierarchical vae @xcite @xcite , which is mainly designed for image generation and not suitable for text. for text generation, hierarchical latent variables are either independent of each other @xcite , or corresponding to different text granularities (sentence or word level), while our della learns conditionally inferred and layer-wise latent variables based on transformer. ## methodology as demonstrated in sec. 3, existing three paradigms make latent variables gradually diminish through layers, be ignored by self-attention or inadequately interact with hidden states, which would not mitigate but even worsen the kl vanishing problem. to deeply fuse latent variables with the whole computation path of transformer, we propose della to learn a series of layer-wise posterior latent variables which are conditionally inferred in encoder, and injected into hidden states in decoder by low-rank tensor product. we present layer-wise latent variables in sec. 4.1, describe the tensor product fusion in sec. 4.2, give the theoretical verification of della's effectiveness for ameliorating kl vanishing in sec. 4.3, and then extend della to conditional vae (cvae) in sec. 4.4. ## conclusion in this paper, we propose a novel variational transformer framework della. our framework learns a series of layer-wise latent variables with iterative dependence. these latent variables are conditionally inferred and injected into corresponding decoder layers by low-rank tensor product for deeper fusion. the experiments on both unconditional and conditional generation tasks demonstrate della's ability to significantly mitigate kl vanishing and improve generated text's quality and diversity. in the future, we plan to explore further the potential of della in larger pretrained language models.