Nonparametric Variational Regularisation of Pretrained Transformers
Authors: Fabio James Fehr
James Henderson
Venue: COLM (2024)
Keywords: Transformers, Nonparametric VIB, Reinterpretation, Post-training regularisation, Out-of-domain generalisation
TL;DR: We propose a novel reinterpretation of pretrained transformers which allows for an information-theoretic regularisation that improves out-of-domain generalisation on NLP tasks without retraining.

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.

Published as a conference paper at COLM 2024
Nonparametric Variational Regularisation of Pretrained Transformers
Fabio Fehr't & James Henderson"
t Idiap Research Institute, Switzerland
$ Ecole Polytechnique Fédérale de Lausanne, Switzerland [fabio.fefl james. h @idiaap.ch
Abstract
Pretrained transformers have demonstrated impressive abilities, but tend not to generalise well out-of-domain and are very expensive to fine-tune 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 func- tions transformers. We show that existing pretrained transformers can be reinterpreted as nonparametric variational models using an empirical prior distribution and identity initialisation with controllable hyperparam- eters. We then show that changing the initialisation introduces novel, information-theoret post-training regularisation the attention mecha- nism, 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.
1 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 successtul, improving performance in virtually every NLP task. This indicates that the inductive bias of sieention-baeed 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 (023)) 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-thcoretiii sparsity-inducing regulariser over attention-based representations. In this paper, we investigate the possibility that NVIB also provides accurate theoretical model of existing attention-based models which have been trained without NVIB regulari- sation. 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 pretrained transformer into the weights of an equivalent nonparametric variational (NV)
Bayesian model (NV-Tranfformer)
This NV-Transformer is Bayesian in that it embeds text into a probability distribution over transformer embeddings, but remains equivalent by adding uncertainty around the em- bedding computed by the pretrained transformer, as illustrated in Figure [ Theoretically, the exact equivalence only occurs when this uncertainty is exactly zero, but empirically
httpps: http://tttuuu..com/iii////// transformers and https: wttpttb.cmmmidiaa//iiiiiiiaia

---

Published as a conference paper at COLM 2024
find practical range of non-zero uncertainty levels where the model's predictions 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 hyperparam- eters 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 does not degrade the model's task accuracy. Instead, it adds 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 Varia- tional Information Bottleneck regularisa- tion, and to the understanding of pretrained transformers. 1. We extend NVIB regulari- sation beyond single-head, cross-attention to all forms of multi-head attention in a transformer, resulting in our proposed NV-
Denoising Multi-Head Attention
Foed
Fonward
Reqularised NiB Emtoedking
Exuisalent NVI Enbadding
Add & Narm
Dencising Mi-i 4 NYB
NVIB projection
Transformer architecture. 2. We define a novel NVIB prior using the empirical dis- tribution of model's embeddings given small amount of data (Section 3.2). We propose reinterpretation of pretrained transformers as variational Bayesian mod- els by defining novel identity initialisation for NVIB with its controllable hyperparam- eters (Section 3.11. Adjusting these hyperpa- rameters regularises the embeddings of this equivalent NV-Transformer without requir- ing retraining. 4. Empirically, we show the usefulness of these proposals by smoothly varying the amount of this post-training reg- ularisation (Section 4.1 and achieving im- proved performance in out-of-domain text summarisation and translation (Section 4.2 This successful Bayesian interpretation helps characterise the way that pretrained trans- formers represent reliable information in their embeddings (Section 5).
Transformer
Empricall
Transforerr
Embedding
Prior
Embedding
Figure 1: Left: Reinterpreted transformer en- coder layer with NVIB and denoising multi- head attention (MHA). Right: A NVIB pro- jection layer which combines the empirical prior with the transformer embeddings and produces a representation for denoising MHA for either post-training regularisation exact equivalence.
Background
Henderson & Fehr (2023) propose NVIB as an information-theorrtii regulariser for attention
layers. It extends vector-space Variational Information Bottleneck (VIB) (Alemi al. 2017
to the unboundedly large oet-of-vectors representations supported by attention. Henderson
Fehr (2023) define Variational Auto-Encoder (Kingma & Welling 2014 for transtormers with NVIB applied to its cross-attention layers. [Bejjati 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 specified with distribution over latent representations.
2

---

Published as a conference paper at COLM 2024
For attention-based representations, these distributions need to capture the essential struc- ture imposed by the attention function on the sequence of vectors being accessed by atten- tion, namely that these vectors are permutation invariant, have normalised weighting, and can be arbitrarily many. 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 dis- tributions (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 ss 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 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 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 nonparametries (Equation (3) below). In the remainder of this section we provide more details about denoising attention and NVIB's distributions over mixture distributions.
Denoising attention To show that the denoising attention function is a generalisation of the standard attention function, Henderson & Fehr (2023 provide constructive proof of exact equivalence. The standard attention function projects the accessed set of vectors Z RRndd via weight matrices W WV Rdxd to keys and values, respectively, and projects the accessing input vector u' E R1xd via the weight matrix Wº E RR 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 (n', ; W WK WV) Attn(n'/Q((W)
Z) WV Attn Z) W
where (w'ko(wk)) T) E R The function Attn Z) can then be defined both in terms of a sum over the vectors z; in Z, or in terms of an integral over a distribution which is only non-zero at the zj:
Attn(u, Z) = softmax
uz T
Z = (Attn Fz)
exp(
2 (|z;|12)
Fz =
8z,
(1)
exp( 2vd (2:
DAtnn(u; F) !.
f(v) g(u; v, vál)
v do
(2)
fof(o) g(u v, Vdl du
where 82, is an impulse distribution at z,, f(-) is the probability density function for distri- bution F, and g(u; v, v VdI) is the multivariate Gaussian function with diagonal variance of Vd. DAttn( Fz) can be interpreted as Bayesian query denoising where Fz is the prior and g(u; v, vdT) is noisy observation of some true vector This construction shows that any transformer embedding Z has an equivalent mixture of impulse distributions, namely Fz, where denoising attention DAttn(u; Fz) gives us exactly the same result as attention Attn(u, Z), for all queries u.
Distributions over mixture distributions Given the generalisation of attention-based representations to nonparametric mixture distributions, Henderson & Fehr (0223) 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 -Byenal model computes a DP estimate of the posterior distribution over mixture distributions given the
3

---

Published as a conference paper at COLM 2024
input. More specifically, the model maps the input to set of pseudo-observations, and then the posterior DP is computed with exact Bayesian inference given prior DP and these observations. This posterior "q" distribution DP(C) a²) consists of a base distribution Gj for generating the vectors of the mixture distribution and pseudo-count ag for generating their mixture weights. The base distribution Gj is a mixture with one Gaussian component for the prior "p" distribution DP(N(H², T(or))) a²² plus one Gaussian component for each 7th pseudo-observation (11, oi,aaa computed by the model:
c/ = 1=nn(oin I(OP))) + (NP,1((P))) N (M) ((FP,
ap
ag
ag = a Σ
(3)
Denoising attention during evaluation As for other VAEs (Kingma & Welling 2014), during evaluation Henderson & Fehr (2023) do not sample a mixture distribution from the posterior DP, but instead use the posterior's 's mean, which is its base distribution Go
Denoising attention applied GJ can be computed straightforwardly because it a mixture of Gaussians. This function differs from standard attention in that the component weights x" /aj replace the key biases implicit in the query-key dot product, and an interpolation between queries and values replaces the values in the attention- weightt average.
119
(09)2
Vd
DAttn(u; G)) softmax
((rrr)
b
(1
(4)
Bias
Query-Value interpolation
a b log( 9
109
2
T
1 (log(o"))
(5)
2
20
where (Vā+(oo))) and 1 is vector of d ones. When the attention bias term 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 (or9) is approxi- mately 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 Related work in post-training regularisation such as quantisation (Dettmers) et al. 2022 Yao et al. 2022 Xiao et al.) 2023; Frantar et al. (023b) and sparsity (Hubara et al. 2021 Frantar et al. 2023a 2022) locus on regularising the model weights and not the latent embeddings. Similar to Frantar et al. (2023a), who propose a data-driven way to sparsify model 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
Our work shares similarity to the out-of-distribution generalisation literature in learning variational models with invariant features (Isse 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 Bayesian approach and detine our prior distribution based on empirical dataset statistics that allows for properties which are robust to domain shift. Park & Lee (021) proposes way to define pretrained language models as variational models by fine- tunettnnng.. Our work does not fine-tune or update the original model weights, nor does train additional weights as in adapters (Houlsby al., 2019 Hu et al. 2022). We keep the model weights frozen and map attention layers to NVIB layers, setting just few hyperparameters based on validation performance.
3 A Nonparametric Variational Reinterpretation of Transformers
We propose a novel method for reinterpreting pretrained transformers as nonparametric variational Bayesian models. To support this construction, 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).
4

---

Published as a conference paper at COLM 2024
This allows us to apply NVIB to every form of attention used in standard transformers, resulting in our proposed NV-Iransformer reinterpretation of pretrained transformers. Using these extensions, 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 way that captures the implicit uncertainty in pretrained transformers (Section 3.2)
Pseudocode for denoising attention is given in Appendix B.5
3.1 Identity Initialisation
We define an identity initialisation for NVIB such that denoising attention is effectively equivalent to standard attention, inspired by Equation [ Given set of vectors Z input to the attention layer, our proposed NVIB layer converts it to the parameters (j, 0,, a) which specify a posterior "q" DP distribution, and then applies the evaluation denoising attention
function (Equation 4 to the resulting base distribution Cg. Excluding the prior component
(discussed in Section 3.2, we define the projections:
p μ(Z) ZW" + b
(6)
a² = d2 = exp (ZW" + b²)
(7)
a = a (Z) = exp(Z'uf Zw + b*)
(8)
where Z2 is the component-wise square. We choose these forms of projections because we want to set its parameters, W W E R b", b" E Rd, w w/ E Rd and Ra E I, so that the evaluation denoising attention function 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 query- value interpolation to the value by reducing the uncertainty of the non-prior components.
Alternatively, we can this requirement making the mixture of Gaussians Gj equivalent See Figure for a visual depiction.
to the mixture of impulse distributions Fz in Equation [ We can down-weight the prior component by making the pseudo-counts (Z) for the non- prior components be large. We can make the non-prior component Gaussians look like the impulse distributions of Fz by setting the mean projection ((.) to be the identity and the variance projection (.) be approximately zero (exactly zero is not possible with the exponential activation function). Then to make the mixture weights the same as for Fz in Equation 1 the pseudo-count projection a.)) must be proportional to the exponent of the scaled L3 norm of Z, which is why it has log-quadratic form in Equation g
Initialisation hyperparameters 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 smooth regularisation of the embeddings (illus- trated in Figure (1). The hyperparameters Tę and T: control the level of uncertainty of the mixture distributions by changing the initialisation of the pseudo-counts and variance, respectively. The indicator is for the different sections of the model, which allows for independent control of the encoder's self-attention (e), decoder's 's rross-attention (c) and decoder's 's causal self-attention (d). Empirically, this allows for more flexibility than single hyperparameter and is practically easier to tune than defining a hyperparameter per layer. In general we define the layer projection weights as follows:
W = I;
b² = 0
W 0:
b" log( (o" }
(9)
1
w
11;
1 0; ba = e/ c
(10)
2/d/h
where d and ht denote the model projection size and number of attention heads. As dis- cussed in the next section, the empirical distribution is used to define the prior standard
deviation oP, and the constant ea, which denotes the empirical standard deviation of the
5

---

Published as a conference paper at COLM 2024
scaled L3-norm in log space per layer. The hyperparameter & in Equation 10 controls the relative weight the prior component to the non-prior components in the mixture distribution C/. Since ex is a standard deviation, it reflects the normal range of values for use Tz to control the weight given to
the non-prior-component pseudo-counts, the prior component with respect to this range. When र. 0, the non-prior pseudo-counts are their scaled Lī-norms and the prior pseudo-count the expected scaled L3-norm. When we increase or decrease the Th, it adjusts the magnitude of the non-prior L3-norms, which relatively decreases or increases the weight on the prior proportionately to the standard deviation of the scaled I.2norm. The hyperparameter To (Equation 9) controls the inter- polation between the query and value (Equation 4 which we set proportionately to the variance of the prior distribution. When To O, there is effectively no interpolation, as with standard attention. When To 1, the uncertainty is increased to the level of the empirical prior distribution.
3.2 Empirical Priors
The prior distribution is Dirichlet Process, so it views all vectors in all transformer embed- dings 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 for- ward 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.
Data-informed prior We can estimate the empirical parameters for an isotropic Gaussian G6 ~ NNH,, ((o)) and the prior's 's pseudo-count a by using the latent vectors Z as follows: n²² -
IN N (oP)² I N-1 n²)
(z, "z'
and log(a = EN Σ (Z))) )/N, where N is the 2vd/h
total number of tokens in the training corpus, the dimension the embedding and 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 n2-noom of the latent vectors.
Analysis of the empirical prior To further understand the distributions of our latent embeddings, calculate the distribution of the empirical priors across all layers of the encoder and decoder attention mechanisms (Appendix F.I. Given a model 0 that has been trained on data x, we analyse priors generated from both in-domain 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 TI and TÓ 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 subsample of the training data. We found that the empirical prior can be created from as few as 0.1% 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- Tost-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-theorett post- training regularisation by changing the initialisation of the NVIB layers. Once the empirical

---

Published as a conference paper at COLM 2024
prior is estimated, adapting pretrained transformer to new domain simply requires hyperparameter selection on data from the new domain, which can be done solely with Norward-passer of the model. This is important contribution as the cost of backward- passes, 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 information- theoretic post-training regularisation.
Models We consider models that are already pretrained and fine-tuned, available on Hug- gingFace|Wol et al. (2020). In our experiments, 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 totpleennnc---suuennee transformer mod- els, 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/DailyMaii data (See et al., 2017), or Xsum data et al. 2018 encoder-decoder translation models already fine-tuned on the OPŪST00 (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
(Narayan et al. 2018). For translation we consider Marian Junccys-Dowmen
4.1 Identity Initialisation Equivalence
Validation Rauge-L
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 =10 and of x 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 Equa- tion (4) Appendix Table shows that the NV-Transformer initialisation has the same accuracy as the pretrained trans- former. More precisely, the leftmost points in Figure being at 100% on the x-axis shows that these two models generate the same outputs.
30-
X9000666666 -X.
20-
10
100
80
60
40
20
Original prodiction Rouge-L
40
DPUS Ensfr
-do
OPUS En De
Validatien BLEU
30
20
10.
Analysis of Increasing Uncertainty show that increasing the uncertainty within our regularisation allows access 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 line ordered by their linear interpolations between the innaiiallnntiiiialiiattoo and the everything regularised corners of the sample space. The Ni-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 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.
100
80 Original prediction BLEU
Figure 2: initialisation, the NV- Transformer's output is compared against the non-regularised original model (x-axis) and against the gold prediction (y-axis). Summarisation Translation (Below).
For each NVIB
(Top),
Attention maps Examining the attention maps (Appendix F3 of the summarisation encoder and decoder, we see that the prior componens is included in the keys of attention,
7

---

Published as a conference paper at COLM 2024
Table 1: Post-training regularisation on OOD summarisation. We report test set Rouge-I.
Out-of-Domain SAMsum
Model
CNN/DM
Xsum
CC
WikiHow
BART (CNN/DM)
29.99
13.12
24.99
22.42
9.26
BART-16bit (CNN/DM)
29.97 [-0.02]
13.13 [-0.01]
24.99 [0.00]
22.36 [-0.06]
9.27 [+0.01]
BART-sbit (CNN/DM)
29.55 [-0.44]
13.13 [+0.01]
24.67 [-0.32]
22.13 [-0.29]
9.36 [+0.10]
NV-BART (CNN/DM)
29.33 [-0.66]
13.99 [+0.87]
25.04 [+0.05]
22.60 [+0.18]
9.41 [+0.25]
BART (Xsum)
16.61
36.42
14.37
18.33
13.43
BART-16bit (Xsum)
16.60 [+0.01]
36.44 [+0.02]
14.38 [+0.01]
18.24 [-0.09]
13.43 [0.00]
BART-8bit (Xsum)
16.56 [-0.05]
36,25 [-0.1]]
14.33 [-0.04]
18.05 [-0:28]
13.53 [+0.10]
NV-BART (Xsum)
19.42 [+2.81]
36.25 [-0.17]
17.61 [+3.24]
21.94 [+3.61]
15.25 [+182]
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.
4.2 Out-of-Domain Generalisation
We evaluate whether our post-training information-theoreee regulariser improves out-of- domain 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 shift in the data domain. To evaluate our regulariser' OOD generalisation we define the following experiment: given model 8r that has been trained evaluate the model on the same task in the presence of a domain-shift to the text domain For example, given pretrained news summarisation or translation model 0 we evaluate the OOD generalisation by shifting the domain from news x to informal dialogues y without further training.
text domain x, we
For baselines we consider the original model (32-bit) and quantisation of 16-bit and 8-bit (Dettmers et al. 2022). Quantisation an alternative form of post-training regularisation which has been shown to improve performance (Dettmers & Zettlemoyer (2022). We also considered 4-bit quantisation model and combining NVIB regularisation with 16-bit quantisation (Appendix Table 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 tossseeee----eeeencce models trained on summarisation and machine translation.
Summarisation For our summarisation experiments evaluate the models on the vali- dation sets of out-of-domain summarisation corpora from either Xsum or CNN/ Daily 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 Trblerrepotts the test set Rouge-L on the OOD text summarisation datasets for the
post-training regularisation methods. Considering the original model trained on CNN/- Maily 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/ CNN/Daily) MailyMail is more extractive. This shows that our information-theone can improve OOD generalisation, both over the original model and over quantisation.
post-training regularisation
Analysis of summarisation improvement We conjecture that the length of the summary is strong proxy for the information within the document. Since the baseline model has been trained to produce a summary of certain length, the summaries on OOD data can be improved by adapting the length of the summary to the information content. We
8

---

Published as a conference paper at COLM 2024
Bosofie 3.58
0.50
125
Densiy
Derwty
1.3
Pretreed Summeny
L.0-
375
111
125
075
22
044
Rouge:
Rouge2
Rouge
Emotrical
Figure 3: A BART baseline model trained on Xsum compared to its reinterpretation with
NVIB regularisation on the SAMsum test dataset. Left-Right: Score distributions for Rouge- 1, Rouge-2, Rouge-L, and empirical vs predicted summary length with its Spearman's correlation.
see this pattern in the mean sentence lengths in Table 2 The baseline pro- duces summaries slightly shorter than its empirical Xsum training data for all datasets, while the reinterpreted NV- Transformer model matches this length in-domain but produces shorter sum- maries for shorter-summary datasets and longer summaries for longer-summary datasets. Looking within the distribution of a single OOD dataset, we compare the output summaries of the original BART model and a reinterpreted NV-Transformer model on the OOD SAMsum test dataset. Figure (3 left) shows consistent improvements across the Rouge-1, Rouge-2 and Louge-L distributions, which measures overlap metrics against the gold summaries. Figure 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 R4
Table 2:
Mean number of words in validation set summaries for the original BART model trained on Xsum data and a reinterpreted NV- Transformer model.
Model
cc CNN/DM SAMsum WikiHow Xsum
Empirical
85.2
57.9
20.3
6.5
21.1
BART (Xsum)
19.3
20.7
17.4
16.4
18.9
NV-BART (Xsum)
28.2
30.8
16.6
14.4
18.6
Translation 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 (Christodoulopou los & Steedman 2014), IWSLT (Cettolo et al., 2017) and TedTalks (Cettolo et al., 2012). See Appendix F. for further dataset details. We find the best model on the validation set (Appendix TT) and then evaluate this model on the test set. Table reports the test set BLEU on the OOD translation datasets. We notice fairly consistent minor improvements out of 6 cases) across all OOD translation datasets. We speculate that the improvements in translation have lower magnitude than summarisation, as this task benefits less from re- duction in information. These results further support our hypothesis that our post-training, information-theoreti regulariser improves OOD generalisation of pretrained transformers.
Discussion
Since regularisation does not guarantee improved performance, 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 selecting a summary from the document.
The improvements from post-training regularisation suggest that our variational Bayesian reinterpretation of pretrained transformers is accurate model of how information is

---

Published as a conference paper at COLM 2024
Table 3: Post-training regularisation on text translation. We report test set BLEU.
Out-of-Domaii IWSLT
Model
OPUS100
Bible
TedTalks
Marian (OPUS En-De)
24.78
23.02
27.16
24.80
Marian-16bit (OPUS En-De)
24.70 [-0.08]
23.02 [0.00]
27.18 [+0.02]
24.79 [-0.01]
Marian-8bit (OPUS En-De)
24.65 [-0.13]
23.11 [+0.09)
27.05 [-0.11]
24.67 [-0.13]
NV-Marian (OPUS En-De)
24.84 [+0.06]
22.95 [-0.07]
27.30 [+0.14]
25.06 [+0.26]
Marian (OPUS En-Fr)
35.22
27.23
39.03
31.92
Marian-16bit (OPUS En-Fr)
35.22 [0.00]
27.21 [-0.02]
39.05 [+0.02]
31.97 [+0.05]
Marian-8bit (OPUS En-Fr)
35.23 [+0.01]
27.19 [-0.04]
39.03 [0.00]
31.89 [-0.03]
NV-Marian (OPUS En-Pr)
35.22 [0.00]
27.41 [+0.18]
39.28 [+0.25]
32.43 [+0.51]
being captured in transformer embeddings. When we map an embedding into probability distribution, the variance specifies what differences between embeddings constitute reliable information, and what differences should be unreliable or unimportant. The denoising attention function then ignores this unreliable information, which should not affect accuracy in-domain, and could reduce overfitting out-of-domain. The above experiments succeed in finding such empirically effective patterns of variance for our proposed DP distributions (Figures 2 I0 11 and Tables 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 L2-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 ar- chitecture allows us to propose a Bayesian reinterpretation 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.
Future work In future work, we plan to evaluate our post-training regularisation on dif- ferent 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 generatinn,,, 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 flash-attention (Dao et al. 2022).
10

---

Published as a conference paper at COLM 2024
References
Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep variational information bottleneck. In 5th International Conference on Learning Representations, ICLR 2017, Conference Track Proceedings, Toulon, France, 2017. OpenReview.net. URL https //
Melika Behjati, Fabio James Fehr, and James Henderson. Learning to abstract with nonpara- metric variational information bottleneck. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023. http::///ppeeeeeeeeeeee////e//os
Gilles Blanchard, Aniket Anand Deshmukh, Ürun Dogan, Gyemin Lee, and Clayton Scott. Domain generalization by marginal transfer learning. J. Mach. Learn. Res., 22(1), jan 2021. ISSN 1532-4435.
Mauro Cettolo, Christian Girardi, and Marcello Federico. WIT3:
Web inventory of tran- scribed and translated talks. In Proceedings of the 16th Annual conference of the European Association for Machine Translation, pp. 261-268, Trento, Italy, May 28-30 2012. European Association for Machine Translation. URL https: h/wp:///wwwwwwwo http://wwwwwwwwwbbbbb/ acwweb. wwtp./aag//nooooo//222222
eemt-1160
Mauro Cettolo, Marcello Federico, Luisa Bentivogli, Jan Niehues, Sebastian Stüker, Kat- suhito Sudoh, Koichiro Yoshino, and Christian Federmann. Overview of the IWSLT 2017 evaluation campaign. In Proceedings of the 14th International Conference on Spoken Language Translation, pp. 2-14, Tokyo, Japan, December 14-15 2017. International Workshop on Spoken Language Translation. URL httphtsss/aaccaantooooo.o.o/////
iws1t-1.1
Christos Christodoulopoulos and Mark Steedman. A massively parallel corpus: the bible in 100 languages. Language Resources and Evaluation, 49:1-21, 06 2014. doi: 10.1007/ $10579-014-9287-v
Curation. Curation corpus base, 2020,
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashatten- tion: Fast and memory-efficient exact attention with io-awareness. S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 16344-16359. Curran Asso- ciates, Inc., 2022. URL
S. Koyejo,
neurips wwtppccpppppiiiceeee wwweerrfsssseseeee2...2 pdf]
Tim Dettmers and Luke Zettlemoyer. The case for 4-bit precision: k-bit inference scaling laws. CoRR, 2b12..220097 2022. URL https://6oiiorr//11 aatp wttp:a/wwww..22..722222
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3.int8(: 8-bit matrix multiplication for transformers at scale. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 30318-30000 30332. Curran Associates, Inc., 2022.
Jacob Devlin, Ming-Wa Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019, Association for Computational Linguistics. doi: 10.18653/v /v1 N19-1423, URL httptpss//aaccaaaooooo...rr//r/e
1423
Elias Frantar, Sidak Pal Singh, and Dan Alistarh. Optimal Brain Compression: a framework for accurate post-training quantization and pruning. Advances in Neural Information Processing Systems, 36, 2022.
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. OPTQ: Accurate quanti- zation for generative tre-trained transformers. In The Eleventh International Conference on Learning Representations, 2023a. URL httppp///ppprreeeeeeeeee//otrttt
11

---

Published as a conference paper at COLM 2024
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. OPTQ: Accurate quanti- zation for generative pre-trained transformers. In The Eleventh International Conference on Learning Representations, 2023b. URL URLLhttps://pperrviiee
http:////pppeeeeeeeeeeeet/rrrrrs http://wwwwweeetoonnooooooe net//formmidd-ccbBrnt
Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. SAMSum corpus: A human-annotated dialogue dataset for abstractive summarization. In Proceedings of the 2nd Workshop on New Frontiers in Summarization, pp. 70-79, Hong Kong, China, November /V/19-5599 URL
2019. Association for Computational Linguistics. doi: httpss//accaathoooog.oro////// 5489
James Henderson and Fabio James Fehr. A VAE for transformers with nonparametric variational information bottleneck. In The Eleventh International Conference on Learning Representations, 2023. URL htttps://///ppeeeeeviiio www..eee/oorumiiooooc......s
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Larous- silhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for NLP. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Pro- ceedings of Machine Learning Research, pp. 2790-2799. PMLR, 09-15 Jun 2019. URL https::/pproceedingg
wwt.ppeeeeeee//sseeeeeeeeeeee 19. 9tm
Edward Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. URL https::/ net/
forum2id==ZeWWeeFFf9
Itay Hubara, Brian Chmiel, Moshe Island, Ron Banner, Joseph Naor, and Daniel Soudry. Ac- celerated sparse neural training: A provable and efficient method to find n:m transposable masks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 21099-21111. Cur- ran Associates, Inc., 2021. URL https: http:hh/tt/////ooeeeeeee //wpe//ddcs.ccccec.///ooo.o neuriinggs.nerrppi www.iccpppperrrrreeeecceeeeeeeee files/paprr///e//////
Maximilian Ilse, Jakub M. Tomczak, Christos Louizos, and Max Welling, DIVA: Domain invariant variational autoencoder. In Medical Imaging with Deep Learning, 2020. URL
Marcin Junczys-Dommunt, Roman Grundkiewicz, Tomasz Dwojak, Hieu Hoang, Kenneth Heafield, Tom Neckermann, Frank Seide, Ulrich Germann, Alham Fikri Aji, Nikolay Bogoychev, André F. T. Martins, and Alexandra Birch. Marian: Fast neural machine translation in C++. In Proceedings of ACL 2018, System Demonstrations, pp. 116-121, Melbourne, Australia, July 2018. Association for Computational Linguistics. URL http: //wwe. org/ http//oooooooo/--4444/
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Conference Track Proceedings, Banff, AB, Canada, 2014. URL URL/htpp///arxxvoorr/aa////
6114
Mahnaz Koupaee and William Yang Wang. Wikihow: large scale text summarization dataset. CoRR, 0bs//180..990s 2018. /tthtt/://arrxiv http://arxxi...rrab///11111 httpg//////11103335555
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: Denoising sequence-to- sequence pre-training for natural language generation, translation, and comprehension. In Proceedings of the 58th Anmual Meeting of the Association for Computational Linguistics, pp. 7871-7880, Online, July 2020. Association for Computational Linguistics. doi: 10.18653/ ///2220.aclmmii....
URt. httptps://aaa.aaooooogg....o////
ac1l aollmminn. 703.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales. Deeper, broader and artier domain generalization. 2017 IEEE International Conference on Computer Vision (ICCV), pp. 5543-5551, 2017.
12

---

Published as a conference paper at COLM 2024
K. Muandet, K. Fukumizu, B. Sriperumbudur, and B. Schölkopf. Kernel mean embedding of distributions: A review and beyond. Foundations and Trends in Machine Learning, 10 (1-2):1-141, 2017. doi: 001561/2200000600
Shashi Narayan, Shay B. Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 1797-1807, Brussels, Belgium, Detober-November 2018. Association for Computational Linguistics. doi: 10.18653/ v1/D18-1200 URL httpss//acclatthooo...rrrg 1286
Seongmin Park and Jihwa Lee. Finetuning pretrained transformers into variational au- toencoders. In Proceedings of the Second Workshop on Insights from Negative Results in NLP, pp. 29-35, Online and Punta Cana, Dominican Republic, November 2021. As- sociation for Computational Linguistics. doi: 10.18653/v1/2000 wttp:///....n.ittssl1l.1 URL
https://aclanttooogy..orr htt./ooooo.....r/.........sssss. insights-1.55
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yangi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. 2020. http: http://iwr..rrppppps//////////// -000 074. html
Abigail See, Peter Liu, and Christopher D. Manning- Get the point: Summarization with pointer-generator networks. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1073-1083, Vancouver, Canada, URL
July 2017. Association for Computational Linguistics. doi: wwtp://.www....///oooooeeooeeeee
https://ww/ ac wwwp///wwwwwwwb.b........... web. orgo httg///hhoooog//111e9999
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomput, 568(C), mar 2024. ISSN 0925-2312. doi: htloo//.j.eucom.2222. 20t...2222.... URL httpt://////o1. http://doooorgg/////// org/10. 1016/j neucom. 2023. 127063
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran As- sociates, Inc., 2017. URL f
www.:nie.rr..ppppppep files/paper/2911/
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38 38-45, Online, October 2020. Association for Computational Linguistics. URL hhttps://www ac bewwwb..c org/ anthology/200 2020. emnlp- 1amnp-dde pepp--deottt
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. SmoothQuant: Accurate and efficient post-training quantization for large language mod- els. In Proceedings of the 40th International Conference on Machine Learning, 2023.
Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yux- iong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 27168-27183 Curran Associates, Inc., 2022.
Biao Zhang, Philip Williams, Ivan Titov, and Rico Sennrich. Improving massively mul- tilingual neural machine translation and zero-shot translation. In Dan Jurafsky, Joyce
13

---

Published as a conference paper at COLM 2024
Chai, Natalie Schluter, and Joel Tetreault (eds.), Proceedings of the 58th Anmual Meet- ing of the Association for Computational Linguistics, pp. 1628-1639, Online, July 2020. As- sociation for Computational Linguistics. doi: 10.18653/v http:///////22...ammmm..eeeee URL https://accannthooogy.orr///////
amlln 1
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.
A Appendix
Ethics Statement
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.
Limitations
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 encoder-decoder transformers.
Moreover, we limit our
choice of models to only consider seguence-to-sequeee These models were selected they comprehensively account for all forms of attention: encoder self-attention, cross-attention and decoder causal-atten decoder-only models, such as many large language models (LLMs). We limit our experimen- tation to Natural Language Processing (NLP) tasks such as sunmarisation and translation. Further experimentation
Hence, we exclude
the domain of vision and speech beyond the scope of this work. We plan to address these limitations in future work.
B Denoising Multihead Attention
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 E Roxd and set of pre-projected queries as u' Raxd We assume the latent vectors are square such that Wº, W w E Idddd and biases 6 6, b, E Rd are used to linearly project to the queries, keys and values, respectively. We define the standard attention weights before the softmax as follows:
A (u'we + 69) (ZW + b E R
(11)
KT
Typically, for multihead attention the projected query Q and keys K are split into heads. In this definition, we split the linear projections by divisible number of heads ht such that Wº, W W E _ and biases bº, ,,, bV E R so that Q E Rhxmx ş and
K E R bxxxxx R. We can then specify multihead attention by defining a matrix of attention scores A E Rlxmxn, for each head i:
A
Va/h (W,) + ((((TTTT)
where the bias term (b E R" 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 Vd/h. For denoising attention, each head's query is projected into the space of the original set of vectors Z, namely U, (W=), and so is still in Ru Thus, each head can be viewed as
14

---

Published as a conference paper at COLM 2024
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 vd/hl.
B.1 Training denoising attention
In this work we do not do any training, but for completeness we include the equations for multihead attention training time. The NVIB layer outputs the isotropic Gaussian parameters и E [R(++)xdd r E R( and Dirichlet parameters a E R((++), which include the 1t component for the prior. During training, these are used for sampling such that π Dir(a) and Z (N, σ). Given these sampled weights and vectors, the training-time timen denoising attention function is the same as the standard attention function with two changes: (1) the keys come from the sampled vectors
((++1)xd which include vector sampled from the prior component, and (2) each key has an attention bias b (1+1)) which is determined by its weight E ((+1)), Summing over heads i, the training-time denoising attention function is:
DAttn(.) = Σ Softmax(A; + log πτ) 2/J7h | ( + +bK))
V:
The biases are defined by adding the log of the sampled weights log πT) E R(++)) from the
NVIB layer and subtracting the scaled 12-norms of the sampled vectors 2vd/h / R
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.
B.2 Evaluation denoising attention
During evaluation, as for training, the NVIB layer outputs the isotropic Gaussian parameters
E Rx+1)xd E (n+1)xd and Dirichlet parameters a (n-1), For evaluation, the denoising attention function is similar the training-time denoising attention function with three changes: (1) the parameters are used directly without sampling; (2) there is different attention bias (" 11) for each of the (I+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 E R (++1), for each head as follows:
A, (W!
va TT 1111 (b,)
Va/h
where the bias term (b, E R" is added across all n keys, and thus normalised out in the softmax below.
To this attention score matrix A, multihead evaluation denoising attention adds the same key biases e E (hx(n+1) 1) across all queries and I heads. For ease of notation we define of (V Vd/h + o2) and note that 0g
- aj.
= log _ 211 1 o,
2 (log((())))
ao
where 14 is a row vector of d ones.
Now, we define the interpolation of each head's projected query U, = (W e Raxd with the inference time value vectors E R (n=11)xd The interpolation weights, (Vd/h oJ and va/h Softmax A; + 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.
/dd/hooyy specific to values but not to queries, and the attention matrix,
D.ttn(.) = ([(( ((Sottmax(AA) +c)
(aahh (o7 O )) w/ + +b²
3 U, Softmax(A, (c;
o7
15

---

Published as a conference paper at COLM 2024
B.3 Denoising Self-Attention
We 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 transformer encoder. The queries for denoising self-attention are computed from the original vectors of the transformer before they are projected to the (H, O,) parameters of NVIB, but the keys and values are computed from the vectors Z [[(++])xd 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, remove the skip connection between pseudo-counts, because it is not part of the pretrained transformer.
B.4 Denoising Causal Attention
causal mask applied self-attention 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 new additional form of start-of-sequence token, without positional embedding. The bias terms do not depend on which other keys are masked, since the only term which depends other keys is ag, and this term normalises out in the softmax. Thus, the implementation is identical to denoising self-attention with the inclusion of a diagonal mask.
B.5 Pseudocode
We provide pythonic pseudocode for the following functions: scaled-dot- product attention (Algorithm [ denoising attention during training (Algorithm 2), and denoising attention during evaluation (Algorithm 3). The purple colour defines the differences between standard scaled-dot-p proddut attention and denoising attention. The and symbols are for matrix multiplication and element wise multiplication, respectively. The ** used to denote an element-wiee power.
16

---

Published as a conference paper at COLM 2024
Algorithm 1 Scaled-dot-product Attention
Algorithm 2 Denoising Attention (training)
class Attention():
class DemoisingAttentiion(]
def ...init_. d, h):
def (se1f, d, h);
Projecctions
reshaped
Projections Q. reshaped by
[d,d] [h,d,
[d.d] +> [h.d.d/h]
self.q linear d). (eshape d, d/h)
self.q = linear(d, d) re.rhphape d, d/h)
self.k linear reshapepe d. d/h)
self.k linear(d. reshappp, d/h)
self. linear(d, reshappe,,,, d,
re.happpp, d/h)
def (self, u. 2): [B, d]
def forlf, 2, pi):
queries
queries
[8, H.
keys
keys
2: [8, N41
scale d/h
scale d/h
scale shapeehape(2) / self k. shape) (b)
scale keys. (2) self.k. shape
Project Q. self.q(u)
Project
q self.q(u)
k self-k(2) sart(scale) self.v(z)
k self.k(z) sart(scale) self.v(z)
Actention scores [h, B, M, 8]
MIB bias [1, B, N+1] b log(pi)
attn (ranspspo
probabilities H, attn (attna
scores [h, B, N11
attn = =qQ k. krtaassposee) +b
Value projection [h. d/h] attn
Attention probabilities [h, B, N91 softwax(attn))
Reshape [8, H. d] (u.
# Value projection [h, B, d/h]
return out
# Reshape [8. N. d]
out rw.hape ())
17

---

Published as a conference paper at COLM 2024
Algorithm 3 Denoising Attention (evaluation)
Denoisingtttennion)):
def
(seir d, h3:
Projections to Q, K, v reshaped by heads [h,d.6/h]
self. linear rh, d,
self.k linear re.hhppap, d, d/h)
selr. ld,ear(d, d). (hshappe,, d, d/h)
fsel,, u, su, sigss2, alpha):
queries u: [8, N, d] keys values au: [8. N+1, d] (2)shapa sh)
Project Q. K,
self-a(u)
self.k(mu (sor((caalo)ssim2)))
interpolation
NIB blas (1, B, 1, N+1]
sun(alpha))
Actention
attn transpose()
Actention probabilities [h, H, N+13
softnax
Query projection to kay-space
[h, H, d/h] >> B,
Value interpolation [h, B, H, out (attn 9 (sigaa2/<sQrt(0)) (s(rt(())++siaaa))) (EEer(cc)/(arrtt))
Project into self.v space
[h, N, d] v(out)(
d/h3
Reshape [6, H, d]
out. (u.
return
C Empirical Equivalence
Table 4: NV-Transformer equivalence to pretrained transformers. Validation cross-entropy (CE), Rouge-L scores for in-domain text summarisation (above) and machine translation BLEU (below).
Data
CE
Rouge-L / BLEU
BART (CNN/DM)
2.71
30.56
NV-BART (CNN/DM)
2.71 [0.00]
30.56 0.00]
BART (Xsum)
2.30
36.47
NV-BART (Xsum)
2.30 [0.00]
36.47 [0.00]
Marian (OPUS En-De)
1.87
24.43
NV-Marian (OPUS En-De)
1.87 [0.00]
24.43 [0.00]
Marian (OPUS Ea-Pr)
1.45
33.25
NV-Marian (OPUS En-Fr)
1.45 [.000
33.25 0.000
18

---

Published as a conference paper at COLM 2024
Table Dataset statistics.
Examples
Mean words
Dataset
Train
Val
Test
Document
Summary
CNN/ CNN/Daily MailyMail
287K
13.4K
11.5K
685
52
Xsum
204K
11.3K
11.3K
431
23
Curation Corpus
15K
7.5K
7.5K
504
83
SAMsum
14.7K
0.8K
0.8K
94
20
Wikihow
198K
6K
6K
580
62
D Summarisation Experimental Setup
D.1 Data
In this work we use commonly available summarisation datasets from HuggingFace 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 concatenation of human-genneated "highlights" and bullet points. Xsum [araaan et al. 2018 abbreviated for the "extreme summarization" dataset and is created from BBC news articles. The summaries to be the first sentence of the article and the source document is the rest of the article. Curation Corpus (CC) (Curation, 2020) is 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 the version available on HuggingFace and split it manually into train, validation and test by 50% / 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 written by 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 5able5prooe provides dataset statistics.
taken
D.2 Models
In this work we consider the BART encoder-decoder model (Lewis et al. 2020), which is transformer-based sequence-to-sequeece model and is pretrained as a denoising autoen- coder. This model has shown efficacy in wide range of tasks including summarisation. In our work we consider BART (large) summarisation models that are already pretrained and fine-tuned on CNN/DailyMai and Xsum available on Hugging We do not fine-ture update the original model weights; we only initialise NVIB layer form of post-training regularisation.
The BART (large) model uses 12 layer transformer encoder and decoder with 16 attention- heads. The size for the word embedding vectors and model projections are 1024, and the size of the feed forward dimensions are 4096, which leads to models of approximately 406 million parameters. The inclusion of the NVIB projection layers, per attention mechanism, results in an increase of approximately 11% parameters. These projections
not trained and only initialised, which results in 459 million parameters in total. It has an input context length of 1024 tokens. Where the tokens are created from a Pite-PPar-Encodiin (BPE) tokenizer. Tabvedprovides specific autoregressive generation details for each of the models.
19

---

Published as a conference paper at COLM 2024
Table BART models specific autoregressive generation details.
BART (CNN/DM)
BART (Xsum)
Number of beams
4
Length penalty
2
Max length
142
62
Min length
56
11
D.3 Validation Hyperparameters
To get final evaluation scores, we first decrease the search space of NVIB hyperparmeters by finding the points which each hyperparameter individually has full equivalence and has degradation in performance. We record this space of parameters in Table
Table 7: BART model's NVIB hyperparmeter selection space for random search.
Tx
7/
T/
min
-10
-55
1
1e
38
le 1e-38
le 38
max
0.5
0.5
0.5
We notice that of and Tf can be decreased by several standard deviations before the noise affects the performance. We also notice that the Tf range shows that the decoder is more sensitive to this parameter. The interpolation parameters Ty 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
T5
15
0.40
0.45
-0.40
0.40
-0.00
0:35
0.35
0.25
66
0.25
.220
4:25
-0.20
0.10
-8
-0.15
0.10
-6.10
CHNVDM
Xsum
CC
SAM
WiliHow
--
WV WW-BBATTWWEM
NV -VARToun
Figure 4: Parallel coordinate plots of best hyperparemeters across models and validation datasets.
After finding the hyperparameter range, we perform random search of 50 trials for each dataset to find the best regularised models. Table 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 we compare each predicted validation summary 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 measure of information in the document such that a longer summary means more information in summary, given the compression ratio of the dataset. We compare the predicted and
20

---

Published as a conference paper at COLM 2024
3.2
0.00
32
cred Summery Length
2.3
12.
Racdine
1.7-
12-
Pougel
Rouge2
Rougel
Emprniall Summmry Length
Seadline 0:15
Predk Summary Length
0.27
2.3
Iandine
W
1.4-
0.5-
6.3
0.6
6-6
0.6
6.3
22
0.5
2.
1.0
0.0
0.2
0.4
08
16
500
1903
1995
Rougel
hougeL
Empriiall Summary Langth
4.0
15
3.2
Meoteree 5 Cammn Sode
20:
Eson the
feadine
ascol
6.2-
6.5
24
Empittall Summary Length
Buaeling
Wy
0.36
Redized Summany ichech
1.90
2.20
Basckie NY
Buselre
6.50
6.30
0.30
02%
0.50 6.75
1.00
0.2
Summary
Figure 5: baseline BART (Xsum) is compared against the same reinterpreted model with NVIB regularisation Top-Bottom Curation Corpus, CNN/DailyMail, SAMsum and WikiHow validation datasets. Left-Right: Rouge-1, Rouge-2, and Rouge-L score distributions; Empirical predicted summary length with Spearman's correlation.
21

---

Published as a conference paper at COLM 2024
Table 8: Post-training regularisation on OOD text summarisation. We report validation set Rouge-L.
Out-of-Domain SAMsum WikiHow
Model
CNN/DM Xsum
CC
BART (CNN/DM)
30.56
13.12
25.24
23.11
9.19
BART-16bit (CNN/DM)
30.55 [-0.01]
13.13 [-001]]
25.25
[+0.01]
23.11
[0.00]
9.20 [+0.0]]
BART-Sbit (CN
30.47 [-0.09]
13.05 [-0.07]
25.01
[-0.23]
23.22 [+0.11]
9.17 [-0.22]
BART-4bit (CNN/DM)
30.33 [-0.23]
13.14
[-0.02]
24.78
[-0466]
22.51 [-0600] 9.17 [-0.2]
NV-BART-16bit (CNN/DM)
29.70 [-0.86]
14.05
I+0.931
25.38
[+0.14]
23.38
[+0.27] 9.40 1+0.21]
NV-BART (CNN/DM)
30.80 |+0.241
14.00
[-0.88]
25.46
1+0.221
23.57
[+0.66]
9.37 [+0.8]]
BART (Xsum)
16.57
36.47
14.41
18.68
13.35
BART-16bil (Xsum)
16.57 [0.00
36.48 [+0.01]
14.42
[+0.01]
18.67 [-0.1]]
13.35 [000]
BART-8bit (Xsum)
16.53 [-0.04]
35.78 [-0.69]
14.42
[+0.01]
17.84 [-0.84]
13.14 [-0.1]]
BART-4bit (Xsum)
16.39 0-018) 35.05
[-1.22]
14.46 [+0.04]
16.45 [-2.23] 13.05 [-0:20]
NV-BART-11bbl (Xsum)
18.96 [+2.39]
36.22
[-0.25]
17.43
[+3.02]
23.31 [+4.63]
15.06 [+171]]
NV-BART (Xsum)
19.43 [42.86]
36.45 [-0.02]
17.70 [+3.39]
23.29
[+4.61] 14.96 [+1661]
true summary length and compute their Spearman's correlation coefficent. If the summary lengths are more correlated than the baseline, it shows the regularisation adjusting the length according to the information in the article.
Showw plots for the model trained on Xsum. The first two rows show that NVIB regularisation increases the performance on the Curation Corpus and CNN CNN/DailyMail datasets and is producing longer, more accurate summaries, but the adaptation to the information in the document 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.
E
Translation Experimental Setup
E.1
Data
In this work we use commonly available translation datasets from HuggingFace We consider high-resource language pairs English-German Gelish (En-De), English to French (En-Fr) for our translation experiment with the in-domain training data being the OPUS100 dataset (Zhng et al. 2020). This is a publicly available parallel corpus with a wide domain. For out-ot-domain translation corpora
consider the Bible (Christodoulopollss & Steedman 2014), IWSLT translation dataset from 2017 (Cettolo et al. (017)) and Ted (Cettolo et al. 2012) datasets. For consistency we randomly split the 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 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 below.
E.2 Models
In this work we consider the Marian encoder-decoder model (Junczys-Dowmunt et al., 2018), which is transformer-based sequence-to-sequene model and is fine-tuned on the
httpps//lugginffaaccccooaaaats
22

---

Published as a conference paper at COLM 2024
Table Dataset statistics.
Train
Validation
Test
Dataset
En-De
En-Fr
En-De
En-Fr
En-De
En-Fr
OPUS100
1000K
1000K
2K
2K
1K
1K
Bible
42K
42K
10K
10K
10K
10K
IWSLT
197K
224K
10K
10K
8K
8K
TedTalks
1K
1K
1K
1K
1K
1K
OPUS dataset in specific monolingual language pairs En and En-Fr available on HuggingFace. We do not fine-tune or update the original model weights; we only initialise our NVTB layer form of post-training regularisation.
The Marian models use 6 layer transformer encoder and decoder with 8 attention-heade The size for the word embedding vectors and model projections are 512, and the size of the feed forward dimensions are 2048, which leads to models of approximately 74 million parameters. The inclusion of the NVIB projection layers, per attention mechanism, results in an increase of approximately 9% parameters. These projections are not trained and only initialised, which results in 81 million parameters in total. It has input context length of 512 tokens. Where the tokens are created from a SentencePiece tokenizer. For autoregressive generation these models use beam search with 4 beams, and max length of 512.
E.3 Validation Hyperparameters
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.
Table Marian model's NVIB hyperparmeter selection space for random search.
T2
T
T2
d
min
-2
-7
0
13-
le 1
le 1 38
max
5
10
5
0.05
0.8
0.3
We notice that TH and Th can be decreased by several standard deviations before the noise affects the performance. We also notice that the Tá range shows that the decoder is more sensitive to this parameter. The interpolation parameters Tf vary a lot with their sensitivity for translation, especially for the encoder. After finding suitable range for the hyperparam- eters, 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 66gure
Table 11 reports the validation results the out-of-domain translation with original base- lines and NVIB regularisation.
hnttps://huggingfacce http://..............././//e///a
23

---

Published as a conference paper at COLM 2024
0.7
4.0
6,023
0.6
0.25
-3.5
0.20
0,015
0.15
2.0
-0.010
-0,005
0.000
0.0
Figure 6: Parallel coordinate plots of best hyperparemeters across models and validation datasets.
Table 11: Post-training regularisation on text translation. We report validation set BLEU.
Out-of-Domain IWSLT
Model
OPUS100
Bible
Ted
Marian (OPUS En-De)
24.43
23.61
24.34
26.90
Marian-16bit (OPUS
24.44 [=0.01]
23.64 [-0.03]
24.36 [=0.02]
26.85 [-0.05]
Marian-Sbit (OPUS En-De)
24.33 [-0.10]
23.69 [-0.08]
24.37 [+0.03]
27.01 [+0.11]
Marian-4bit (OPUS En-De)
23.89 [-054]
22.57 [-L04]
23.79
[-0.55]
26.77 [-0.3]]
NV-Mariann-6bit (OPUS En-De)
24.11 [-032]
23.56 [-0.05]
24.50
[+0.16]
27.26 [+036]]
NV-Marian (OPUS En-De)
24.67 [+0.24]
23.64 [-0.03]
24.52
[+0.8]]
27.28 [+038]]
Marian (OPUSEn-l)
33.25
27.15
34.04
31.38
Marian-16bit En-FS
33.23 [-0.02]
27.16 [-0.01]
34.04 0.00
31.35 [-0.03]
Marian-8bit (OPUS Fr-
33.26 [+0.01]
27.16 [-0.01]
34.00 [-0.44]
31.38 [0.00]
Marian-4bit (OPUS Fr)
32.89 [-0.36]
25.09 [-206]]
33.58 [-0.46]
30.84 [-0.4]]
NV-Marian-16bit (OPUS En-Fr)
33.27 [+0.02]
27.20 [=0.05]
34.21 [+0.17]
31.66 [+028]
NV-Marian (OPUS En-Fr)
33.30 [+0.05]
27.39 [-0.24]
34.19 [-0.15]
31.89 [+0.51]
Additional Analyses
E1 Empirical Prior Analysis
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/ DDilly Mail and Xsum and calculate the em- pirical priors (n²², o², a given different in- and out-of-domatt 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) the embedding for cross-attention.
Figure 7 displays the average empirical embedding across all encoder layers for model fine- tuned on CNN/DailyMail and Xsum, respectively. We notice that the mean is approximately zero and very similar across datasets and across models. The variance close 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 for cross attention (Encoder layer 13) with consistent variation across datasets. Its clear the distribution of embeddings between fine-tuned BART encoders similar.
Figure 8 shows the average empirical embedding across all decoder layers for model fine-tuned on CNN/ Mail and Xsum, respectively. The decoder means have larger values in comparison to the encoder and similar low variance. Considering the logged
24

---

Published as a conference paper at COLM 2024
1.0
CNNION Xaum
CC
Encoder moan tony?
SANSum Vil
Enceder loglaş1
8
8.0000
Causs
-0.0025
SAMSum
-2.090
WildHow
Xaum
0.0275
CC
0.2
-2.0100
WikHoW
Gaussian
Encoder layers
11
11
Encoder layers
Encoder layers
1.0
8.0050
CNNIRM
Xeum
Encoder (dP)
0.8
SANsUm
Encoder logiag)
CNN/DM
Causs
0.6
SANsum
WkkHow
Xsum
CC
0.2
10
0.0
Goussian
Encoder layers
6 16 12
layers
Encoder layers
Figure 7: Top: BART fine-tuned on CNN/DailyMail. Bottom: BART fine-tuned on Xsum. Averaged empirical embeddings per layer Left: encoder mean component z². Middle: encoder variance (oP)2. Right: encoder logged pseudo-count log(a)).
1.0
ChINDDM
18.
0.055
Decoder mean (ant?
SAHsum
WikHHw Goussian
Decoder lagles)
WiiHoww
Gaussian
0.005
2.8
SAMsum WildHow
0.4
Gousston
layers
Decoder layers
10
11 12
Decader layers
30 12
14
0.055
0.010
CNNJDM
20
0.055
0.6
CE
SAMsum Wikitow
Decoter
-0.055
SAMsuum
WikiHow
Decoder
10
layers
Decoder layers
Figure 8: Top: BART fine-tuned on CNN/DailyMail Bottom: BART fine-tuned on Xsum. Averaged empirical embeddings per layer Left: decoder mean component uP, Middle: decoder variance (oP)?, Right: decoder logged pseudo-count log(a"). "Gaussian" is a unit Gaussian, for reference.
25

---

Published as a conference paper at COLM 2024
pseudo-count we notice that they get exceedingly large through the layers which are near double the magnitude to the encoder in log-space.
F2 Empirical prior data requirements
In this section we restrict the amount of data used to create the empirical prior and see if 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 Ihowss 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.
1.02
Relativa validation ROUGE-L
1.01
1.00
0.99
CC
CNNVOM
0.98
SAMsum
WidHow
0.97
Xsum
Model
NV BART
NV BART
0.95
100.0
75.0
50.0
25.0 10.0 10 0.1
(%) of training data
Figure The relative validation performance of Rouge-L (y-axis) is compared for different amounts of training data used to create the empirical prior (x-axis).
E3 Attention plots
In this section we consider the attention plots of our reinterpreted model in the presence of different levels of regularisation. We selected 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 Ny-Transformers. We consider three cases: firstly the identity initialisation of
and र = 10; secondly with regularisation where the tf and T 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 the prior component [P), where the initialisation parameters are &e 0 and रę -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 atten- tion maps completely ignore the prior component representation [P] and has equivalence in attention scores across attention functions. In contrast, when Ti = -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, 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.
notice the attention patterns generally shift attention away
26

---

Published as a conference paper at COLM 2024
Figure 10: Encoder at layer 10 averaged over all heads self-attention maps. Top: CNN/Dai- lymail. Bottom: Xsum. Left: Equivalence initialisation. Middle: Best post-training NVIB regularisation from validation set Right: Prior collapse.
Figure 11: Decoder cross-attention maps layer 3 averaged over all heads. Top: CNN/Dai- lymail. Bottom: Xsum. Left: Equivalence initialisation. Middle: Best post-training NVIB regularisation from validation set Right: Prior collapse.
27

---

Published as a conference paper at COLM 2024
Figure 12: Decoder causal self-attention maps averaged over heads. Top: CNN/- Dailymail layer 6. Bottom: Xsum layer 6. Left: Equivalence initialisation. Middle: Best post-training NVIB regularisation from validation set Right: Prior collapse.
28

---

Published as a conference paper at COLM 2024
E4
Generated outputs
In this section we provide example output summaries from our original baseline and reinterpreted NV models for the first example of each validation dataset. Each model only has a context of 1024 tokens in the BART model. We see that the models with post-training regularisation produce similar but different and sometimes improved predictions.
SAMsum
Article A: Hi Tom, are you busy tomorrow's 's afternoon? B: I'm pretty sure I am. What's up?
A: Can you go with to the animal shelter?. B: What do you want to do? A: want get a puppy for my son. B: That will make him happy. A: Yeah, we' ve discussed many times. think he's ready now. B: That's good. Raising a dog is a tough issue. Like having a baby 7) A: I'll get him one of those little dogs. B: One that won't grow up too big-)) A: And eat too much;-) B: Do you know which one he would like? A: Oh, yes, I took him there last Monday. He showed me one that he really liked. B: bet you had to drag him away. A: He wanted to take it home right away 7-). B: wonder what he'll name it. A: He said he' name it after his dead hamster Lemmy- he's a great Motorhead fan :-))
Summary will go the animal shelter tomorrow to get puppy for her son. They already
visited the shelter last Monday and the son chose the puppy.
BART (CNN/DM) A mother wants to get puppy for her son. She asks her son to accompany her to the animal shelter to buy him a dog. The boy takes the dog home and names it Lemmy after his dead hamster - Lemmy-- Motorhead fan. The dog is taken back to the shelter.
NV-BART (CNN/DM) A mother wants to get puppy for her son. She asks her son to take her the animal shelter. He wants a little dog. She says he wanted to take it home right away. He said he'd d name it after his dead hamster- - Lemmy.
BART (Xsum) A conversation between a mother and her son has been caught on camera, which the mother says she wants to get puppy for her son.
NV-BART (Xsum) I've been chatting to my friend Tom and he's got a puppy for his son and we're planning to get him dog.
XSum
Article The ex ex-Reading defender denied fraudulent trading charges relating to the Sodje Sports Foundation charity to raise money for Nigerian sport. Mr Sodje, 37, is jointly charged with elder brothers Efe, 44, Bright, 50 and Stephen, 42. Appearing the Old Bailey earlier, all four denied the offence. The charge relates to offences which allegedly took place between 2008 and 2014. Sam, from Kent, Efe and Bright, of Greater Manchester, and Stephen, from Bexley, are due to stand trial in July. They were all released on bail.
Summary Former Premier League footballer Sam Sodje has appeared in court alongside three brothers accused of charity fraud.
BART (CNN/DM) Sam Sodje, 37, is jointly charged with elder brothers Efe, 44, Bright, 50 and Stephen, 42. The charge relates to offences which allegedly took place between 2008 and 2014. Sam, from Kent, Efe and Bright, Greater Manchester, and Stephen, from Bexley, are due stand trial in July.
NV-BART (CNN/DM) the Sodje Sports Foundation charity to raise money for Nigerian sport.
Mr Sodje, 37, is jointly charged with elder brothers Efe, 44, Bright, 50, and Stephen, 42. The four deny fraudulent trading charges relating to the alleged offences between 2008 and 2014.
BART (Xsum) Former Premier League footballer Sam Sodje has appeared in court charged with fraud.
NV-BART (Xsum) Former Premier League footballer Sam Sodje has appeared in court charged with fraud.
29

---

Published as a conference paper at COLM 2024
Curation Corpus
Article China will allow self-driving car tests on public roads Companies can apply to test their technology on an approved set of roads.China is opening up its roads to self-driving cars. The Beijing Municipal Transport Commission released a statement today saying that on certain roads and under certain conditions, companies registered in China will be able to test their autonomous vehicles. Last year, Chinese authorities banned self-driving vehicles from the country's 's highways until new regulations could be created and approved. But those have been slow to arrive, which is why Chinese company Baidu and its CEO Robin Li came under fire this summer for conducting an apparently illegal demonstration of its driverless technology. The new regulations will allow companies to apply for temporary permission to test their autonomous vehicles on ge-tt-bbe-detemmined group of approved roads. The cars will have to have traffic accident liability insurance and human behind the wheel to take over if anything goes wrong- With this development, China now joins a number of other countries allowing self-driving technology developers to test their products in real-life scenarios.
Summary iThe Beijing Municipal Transport Commission will license companies based China to conduct tests of self-driving cars on some roads under certain conditions, set out in new regulations. These include requirement for traffic accident insurance and for a human back-up driver to be behind the wheel to take over if necessary.
BART (CNN/DM) China will allow self-driving car tests on public roads. Companies can apply test their technology on an approved set of roads. The cars will have to have traffic accident liability insurance and human behind the wheel to take over if anything goes wrong. China now joins a number of other countries allowingSelf-driinn technology developers totest their products in real-life scenarios.
NV-BART (CNN/DM) China will allow self-driving car tests on public roads. Companies can apply to test their technology on an approved set of roads. The cars will have to have traffic accident liability insurance and human behind the wheel to take over if anything goes wrong- China joins number of other countries allowing technology developers to test products in real-life scenarios.
BART (Xsum) China is opening up its roads to self-driving cars.
NV-BART (Xsum) China is opening up its roads to self-driving car tests on public roads and under certain conditions.
CNNiDailyMail
Article (CNN)Share, and your gift will be multiplied. That may sound like esoteric adage, but when Zully Broussard selflessly decided to give one of her kidneys to stranger, her generosity paired up with big data. It resulted six patients receiving transplants. That surprised and wowed her. "I thought I was going to help this one person who I don't know, but the fact that so many people can have life extension, that's pretty big," Broussard told CNN affiliate KGO. She, may feel guided in her generosity by higher power. "Thanks for all the support and prayers," a comment on a Facebook page in her name read. "I know this entire journey much bigger than all of us. I also know I'm just the messenger." " CNN cannot verify the authenticity of the page.... (Beyond 1024 tokens input)
Summary Zully Broussard decided to give a kidney to stranger.
BART (CNN/DM) Zully Broussard gave one her kidneys to a stranger. Her generosity paired up with big data. It resulted in six patients receiving transplants. "I thought was going to help this one person who I don't know," she says. The chain of surgeries is to be wrapped up Friday. NV-BART (CNN/DM) Zully Broussard gave one of her kidneys to stranger. Her generosity paired up with big data. It resulted in six patients receiving transplants. "I thought was going to help this one person who I don't 't know, but the fact that so many people can have a life extension, that's pretty big"
BART (Xsum) A San Francisco woman has become the first person in the world to donate a kidney to six people at once.
NV-BART (Xsum) A woman California has become the first person in the world to donate a kidney to a stranger in a "super-donor chain" by giving one of her own organs to another person in exchange for another person's 's kidney, in what is being hailed as an act of generosity.
30

---

Published as a conference paper at COLM 2024
WikiHow
Article Adult humans should drink 2-3 litres((rr roughly eight 8 oz glasses) of water per day whilst children should drink 1-2 litres ( (r roughly five oz glasses. Thass is in addition to things like tea and coffee. Water keeps bodies at the correct temperature and removes toxins that are the inevitable result of metabolism and industrial life. You'll be healthier automatically. Summary Drink more water.
BART (CNN/DM) Adult humans should drink 2-3 litres(or roughly eight 8 oz glasses) water per day. Water keeps bodies at the correct temperature and removes toxins. Water also clears your skin, helps your kidneys, helps to control your appetite, and keeps you energized. It also keeps you from drinking unhealthy beverages like soda and juice.
NV-BART (CNN/DM) Adult humans should drink 2-3 litres of water per day. Children should drink 1-2 litres (ooughly five 8 oz glasses) Water keeps bodies at the correct temperature and removes toxins. Water also clears your skin, helps your kidneys, and keeps you energized. It also keeps you from drinking unhealthy beverages like soda and juice.
BART (Xsu)) If you want to be healthier, you need to drink more water. NV-BART (Xsum) If you want to be healthier, and have better skin, you should drink more water.
31


Reviews:
[{"review_id": "tLljbTyOgs", "reviewer": "Reviewer_59aH", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank Reviewer **59aH** for their time, effort and clarifying questions. Furthermore, we are pleased they highlighted the strengths of our contribution: **model understanding that can benefit the community**; Novelty of our **new research direction**; and rigour by being **solid**.   \n\n*Writing issue. The paper starts with the NVIB, however, without sufficient background, section 2 is a bit hard to follow.*\n\nWe appreciate that section 2 is a very compressed presentation of the relevant aspects of prior work, and we agree that we need to use the extra page available for any camera-ready version to make this explanation more accessible.  However, for scientific diversity and developing new research directions, it must be possible to build on top of previously published work without reiteration. We believe that section 2 adequately introduces the concepts which are needed to understand our contribution.\n\n*Advantage over just keeping top-k weights?*\n\nThis question does not directly apply to the submitted paper because we evaluate NVIB entirely post-training.  For training, NVIB encourages a continuous transition from not-sparse to sparse models, whereas top-k pruning of vectors introduces a discrete discontinuity between not-pruning and pruning vectors, which is potentially an advantage for NVIB, as well as its strong information-theoretic foundations. \n\n*Understand the concept of denoising?*\n\nIn this paper, denoising attention just uses Bayesian denoising as a mathematical function used to define the generalisation of attention.  In the previous work, Nonparametric VAEs (like all VAEs) are closely related to diffusion in that they introduce noise to the embeddings during training, which the model needs to learn to remove to reconstruct the input. \n\n*Is it possible to link this research to LLMs?*\n\nCertainly; this is a topic we are currently working on, and we would be happy to discuss it more under future work. However, there are a number of complexities in LLMs (e.g. relative position embeddings, flash-attention and scaling for compute) which make this work outside the scope of the submitted paper.  The submitted paper is an essential step towards NVIB for LLMs. As LLMs get larger, the value of regularisation should increase.  Also, during LLM generation, NVIB would allow sampling to be done for the embeddings of the model, resulting in more diversity in generation.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "Model understanding can benefit the community. Overall, I think the paper is solid in some sense, but needs modification. It seems to be a new research direction without too many studies.\n\n-------\n\nUpdate after rebuttal: I raised the score --> 6.", "reasons_to_reject": "- Writing issue. The paper starts with the NVIB, however, without sufficient background, section 2 is a bit hard to follow.", "questions_to_authors": "- 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?\n\n- 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?\n\n- Is it possible to link this research to LLMs?", "extra_scores": {}}, {"review_id": "oy1EwhPuud", "reviewer": "Reviewer_SbeV", "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.\n\nIn 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. \n\nI 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.\n\nFrom 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.\n\nIf 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.\n\nI think it is borderline paper. I enjoy the general writing and presentation of this work.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank Reviewer **SbeV** for their thorough review, and for taking the time to read the prior work, which clearly helped in understanding our contribution. We are delighted they: **enjoy the general writing and presentation**; emphasized the novelty of our **new and preliminary direction** and highlighted our potential by performing **significantly better than post-training methods**.\n\n*If the authors could add some studies on a decoder-only language model with NVIB*\n\nWe certainly agree that leveraging the submitted work to apply NVIB to pretrained LLMs would improve the impact of this work, but this is a complex topic which cannot be adequately addressed in the length of this paper.  This topic is part of our current work, and we would be happy to include it in an extended discussion of future work.  \n\n*The performance of NV-based LM seems to be little bit inconsistent and dataset dependent. Some in-depth analysis is encouraged.*\n\nWe consider any improvement as an achievement for a method which is applied post-training without backpropagation or parameter updates, as shown by the comparison to baselines.\nSecondly, we believe that the variance in performance does not reflect inconsistent performance of NVIB but rather is due to the nature of the different tasks. Summarisation is fundamentally an information compression task, whereas translation is a less compressive transformation between languages (mentioned top pg9). We speculate the larger improvement in NV-BART(Xsum) is due to the abstractive nature of the Xsum dataset (mentioned mid pg8). This requires more generalisation than the extractive CNN/DM dataset which relies on selecting a summary from the document. In general, regularisation does not guarantee improved performance. Some discussion of this issue is already in the paper, but we will add further discussion if space allows.\n\n*add some min-max perturbation to simulate severe OOD cases.*\n\nThis is a very interesting and informative suggestion, and we will definitely investigate how we can construct such an evaluation.\n\n*Could adding a causal graph over the VAE framework further help with some counterfactual OOD evaluation cases?*\n\nThis suggestion and the cited papers look very interesting, but it will require further study.  Causal graphs and Counterfactual OOD are very relevant, and we think this could be a strength of our Bayesian approach, but it is out of the scope for the current paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The work aims to extend the nonparametric mixture distributions of transformer to applying regularization post-training.\n\n2. On OOD sub-tasks of summarisation, the performance NV-BART is significantly better than post-training methods.", "reasons_to_reject": "1. the performance of NV-based LM seems to be little bit inconsistent and dataset dependent. Some in-depth analysis is encouraged.", "questions_to_authors": "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.\n\n\n2. 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?\n\n***\n\nA. Large Language Models and Causal Inference in Collaboration: A Comprehensive Survey\n\nB. Training a Resilient Q-Network against Observational Interference\n\nC. Lipschitz Parametrization of Probabilistic Graphical Models", "extra_scores": {}}, {"review_id": "afKhv6CgJ7", "reviewer": "Reviewer_nToS", "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.\n\nMinor issues:\n* \"HuggingFace Wolf et al. (2020).\" => \"HuggingFace (Wolf et al. 2020).\"\n* \"BART Lewis et al. (2020)\" => \"BART (Lewis et al. 2020)\"", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate Reviewer **nToS** for their analytical view of our experiments section. We are pleased they emphasized the contributions of: our **useful** identity initialisation, which allows access to pretrained models; Post-training regularisation **without decreasing performance** and  **without backpropagation**; validating our results on **summarization** and **translation** across **in domain** and  **out of domain**.\n\n*The regularization doesn't generally yield benefits, even in out-of-domain settings.*\n\nWe appreciate that Reviewer **nToS** found our paper **theoretically interesting**.  We believe that the empirical results are sufficient to support these theoretical claims about the nature of pretrained transformers, given that there are small improvements in most cases. In general, regularisation does not guarantee improved performance.  It is surprising there exists any improvements that generalise to test set performance given that this method is applied post-training without backpropagation or parameter updates. We believe that the variance in performance is due to the nature of the different tasks. Summarisation is fundamentally an information compression task, whereas translation is a less compressive transformation between languages (mentioned top pg9). We conjecture that the larger improvement in NV-BART(Xsum) is due to the abstractive nature of the Xsum dataset (mentioned mid pg8). This requires more generalisation than the extractive CNN/DM dataset which relies on selecting a summary from the document. Some discussion of this issue is already in the paper, but we will add further discussion if space allows.\n\n*The paper is hard to read in several places because it offloads too much to the appendix*\n\nWe agree that the analysis paragraph in section 4.2 would be much improved by moving some representative plots from the appendix into the body of the paper, and we will do this for any camera-ready version.  This was necessary given space constraints.  We will expand the discussion in the body of the paper, as space allows, but given the content quantity, the appendix will still be necessary.\n\n*Is there prior research that would have suggested that quantization would improve performance?*\n\nYes. \"We find that reducing the precision steadily from 16 to 4 bits increases the zero-shot performance\" - The case for 4-bit precision: k-bit Inference Scaling Laws - Dettmers - ICML2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* Showing that NVIB attention can be initialized such that it matches standard attention seems like a useful result.\n* The post-training regularization approach allows model uncertainty to be increased without decreasing performance and without back propagation or parameter updates.\n* Summarization and translation models are evaluated with the proposed regularization both in domain and out of domain", "reasons_to_reject": "* 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.)\n* 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_to_authors": "* 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.", "extra_scores": {}}, {"review_id": "1Ff7bgVN8N", "reviewer": "Reviewer_nwsJ", "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.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer **nwsJ** for their clarifying questions. We are pleased they found our paper **interesting and elegant**.\n\n*The manuscript is extremely hard to follow*\n\nWe appreciate that the presentation of the relevant aspects of prior work is very compressed, and we agree that we need to make this explanation more accessible in any future version.  This will make the Bayesian nature of our model clearer, in addition to the answers below. \n\n*Make the prior distribution, likelihood and posterior distribution more clear?*\n\nWe had hoped that the contribution of this paper could be understood without this explanation, but we are happy to add it as it clarifies the Bayesian nature of our contribution.  In section 3, we are inferring DP representations (distributions over mixtures) given the input.  The prior is a DP with a pseudocount $\\alpha^p_0$ and mono-modal base distribution $G^p_0 = N(\\mu^p,\\sigma^p),$ which are set from the empirical distribution.  The NN converts the input into a set of pseudo-observations $(\\mu^q,\\sigma^q,\\alpha^q)$.  Using exact inference of DPs, each of these pseudo-observations $(\\mu^q_i,\\sigma^q_i,\\alpha^q_i)$ becomes one component $N(\\mu^q_i,\\sigma^q_i)$ of the posterior's base distribution $G^q_0$, which are added to the one component $N(\\mu^p,\\sigma^p)$ coming from the prior.  These components are weighted by their respective alphas (bottom pg3).  The likelihood is only used to justify the KL-divergence regulariser via the ELBo.  We will add a more explicit specification with precise equations to the paper. \n\n*What’s the meaning of $Z^2$ in Equation 7?*\n\nThis and other algebraic computations are meant to be interpreted component-wise.  We will make this explicit.\n\n*Could you provide an algorithm box showing your modification to attention?*\n\nYes. An algorithm or pseudocode box is a great idea. We can provide this in the appendix.\n\n*What’s the motivation to set the parameters as in Equation 7?*\n\nIntuitively, $\\alpha_i$ determines the weight of the component for $z_i$, which we can see in equation (1) needs to be a function of the $L^2_2$ norm of $z_i$.  For this reason we allow a log-quadratic relationship between $z_i$ and $\\alpha_i$ in equation (7).  We can make this explanation in the paper clearer.  More precisely, the bias term in equation (4) needs to go to a constant to get equivalence, which again requires a log-quadratic relationship between $\\alpha_i$ and $z_i$ given the linear relationship between $z_i$ and $\\mu_i$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* The idea is interesting and elegant.", "reasons_to_reject": "* The manuscript is extremely hard to follow, even for a Bayesian expert.\n* There are several points not clear to me even after I go through the manuscript for many times. See the Questions part below.", "questions_to_authors": "* 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.\n* What’s the meaning of $Z^2$ given $Z$ is a vector (e.g. in Equation 7)?\n* Could you provide an algorithm box showing what’s your modification compared with the standard attention mechanism?\n* What’s the motivation to set the parameters as in Equation 7? I don’t get any motivation on the necessity of $Z^2$.", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept