Venue: NAACL 2022

# Fuse It More Deeply! A Variational Transformer with Layer-Wise

Latent Variable Inference for Text Generation

Jinyi Hu\({}^{1,2,3}\), Xiaoyuan Yi\({}^{5}\), Wenhao Li\({}^{1,2,3}\), Maosong Sun\({}^{1,2,3,4}\), Xing Xie\({}^{5}\)

\({}^{1}\) Department of Computer Science and Technology, Tsinghua University, Beijing, China

\({}^{2}\) Beijing National Research Center for Information Science and Technology

\({}^{3}\) Institute for Artificial Intelligence, Tsinghua University, Beijing, China

\({}^{4}\) International Innovation Center of Tsinghua University, Shanghai, China

\({}^{5}\) Microsoft Research Asia

hu-jy21@mails.tsinghua.edu.cn,xiaoyuanyi@microsoft.com

Corresponding author. Email: sms@tsinghua.edu.cn _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_, pages 697 - 716 July 10-15, 2022 0202 Association for Computational Linguistics

###### 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 _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.

+
Footnote †: _Keywords:_ TensorVAE, Generative Adversarial

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_) Li et al. (2020); (b) using latent variables as a separate memory token vector to be attended by self-attention in each layer (abbr. _Memory_) Fang et al. (2021); (c) combining latent variables with the last-layer decoder states before output softmax (abbr. _Softmax_) Wang and Wan (2019). 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 Della1, 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.

Footnote 1: Della: **DE**eply Fused **L**ayer-wise **LA**tent Variables

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](https://github.com/OpenVLG/DELLA.git).

## 2 Related Work

Thanks to the representation capacity of latent space, VAE has been widely adopted for both image generation van den Oord et al. (2017); Vahdat and Kautz (2020) and text generation Bowman et al. (2016); Hu et al. (2017). In the early stage, VAE was combined with RNN decoders for generating a broad range of text, varying from dialogue Serban et al. (2016), image caption Wang et al. (2017), text summarization Gupta et al. (2017) to story Yu et al. (2020) and poetry Yi et al. (2020). In this case, latent variables are usually utilized as either the initial decoder state Li et al. (2018) or input at each time step Gupta et al. (2017).

In spite of extensive applications, VAE suffered from _KL vanishing_ in the scenario of text generation Bowman et al. (2016). 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 Bowman et al. (2016) linearly increased the weight of KL term from 0 to 1 during the warm-up period. Fu et al. (2019) 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 Kingma et al. (2016) achieved a fixed minimum by combining a hinge loss, while BN-VAE Zhu et al. (2020) learned more flexible ones via batch normalization. \(\delta\)-VAE Razavi et al. (2019) chose to restrain the family of posterior distributions. The third line aims to constraint decoders to force a more informative latent variable. Wang et al. (2017) introduced an auxiliary BOW bag-of-words loss. He et al. (2019) added additional training loops for the encoder. Yang et al. (2017) adopted dilated CNN as decoder, and Dieng et al. (2019) added skip connections to the decoder. Although the above methods mitigate _KL vanishing_ to some extent, it is still challenging for either tuning or optimization.

In these years, the powerful Transformer has been integrated with VAE to benefit diverse tasks, including text classification Gururangan et al. (2019), story generation Wang and Wan (2019); Fang et al. (2021) and dialogue generation Lin et al. (2020). Optimus Li et al. (2020) further bridged the pre-trained BERT Devlin et al. (2019) and GPT-2 Radford et al. (2019) with VAE for pre-training. Most existing works inject latent variables into the Transformer decoder by the three paradigms, Embedding Li et al. (2020), Memory Li et al. (2020); Fang et al. (2021) and SoftmaxWang and Wan (2019), 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 Sonderby et al. (2016); Klushyn 

et al., 2019; Vahdat and Kautz, 2020; Child, 2020), which is mainly designed for image generation and not suitable for text. For text generation, hierarchical latent variables are either independent of each other Serban et al. (2016), or corresponding to different text granularities (sentence or word level), while our Della learns conditionally inferred and layer-wise latent variables based on Transformer.

## 3 Preliminaries

### Transformer

Transformer Vaswani et al. (2017) represents an input sequence \({\mathbf{x}}=\{x_{1},\ldots,x_{i},\ldots,x_{n}\}\) as contextualized distributed hidden states \({\mathbf{h}}=\{h_{1},\ldots,h_{i},\ldots,h_{n}\}\) by a series of stacked layers, and states in the \(l\)-th layer, \({\mathbf{h}}^{(l)}\), are calculated with scaled dot-product attention:

\[\mathrm{Attention}(Q,K,V)=\operatorname*{softmax}\left(\frac{Q^{T}K}{\sqrt{d }}\right)V^{T}, \tag{1}\]

where \(Q,K,V\) stand for Query, Key, Value, respectively, which are projected from outputs of the previous layer: \(Q=W^{q}{\mathbf{h}}^{(l-1)}\), \(K=W^{k}{\mathbf{h}}^{(l-1)}\), \(V=W^{v}{\mathbf{h}}^{(l-1)}\). \(d\) is the dimension of hidden states. In practice, multiple groups of states are calculated with different attention parameters and then concatenated, known as multi-head attention.

### VAE

As a kind of generative model, VAE estimates the intractable data distribution \(p({\mathbf{x}})\) by deriving and maximizing its lower bound as:

\[\begin{split}&\log p({\mathbf{x}})\geq\mathcal{L}_{ELBO}({\mathbf{x}};{ \mathbf{\theta}},{\mathbf{\phi}})=\\ &\mathbb{E}_{q_{{\mathbf{\phi}}}({\mathbf{z}}|{\mathbf{x}})}[\log p_{{\mathbf{ \theta}}}({\mathbf{x}}|{\mathbf{z}})]-\mathrm{KL}(q_{{\mathbf{\phi}}}({\mathbf{z}}|{\mathbf{x}})||p ({\mathbf{z}})),\end{split} \tag{2}\]

where \({\mathbf{z}}\) is the latent variable and \(p({\mathbf{z}})\) is the prior distribution of latent variable which is commonly assumed as standard Gaussian; the posterior distribution \(p({\mathbf{z}}|{\mathbf{x}})\) is approximated by an inference network (encoder) \(q_{{\mathbf{\phi}}}({\mathbf{z}}|{\mathbf{x}})\); \(p_{{\mathbf{\theta}}}({\mathbf{x}}|{\mathbf{z}})\) is a generator (decoder) to generate text \({\mathbf{x}}\) from the latent variable \({\mathbf{z}}\); \({\mathbf{\theta}}\) and \({\mathbf{\phi}}\) are corresponding parameters.

The whole lower bound in Eq.(2), called Evidence Lower BOund (ELBO), consists of two terms: the reconstruction loss,

\[\mathcal{L}_{E}=-\mathbb{E}_{q_{{\mathbf{\phi}}}({\mathbf{z}}|{\mathbf{x}})}\left[\log p_ {{\mathbf{\theta}}}({\mathbf{x}}|{\mathbf{z}})\right], \tag{3}\]

which helps reconstruct the input given the posterior latent variable \({\mathbf{z}}\), and the KL divergence,

\[\mathcal{L}_{R}=\mathrm{KL}\left(q_{{\mathbf{\phi}}}({\mathbf{z}}|{\mathbf{x}})||p({\mathbf{z} })\right). \tag{4}\]

In practice, VAE is considered as a regularized Auto-encoder, and a hyper-parameter \(\beta\) is introduced to control the strength of KL, \(\beta\mathcal{L}_{R}\), usually used in KL annealing methods Fu et al. (2019).

### Incorporate Transformer into VAE

For Transformer encoder, the posterior \({\mathbf{z}}\) is mapped from the text representation, which can be the pooling of all hidden states in the last layer Fang et al. (2021), or state of a special token Li et al. (2020a), _e.g._, [CLS]. Then \({\mathbf{z}}\) is injected into Transformer decoder by the paradigms discussed in Sec. 1.

Now we take a further step and investigate why intrinsically these three paradigms, namely Embedding, Memory and Softmax, would perform poorly.

**Embedding**: Define \({\mathbf{e}}_{i},{\mathbf{e}}_{j}\) as two token embeddings and \(\alpha_{i,j}\) as the attention weight of \(i\)-th and \(j\)-th tokens. From Eq.(1), we have \(\alpha_{i,j}=(W^{q}{\mathbf{e}}_{i})^{T}(W^{k}{\mathbf{e}}_{j})={\mathbf{e}}_{i}^{T}(W^{ q})^{T}W^{k}{\mathbf{e}}_{j}\), which is further abbreviated as \(\langle{\mathbf{e}}_{i},{\mathbf{e}}_{j}\rangle\). Such Embedding paradigm directly adds \({\mathbf{z}}\) to token embeddings as:

\[\begin{split}\alpha^{\prime}_{i,j}&=\big{[}W^{q}({ \mathbf{e}}_{i}+{\mathbf{z}})\big{]}^{T}\big{[}W^{k}({\mathbf{e}}_{j}+{\mathbf{z}})\big{]}\\ &=\langle{\mathbf{e}}_{i},{\mathbf{e}}_{j}\rangle+\langle{\mathbf{e}}_{i},{ \mathbf{z}}\rangle+\langle{\mathbf{z}},{\mathbf{e}}_{j}\rangle+\langle{\mathbf{z}},{\mathbf{z}} \rangle,\end{split} \tag{5}\]

