Nonparametric Variational Regularisation of Pretrained Transformers

Authors: fabio  fehr
james  henderson
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 12:07:25

Abstract: Pretrained transformers have demonstrated impressive abilities, but tend not to generalise well out-of-domain and are very expensive to fine-tune on new domain data. Nonparametric Variational Information Bottleneck (NVIB) has been proposed as a regulariser for training cross-attention in transformers, potentially addressing this domain overfitting problem. We extend the NVIB framework to replace all types of attention functions in transformers. We show that existing pretrained transformers can be reinterpreted as nonparametric variational models using an empirical prior distribution and identity initialisation with controllable hyperparameters. We then show that changing the initialisation introduces a novel, information-theoretic post-training regularisation in the attention mechanism, which improves out-of-domain generalisation on NLP tasks without any additional training. This success supports the hypothesis that the way pretrained transformer embeddings represent information is accurately characterised by nonparametric variational Bayesian models.

Main paper: 
 introduction: 
 Self-supervised pretraining of transformer models (Devlin et al., 2019; Lewis et al., 2020; Raffel et al., 2020; Zhang et al., 2022) has been hugely successful, improving performance in virtually every NLP task. This indicates that the inductive bias of attention-based representations is extremely effective for language, but it is not clear which characteristics of transformers are essential to this inductive bias and which are implementation details. We shed light on this issue by developing a variational Bayesian reinterpretation of pretrained transformers, which more explicitly characterises how information about the real data distribution is represented. This insight allows us to improve the models' out-of-distribution generalisation without any additional training.In previous work, Henderson & Fehr (2023) derive a variational Bayesian generalisation of a single-head cross-attention layer of transformers, called Nonparametric Variational Information Bottleneck (NVIB). When used during training, the NVIB layer provides an information-theoretic, sparsity-inducing regulariser over attention-based representations. In this paper, we investigate the possibility that NVIB also provides an accurate theoretical model of existing attention-based models which have been trained without NVIB regularisation. We extend NVIB to all the uses of attention in transformers (single-and multi-head; cross-and self-attention; in encoders and decoders), and propose a method for converting a pretrained transformer into the weights of an equivalent nonparametric variational (NV) Bayesian model (NV-Transformer) . 1This NV-Transformer is Bayesian in that it embeds text into a probability distribution over transformer embeddings, but remains equivalent by adding uncertainty around the embedding computed by the pretrained transformer, as illustrated in Figure 1 . Theoretically, the exact equivalence only occurs when this uncertainty is exactly zero, but empirically we find a practical range of non-zero uncertainty levels where the model's predictions are unchanged (Section 4.1). Continuing training of this initial model with NVIB regularisation would push this uncertainty higher, but training large models is computationally expensive and requires significant amounts of in-domain data. Instead, we only adjust hyperparameters of the initialisation to better satisfy the NVIB regulariser by increasing uncertainty, without any backward-passes or parameter updates.Increasing the uncertainty level does change the model's predictions, but interestingly it does not degrade the model's task accuracy. Instead, it adds a form of post-training regularisation, which accesses a space of models which generate different outputs with the same level of accuracy (Section 4.1). Surprisingly, this post-training regularisation even improves out-of-domain generalisation (Section 4.2). These results show that adding smaller amounts of uncertainty removes unreliable information which does not generalise in-domain and can cause overfitting out-of-domain. This agreement between uncertainty in the Bayesian model and unreliability in pretrained transformer embeddings shows that our Bayesian reinterpretation helps characterise how transformers represent information. Contributions In this work we contribute technically to the development of novel models which use Nonparametric Variational Information Bottleneck regularisation, and to the understanding of pretrained transformers. 1. We extend NVIB regularisation beyond single-head, cross-attention to all forms of multi-head attention in a transformer, resulting in our proposed NV-Transformer architecture. 2. We define a novel NVIB prior using the empirical distribution of a model's embeddings given a small amount of data (Section 3.2). 3. We propose a reinterpretation of pretrained transformers as variational Bayesian models by defining a novel identity initialisation for NVIB with its controllable hyperparameters (Section 3.1). Adjusting these hyperparameters regularises the embeddings of this equivalent NV-Transformer without requiring retraining. 4. Empirically, we show the usefulness of these proposals by smoothly varying the amount of this post-training regularisation (Section 4.1) and achieving improved performance in out-of-domain text summarisation and translation (Section 4.2). 5. This successful Bayesian interpretation helps characterise the way that pretrained transformers represent reliable information in their embeddings (Section 5).2 Background Henderson & Fehr (2023) propose NVIB as an information-theoretic regulariser for attention layers. It extends vector-space Variational Information Bottleneck (VIB) (Alemi et al., 2017) to the unboundedly large set-of-vectors representations supported by attention. Henderson & Fehr (2023) define a Variational Auto-Encoder (Kingma & Welling, 2014) for transformers with NVIB applied to its cross-attention layers. Behjati et al. (2023) apply NVIB to stacked transformer self-attention layers in the encoder, which when trained learns increasingly sparse representations across the layers of the model. They found that the stacked NVIB layers abstracted the character-level inputs into word-like units, when applied to the final layers of the encoder.In variational Bayesian models like NVIB layers, the information which the input conveys about the latent representation is specified with a distribution over latent representations.For attention-based representations, these distributions need to capture the essential structure imposed by the attention function on the sequence of vectors being accessed by attention, namely that these vectors are permutation invariant, have a normalised weighting, and can be arbitrarily many. A key insight of Henderson & Fehr (2023) is that these properties can be captured by interpreting attention-based representations as nonparametric mixture distributions. The accessed vectors are interpreted as specifying a mixture of impulse distributions (Equation (1) below). This is an equivalent representation in that we can define an exactly equivalent attention function in the form of Bayesian query denoising using this distribution as the prior. The query is interpreted as being corrupted by Gaussian noise, the attention weights are the query posterior's mixture weights given the query observation, and the output of the attention function is the expected value of this posterior distribution. The authors define a generalisation of attention known as denoising attention (Equation (2) below), which generalises the attention function to query denoising with any mixture distribution as the representation being accessed.Since the number of vectors in a transformer embedding grows to accommodate the length of the input text, the latent space of equivalent mixture distributions is nonparametric in nature (i.e. there is no fixed set of parameters which can specify the entire space of mixture distributions). Thus, Henderson & Fehr (2023) define distributions over these generalised attention-based representations using Bayesian nonparametrics (Equation (3) below). In the remainder of this section we provide more details about denoising attention and NVIB's distributions over mixture distributions.To show that the denoising attention function is a generalisation of the standard attention function, Henderson & Fehr (2023) provide a constructive proof of exact equivalence. The standard attention function projects the accessed set of vectors Z ∈ R n×d via weight matrices W K , W V ∈ R d×d to keys and values, respectively, and projects the accessing input vector u ′ ∈ R 1×d via the weight matrix W Q ∈ R d×d to a query. It uses the keys' dimensionality d for scaling. This scaled dot product attention function can be regrouped into a core dot product attention function Attn(u, Z) in which all operations are done in the space of Z.Attention(u ′ , Z ; W Q , W K , W V ) = Attn(u ′ W Q (W K ) ⊤ , Z) W V = Attn(u, Z) W V where u = (u ′ W Q (W K ) ⊤ ) ∈ R 1×d .The function Attn(u, Z) can then be defined both in terms of a sum over the vectors z i in Z, or in terms of an integral over a distribution which is only non-zero at the z i :Attn(u, Z) = softmax 1 √ d uZ ⊤ Z = DAttn(u; F Z ) F Z = n ∑ i=1 exp( 1 2 √ d ||z i || 2 ) ∑ n i=1 exp( 1 2 √ d ||z i || 2 ) δ z i (1) DAttn(u; F) = v f (v) g(u; v, √ dI) v f (v) g(u; v, √ dI) dv v dv (2)where δ z i is an impulse distribution at z i , f (•)dI) is a noisy observation of some true vector v. This construction shows that any transformer embedding Z has an equivalent mixture of impulse distributions, namely F Z , where denoising attention DAttn(u; F Z ) gives us exactly the same result as attention Attn(u, Z), for all queries u.Given the generalisation of attention-based representations to nonparametric mixture distributions, Henderson & Fehr (2023) propose to use Dirichlet Processes (DPs) to define distributions over this latent space of mixture distributions. Instead of directly computing the latent Z, their variational-Bayesian model computes a DP estimate of the posterior distribution over mixture distributions given the input. More specifically, the model maps the input to a set of pseudo-observations, and then the posterior DP is computed with exact Bayesian inference given a prior DP and these observations. This posterior "q" distribution DP(G q 0 , α q 0 ) consists of a base distribution G q 0 for generating the vectors of the mixture distribution and a pseudo-count α q 0 for generating their mixture weights. The base distribution G q 0 is a mixture with one Gaussian component for the prior "p" distribution DP(N (µ p , I(σ p ) 2 ), α p ) plus one Gaussian component for each i th pseudo-observation (µ q i , σ q i , α q i ) computed by the model:G q 0 = α p α q 0 N (µ p , I(σ p ) 2 ) + ∑ i α q i α q 0 N (µ q i , I(σ q i ) 2 ) ; α q 0 = α p + ∑ i α q i(3)As for other VAEs (Kingma & Welling, 2014), during evaluation Henderson & Fehr (2023) do not sample a mixture distribution F from the posterior DP, but instead use the posterior's mean, which is its base distribution G q 0 . Denoising attention applied to G q 0 can be computed straightforwardly because it is a mixture of Gaussians. This function differs from standard attention in that the component weights α q /α q 0 replace the key biases implicit in the query-key dot product, and an interpolation between queries and values replaces the values in the attention-weighted average.EQUATIONwhere (σ r i ) 2 = (√ d + (σ q i ) 2 ), and 1 d is a row vector of d ones. When the attention bias term b is large across the non-prior components and thus down-weights the prior component, we can have an identical softmax to attention. When the posterior variance (σ q ) 2 is approximately zero, we get no interpolation between the query and value and get an identical value computation to attention. Henderson & Fehr (2023) propose this formulation of attention only for the case of single-head cross-attention during evaluation.Related work in post-training regularisation such as quantisation (Dettmers et al., 2022; Yao et al., 2022; Xiao et al., 2023; Frantar et al., 2023b) and sparsity (Hubara et al., 2021; Frantar et al., 2023a; 2022) focus on regularising the model weights and not the latent embeddings. Similar to Frantar et al. (2023a) , who propose a data-driven way to sparsify a model in one-shot without any retraining, we propose a data-driven form of soft sparsity for attention, without any retraining. Moreover, our method's regularisation uses information theory to regularise post-training. Our work shares similarity to the out-of-distribution generalisation literature in learning variational models with invariant features (Ilse et al., 2020) and parameter sharing in the embedding space to adapt to domain-shifts (Muandet et al., 2017; Li et al., 2017; Blanchard et al., 2021) . We take a Bayesian approach and define our prior distribution based on empirical dataset statistics that allows for properties which are robust to domain shift. Park & Lee (2021) proposes a way to define pretrained language models as variational models by fine-tuning. Our work does not fine-tune or update the original model weights, nor does it train additional weights as in adapters (Houlsby et al., 2019; Hu et al., 2022) . We keep the model weights frozen and map attention layers to NVIB layers, setting just a few hyperparameters based on validation performance. 
