A Variational AutoEncoder for Transformers with Nonparametric Variational Information Bottleneck

Attention-based deep learning models, such as Transformers (Vaswani et al., 2017; Devlin et al., 2019), have achieved unprecedented empirical success in a wide range of cognitive tasks, in particular in natural language processing (NLP). On the other hand, deep variational Bayesian approaches to representation learning, such as variational autoencoders (VAEs) (Kingma and Welling, 2014), have also been very influential, especially due to their variational information bottleneck (VIB) (Alemi et al., 2017; Kingma and Welling, 2014) for regularising the induced latent representations. Previous VIB methods only apply to a vector space, and Transformers crucially do not use a single vector as their latent representation, instead using a set of vectors (Lin et al., 2020; Fang et al., 2021; Park and Lee, 2021). This allows the number of vectors in a Transformer embedding to grow with the size of the input, which is essential for embedding natural language text (Bahdanau et al., 2015), where the size of the input can range from a single word to thousands of words. In this paper, we propose a variational information bottleneck regulariser for set-of-vector latent representations, and use it to regularise the induced latent representation of a Transformer encoder-decoder variational autoencoder.