where we can find that a redundant term, \(\langle{\mathbf{z}},{\mathbf{z}}\rangle\), is introduced, bringing extra noise for attention mechanism. Moreover, information in \({\mathbf{z}}\) could diminish with propagation through layers (Fig. 2), aggravating KL vanishing.

**Memory**: This paradigm treats \({\mathbf{z}}\) as an additional memory token and places it at the beginning of \({\mathbf{x}}\) to be attended by other tokens via attention. Nevertheless, as mentioned in Sec. 1, the powerful Transformer decoder may only rely on preceding decoded tokens. Consequently, with no explicit constraints (_e.g._, auxiliary loss), such a memory token is more likely to be ignored by self-attention (Fig. 6 & 7), even exacerbating KL vanishing.

**Softmax**: This paradigm first adds \({\mathbf{z}}\) to the last-layer hidden states \({\mathbf{h}}\), and then projects \({\mathbf{z}}+{\mathbf{h}}\) into a logit vector \({\mathbf{p}}\in\mathbb{R}^{v}\) over the vocabulary, where \(v\) is vocab size. In this method, latent variables do not interact with hidden states until the last layer, which erodes the effect of latent variables (see Fig. 2).

## 4 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.

### Layer-wise Latent Variables

Different from previous work where only one latent variable is calculated and shared by Li et al. (2020) or projected to Fang et al. (2021) decoder layers, we involve a series of latent variables , where is the number of Transformer layers. Then we reformulate the prior and posterior distributions as , respectively, with each still following Gaussian distribution. Then we rewrite in Eq.(4) similar to Vahdat and Kautz (2020):

(6)

When is the standard Gaussian distribution, . We give detailed derivations in Appendix B.1.

These latent variables are calculated (inferred) layer by layer using representations of the corresponding layer. Concretely, in -th layer, we use the hidden state of the first token in text , as its -layer representation, denoted as , where is hidden size. Then we represent latent variables in lower layers as and obtain it by:

(7)

where , so is an -dimensional of latent variable and are set as zero vector. We calculate the mean and variance vectors of and by:

(8)

where , where , .

The latent variable is sampled from the posterior distribution for training, and from the prior one , is an -dimensional for testing. Since hidden states in each layer belong to different vector spaces, the parameters to calculate each , , , and , _do not_ share throughout different layers.

### Low-rank Tensor Product

We inject the latent variable , which is obtained based on -th encoder layer, into the corresponding -th decoder layer. Instead of simply using as a memory token as discussed in Sec. 3.3, we resort to low-rank tensor product, which has been successfully utilized for fusing multimodal representations Liu et al. (2018), to deeply fuse latent variables with hidden states in the decoder.

In detail, we conduct low-rank tensor product on and -hyper value vector as:

(9)

where is a hyper-parameter, means element-wise multiplication, are shared across all positions but not shared with layers (), considering distinct vector spaces in different layers, as mentioned in Sec. 4.1. Then the fused Value is used in Eq.(1)

In this way, layer-wise is conditionally inferred from latent variables in previous encoder layers, together with -th-layer text representation, and then explicitly fused with the corresponding decoder layer, yielding a deeper intervention throughout the whole computation path of Transformer.

### Why Could Della Work Well?

To theoretically interpret the advantage of layer-wise latent variables which contributes most to Della (Table 4), we give the following theorem:

**Theorem 1**: _For an observation and a sequence of latent variables , satisfying , and , then the expectation of the KL term is an upper bound of:_ 

_where \(I\) is the interaction information2._

