Nonparametric Variational Regularisation of Pretrained Transformers

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.

Rating

4

Confidence

4

Ethics flag

1

Reasons to accept

* The idea is interesting and elegant.

Reasons to reject

* The manuscript is extremely hard to follow, even for a Bayesian expert. * There are several points not clear to me even after I go through the manuscript for many times. See the Questions part below.

Questions 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. * What’s the meaning of $Z^2$ given $Z$ is a vector (e.g. in Equation 7)? * Could you provide an algorithm box showing what’s your modification compared with the standard attention mechanism? * What’s the motivation to set the parameters as in Equation 7? I don’t get any motivation on the necessity of $Z^2$.

Summary

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

Rating

6

Confidence

3

Ethics flag

1

Reasons to accept

* Showing that NVIB attention can be initialized such that it matches standard attention seems like a useful result. * The post-training regularization approach allows model uncertainty to be increased without decreasing performance and without back propagation or parameter updates. * 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.) * 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.

Summary

This paper explores a method based on NVIB (Henderson and Fehr, ICLR 2023) to apply as regularization over multi-head transformers for OOD text summarization and machine translation tasks. In an encoder-decoder model, BART, the NVIB projection layer combines the Dirichlet Process-based empirical prior with the transformer embeddings to produce a representation for denoising MHA. In the NV-Marian case, the improvement of NVIB is less significant. I spent some time reading NVIB in two of its original references; it appears to be a relatively new and preliminary direction with only a few studies. From my perspective, it is acceptable if NVIB does not serve as the best regularization/post-training boosting method for some applications. For example, motivating and explaining the limitations of this form of non-parametric variational Bayesian application in language modeling would be interesting to general audiences. If the authors could add some studies on a decoder-only language model with NVIB (e.g., small-scale GPT-2 or LLama), it could potentially increase the general impact and interest of this work across more language modeling applications. I think it is borderline paper. I enjoy the general writing and presentation of this work.

Rating

6

Confidence

3

Ethics flag

1

Reasons to accept

1. The work aims to extend the nonparametric mixture distributions of transformer to applying regularization post-training. 2. 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. 2. This might be a bit open. There are some causal variational Bayesian works [A,B] in robust RL and sequence modeling that improve OOD robustness via local Lipschitz constant measurement. Could adding a causal graph over the VAE framework further help with some counterfactual OOD evaluation cases? *** A. Large Language Models and Causal Inference in Collaboration: A Comprehensive Survey B. Training a Resilient Q-Network against Observational Interference C. Lipschitz Parametrization of Probabilistic Graphical Models

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.

Rating

6

Confidence

3

Ethics flag

1

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. ------- Update 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? - For denoising attention, how should we understand the concept of denoising? For example, in diffusion generative modeling, we want to recover clean images from noisy ones, involving a denoising process. How does denoising work in your context? - Is it possible to link this research to LLMs?

Thanks for your clarification.

Thanks for your clarification. Indeed I think this manuscript can be interesting and elegant with some potential impacts. However the overall writing needs to be improved to make sure the core component can be well understood by the general audience. I would rather encourage the authors to polish the paper and re-submit.

Update after rebuttal

Sorry for the late reply and thank you for your clarification. I think the paper could be sufficiently clarified given an extra page in the camera-ready version. I raised my score.

Re: Rebuttal

Thank the authors for the response. Given the current scope is still fixed (e.g., decoder lm exclusive & evaluation on OOD), I will keep my current score. If the authors will revise the presentation on mentioning the challenges and related directions discussed here, the final version would be with boarder impacts and audiences. It is a borderline and interesting paper to me but I won’t strongly argue for acceptance.