a nonparametric variational reinterpretation of transformers: 
 We propose a novel method for reinterpreting pretrained transformers as nonparametric variational Bayesian models. To support this construction, we implemented the following applications of NVIB: multihead denoising attention (Appendix B); encoder denoising self-attention (Appendix B.3); and decoder denoising causal attention (Appendix B.4).Pseudocode for denoising attention is given in Appendix B.5. This allows us to apply NVIB to every form of attention used in standard transformers, resulting in our proposed NV-Transformer reinterpretation of pretrained transformers. Using these extensions, we propose an identity initialisation with controllable hyperparameters, which allows us to achieve an equivalence with the standard attention mechanisms in transformers (Section 3.1). We then define a novel empirical prior distribution, which introduces uncertainty in the latent representations in a way that captures the implicit uncertainty in pretrained transformers (Section 3.2).We define an identity initialisation for NVIB such that denoising attention is effectively equivalent to standard attention, inspired by Equation 1. Given a set of vectors Z input to the attention layer, our proposed NVIB layer converts it to the parameters (µ, σ, α) which specify a posterior "q" DP distribution, and then applies the evaluation denoising attention function (Equation 4) to the resulting base distribution G q 0 . Excluding the prior component (discussed in Section 3.2), we define the projections:µ = µ(Z) = ZW µ + b µ (6) σ 2 = σ 2 (Z) = exp(ZW σ + b σ ) (7) α = α(Z) = exp(Z 2 w α 1 + Zw α 2 + b α ) (8)where Z 2 is the component-wise square. We choose these forms of projections because we want to set its parameters,W µ , W σ ∈ R d×d , b µ , b σ ∈ R d , w α 1 , w α 2 ∈ R d and b α ∈ R,so that the evaluation denoising attention function is effectively equivalent to standard attention over Z. For equivalence between standard attention and denoising attention, we must remove the influence of the prior component by down-weighting it, and skew the queryvalue interpolation to the value by reducing the uncertainty of the non-prior components. Alternatively, we can see this requirement as making the mixture of Gaussians G q 0 equivalent to the mixture of impulse distributions F Z in Equation 1. See Figure 1 for a visual depiction. We can down-weight the prior component by making the pseudo-counts α(Z) for the nonprior components be large. We can make the non-prior component Gaussians look like the impulse distributions of F Z by setting the mean projection µ(.) to be the identity and the variance projection σ 2 (.) to be approximately zero (exactly zero is not possible with the exponential activation function). Then to make the mixture weights the same as for F Z in Equation 1, the pseudo-count projection α(.) must be proportional to the exponent of the scaled L 2 2 norm of Z, which is why it has a log-quadratic form in Equation 8.We define the parameters of these projections to achieve these requirements for equivalence but still allow some uncertainty. Based on preliminary experiments, we propose a small number of initialisation hyperparameters which controls the transition between equivalence and a smooth regularisation of the embeddings (illustrated in Figure 1 ). The hyperparameters τ i α and τ i σ control the level of uncertainty of the mixture distributions by changing the initialisation of the pseudo-counts and variance, respectively. The indicator i is for the different sections of the model, which allows for independent control of the encoder's self-attention (e), decoder's cross-attention (c) and decoder's causal self-attention (d). Empirically, this allows for more flexibility than a single hyperparameter and is practically easier to tune than defining a hyperparameter per layer. In general we define the layer projection weights as follows:EQUATION)where d and h denote the model projection size and number of attention heads. As discussed in the next section, the empirical distribution is used to define the prior standard deviation σ p , and the constant ϵ α , which denotes the empirical standard deviation of the scaled L 2 2 -norm in log space per layer. The hyperparameter τ i α in Equation 10 controls the relative weight of the prior component to the non-prior components in the mixture distribution G q 0 . Since ϵ α is a standard deviation, it reflects the normal range of values for the non-prior-component pseudo-counts, so we can use τ i α to control the weight given to the prior component with respect to this range. When τ i α = 0, the non-prior pseudo-counts are their scaled L 2 2 -norms and the prior pseudo-count is the expected scaled L 2 2 -norm. When we increase or decrease the τ i α , it adjusts the magnitude of the non-prior L 2 2 -norms, which relatively decreases or increases the weight on the prior proportionately to the standard deviation of the scaled L 2 2 -norm. The hyperparameter τ σ (Equation 9) controls the interpolation between the query and value (Equation 4), which we set proportionately to the variance of the prior distribution. When τ σ ≈ 0, there is effectively no interpolation, as with standard attention. When τ σ = 1, the uncertainty is increased to the level of the empirical prior distribution.The prior distribution is a Dirichlet Process, so it views all vectors in all transformer embeddings as impulses generated from its own base distribution. Therefore, we can estimate the base distribution of the prior empirically by observing the distribution of vectors given forward passes of the transformer model. Taking a Bayesian approach, the NV-Transformer's prior should reflect the distribution over vectors which the pretrained transformer knows before seeing the input text. This is the distribution observed during training. We compute statistics from the fine-tuned corpora and use them to define our priors. We estimate the prior as the best fit of an isotropic Gaussian distribution to the empirical distribution over latent vectors computed when embedding this training corpus.We can estimate the empirical parameters for an isotropic Gaussian G p 0 ∼ N (µ p , I(σ p )) and the prior's pseudo-count α p 0 by using the latent vectors Z as follows:µ p = ∑ N i Z i N , (σ p ) 2 = ∑ N i (Z i -µ p ) 2 N-1and log(αp 0 ) = ∑ N i ( ∑ d j (Z ij ) 2 2 √ d/h )/N,where N is the total number of tokens in the training corpus, d is the dimension of the embedding and h is the number of attention heads. This allows the prior mean to be the least informative representation in the center of the latent embeddings vector space. The variance is estimated from this mean. The empirical pseudo-count is kept in log space for numerical stability and is the expected scaled L 2 2 -norm of the latent vectors.To further understand the distributions of our latent embeddings, we calculate the distribution of the empirical priors across all layers of the encoder and decoder attention mechanisms (Appendix F.1). Given a model θ *x that has been trained on data x, we analyse priors generated from both in-domain x and out-of-domain summarisation distributions. We find that the distribution of embeddings for the groups: encoder self-attention, decoder cross-attention and decoder causal self-attention, behave differently between and similarly within these attention groups. This observation motivates the choice of grouping the hyperparameters τ i α and τ i σ where i is an indicator for encoder, cross or decoder attention mechanisms. We also consider the data requirements of the empirical prior by calculating the prior from a subsample of the training data. We found that the empirical prior can be created from as few as 0.1% of the training data (≈ 200 examples) and achieve similar performance (Appendix F.2). Hence, the empirical prior is data-efficient as it requires a minimal amount of training data for creation. 