Footnote 2: [https://en.wikipedia.org/wiki/Interaction](https://en.wikipedia.org/wiki/Interaction) ̇information

See Appendix B.2 for proof. Based on Theorem 1, minimizing \({\cal L}_{R}\) approximatively means maximizing each interaction information term in Eq.(10), which forces the entanglement of all latent variables \(\mathbf{z}_{1};\ldots;\mathbf{z}_{L}\) given the observation \(\mathbf{x}\), alleviating the diminishing of information encoded in latent variables when propagating through layers.

### Extension to CVAE

Della could also be applied to CVAE for conditional generation tasks like storytelling. Given an observation \(\mathbf{x}\) and its condition \(\mathbf{c}\), we can optimize:

\[\log p(\mathbf{x}|\mathbf{c}) \geq\mathbb{E}_{q_{\mathbf{\phi}}(\mathbf{x}|\mathbf{x},\mathbf{c})}[\log p_{\bm {\theta}}(\mathbf{x}|\mathbf{z},\mathbf{c})] \tag{11}\] \[-\mathrm{KL}(q_{\mathbf{\phi}}(\mathbf{z}|\mathbf{x},\mathbf{c})||p(\mathbf{z}|\mathbf{ c})),\]

and then replace the prior distribution \(q(\mathbf{z}_{l}|\mathbf{x},\mathbf{z}_{<l})\) and posterior distribution \(p(\mathbf{z}_{l}|\mathbf{z}_{<l})\) in Eq.(6) with \(q(\mathbf{z}_{l}|\mathbf{x},\mathbf{c},\mathbf{z}_{<l})\) and \(p(\mathbf{z}_{l}|\mathbf{z}_{<l},\mathbf{c})\), respectively.

In this case, we encode the condition \(\mathbf{c}\) with the same encoder. Similarly, we can obtain the representation of \(\mathbf{c}\) at \(l\)-th layer, denoted as \(\mathbf{c}^{(l)}\in\mathbb{R}^{d}\), and then calculate the mean and log variance of \(p(\mathbf{z}_{l}|\mathbf{z}_{<l},\mathbf{c})\) and \(q(\mathbf{z}_{l}|\mathbf{z}_{<l},\mathbf{x},\mathbf{c})\) by:

\[\begin{split}&\left(\begin{matrix}\mathbf{\mu}_{p}\\ \log(\mathbf{\sigma}_{p}^{2})\end{matrix}\right)=\hat{\mathbf{W}}_{p}^{(l)}\left(\begin {matrix}\mathbf{z}_{<l}\\ \mathbf{c}^{(l)}\end{matrix}\right),\\ &\left(\begin{matrix}\mathbf{\mu}_{q}\\ \log(\mathbf{\sigma}_{q}^{2})\end{matrix}\right)=\hat{\mathbf{W}}_{q}^{(l)}\left( \begin{matrix}\mathbf{z}_{<l}\\ \mathbf{x}^{(l)}\end{matrix}\right),\end{split} \tag{12}\]

where \(\hat{\mathbf{W}}_{p}^{(l)}\in\mathbb{R}^{(p+d)\times 2p}\), \(\hat{\mathbf{W}}_{q}^{(l)}\mathbb{R}^{(p+2d)\times 2p}\).

## 5 Experiment

### Dataset

We consider four datasets for language modelling and unconditional generation, including the Yelp, and Yahoo Yang et al. (2017); He et al. (2019), Penn Treebank (PTB) Marcus et al. (1993), and SNLI Bowman et al. (2015), and three datasets for conditional generation tasks, including summarization generation with CNN/DailyMail (CNN/DM) See et al. (2017), story generation with WritingPrompts (WP) Fan et al. (2018) and paraphrase generation with Quora 3. Detailed data statistics are listed in Table 7. Due to the limited computation capability, we use 165,157 samples in CNN/DM and 22,2614 in WP with the max length of 900 for training.

Footnote 3: [https://quoradata.quora.com/First-Quora-Dataset-Release-Question-Pairs](https://quoradata.quora.com/First-Quora-Dataset-Release-Question-Pairs)

### Implementation Details

We use pretrained language models as the backbone and fine-tune them on each task mentioned above with our Della as in Li et al. (2020). For unconditional generation and story generation, encoder and decoder **shared** the same parameters initialized with 12-layer GPT-2 Radford et al. (2019). For summarization and paraphrase generation, parameters are **not shared** and initialized with BART-base Lewis et al. (2020). We set the dimension of latent variable as 32 for all VAE-based models and use cyclical annealing for training, following Li et al. (2020). More details are given in Appendix A.1.

### Baseline

We make a comprehensive comparison with strong Transformer-based baselines. We do not consider RNN-based models that are inferior to Transformer for text generation as shown in Li et al. (2020).

**Finetuned Pretrained Models**. To manifest the suitability of Della for different pretrained language models, we compare it with fine-tuned GPT2 on unconditional generation and story generation, and with fine-tuned BART-base on summarization generation and paraphrase generation.

**Optimus**Li et al. (2020): a large-scale VAE model which takes a pre-trained BERT as encoder and pretrained GPT-2 as decoder. This model is first pretrained as a VAE, which simultaneously utilizes the two paradigms, Embedding and Memory as introduced in Sec. 3.3, for injecting latent variables, with both KL annealing and KL threshold tricks, and then fine-tuned on downstream tasks.

**Transformer-based VAE**. Besides Optimus, we also compre the three paradigms, namely **Embedding**Li et al. (2020), **Memory**Fang et al. (2021) and **Softmax**Wang and Wan (2019), and incorporate each paradigm into the same pre-trained model as Della on each dataset for fair comparison.

### Metrics

_For unconditional generation tasks,_ we consider three types of metrics. **(a) Representation Learning Capability**: we report PPL, ELBO, KL, mutual information (MI) Alemi et al. (2016) and activate units (AU) Burda et al. (2016). These metrics measure VAE's ability to mitigate KL vanishing and learn meaningful representations. Different from traditional language models like GPT-2, VAEBased models could not produce exact PPL due to randomness, so we use importance-weighted sam 

ples to estimate PPL, following He et al. (2019). We set the threshold in AU to 0.2 to further distinguish different models. **(b) Generation Quality**: we report BLEU Papineni et al. (2002), CND Li et al. (2020) and MAUVE Pillutla et al. (2021). CND and MAUVE measure the divergence between human-authored text and the generated one. **(c) Generation Diversity**: we report Self-BLEU Zhu et al. (2018), Dist Li et al. (2016) and JS Jaccard similarity Wang and Wan (2018) to assess the diversity and novelty of generated text.

_For conditional generation tasks_, we report BLEU, Rouge-1, Rouge-2, Rouge-L Lin and Hovy (2002), and BERTScore Zhang et al. (2020) to evaluate the quality of generated texts, as well as the same diversity metrics used in unconditional generation. We also report KL and AU value to present representation learning capability. More details of metrics are provided in Appendix A.3.

### Results

#### 5.5.1 Unconditional Generation

We present results on Yelp and Yahoo in Table 1 and leave the those on PTB and SNLI in the Appendix A.5 due to space limitations. We also show the learning curves of ELBO and KL in Fig. 5.

As shown in Table 1, Della achieves notably improvement on almost all the metrics, especially superior on representation learning metrics. Much higher KL, MI and AU, and a big gap in PPL obtained by Della indicate the latent variables encode more meaningful text information and won't diminish when propagating through Transformer layers, which strongly supports our motivation that fusing latent variables with hidden states more deeply could effectively alleviate the KL vanishing problem. Such results also empirically verify the theoretical advantage of our model (Theorem 1), demonstrating entangled layer-wise latent variables can preserve more encoded knowledge for decoder. We will show that \(z\) can involve more information when injected into more layers in Sec. 5.8.

Besides, Della also gets good performance (comparable BLEU and much better CND and MAUVE) on generation quality. With more informative latent variables, Della could achieve a better ELBO and hence further boost the learning of data distribution \(p({\mathbf{x}})\) in Eq.(2), leading to satisfactory quality of generated texts.

Generally, Della also outperforms baseline models on generation diversity. The reason is twofold: randomly sampled latent variables \({\mathbf{z}}\) should bring diversity, while the VAE-based baselines tend to ignore \({\mathbf{z}}\) as mentioned before, losing some randomness. In contrast, latent variables are deeply fused in Della, maintaining enough randomness. Besides, each latent variable is sampled in corresponding layer, and thus such a sampling process accumulates and enhances randomness, further benefiting diversity while keeping good quality.

#### 5.5.2 Conditional Generation

We report the results of WP and CNN/DM in Table 2, and leave those of Quora in Appendix A.5. As we can see, Della performs better on most quality metrics, but gets a little worse on diversity

\begin{table}
\begin{tabular}{l|c|c|c|c|c|c|c|c|c|c} \hline \multirow{2}{*}{Model} & \multicolumn{4}{c|}{Representation Learning} & \multicolumn{4}{c|}{Generation Quality} & \multicolumn{2}{c}{Generation Diversity} \\ \cline{2-11}  & PPL\(\downarrow\) & ELBO\(\downarrow\) & KL\(\uparrow\) & M\(\uparrow\) & AU\(\uparrow\) & BLEU\(\uparrow\) & CND\(\downarrow\) & MAUVE\(\uparrow\) & SB\(\downarrow\) & Dist\(\uparrow\) & JS \(\downarrow\) \\ \hline \multicolumn{11}{c}{Dataset: Yelp} \\ \hline GPT-2 & 22.13 & - & - & - & - & 56.92 & 0.68 & 0.12 & 65.90 & **17.96** & 0.51 \\ Optimus & 22.79 & 344.10 & 15.09 & 7.67 & - & - & - & - & - & - \\ Embed & 19.98 & 327.28 & 4.77 & 4.14 & 6 & 56.34 & 0.31 & 0.42 & 65.27 & 15.59 & 0.44 \\ Memory & 19.95 & 326.60 & 5.70 & 5.30 & 11 & **57.37** & 0.27 & 0.46 & 63.90 & 16.91 & **0.39** \\ Softmax & 20.14 & 328.13 & 7.50 & 6.29 & 13 & 56.83 & 0.30 & 0.45 & 64.26 & 16.51 & 0.40 \\ Della & **12.35** & **239.83** & **29.47** & **10.78** & **23** & 57.15 & **0.13** & **0.55** & **60.02** & 17.63 & 0.43 \\ \hline \multicolumn{11}{c}{Dataset: Yahoo} \\ \hline GPT-2 & 24.17 & - & - & - & - & 44.25 & 0.55 & 0.15 & 54.06 & 21.07 & **0.28** \\ Optimus & 23.11 & 293.34 & 17.45 & 8.85 & - & - & - & - & - & - \\ Embed & 22.18 & 286.85 & 3.63 & 3.03 & 3 & 42.27 & 0.45 & 0.31 & 54.15 & 20.80 & 0.32 \\ Memory & 22.03 & 285.47 & 4.87 & 4.62 & 18 & **45.20** & 0.46 & 0.37 & 54.59 & 21.87 & 0.33 \\ Softmax & 22.35 & 287.44 & 6.35 & 5.52 & 19 & 44.28 & 0.44 & 0.34 & 54.49 & 21.65 & 0.32 \\ Della & **11.49** & **201.34** & **27.84** & **12.31** & **21** & 44.67 & **0.19** & **0.38** & **48.53** & **21.88** & 0.31 \\ \hline \end{tabular}
\end{table}
Table 1: Evaluation results for language modelling and unconditional generation. Results of Optimus are directly copied from the original paper with \(\lambda=0.5\). SB means Self-BLEU.

 

compared to GPT-2. This is because GPT-2 may produce some ill-formed contents which 'improve' diversity by cheating the metrics but also lead to much worse quality (lower BLEU and Rouge). Even so, on both WP and CNN/DM, Della still beats all previous VAE paradigms in diversity, manifesting the effectiveness of our Della.

In addition, all baselines methods suffer from severer KL vanishing problems on conditional generation tasks than on the unconditional ones. This is because the given condition text could aggravate the reliance of these models on preceding generated tokens and the condition, and therefore bypass latent variables. By contrast, Della could learn more informative \({\mathbf{z}}\) and hence keep a relatively higher KL value even given the condition text.

### Human Evaluation

To better verify the effectiveness of Della, we also conduct human evaluation on the two conditional generation tasks. For each model, we generated 30 samples on each task, and invite 5 competent annotators to score these samples in terms of three criteria, **Fluency**, **Coherence** and **Novelty** for story generation, and **Informativeness**, **Coherence** and **Novelty** for summarization generation.

As shown in Table 3, Della obtains satisfactory performance in quality, and is consistently superior to all baselines on diversity and novelty. See Appendix A.4 for more detailed evaluation protocols.

### Ablation Study

Table 4 shows the results of ablation study on Yelp. We can find both tensor product and the layer-wise latent variables benefit the learning of informative

\begin{table}
\begin{tabular}{l|c|c|c|c|c} \hline \hline Model & PPL\(\downarrow\) & ELBO\(\downarrow\) & KL\(\uparrow\) & MI \(\uparrow\) & AU\(\uparrow\) \\ \hline Della & **12.35** & **239.83** & **29.47** & **10.78** & **23** \\ -LTP & 12.68 & 249.32 & 28.52 & 9.77 & 21 \\ -LW & 19.88 & 324.45 & 20.12 & 7.23 & 18 \\ \hline Separate & 14.17 & 286.30 & 28.82 & 9.88 & 16 \\ \(l=1\) KL & 12.55 & 266.97 & 0.15 & 0.15 & 0 \\ \(l=12\) KL & 12.48 & 263.38 & 0.73 & 0.61 & 0 \\ \hline Embed(384) & 20.11 & 327.29 & 0.55 & 0.38 & 0 \\ Memory(384) & 20.09 & 326.24 & 0.46 & 0.25 & 0 \\ Softmax(384) & 20.15 & 330.24 & 5.04 & 7.15 & 0 \\ \hline \hline \end{tabular}
\end{table}
Table 4: Ablation study on Yelp dataset. LTP: low-rank tensor product. LW: layer-wise latent variables. Separate: latent variables in each layer are independent. \(l=1\) or \(4\) KL means we only compute KL loss on \({\mathbf{z}}_{1}\) or \({\mathbf{z}}_{L}\), respectively. 384 means the dimension of latent variable used in baseline are \(12\times 32=384\).

\begin{table}
\begin{tabular}{c|c|c|c|c|c|c|c|c|c} \hline \hline \multirow{2}{*}{Model} & \multicolumn{6}{c|}{Quality} & \multicolumn{3}{c|}{Diversity} & \multirow{2}{*}{KL\(\uparrow\)} & \multirow{2}{*}{AU\(\uparrow\)} \\ \cline{2-2} \cline{6-10}  & BLEU\(\uparrow\) & Rouge-1\(\uparrow\) & & Rouge-2\(\uparrow\) & & Rouge-L\(\uparrow\) & BERTScore\(\uparrow\) & SB\(\downarrow\) & Dist\(\uparrow\) & JS\(\downarrow\) \\ \hline \multicolumn{10}{c}{Dataset: WritingPrompts} \\ \hline GPT-2 & 27.89 & 27.72 & 7.96 & 14.30 & 78.12 & **53.78** & **22.99** & **0.51** & - & - \\ Embed & 39.67 & **36.17** & 7.96 & 15.78 & 81.64 & 64.55 & 14.31 & 0.73 & 2.35 & 3 \\ Memory & 40.79 & 36.13 & 8.04 & 16.16 & 81.68 & 67.56 & 12.90 & 0.80 & 0.07 & 0 \\ Softmax & 41.04 & 36.14 & 8.12 & 16.30 & 81.75 & 67.02 & 13.08 & 0.78 & 0.32 & 0 \\ Della & **41.39** & 35.46 & **8.78** & **17.20** & **81.77** & 56.28 & 20.91 & 0.60 & **28.14** & **8** \\ \hline \multicolumn{10}{c}{Dataset: CNN/DM} \\ \hline Bart-base & 48.74 & **41.33** & 19.82 & 29.63 & 87.75 & 29.94 & 43.68 & 0.10 & - & - \\ Embed & 44.10 & 40.43 & 19.41 & 29.43 & 87.60 & 29.60 & 44.04 & 0.10 & 0.0 & 0 \\ Memory & 46.02 & 41.18 & 19.74 & 29.64 & 87.78 & 29.79 & 43.92 & 0.11 & 0.0 & 0 \\ Softmax & 44.40 & 40.94 & 19.63 & 29.61 & 87.00 & 29.64 & 44.11 & 0.10 & 0.0 & 0 \\ Della & **49.18** & 41.27 & **19.85** & **29.84** & **88.09** & **29.07** & **44.24** & **0.09** & **0.91** & **1** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Evaluation results for conditional generation.

\begin{table}
\begin{tabular}{c|c|c|c} \hline \hline Model & Fluency & Coherence & Novelty \\ \hline GPT2 & 1.83 & 2.12 & 2.50 \\ Embed & 2.16 & 2.33 & 2.67 \\ Memory & 2.45 & 2.28 & 2.78 \\ Softmax & 2.48 & **2.42** & 2.85 \\ Della & **2.51** & 2.38 & **2.89** \\ \hline \multicolumn{10}{c}{Dataset: CNN/DM} \\ \hline Model & Informativeness & Coherence & Novelty \\ \hline Bart-base & **3.12** & 4.32 & 3.52 \\ Embed & 2.88 & 4.08 & 3.50 \\ Memory & 2.95 & 4.23 & 3.48 \\ Softmax & 2.91 & **4.33** & 3.50 \\ Della & 3.05 & **4.33** & **3.56** \\ \hline \hline \end{tabular}
\end{table}
Table 3: Human evaluation results on conditional generation. The scores range from 1 (worst) to 5 (best). The p-value is 0.002 and Kappa score is 0.64 which indicates acceptable inter-annotator agreement.

 

latent variables, while the latter contributes the most to Della. To further verify the performance gain originating from Theorem 1 instead of simply increasing the number or the dimension of latent variables, we conduct two groups of experiments.

First, we remove the conditional dependence between layer-wise latent variables by independently sampling each in both training and testing. We can see that removing dependence causes a significant performance drop. Besides, we keep the dependence between but optimize only one of the KL terms in Eq.(6), and find all representation capability metrics deteriorate, especially KL, MI and AU. Such results effectively demonstrate the necessity of using and optimizing the conditional inference of layer-wise latent variables, supporting our theoretical interpretation of Della.

Second, we enlarge the dimension of used in the three paradigms to \(384\) (\(12\times 32\)), equal to the total latent dimension used in Della. The results show that simply increasing the dimension of latent variables brings a more sparse latent space, even exacerbating the KL vanishing problem.

### Analysis

**Training Tricks** To reveal the robustness of our model, we evaluate the influence of three commonly used training tricks to relieve KL vanishing, BOW (bag-of-words) loss Wang et al. (2017), batch normalization Zhu et al. (2020) and KL annealing Fu et al. (2019), to the performance of Della and the three paradigms. As shown in Table 5, previous methods suffer KL vanishing seriously without annealing or BOW loss, getting KL, MI and AU almost 0. Though not good as using annealing, Della still maintains acceptable performance and mitigates KL vanishing even without any training tricks. Bow and batch normalization dramatically prevent low KL divergence, but obstruct the optimization and thus cause higher PPL.

**Number of Latent Variables** We observe the change of PPL, KL and ELBO with different numbers of latent variables. We conduct two groups of experiments where we produce and utilize layer-wise latent variables starting from and ending at different layers. As shown in Fig. 2, incorporating more latent variables could continuously improve performance, consistent to our claim in Sec. 4. With the same number of latent variables, starting from a higher layer is better than ending at a lower layer, which indicates that latent variables

Figure 2: PPL, ELBO end KL on Yelp with different numbers of latent variables. The values _start layer_ and _end layer_ means latent variables are produces and utilized only from \(i\)-th layer to the last layer, or from the first layer to \(j\)-th layer of the encoder respectively.

 

generated from higher layers encode more helpful information compared to those from lower layers, manifesting disadvantages of the two previous paradigms, Softmax (starting from the last layer) and Embedding (ending at the first layer).

**Model size** We compare the performance of Della and three paradigms with 24-layer GPT2-medium as backbone. As shown in Table 5, with the increasing of model size, Della consistently achieves better performance than baselines.

### Case Study

VAE captures text representations in a smooth latent space. We take two sentences \({\mathbf{x}}_{1}\) and \({\mathbf{x}}_{2}\) and sample two posterior latent variables \({\mathbf{z}}^{(1)}\) and \({\mathbf{z}}^{(2)}\) from \(p({\mathbf{z}}^{(1)}|{\mathbf{x}}_{1})\) and \(p({\mathbf{z}}^{(2)}|{\mathbf{x}}_{2})\), and get interpolated latent variables with \({\mathbf{z}}=\tau{\mathbf{z}}^{(1)}+(1-\tau){\mathbf{z}}^{(2)}\). We generate multiple sentences with a continuously changed \(\tau\) from 0 to 1. As shown in Fig. 3, sentences generated from interpolated \({\mathbf{z}}\) mix the semantics of the two initial sentences and smoothly change from \({\mathbf{x}}_{1}\) to \({\mathbf{x}}_{2}\), showing Della's ability of learning a flexible latent space.

Fig. 4 shows the generation examples of Della and one of baseline, Memory, given the same prompt WritingPrompts. We observe that the generated text of Memory is irrelevant to the prompt, while Della generates coherent and vivid text.

## 6 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.

## Acknowledgement

Thanks for the anonymous reviewers for their comments. This work is supported by the National Key R&D Program of China (No. 2020AAA0106502) and International Innovation Center of Tsinghua University, Shanghai, China.

Figure 4: Generation examples of Memory and Della based on the prompt from test set of WritingPrompts.

Figure 5: Reconstruction loss and KL Divergence throughout training process.

 



## References

* A. Alemi, I. Fischer, J. V. Dillon, and K. Murphy (2016)Deep variational information bottleneck. In International Conference on Learning Representations, Cited by: SS3.4.
* S. R. Bowman, G. Angeli, C. Potts, and C. D. Manning (2015)A large annotated corpus for learning natural language inference. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, Lisbon, Portugal, pp. 632-642. External Links: Link, Document Cited by: SS1, SS2.
* S. R. Bowman, L. Vilnis, O. Vinyals, A. Dai, R. Jozefowicz, and S. Bengio (2016)Generating sentences from a continuous space. In Proceedings of The 20th SIGNLL Conference on Computational Natural Language Learning, Berlin, Germany, pp. 10-21. External Links: ISBN 978-3

Bohan Li, Junxian He, Graham Neubig, Taylor Berg-Kirkpatrick, and Yiming Yang. 2019. A surprisingly effective fix for deep latent variable modeling of text. In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)_, pages 3603-3614, Hong Kong, China. Association for Computational Linguistics.
* Li et al. (2020a) Chunyuan Li, Xiang Gao, Yuan Li, Baolin Peng, Xiujun Li, Yizhe Zhang, and Jianfeng Gao. 2020a. Optimus: Organizing sentences via pre-trained modeling of a latent space. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_, pages 4678-4699, Online. Association for Computational Linguistics.
* Li et al. (2020b) Jaining Li, Yanyan Lan, Jiafeng Guo, and Xueqi Cheng. 2020b. On the relation between quality-diversity evaluation and distribution-fitting goal in text generation. In _International Conference on Machine Learning_.
* Li et al. (2016) Jiwei Li, Michel Galley, Chris Brockett, Jianfeng Gao, and Bill Dolan. 2016. A diversity-promoting objective function for neural conversation models. In _Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_, pages 110-119, San Diego, California. Association for Computational Linguistics.
* Li et al. (2018) Juntao Li, Yan Song, Haisong Zhang, Dongmin Chen, Shuming Shi, Dongyan Zhao, and Rui Yan. 2018. Generating classical Chinese poems via conditional variational autoencoder and adversarial training. In _Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing_, pages 3890-3900, Brussels, Belgium. Association for Computational Linguistics.
* Lin and Hovy (2002) Chin-Yew Lin and Eduard Hovy. 2002. Manual and automatic evaluation of summaries. In _Proceedings of the ACL-02 Workshop on Automatic Summarization_, pages 45-51, Philadelphia, Pennsylvania, USA. Association for Computational Linguistics.
* Lin et al. (2020) Zhaojiang Lin, Genta Indra Winata, Peng Xu, Zihan Liu, and Pascale Fung. 2020. Variational transformers for diverse response generation. _arXiv: Computation and Language_.
* Liu et al. (2018) Zhun Liu, Ying Shen, Varun Bharadhwaj Lakshmiarasimhan, Paul Pu Liang, AmirAli Bagher Zadeh, and Louis-Philippe Morency. 2018. Efficient low-rank multimodal fusion with modality-specific factors. In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_, pages 2247-2256, Melbourne, Australia. Association for Computational Linguistics.
* Marcus et al. (1993) Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. 1993. Building a large annotated corpus of English: The Penn Treebank. _Computational Linguistics_, 19(2):313-330.
* Papineni et al. (2002) Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002. Bleu: a method for automatic evaluation of machine translation. In _Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics_, pages 311-318, Philadelphia, Pennsylvania, USA. Association for Computational Linguistics.
* Pillutla et al. (2021) Krishna Pillutla, Swabha Swayamdipta, Rowan Zellers, John Thicktun, Sean Welleck, Yejin Choi, and Zaid Harchaoui. 2021. Mauve: Measuring the gap between neural text and human text using divergence frontiers. In _Neural Information Processing Systems_.
* Radford et al. (2019) Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners.
* Razavi et al. (2019) Ali Razavi, Aaron van den Oord, Ben Poole, and Oriol Vinyals. 2019. Preventing posterior collapse with delta-vaes. In _International Conference on Learning Representations_.
* Rezende et al. (2014) Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. 2014. Stochastic backpropagation and approximate inference in deep generative models. In _International conference on machine learning_, pages 1278-1286. PMLR.
* See et al. (2017)bigail See, Peter J. Liu, and Christopher D. Manning. 2017. Get to the point: Summarization with pointer-generator networks. In _Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_, pages 1073-1083, Vancouver, Canada. Association for Computational Linguistics.
* Semeniuta et al. (2017) Stanislau Semeniuta, Aliaksei Severyn, and Erhardt Barth. 2017. A hybrid convolutional variational autoencoder for text generation. In _Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing_, pages 627-637, Copenhagen, Denmark. Association for Computational Linguistics.
* Serban et al. (2016) Iulian Vlad Serban, Alessandro Sordoni, Ryan Lowe, Laurent Charlin, Joelle Pineau, Aaron Courville, and Yoshua Bengio. 2016. A hierarchical latent variable encoder-decoder model for generating dialogues. _arXiv: Computation and Language_.
* Sonderby et al. (2016) Casper Kaae Sonderby, Tapani Raiko, Lars Maaloe, Soren Kaae Sonderby, and Ole Winther. 2016. Ladder variational autoencoders. In _Neural Information Processing Systems_.
* Vahdat and Kautz (2020) Arash Vahdat and Jan Kautz. 2020. Nvae: A deep hierarchical variational autoencoder. In _Neural Information Processing Systems_.

 

Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. 2017. Neural discrete representation learning. In _Neural Information Processing Systems_.
* Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In _Advances in neural information processing systems_, pages 5998-6008.
* Wang and Wan (2018) Ke Wang and Xiaojun Wan. 2018. Sentigan: Generating sentimental texts via mixture adversarial networks. In _International Joint Conference on Artificial Intelligence_.
* Wang et al. (2017) Liwei Wang, Alexander G. Schwing, and Svetlana Lazebnik. 2017. Diverse and accurate image description using a variational auto-encoder with an additive gaussian encoding space. In _Neural Information Processing Systems_.
* Wang and Wan (2019) Tianming Wang and Xiaojun Wan. 2019. T-cvae: transformer-based conditioned variational autoencoder for story completion. In _International Joint Conference on Artificial Intelligence_.
* Wolf et al. (2020) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Phu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. 2020. Transformers: State-of-the-art natural language processing. In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_, pages 38-45, Online. Association for Computational Linguistics.
* Yang et al. (2017) Zichao Yang, Zhiting Hu, Ruslan Salakhutdinov, and Taylor Berg-Kirkpatrick. 2017. Improved variational autoencoders for text modeling using dilated convolutions. In _International conference on machine learning_, pages 3881-3890. PMLR.
* Yi et al. (2020) Xiaoyuan Yi, Ruoyu Li, Cheng Yang, Wenhao Li, and Maosong Sun. 2020. Mixpoet: Diverse poetry generation via learning controllable mixed latent space. In _Proceedings of The Thirty-Fourth AAAI Conference on Artificial Intelligence_, New York, USA.
* Yu et al. (2020) Meng-Hsuan Yu, Juntao Li, Danyang Liu, Dongyan Zhao, Rui Yan, Bo Tang, and Haisong Zhang. 2020. Draft and edit: Automatic storytelling through multi-pass hierarchical conditional variational autoencoder. In _National Conference on Artificial Intelligence_.
* Zhang et al. (2020) Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, and Yoav Artzi. 2020. Bertscore: Evaluating text generation with bert. In _International Conference on Learning Representations_.
* Zhao et al. (2017) Tiancheng Zhao, Ran Zhao, and Maxine Eskenazi. 2017. Learning discourse-level diversity for neural dialog models using conditional variational autoencoders. In _Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_, pages 654-664, Vancouver, Canada. Association for Computational Linguistics.
* Zhu et al. (2020) Qile Zhu, Wei Bi, Xiaojiang Liu, Xiyao Ma, Xiaolin Li, and Dapeng Wu. 2020. A batch normalized inference network keeps the KL vanishing away. In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_, pages 2636-2649, Online. Association for Computational Linguistics.
* Zhu et al. (2018) Yaoming Zhu, Sidi Lu, Lei Zheng, Jiaxian Guo, Weinan Zhang, Jun Wang, and Yong Yu. 2018. Texygen: A benchmarking platform for text generation models. In _International ACM SIGIR Conference on Research and Development in Information Retrieval_ 

Experiment Details

### Implementation Details

We load pretrained model GPT-2 (Radford et al., 2019) as initial parameters for unconditional generation and story generation, and pretrained BARTbase Lewis et al. (2020) for summarization and paraphrasing generation tasks. For the summarization and paraphrasing generation, we keep the encoder-decoder attention block. No encoder-decoder attention is used in unconditional generation and story generation tasks. The number of layers and dimensions of hidden states in Della is consistent with the configurations of corresponding pretrained models (GPT-2 has 12 layers and Bart-base has 6-layer encoder and 6-layer decoder. The hidden size of both is 768). We use the state of a special token to obtain the representation in the encoder. We utilize cyclical annealing tricks to train Della and other VAE baselines. Specifically, two epochs are one annealing period. In one period, \(\beta\) (the weight of KL term in ELBO) keeps 1e-5 in the first half, then linearly increases to 1 in the next quarter, then keeps at 1 for the last quarter. We select batch size over \(\{16,32\}\) and learning rate over {5e-5, 7e-5}. We use beam search for Della and top-k sampling for compared baseline models for the unconditional generation and story generation. For the summarization and paraphrasing generation, we use beam search in all the models.

We implement Della and other VAE baselines based on Huggingface Transformers Wolf et al. (2020) library of v4.10.0 and use NVIDIA GeForce RTX 3090 to train our model. The total number of training GPU hours on different datasets is in Table 6. The number of parameters for our model is 193,353,984 in the unconditional generation setting and 195,180,114 in the conditional generation one. All experimental results are trained and tested in a single run.

### Datasets Details

The detailed dataset statistics are in Table 7. For the licenses of the datasets we use, CNN/DM and WritingPrompts use MIT License, while SNLI uses CC BY-SA 4.0. Meanwhile, PTB, Quora, and Yelp use their own license: LDC User Agreement, Yelp Data Agreement, and Quora's Terms of Service, respectively. All of these licenses and agreements allow their data for academic use. Unfortunately, we did not find the license for the Yahoo Dataset.

### Metrics Details

Here we provide more details of the metrics used in our experiments.

**Perplexity (PPL)**. PPL \(=p(\mathbf{x})^{-1/n}\) is commonly used to evaluate the performance of language models, where \(n\) is number of tokens \(\mathbf{x}\) contains. For VAE-based model, we can only obtain the lower bound of \(\log p(\mathbf{x})\). We consider \(k\) latent variables \(z_{1},z_{2},\ldots,z_{k}\) sampled from the posterior distribution \(q(\mathbf{z}_{i}|\mathbf{x})\). Based on the fact that average importance weights are an unbiased estimator of \(\log p(\mathbf{x})\)Burda et al. (2016) and Jensen's Inequality, we have:

\[\mathcal{L}_{k} =\mathbb{E}\left[\log\frac{1}{k}\sum_{i=1}^{k}\frac{p(\mathbf{x},\mathbf{ z}_{i})}{q(\mathbf{z}_{i}|\mathbf{x})}\right] \tag{13}\] \[\leq\log\mathbb{E}\left[\frac{1}{k}\sum_{i=1}^{k}\frac{p(\mathbf{ x},\mathbf{z}_{i})}{q(\mathbf{z}_{i}|\mathbf{x})}\right]=\log p(\mathbf{x}).\]

We use \(\mathcal{L}_{k}\) to estimate \(\log p(\mathbf{x})\) and calculate PPL.

**ELBO**. The ELBO is the sum of reconstruction loss and KL divergence.

**KL**. The KL divergence of the posterior and prior distribution.

**Mutual Information(MI)**Alemi et al. (2016).

\begin{table}
\begin{tabular}{l|c} \hline \hline Dataset & Training Time \\ \hline Yelp & 20h \\ Yahoo & 20h \\ PTB & 6h \\ SNLI & 12h \\ CNN/DM & 40h \\ WP & 170h \\ Quora & 5h \\ \hline \hline \end{tabular}
\end{table}
Table 6: GPU hours of training Della with RTX3090

\begin{table}
\begin{tabular}{l|c|c|c|c} \hline \hline Dataset & \# Train & \# Dev & \# Test & Average Length \\ \hline Yelp & 100k & 10k & 10k & 96 \\ Yahoo & 100k & 10k & 10k & 79 \\ PTB & 42k & 3k & 3k & 21 \\ SNLI & 100k & 10k & 10k & 10 \\ \hline CNN/DM & 287k & 13k & 11k & S: 790 & T: 61 \\ WP & 272k & 15k & 15k & S: 28 & T: 674 \\ Quora & 134k & 5k & 10k & S: 10 & T: 10 \\ \hline \hline \end{tabular}
\end{table}
Table 7: **Statistics of datasets**. We present the size of train/dev/test sets and the average length for 7 datasets. S means source text and T means target text.

 

Mutual Information \({\cal I}(\mathbf{x},\mathbf{z})\) is defined as:

\[{\cal I}_{q}(\mathbf{x},\mathbf{z}) \tag{14}\] \[= \mathbb{E}_{p(\mathbf{x})}\mathbb{E}_{q(\mathbf{z}|\mathbf{x})}\log q(\mathbf{z}| \mathbf{x})-\mathbb{E}_{q(\mathbf{z})}\log q(\mathbf{z})\]

where \(q_{(}\mathbf{z})=\mathbb{E}_{p(\mathbf{x})}q(\mathbf{z}|\mathbf{x})\) is called the aggregated posterior.

**Activate Units(AU)**Burda et al. (2016). AU is the active units in latent varibles, defined as \(A_{\mathbf{z}}=\text{Cov}_{\mathbf{x}}(\mathbb{E}_{z\sim q(\mathbf{z}|\mathbf{x})}[z])>\delta\), where \(\delta\) is a threshold, commonly set as 0.01. However, we find that with \(\delta=0.01\), all VAE models in our experiments have full active unit. So we increase the threshold to 0.2 to distinguish the performance of different models on this metric. **Please note** that Della incorporates latent variables in all layers, and hence we calculate AU for the latent variable in each layer and then report the average.

**BLEU**Papineni et al. (2002). BLEU measures the n-gram overlap of generated sequences and the reference ones. For unconditional setting, we regard all samples in the test set as references to each generated example.

**CND**Li et al. (2020b). CND approximates the divergence of the empirical reference distribution and generated text distribution in n-gram spaces.

**MAUVE**Piillutla et al. (2021). MAUVE measures the gap between reference text and generated text using divergence frontiers.

**Self-BLEU**Zhu et al. (2018). Self-Bleu calculates the BLEU score on the generated samples, which averages the BLEU score of each generated sequence calculated with other generated ones as references. This metric measures the diversity of a set of generated sequences. Higher Self-BLEU means these generated sequences are more distinguishable from each other.

**Dist**Li et al. (2016). Dist measures the proportion of distinct n-grams on generated samples.

**Jaccard Similarity(JS)**Wang and Wan (2018). JS calculates the average n-gram Jaccard similarity between every two generated sequences.

**Rouge**Lin and Hovy (2002). Rouge computes n-gram overlap of generated examples with given target samples. We use rouge-score v0.0.4 to evaluate the rouge score of our model and the baselines.

**BERTScore**Zhang et al. (2020). BERTScore uses pre-trained BERT Devlin et al. (2019) to obtain the vector representations of generated and reference text and calculates their cosine similarity. We use bert-score v0.3.10 to calculate the BERTScore of our model and the baselines.

### Human Evaluation Details

Due to the relatively long length of generated text, we randomly sample 30 examples in the test set of WP and CNN/DM as input to Della and other compared baseline models to generate the target. We invite five graduate students proficient in English to score the generated text. The criteria for story generation include fluency, coherence, and novelty, and the criteria for summarization generation include informativeness, consistency, and novelty. Specifically, fluency measures whether the generated sentences are syntactically fluent; coherence measures whether the generated text is logically structured and consistent with the input text; novelty measures whether the content is novel and attractive; informativeness measures to what extent the generated summarization summarizes the general idea of the article.

When conducting the human evaluation, we informed the participants as follows:

* The following contents are generated by the automatic models. Some of them may be offensive or contain improper arguments. Please be conscious of these risks and evaluate these contents equitably and adequately.
* The evaluation you provide will be used only for academic use and will never be used commercially.

Every evaluator will sign their signature below these warnings to confirm that they have read those words. After finishing the annotation, they will receive $25. This amount is determined by the time of the whole annotation process and the estimation of average hourly income. The ethics review board for data collection protocol is not essential in our country, so we did not conduct this review for our data collection protocol.

### Additional Experimental Result

Table 8 and Table 9 report the results on PTB, SNLI and Quora dataset.

### Case Study Details

We take two sentences \(\mathbf{x}_{1}\) and \(\mathbf{x}_{2}\) and sample two groups of latent variables \(\mathbf{z}^{(1)}=\{\mathbf{z}^{(1)}_{1},\mathbf{z}^{(1)}_{2},\ldots,\mathbf{z}^{(1)}_{L}\}\) and \(\mathbf{z}^{(2)}=\{\mathbf{z}^{(2)}_{1},\mathbf{z}^{(2)}_{2},\ldots,\mathbf{z}^{(2)}_{L}\}\) from posterior distributions \(p(\mathbf{z}^{(1)}|\mathbf{x}_{1})\) and \(p(\mathbf{z}^{(2)}|\mathbf{x}_{2})\). We obtain the weighted latent variables \(\hat{\mathbf{z}}=\{\hat{\mathbf{z}}_{1},\hat{\mathbf{z}}_{2},\ldots,\hat{\mathbf{z}}_{L}\}\) by 

taking weighted sum at each corresponding element in two groups, i.e. \(\hat{\mathbf{z}}_{i}=\tau*\mathbf{z}^{(1)}_{i}+(1-\tau)*\mathbf{z}^{(2)}_{i}\). The mixed sentence \(\hat{\mathbf{x}}\) is generated conditioned on \(p(\hat{\mathbf{x}}|\hat{\mathbf{z}})\) by the decoder.

### Potential Risks and Limitations of our work

Due to the unclean corpus (especially in the WP dataset) we use where slang repeatedly appears, the model training on this corpus may also output some rude expressions during generation. Also, the text generated in the unconditional generation task is not controllable, which may contain some bias or politically sensitive expression. Besides, since our model significantly improves the quality and diversity of generated, it can produce more plausible texts like news, which could be possibly utilized to create fake news or disinformation. However, on the other hand, our model could benefit fairness in language generation. Previous text generation models tend to produce biases like gender or nationality biases, which means only the majority would be appropriately described while the minority may be ignored. These biases are mainly caused by the biased training corpus. With the same data, our model can improve the diversity of generated text, which is also potential for mitigating these biased. We will try to develop debiased language generation systems in future work to avoid these risks harming society.

While Della shows good performance on text generation, it has one limitation: training efficiency. Della brings more parameters compared with three baseline methods. Training efficiency needs to be considered if we further explore the performance of Della on the large pretrained model 



\begin{table}
\begin{tabular}{c|c|c|c|c|c|c|c|c|c|c|c} \hline \multirow{2}{*}{Model} & \multicolumn{6}{c|}{Representation Learning} & \multicolumn{4}{c|}{Generation Quality} & \multicolumn{3}{c}{Generation Diversity} \\ \cline{2-15}  & PPL\(\downarrow\) & ELBO\(\downarrow\) & KL\(\uparrow\) & MI\(\uparrow\) & AU\(\uparrow\) & BLEU\(\uparrow\) & CND\(\downarrow\) & MAUVE\(\uparrow\) & SB\(\downarrow\) & Dist\(\uparrow\) & JS \(\downarrow\) \\ \hline \multicolumn{10}{c}{Dataset: PTB} \\ \hline GPT-2 & 25.80 & - & - & - & - & 27.91 & 1.12 & **0.73** & 41.55 & 37.79 & 0.30 \\ Optimus & 22.79 & 344.10 & 15.09 & 7.67 & - & - & - & - & - & - \\ Embed & 19.98 & 327.28 & 4.77 & 4.14 & 6 & 28.04 & 1.38 & 0.69 & 41.32 & 34.46 & 0.33 \\ Memory & 24.41 & 90.25 & 1.22 & 1.17 & 4 & 21.31 & 1.21 & 0.58 & 26.58 & 38.28 & **0.08** \\ Softmax & 24.04 & 90.63 & 2.13 & 1.89 & 21 & **28.59** & 1.39 & 0.72 & 42.15 & 33.91 & 0.30 \\ Della & **10.28** & **58.43** & **12.46** & **12.35** & **22** & 28.15 & **0.63** & 0.68 & **24.87** & **41.84** & 0.17 \\ \hline \multicolumn{10}{c}{Dataset: SNLI} \\ \hline GPT-2 & 20.19 & - & - & - & - & **63.57** & 1.95 & 0.71 & 75.34 & 19.11 & 0.58 \\ Optimus & 16.67 & 38.50 & 16.35 & 8.89 & - & - & - & - & - & - \\ Embed & 13.79 & 32.97 & 3.24 & 3.16 & 20 & 59.26 & 0.98 & **0.72** & 65.59 & 20.89 & 0.44 \\ Memory & 13.78 & 32.62 & 2.13 & 2.08 & 10 & 62.80 & 1.24 & 0.67 & 54.59 & 21.87 & 0.33 \\ Softmax & 14.21 & 33.18 & 2.70 & 2.65 & 16 & 60.51 & 1.94 & 0.71 & 71.84 & 18.59 & 0.57 \\ Della & **5.13** & **10.23** & **5.86** & **16.58** & **23** & 62.94 & **0.85** & 0.69 & **36.85** & **32.61** & **0.21** \\ \hline \end{tabular}
\end{table}
Table 8: Additional results for language model and unconditional generation task. The results of Optimus are copied from original paper with \(\lambda=0.5\).

\begin{table}
\begin{tabular}{l|c|c|c|c|c|c|c|c} \hline Model & BLEU\(\uparrow\) & Rouge-1\(\uparrow\) & Rouge-2\(\uparrow\) & Rouge-L\(\uparrow\) & Bertscore\(\uparrow\) & KL\(\uparrow\) \\ \hline Bart-base & 64.34 & 63.27 & 39.83 & 60.28 & 94.72 & - \\ Embed & 63.94 & 63.12 & 39.42 & 60.22 & 94.66 & 0.0 \\ Mem & 63.78 & 62.86 & 39.18 & 59.96 & 94.65 & 0.0 \\ Softmax & 64.30 & 63.25 & 39.92 & 60.39 & 94.71 & 0.0 \\ Della & **64.40** & **63.80** & **40.58** & **61.03** & **94.84** & **3.88** \\ \hline \end{tabular}
\end{table}
Table 9: Results on Quora dataset. Because the sentences in Quora are quite short and constrained, the results of the three diversity metrics on all baselines are almost the same. So we omit them here.



Additional Proof

### Derivation of KL Divergence of Layer-Wise Latent Variables

KL divergence of layer-wise latent variables

\[KL(q({\mathbf{z}}|{\mathbf{x}})||p({\mathbf{z}})) \tag{15}\] \[= \int q({\mathbf{z}}|{\mathbf{x}})\log\frac{q({\mathbf{z}}|{\mathbf{x}})}{p({\mathbf{z} })}\,\mathrm{d}{\mathbf{z}}\] \[= \int\prod_{l=1}^{L}q({\mathbf{z}}_{l}|{\mathbf{x}},{\mathbf{z}}_{<l})\log \frac{\prod_{l=1}^{L}q({\mathbf{z}}_{l}|{\mathbf{x}},{\mathbf{z}}_{<l})}{\prod_{l=1}^{L}p({ \mathbf{z}}_{l}|{\mathbf{z}}_{<l})}\,\mathrm{d}{\mathbf{z}}_{1}\,\mathrm{d}{\mathbf{z}}_{2} \ldots\mathrm{d}{\mathbf{z}}_{L}\] \[= \sum_{i=1}^{L}\int\prod_{l=1}^{L}q({\mathbf{z}}_{l}|{\mathbf{x}},{\mathbf{z}} _{<l})\log\frac{q({\mathbf{z}}_{l}|{\mathbf{x}},{\mathbf{z}}_{<l})}{p({\mathbf{z}}_{l}|{\mathbf{ z}}_{<l})}\,\mathrm{d}{\mathbf{z}}_{1}\,\mathrm{d}{\mathbf{z}}_{2}\ldots\mathrm{ d}{\mathbf{z}}_{L}\] \[= \sum_{l=1}^{L}\int q({\mathbf{z}}_{<l}|{\mathbf{x}})q({\mathbf{z}}_{l}|{\mathbf{ x}},{\mathbf{z}}_{<l})\log\frac{q({\mathbf{z}}_{l}|{\mathbf{x}},{\mathbf{z}}_{<l})}{p({\mathbf{z}}_{ l}|{\mathbf{z}}_{<l})}\,\mathrm{d}{\mathbf{z}}_{1}\,\mathrm{d}{\mathbf{z}}_{2}\ldots \mathrm{d}{\mathbf{z}}_{l}\] \[= \sum_{l=1}^{L}\mathbb{E}_{q({\mathbf{z}}_{<l}|{\mathbf{x}})}\mathrm{KL}( q({\mathbf{z}}_{l}|{\mathbf{x}},{\mathbf{z}}_{<l})||p({\mathbf{z}}_{l}|{\mathbf{z}}_{<l}))\]

### Proof of Theorem 1

First, we consider on term in the summation and can obtain:

\[\begin{split}&\quad\mathbb{E}_{p({\mathbf{x}})}\mathbb{E}_{q({\mathbf{z} }_{<l}|{\mathbf{x}})}[\mathrm{KL}(q({\mathbf{z}}|{\mathbf{x}},{\mathbf{z}}_{<l})||p({\mathbf{z}}_{ l}|{\mathbf{z}}_{<l}))]\\ =&\int q({\mathbf{x}})q({\mathbf{z}}_{<l}|{\mathbf{x}})q({\mathbf{z} }_{l}|{\mathbf{x}},{\mathbf{z}}_{<l})\log\frac{q({\mathbf{z}}_{l}|{\mathbf{x}},{\mathbf{z}}_{l})}{ p({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})}\,\mathrm{d}{\mathbf{x}}\,\mathrm{d}{\mathbf{z}}_{l}\, \mathrm{d}{\mathbf{z}}_{<l}\\ =&\int q({\mathbf{x}},{\mathbf{z}}_{l},{\mathbf{z}}_{<l})\log \frac{q({\mathbf{z}}_{l}|{\mathbf{x}},{\mathbf{z}}_{l})}{p({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})}\, \mathrm{d}{\mathbf{x}}\,\mathrm{d}{\mathbf{z}}_{l}\,\mathrm{d}{\mathbf{z}}_{<l}\\ =&\int q({\mathbf{x}},{\mathbf{z}}_{l},{\mathbf{z}}_{<l})\log \Big{(}\frac{q({\mathbf{z}},{\mathbf{x}}|{\mathbf{z}}_{<l})}{q({\mathbf{z}}|{\mathbf{z}}_{<l})q({ \mathbf{z}}_{l}|{\mathbf{z}}_{<l})}\frac{q({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})}{p({\mathbf{z}}_{l}| {\mathbf{z}}_{<l})}\Big{)}\,\mathrm{d}{\mathbf{x}}\,\mathrm{d}{\mathbf{z}}_{l}\,\mathrm{d }{\mathbf{z}}_{<l}\\ =&\int q({\mathbf{z}}_{<l})q({\mathbf{x}},{\mathbf{z}}_{l}|{\bm {z}}_{<l})\log\frac{q({\mathbf{z}},{\mathbf{x}}|{\mathbf{z}}_{<l})}{q({\mathbf{x}}|{\mathbf{z}}_{< l})q({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})}\,\mathrm{d}{\mathbf{x}}\,\mathrm{d}{\mathbf{z}}_{l}\, \mathrm{d}{\mathbf{z}}_{<l}+\\ &\quad\int q({\mathbf{x}}|{\mathbf{z}}_{l},{\mathbf{z}}_{<l})q({\mathbf{z}}_{l}| {\mathbf{z}}_{<l})q({\mathbf{z}}_{<l})\log\frac{q({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})}{p({\bm {z}}_{l}|{\mathbf{z}}_{<l})}\,\mathrm{d}{\mathbf{x}}\,\mathrm{d}{\mathbf{z}}_{l}\,\mathrm{ d}{\mathbf{z}}_{<l}\\ =&\int q({\mathbf{x}},{\mathbf{z}}_{l}|{\mathbf{z}}_{<l})\log \frac{q({\mathbf{z}},{\mathbf{x}}|{\mathbf{z}}_{<l})}{q({\mathbf{x}}|{\mathbf{z}}_{<l})q({\mathbf{z}}_{ l}|{\mathbf{z}}_{<l})}\,\mathrm{d}{\mathbf{x}}\,\mathrm{d}{\mathbf{z}}_{l}+\\ &\quad\int q({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})q({\mathbf{z}}_{<l})\log \frac{q({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})}{p({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})}\,\mathrm{ d}{\mathbf{z}}_{l}\,\mathrm{d}{\mathbf{z}}_{<l}\\ =& H({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})-H({\mathbf{z}}_{l}|{\mathbf{z}}_{<l},{\bm {x}})\\ \geq& H({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})-H({\mathbf{z}}_{l}|{\mathbf{z}}_{<l},{\bm {x}})\\ \end{split} \tag{16}\]

where \(H\) is the Shannon entropy. Then, the summation has a lower bound:

\[\begin{split}&\quad\sum_{i=1}^{L}\mathbb{E}_{p({\mathbf{x}})}\mathbb{ E}_{q({\mathbf{z}}_{<l}|{\mathbf{x}})}[\mathrm{KL}(q({\mathbf{z}}|{\mathbf{x}},{\mathbf{z}}_{<l})||p({\mathbf{z}}_{ l}|{\mathbf{z}}_{<l}))]\\ \geq&\sum_{i=1}^{L}H({\mathbf{z}}_{l}|{\mathbf{z}}_{<l})-H({ \mathbf{z}}_{l}|{\mathbf{z}}_{<l},{\mathbf{x}})\\ =& H({\mathbf{z}}_{1},\dots,{\mathbf{z}}_{L})-H({\mathbf{z}}_{1}, \dots,{\mathbf{z}}_{L}|{\mathbf{x}})\\ =& I({\mathbf{x}};{\mathbf{z}}_{1},\dots,{\mathbf{z}}_{L})\end{split}\] (17) 

where \(I\) is mutual information. Next, we prove the following inequality with induction:

\[I({\mathbf{x}};{\mathbf{z}}_{1},\ldots,{\mathbf{z}}_{L})\geq I({\mathbf{x}};{\mathbf{ z}}_{1};\ldots;{\mathbf{z}}_{L}) \tag{18}\]

When \(L=2\), we proof \(I({\mathbf{x}};{\mathbf{z}}_{1},{\mathbf{z}}_{2})\geq I({\mathbf{x}};{\mathbf{z}}_{1};{\mathbf{z}}_{2})\). Actually, we have the following facts:

\[I({\mathbf{x}};{\mathbf{z}}_{1},{\mathbf{z}}_{2}) \tag{19}\] \[= H({\mathbf{x}})+H({\mathbf{z}}_{1},{\mathbf{z}}_{2})-H({\mathbf{x}},{\mathbf{z}}_{1} ,{\mathbf{z}}_{2})\] \[I({\mathbf{x}};{\mathbf{z}}_{1};{\mathbf{z}}_{2})\] \[= H({\mathbf{x}})+H({\mathbf{z}}_{1})+H({\mathbf{z}}_{2})+H({\mathbf{x}},{\mathbf{z}} _{1},{\mathbf{z}}_{2})\] (20) \[-H({\mathbf{z}}_{1},{\mathbf{z}}_{2})-H({\mathbf{x}},{\mathbf{z}}_{1})-H({\mathbf{x}} ,{\mathbf{z}}_{2})\]

Based on the facts above, we have:

\[I({\mathbf{x}};{\mathbf{z}}_{1},{\mathbf{z}}_{2})\geq I({\mathbf{x}};{\mathbf{z}}_{1 };{\mathbf{z}}_{2}) \tag{21}\] \[\Leftrightarrow 2H({\mathbf{z}}_{1},{\mathbf{z}}_{2})+H({\mathbf{x}},{\mathbf{z}}_{1 })+H({\mathbf{x}},{\mathbf{z}}_{2})\geq H({\mathbf{z}}_{1})+H({\mathbf{z}}_{2})+2H({\mathbf{x}},{ \mathbf{z}}_{1},{\mathbf{z}}_{2}) \tag{22}\]

It's true because we have:

\[H({\mathbf{z}}_{1},{\mathbf{z}}_{2})+H({\mathbf{x}},{\mathbf{z}}_{1}) \tag{23}\] \[= H({\mathbf{z}}_{2}|{\mathbf{z}}_{1})+H({\mathbf{x}}|{\mathbf{z}}_{1})+2H({\mathbf{z} }_{1})\] \[\geq H({\mathbf{x}},{\mathbf{z}}_{2}|{\mathbf{z}}_{1})+2H({\mathbf{z}}_{1})\] \[= H({\mathbf{x}},{\mathbf{z}}_{1},{\mathbf{z}}_{2})+H({\mathbf{z}}_{1})\]

Similarly, the following inequality also holds true:

\[H({\mathbf{z}}_{1},{\mathbf{z}}_{2})+H({\mathbf{x}},{\mathbf{z}}_{2})\geq H({\mathbf{ x}},{\mathbf{z}}_{1},{\mathbf{z}}_{2})+H({\mathbf{z}}_{2}) \tag{24}\]

Therefore, making sum to Eq.(23) and Eq.(24), we conclude that \(I({\mathbf{x}};{\mathbf{z}}_{1},{\mathbf{z}}_{2})\geq I({\mathbf{x}};{\mathbf{z}}_{1};{\mathbf{z}}_{2})\). Hence, we finish the proof of the \(L=2\) case.

When \(L=k\), suppose \(I({\mathbf{x}};{\mathbf{z}}_{1},\ldots,{\mathbf{z}}_{k})\geq I({\mathbf{x}};{\mathbf{z}}_{1}; \ldots;{\mathbf{z}}_{k})\), we consider \(L=k+1\). In this case, based on the inductive assumption, we have:

\[I({\mathbf{x}};{\mathbf{z}}_{1},\ldots,{\mathbf{z}}_{k+1})\geq I({\mathbf{x}};{ \mathbf{z}}_{1},\ldots,{\mathbf{z}}_{k})\geq I({\mathbf{x}};{\mathbf{z}}_{1};\ldots;{\mathbf{z}}_{ k})\geq I({\mathbf{x}};{\mathbf{z}}_{1};\ldots;{\mathbf{z}}_{k+1}) \tag{25}\]

Hence, the case of \(L=k+1\) also holds true. Therefore, we conclude that \(I({\mathbf{x}};{\mathbf{z}}_{1},\ldots,{\mathbf{z}}_{L})\geq I({\mathbf{x}};{\mathbf{z}}_{1}; \ldots;{\mathbf{z}}_{L})\).

Now, we consider the interaction information and can obtain:

\[I({\mathbf{x}};{\mathbf{z}}_{1};\ldots;{\mathbf{z}}_{L}) \tag{26}\] \[= I({\mathbf{z}}_{L},{\mathbf{z}}_{L-1})-\sum_{i=2}^{L-1}I({\mathbf{z}}_{L}; \ldots;{\mathbf{z}}_{i}|{\mathbf{z}}_{i-1})-I({\mathbf{z}}_{L};\ldots;{\mathbf{z}}_{1}|{\mathbf{x}})\] \[\geq \sum_{i=2}^{L-1}I({\mathbf{z}}_{L};\ldots;{\mathbf{z}}_{i}|{\mathbf{z}}_{i-1} )-I({\mathbf{z}}_{L};\ldots;{\mathbf{z}}_{1}|{\mathbf{x}})\]

Finally, based on Eq.(16), (17), (25), (26), we can conclude:

\[\mathbb{E}_{p({\mathbf{x}})}[\mathcal{L}_{R}] =\sum_{i=1}^{L}\mathbb{E}_{p({\mathbf{x}})}\mathbb{E}_{q({\mathbf{z}}_{<l }|{\mathbf{x}})}[\mathrm{KL}(q({\mathbf{z}}|{\mathbf{x}},{\mathbf{z}}_{<l})||p({\mathbf{z}}_{l}|{ \mathbf{z}}_{<l}))]\] (27) \[\geq I({\mathbf{x}};{\mathbf{z}}_{1},\ldots,{\mathbf{z}}_{L})\] \[\geq\sum_{i=2}^{L-1}I({\mathbf{z}}_{L};\ldots;{\mathbf{z}}_{i}|{\mathbf{z}}_{ i-1})-I({\mathbf{z}}_{L};\ldots;{\mathbf{z}}_{1}|{\mathbf{x}})\] 

Figure 6: Attention weight of the Memory paradigm for layer 0 to layer 5. We plot three heatmaps in each layer. Average means averaging weights through all head. Max and min means we select the head with max and min attention weight on the memory token (latent variable).We can see the memory token tends to be ignored by most heads especially in lower layers.

 

Figure 7: Attention weight of Memory paradigm for layer 6 to layer 11.