post-training evaluation: 
 To evaluate our Bayesian reinterpretation of transformers, we first show that our proposed identity initialisation results in empirical equivalence to pretrained transformers. We then investigate how this reinterpretation allows us to apply an information-theoretic, posttraining regularisation by changing the initialisation of the NVIB layers. Once the empirical prior is estimated, adapting a pretrained transformer to a new domain simply requires a hyperparameter selection on data from the new domain, which can be done solely with forward-passes of the model. This is an important contribution as the cost of backwardpasses, regularisation and parameter updates becomes increasingly expensive as the models scale. Hence, the experimental design for Nonparametric Variational (NV) regularisation is entirely post-training and would not be fair to compare against fine-tuned methods. We evaluate on the validation set of the original dataset the model was trained on. We refer to this as in-domain. We then evaluate out-of-domain (OOD) generalisation, demonstrating that by changing the initialisation of our model we are able to achieve an informationtheoretic post-training regularisation.We consider models that are already pretrained and fine-tuned, available on Hug-gingFace Wolf et al. (2020) . In our experiments, we do not fine-tune or update the original model weights with data, instead only setting hyperparameters in the initialisation of our NVIB layers, which regularise the attention functions post-training. For a comprehensive application of our attention regulariser, we consider sequence-to-sequence transformer models, which accounts for: encoder self-attention, cross-attention and decoder causal-attention. For summarisation we consider BART Lewis et al. (2020) summarisation models that have already been fine-tuned on either CNN/DailyMail data (See et al., 2017) , or Xsum data (Narayan et al., 2018) . For translation we consider Marian (Junczys-Dowmunt et al., 2018) encoder-decoder translation models already fine-tuned on the OPUS100 (Zhang et al., 2020) high resource language pairs: English-German (En-De) and English-French (En-Fr). For all tasks we use empirical priors created from their respective fine-tuning corpora. We provide model implementation details in Appendix D.2 and E.2. We show empirically that pretrained transformer models can be interpreted as NV models by showing that our identity initialisation (Section 3.1) does not change their predictions. We set the NVIB initialisation parameters τ i α =10 and τ i σ ≈ 0 (1e -38 for Float32) for all i. These down-weight the influence of the prior and reduce the interpolation between the query and value in the attention mechanism, respectively (see Equation (4)). Appendix Table 4 shows that the NV-Transformer initialisation has the same accuracy as the pretrained transformer. More precisely, the leftmost points in Figure 2 being at 100% on the x-axis shows that these two models generate the same outputs.We show that increasing the uncertainty within our regularisation allows access to models with different outputs but equivalent accuracy. We evaluate increasing the uncertainty of our NV models on the validation sets for the same datasets they were trained on. To characterise the range of these regularised NV models, Figure 2 plots them in a line ordered by their linear interpolations between the equivalent-initialisation and the everything-overregularised corners of the sample space. The NV-Transformers are plotted by comparing their output using Rouge-L (top) or BLEU (below) against the gold summary or translation (y-axis) and against the original non-NV model's output (x-axis). As we increase the regularisation from the identity initialisation we discover a space of models which are not only different from the original baseline model (x-axis of < 100% overlap), but also equally good (same y-axis performance). The inclusion of our regularisation gives rise to a smooth transition between these models.Examining the attention maps (Appendix F.3) of the summarisation encoder and decoder, we see that the prior component is included in the keys of attention, but the identity initialisation down-weights this prior to the point that there is no attention weight to this component. When regularised, attention weight is redirected from tokens such as punctuation, to the prior component.We evaluate whether our post-training information-theoretic regulariser improves out-ofdomain performance in the presence of a text domain shift. This is an increasingly important problem as pretrained models become larger and more challenging to fine-tune given a shift in the data domain. To evaluate our regulariser's OOD generalisation we define the following experiment: given a model θ *x that has been trained on a text domain x, we evaluate the model on the same task in the presence of a domain-shift to the text domain y. For example, given a pretrained news summarisation or translation model θ *x we evaluate the OOD generalisation by shifting the domain from news x to informal dialogues y without further training.For baselines we consider the original model (32-bit) and quantisation of 16-bit and 8-bit (Dettmers et al., 2022) . Quantisation is an alternative form of post-training regularisation which has been shown to improve performance (Dettmers & Zettlemoyer, 2022) . We also considered a 4-bit quantisation model and combining NVIB regularisation with 11) . When combined, we found improvements over the original models. We did not consider sparsity techniques as they were not easily available without further training. In the following sections we consider sequence-to-sequence models trained on summarisation and machine translation. Summarisation For our summarisation experiments we evaluate the models on the validation sets of out-of-domain summarisation corpora from either Xsum or CNN/DailyMail, respectively, and from Curation Corpus (CC) (Curation, 2020), SAMsum (Gliwa et al., 2019) and WikiHow (Koupaee & Wang, 2018). We provide further dataset details in Appendix D.1. To evaluate on the test set we select the best model on the validation set (Appendix D.3). Table 1 reports the test set Rouge-L on the OOD text summarisation datasets for the post-training regularisation methods. Considering the original model trained on CNN/-DailyMail, we notice that with NVIB regularisation the OOD performance on the test set improves over the baselines consistently, although minorly. In contrast, the original model trained on Xsum is substantially improved when NVIB regularisation is applied to OOD text summarisation. We speculate the larger improvement in the model trained on Xsum is due to the abstractive nature of the training data, whereas the model trained on CNN/DailyMail is more extractive. This shows that our information-theoretic, post-training regularisation can improve OOD generalisation, both over the original model and over quantisation.We conjecture that the length of the summary is a strong proxy for the information within the document. Since the baseline model has been trained to produce a summary of a certain length, the summaries on OOD data can be improved by adapting the length of the summary to the information content. We 3 (3 left) shows consistent improvements across the Rouge-1, Rouge-2 and Rouge-L distributions, which measures overlap metrics against the gold summaries. Figure 3 (1 right) plots the output summary lengths against the true empirical summary lengths, and gives the Spearman's correlation coefficient. This suggests that the NV model is better at generalising to the length of the test summary according to the information variation in the input document. In Appendix Figure 5 we provide these plots for all OOD validation datasets. We provide examples of generated summaries of all our summarisation models and in-and out-of-domain datasets in Appendix F.4.We evaluate the translation models on the English to German (En-De) and English to French (En-Fr) validation sets of out-of-domain corpora Bible (Christodoulopoulos & Steedman, 2014), IWSLT (Cettolo et al., 2017) and TedTalks (Cettolo et al., 2012) . See Appendix E.1 for further dataset details. We find the best model on the validation set (Appendix 11) and then evaluate this model on the test set. Table 3 reports the test set BLEU on the OOD translation datasets. We notice fairly consistent minor improvements (5 out of 6 cases) across all OOD translation datasets. We speculate that the improvements in translation have a lower magnitude than summarisation, as this task benefits less from a reduction in information. These results further support our hypothesis that our post-training, information-theoretic regulariser improves OOD generalisation of pretrained transformers. 
discussion: 
 Since regularisation does not guarantee improved performance, it is surprising that our method, which is applied post-training without backpropagation or parameter updates, can improve performance. We believe the variance in performance is attributed to the nature of the different tasks. Summarisation is fundamentally an information compression task, well suited to an information bottleneck, unlike translation. We speculate the larger improvement in summarisation models trained on Xsum is due to the abstractive nature of the dataset. The Xsum summarisation dataset requires more generalisation than the extractive CNN/DailyMail dataset, which relies on selecting a summary from the document.The improvements from post-training regularisation suggest that our variational Bayesian reinterpretation of pretrained transformers is an accurate model of how information is 1, 3 ). Since this regularisation is done post-training, this argument suggests that the pretrained model has learned what information is unreliable, and that NVIB regularisation exposes this representation of information. In particular, our model suggests that the information in a given embedding dimension value is relative to the distribution over that dimension during training, with the mean value carrying no information and the scale being relative to the variance. They also suggest that the L 2 2 -norm of a vector reflects the importance of the information in a vector. We anticipate that these insights into how transformer embeddings represent information will help in the understanding and development of future improvements to transformer architectures. 
conclusion: 
 This work contributes both to the development of novel models which use Nonparametric Variational Information Bottleneck regularisation, and to the understanding of pretrained transformers. Technically, we contributed an extension of NVIB to all forms of multi-head attention, a novel empirical prior, and an effective identity initialisation with controllable hyperparameters. This development of the Nonparametric Variational Transformer architecture allows us to propose a Bayesian reinterpretation of pretrained transformers. Empirically, we analyse and show the usefulness of this smooth, post-training regularisation by improving the performance of existing pretrained transformers in out-of-domain text summarisation and translation, without any fine-tuning. This success suggests that NVIB provides an accurate characterisation of how pretrained transformers represent information in their embeddings. 
 
 Appendix: In future work, we plan to evaluate our post-training regularisation on different modalities and tasks, such as: encoder-only classification; and decoder-only language modelling. The success of our regularisation in this post-training scenario suggests that even better generalisation could be achieved by fine-tuning with NVIB regularisation, for which our proposed mapping to a regularised NV-Transformer should be an effective initialisation. With Large Language Models (LLMs), their scale makes it more valuable to regularise them post-training, and makes regularisation more relevant during training. Moreover, understanding how transformer embeddings represent information could improve diversity in LLM generation, since sampling can be done directly in the embeddings of the model. We leave application of NVIB regularisation to LLMs for future work because it is non-trivial to connect denoising attention to relative position encodings (Su et al., 2024) and hardware specific implementations such as flash-attention (Dao et al., 2022) .This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.In this work, we limit our claims and evaluation to only consider the post-training paradigm. In this, we exclude sparsity and further fine-tuned baselines as they require considerable computation due to the parameter updates and backward-passes. Moreover, we limit our choice of models to only consider sequence-to-sequence, encoder-decoder transformers. These models were selected as they comprehensively account for all forms of attention: encoder self-attention, cross-attention and decoder causal-attention. Hence, we exclude decoder-only models, such as many large language models (LLMs). We limit our experimentation to Natural Language Processing (NLP) tasks such as summarisation and translation. Further experimentation in the domain of vision and speech is beyond the scope of this work. We plan to address these limitations in future work.In this section, we provide the details for denoising multihead attention at training and evaluation time. We define the set of transformer latent embedding vectors as Z ∈ R n×d and set of pre-projected queries as U ′ ∈ R m×d . We assume the latent vectors are square such that WQ , W K , W V ∈ R d×d and biases b Q , b K , b V ∈ R d areused to linearly project to the queries, keys and values, respectively. We define the standard attention weights before the softmax as follows:A = 1 √ d (U ′ W Q + b Q ) Q (ZW K + b K ) ⊤ K ⊤ ∈ R m×n (11)Typically, for multihead attention the projected query Q and keys K are split into heads.In this definition, we split the linear projections by a divisible number of heads h suchthat W Q , W K , W V ∈ R h×d× d h and biases b Q , b K , b V ∈ R h× d h , so that Q ∈ R h×m× d h and K ∈ R h×n× d h .We can then specify multihead attention by defining a matrix of attention scores A ∈ R h×m×n , for each head i:A i = 1 √ d/h ((Q i (W K i ) ⊤ Z ⊤ + Q i (b K i ) ⊤ )where the bias term Q i (b K i ) ⊤ ∈ R m is added across all n keys, and thus is normalised out in the softmax below. The scaling term also considers the heads and is division by√ d/h.For denoising attention, each head's query is projected into the space of the original set of vectors Z, namely U i =Q i (W K i ) ⊤ , and so is still in R m×d . Thus, each head can be viewed as doing denoising attention in the same way as single-head attention, with the only difference being that the variance of the theoretical query noise is now √ d/hI.In this work we do not do any training, but for completeness we include the equations for multihead attention at training time. The NVIB layer outputs the isotropic Gaussian parameters µ ∈ R (n+1)×d , σ ∈ R (n+1)×d and Dirichlet parameters α ∈ R (n+1) , which include the (n+1) th component for the prior. During training, these are used for sampling such that π ∼ Dir(α) and Z ∼ N (µ, σ). Given these sampled weights and vectors, the training-time denoising attention function is the same as the standard attention function with two changes: (1) the keys come from the sampled vectors Z ∈ R (n+1)×d , which include a vector sampled from the prior component; and ( 2) each key has an attention bias b ∈ R (n+1) which is determined by its weight π ∈ R (n+1) . Summing over heads i, the training-time denoising attention function is:DAttn(.) = ∑ i Softmax(A i + log(π) -1 2 √ d/h ∥Z∥ 2 b ) (ZW V i + b V i ) V i )The biases b are defined by adding the log of the sampled weights log(π) ∈ R (n+1) from the NVIB layer and subtracting the scaled L 2 2 -norms of the sampled vectors 1 ) . For multihead attention we only need to reuse the same biases b for each head, just like we reuse the same vectors Z for each head.2 √ d/h ∥Z∥ 2 ∈ R (n+1During evaluation, as for training, the NVIB layer outputs the isotropic Gaussian parameters µ ∈ R (n+1)×d , σ ∈ R (n+1)×d and Dirichlet parameters α ∈ R (n+1) . For evaluation, the denoising attention function is similar to the training-time denoising attention function with three changes: (1) the parameters are used directly without sampling; (2) there is a different attention bias c ∈ R (n+1) for each of the (n+1) input vectors, including the prior; and (3) the attention values are computed with an interpolation between the query and value. We can write the denoising attention scores A ∈ R h×m×(n+1) , for each head i, as follows:A i = Q i (W K i ) ⊤ ( µ √ d/h + σ 2 ) ⊤ + 1 √ d/h Q i (b K i ) ⊤where the bias term Q i (b K i ) ⊤ ∈ R m is added across all n keys, and thus is normalised out in the softmax below.To this attention score matrix A, multihead evaluation denoising attention adds the same key biases c ∈ R h×(n+1) across all m queries and h heads. For ease of notation we define σ 2 r = (√ d/h + σ 2 ) and note that α 0 = ∑ d j=1 α j . c = log( α α 0 ) -1 2 ∥ µ σ r ∥ 2 -1 d (log(σ r )) ⊤where 1 d is a row vector of d ones.Now, we define the interpolation of each head's projected queryU i = Q i (W K i ) ⊤ ∈ R m×d with the inference time value vectors µ ∈ R (n+1)×d .The interpolation weights,σ 2 ( √ d/h+σ 2 )and√ d/h ( √ d/h+σ 2 ), are specific to values but not to queries, and the attention matrix, Softmax(A i + c), maps values to queries, so we can accommodate the interpolation between queries and values by multiplying in the values before the dot product with the attention matrix and multiplying in the queries after this dot product.DAttn(.) = ∑ i Softmax(A i + c) σ 2 σ 2 r ⊙ U i + Softmax(A i + c) √ d/h σ 2 r ⊙ µ W V i + b V iWe build upon the contribution of Behjati et al. (2023) of including NVIB and denoising attention for encoder self-attention. In that previous work the NVIB regulariser is applied to training single-head self-attention in the stacked layers of a transformer encoder. The queries for denoising self-attention are computed from the original n vectors of the transformer before they are projected to the (µ, σ, α) parameters of NVIB, but the keys and values are computed from the vectors Z ∈ R (n+1)×d of the NVIB layer. This allows every use of the attention function in self-attention to be done with denoising attention. From this previous work we maintain the exponential activation for the pseudo-counts. However, we remove the skip connection between pseudo-counts, because it is not part of the pretrained transformer.A causal mask is applied to self-attention in a transformer decoder to simulate a next token prediction when using teacher-forcing (Vaswani et al., 2017) . The mask may be applied as before over the attention scores before the softmax. The prior component in the keys is never masked. This prior component acts like a new additional form of start-of-sequence token, without a positional embedding. The bias terms do not depend on which other keys are masked, since the only term which depends on other keys is α 0 , and this term normalises out in the softmax. Thus, the implementation is identical to denoising self-attention with the inclusion of a diagonal mask.We provide pythonic pseudocode for the following functions: scaled-dot-product attention (Algorithm 1), denoising attention during training (Algorithm 2), and denoising attention during evaluation (Algorithm 3). The purple colour defines the differences between standard scaled-dot-product attention and denoising attention. The @ and * symbols are for matrix multiplication and element wise multiplication, respectively. The ** is used to denote an element-wise power. In this work we use commonly available summarisation datasets from HuggingFace2 . We use the following datasets: CNN/DailyMail (CNN/DM) (See et al., 2017) which is one of the most widely used summarization corpora. It is based on news articles from the CNN and DailyMail websites. The summary sentences are a concatenation of human-generated "highlights" and bullet points. Xsum (Narayan et al., 2018) is abbreviated for the "extreme summarization" dataset and is created from BBC news articles. The summaries are taken to be the first sentence of the article and the source document is the rest of the article. Curation Corpus (CC) (Curation, 2020) is a dataset of professionally written summaries of news articles. This is the only freely available news summarization dataset with references that were written for the purpose of summarizing the article. For this dataset we use the version available on HuggingFace and split it manually into train, validation and test by a 50%/25%/25% split. SAMsum (Gliwa et al., 2019) is an abstractive dialogue summarization dataset which is constructed to resemble the chats of a mobile messenger app. Each dialogue is written by a single linguist which can be formal or informal, and potentially contains slang, emoticons or typos. WikiHow was constructed from a knowledge base of how-to articles, explaining how to solve a task. We use the text as the document and the headline as the summary. Table 5 provides dataset statistics.In this work we consider the BART encoder-decoder model (Lewis et al., 2020) , which is a transformer-based sequence-to-sequence model and is pretrained as a denoising autoencoder. This model has shown efficacy in a wide range of tasks including summarisation.In our work we consider BART (large) summarisation models that are already pretrained and fine-tuned on CNN/DailyMail3 and Xsum4 , available on HuggingFace. We do not fine-tune or update the original model weights; we only initialise our NVIB layer as a form of post-training regularisation.The To get final evaluation scores, we first decrease the search space of NVIB hyperparmeters by finding the points at which each hyperparameter individually has full equivalence and has degradation in performance. We record this space of parameters in Table 7 . We notice that τ e α and τ c α can be decreased by several standard deviations before the noise affects the performance. We also notice that the τ d α range shows that the decoder is more sensitive to this parameter. The interpolation parameters τ σ have about the same sensitivity across the encoder, cross attention and decoder. The best hyperparamters for each model and each validation dataset are visualised in Figure 4 . After finding the hyperparameter range, we perform a random search of 50 trials for each dataset to find the best regularised models. Table 8 reports the validation results on the out-of-domain text summarisation task for BART models with post-training regularisation methods, including quantisation and NVIB regularisation. For quantisation we consider 16-bit, 8-bit and 4-bit baselines. We also include a combination of NVIB regularisation with quantisation, but the implementation currently only supports 16-bit.In Figure 5 , we compare each predicted validation summary to its corresponding gold empirical summary and report the distribution of Rouge overlap scores and correlations in length. Comparing these distributions for the baseline and for the NV-Transformer tells us in what way we are getting improvements. Comparing the lengths tells us how well the information content of the summary is being modelled. The summary is a compressed and high information version of the original document. The summary length is a measure of information in the document such that a longer summary means more information in a summary, given the compression ratio of the dataset. We compare the predicted and Figure 5 shows plots for the model trained on Xsum. The first two rows show that NVIB regularisation increases the performance on the Curation Corpus and CNN/DailyMail datasets and is producing longer, more accurate summaries, but the adaptation to the information in the document is similar to the baseline. However, NVIB regularisation is resulting in longer summaries than the BART (Xsum) model, which better reflects in absolute terms the much longer summaries in the Curation Corpus dataset (See Table 2 ). The last two rows show that NVIB regularisation increases the performance on the SAMsum and WikiHow datasets and produces shorter, more accurate summaries that are adaptive to the information.In this work we use commonly available translation datasets from HuggingFace5 . We consider high-resource language pairs English-German (En-De), English to French (En-Fr) for our translation experiment with the in-domain training data being the OPUS100 dataset (Zhang et al., 2020) . This is a publicly available parallel corpus with a wide domain. For out-of-domain translation corpora we consider the Bible (Christodoulopoulos & Steedman, 2014) , IWSLT translation dataset from 2017 (Cettolo et al., 2017) and TedTalks (Cettolo et al., 2012) datasets. For consistency we randomly split the OOD data into train, validation and test. For the small TedTalks dataset we use an even split. Since the IWSLT dataset was already split, we increased the validation data from 1K examples to 10K from the training data. Finally, from the Bible data we split the training data into validation (10K) and test (10K) as seen in Table 9 below.In this work we consider the Marian encoder-decoder model (Junczys-Dowmunt et al., 2018) , which is a transformer-based sequence-to-sequence model and is fine-tuned on the To get final evaluation scores we first decrease the search space of NVIB hyperparmeters. We find the points at which each hyperparameter individually has full equivalence and has degradation in performance for each model and each dataset. We notice that τ e α and τ c α can be decreased by several standard deviations before the noise affects the performance. We also notice that the τ d α range shows that the decoder is more sensitive to this parameter. The interpolation parameters τ i σ vary a lot with their sensitivity for translation, especially for the encoder. After finding a suitable range for the hyperparameters, we perform a random search of 100 trials for each dataset to find the best regularised models. The best hyperparamters for each model and each validation dataset are visualised in Figure 6 . We conduct an analysis of how the distribution over vectors changes through the layers and with different datasets, by calculating the empirical prior parameters for each distribution.We consider the BART model trained on CNN/DailyMail and Xsum and calculate the empirical priors (µ p , σ p , α p 0 ) given different in-and out-of-domain datasets. The empirical prior used for NVIB regularisation is the one computed in-domain. We also include a standard normal Gaussian prior in the plots for reference. We plot the average embedding value across all layers, where the last layer of the encoder (Encoder layer 13) is the embedding for cross-attention.Figure 7 displays the average empirical embedding across all encoder layers for a model finetuned on CNN/DailyMail and Xsum, respectively. We notice that the mean is approximately zero and very similar across datasets and across models. The variance is close to 0.1 with the cross-attention (Encoder layer 13) being lower at 0.02. The expectation of the logged pseudo-count is higher and approximately around 7 for encoder layers and lower around 2 for cross attention (Encoder layer 13) with consistent variation across datasets. Its clear the distribution of embeddings between fine-tuned BART encoders is similar.Figure 8 shows the average empirical embedding across all decoder layers for a model fine-tuned on CNN/DailyMail and Xsum, respectively. The decoder means have larger values in comparison to the encoder and similar low variance. Considering the logged pseudo-count we notice that they get exceedingly large through the layers which are near double the magnitude to the encoder in log-space.In this section we restrict the amount of data used to create the empirical prior and see if it is still able to maintain the performance increase across the datasets. We plot the Rouge-L performance as a function of the amount of data used to create the prior. We consider the best selected hyperparameters from the validation set (Section D.3) and show the relative performance for all models and all datasets. Figure 9 shows that the empirical prior is data efficient as it requires few examples (0.1% of training data which is approximately 200) to achieve good stable performance. In this section we consider the attention plots of our reinterpreted model in the presence of different levels of regularisation. We selected a random validation example from the Curation Corpus dataset and shorten the article document for visualisation purposes. We manually selected the layers with the most regularisation towards the prior for best visualisation. The attention plots are averaged over all heads and the attention scores are displayed with light yellow for 1 and dark purple for 0.We show an examples of encoder self attention, decoder cross attention and decoder causal attention for both reinterpreted NV-Transformers. We consider three cases: firstly the identity initialisation τ i σ ≈ 0 and τ i α = 10; secondly with regularisation where the τ i σ and τ i α are set to the best validation hyperparameters (Appendix D.3); and finally an over-regularised example to show what happens in the presence of collapse to the prior component [P] , where the initialisation parameters are τ i σ ≈ 0 and τ i α = -30. The prior collapse case can be interpreted as the non-prior pseudo-counts being 30 standard deviations smaller than the prior.In Figures 10, 11 and 12 we see that when the model uses the identity initialisation the attention maps completely ignore the prior component representation [P] and has equivalence in attention scores across all attention functions. In contrast, when τ i α = -30 the model has over emphasized the importance of the prior component, such that the attention patterns collapse to only considering the prior. When the model is regularised according to the best validation performance, we notice the attention patterns generally shift attention away from the vertical bars at special characters like punctuation and towards the prior. This shows that the models that are getting improvements have attention distributions which are regularised towards the prior.https://github.com/idiap/nvr transformers and https://github.com/idiap/nvib.https://huggingface.co/datasetshttps://huggingface.co/facebook/bart-large-cnnhttps://huggingface.co/facebook/bart-large-xsumhttps://huggingface.co/datasetshttps://huggingface.co/Helsinki-NLP/opus-mt-en-dehttps://huggingface.co/Helsinki-NLP/opus-mt-en-fr 


Decision: Henderson & Fehr (2023) offer a new view of the attention component of transformers through Neural Variational Information Bottleneck. This paper takes the idea further, extending from single-head cross-attention to all forms of attention in the transformer. The reviewers in general find the ideas to be sound and mathematically elegant. Personally I am happy to see work that explores new foundations for the modeling components of LLMs, and would be interested to learn more about the ideas presented in the paper. 

However, as a line of research, NVIB regularization is still somewhat work-in-progress, particularly from the empirical perspective. The reviewers did not find the empirical results to be very persuasive: I agree with the reviewer suggestion to consider more extreme distribution shift; and, given the well-known problems with Rouge- and Bleu-based evals, I wonder whether it would be possible to apply prompted evals, particularly for the summarization task. 

Another area for improvement is in clarity of presentation. The reviewers, some of whom have quite strong mathematical backgrounds, also found the presentation to be difficult to follow, and make useful suggestions for improvement that would likely benefit many CoLM readers. Reviewer `nwsJ` had particularly concrete ideas for improvement, which I endorse. I also found it unfortunate that most of the novel aspects of the work had to be relegated to appendix B.

[comments from the PCs] Please follow up on the AC and reviewer recommendations to improve the presentation. Improving the presentation will help your work have bigger impact.

Reviews: review: paper_summary: paper_summary: The paper extends NVIB regularization to all forms of multi-head attention in a transformer model. Also, it shows that existing pretrained transformers can be reinterpreted as nonparametric variational models using an empirical prior distribution and identity initialization. Experimental results show the effectiveness of the proposed approach.
strength_weakness: strength_weakness: - Writing issue. The paper starts with the NVIB, however, without sufficient background, section 2 is a bit hard to follow.
questions: questions: - When used during training, the NVIB layer provides an information-theoretic, sparsity-inducing regulariser over attention-based representations. Is there any advantage over just keeping top-k weights, though we have to set k in advance?

- For denoising attention, how should we understand the concept of denoising? For example, in diffusion generative modeling, we want to recover clean images from noisy ones, involving a denoising process. How does denoising work in your context?

- Is it possible to link this research to LLMs?
paper_summary: paper_summary: This paper explores a method based on NVIB (Henderson and Fehr, ICLR 2023) to apply as regularization over multi-head transformers for OOD text summarization and machine translation tasks.

In an encoder-decoder model, BART, the NVIB projection layer combines the Dirichlet Process-based empirical prior with the transformer embeddings to produce a representation for denoising MHA. In the NV-Marian case, the improvement of NVIB is less significant. 

I spent some time reading NVIB in two of its original references; it appears to be a relatively new and preliminary direction with only a few studies.

From my perspective, it is acceptable if NVIB does not serve as the best regularization/post-training boosting method for some applications. For example, motivating and explaining the limitations of this form of non-parametric variational Bayesian application in language modeling would be interesting to general audiences.

If the authors could add some studies on a decoder-only language model with NVIB (e.g., small-scale GPT-2 or LLama), it could potentially increase the general impact and interest of this work across more language modeling applications.

I think it is borderline paper. I enjoy the general writing and presentation of this work.
strength_weakness: strength_weakness: 1. the performance of NV-based LM seems to be little bit inconsistent and dataset dependent. Some in-depth analysis is encouraged.
questions: questions: 1. Would it be beneficial to add some min-max perturbation to simulate severe OOD cases for further diving into the properties of NV-based LM? One intuition is that the current performance is relatively diverse and seems to perform better in OOD tasks.


2. This might be a bit open. There are some causal variational Bayesian works [A,B] in robust RL and sequence modeling that improve OOD robustness via local Lipschitz constant measurement. Could adding a causal graph over the VAE framework further help with some counterfactual OOD evaluation cases?

***

A. Large Language Models and Causal Inference in Collaboration: A Comprehensive Survey

B. Training a Resilient Q-Network against Observational Interference

C. Lipschitz Parametrization of Probabilistic Graphical Models
paper_summary: paper_summary: The authors present an initialization for a Bayesian Nonparametric Variational Information Bottleneck (NVIB) interpretation of transformer models, and demonstrate that there exists an identity initialization of the NVIB hyper-parameters such that NVIB attention is equivalent to standard attention. They also show that by tuning the NVIB hyper-parameters away from this identity initialization using priors derived via forward passes through the training data they can produce models with equivalent performance but greater uncertainty. Such tuning yields performance gains for a BART summarization model trained on XSum, but does not yield significant gains for a BART summarization model trained on CNN/Daily Mail nor a Marian translation model trained on OPUS100.

Minor issues:
* "HuggingFace Wolf et al. (2020)." => "HuggingFace (Wolf et al. 2020)."
* "BART Lewis et al. (2020)" => "BART (Lewis et al. 2020)"
strength_weakness: strength_weakness: * The paper is more theoretically interesting than practically useful, as the regularization doesn't generally yield benefits, even in out-of-domain settings. (Gains are only observed in the specific combination of BART trained on XSum; in all other settings performance is similar to the original model.)
* The paper is hard to read in several places because it offloads too much to the appendix, starting a discussion of an analysis but then finishing the explanation of that analysis and the figures showing the results in the appendix.
questions: questions: * Is there prior research that would have suggested that quantization would improve performance? That wouldn't have been my intuition. I'm trying to understand the decision to use quantization as a baseline in this paper. The results certainly show that in the datasets selected quantization does not improve performance.
paper_summary: paper_summary: This manuscript proposed to leverage the idea of nonparametric variational information bottleneck method that modify the original attention module to perform post-training regularization on the pre-trained transformer model. The authors demonstrate the effectiveness on keeping the identity with specific initialization and out-of-domain generalization with additional regularization on some transformer architectures.
strength_weakness: strength_weakness: * The manuscript is extremely hard to follow, even for a Bayesian expert.
* There are several points not clear to me even after I go through the manuscript for many times. See the Questions part below.
questions: questions: * Could you revise Section 3 to make the prior distribution, likelihood and posterior distribution more clear? Specifically, I think it is necessary to show them in some equations, instead of some descriptive texts.
* What’s the meaning of $Z^2$ given $Z$ is a vector (e.g. in Equation 7)?
* Could you provide an algorithm box showing what’s your modification compared with the standard attention mechanism?
* What’s the motivation to set the parameters as in Equation 7? I don’t get any motivation on the necessity of $Z^2$.

Comments: title: Re: Rebuttal
Update after rebuttal
Thanks for your clarification.
comment: Thank the authors for the response.

Given the current scope is still fixed (e.g., decoder lm exclusive & evaluation on OOD), I will keep my current score. If the authors will revise the presentation on mentioning the challenges and related directions discussed here, the final version would be with boarder impacts and audiences.

It is a borderline and interesting paper to me but I won’t strongly argue for acceptance.
Sorry for the late reply and thank you for your clarification. 

I think the paper could be sufficiently clarified given an extra page in the camera-ready version. I raised my score.
Thanks for your clarification. Indeed I think this manuscript can be interesting and elegant with some potential impacts. However the overall writing needs to be improved to make sure the core component can be well understood by the general audience. I would rather encourage the authors to polish the paper and re-submit.