A Variational AutoEncoder for Transformers with Nonparametric Variational Information Bottleneck

TL;DR

This paper introduces a novel nonparametric variational information bottleneck for Transformer models, enabling flexible regularization of attention-based embeddings by capturing permutation invariance and variable vector support.

Contribution

It develops a nonparametric variational information bottleneck (NVIB) for Transformer embeddings using Bayesian nonparametrics, and applies it to create a nonparametric variational autoencoder (NVAE).

Findings

NVAE regularizes the number of vectors in attention.

Embedding space exhibits VAE properties for Transformers.

Initial experiments show desired embedding properties.

Abstract

We propose a VAE for Transformers by developing a variational information bottleneck regulariser for Transformer embeddings. We formalise the embedding space of Transformer encoders as mixture probability distributions, and use Bayesian nonparametrics to derive a nonparametric variational information bottleneck (NVIB) for such attention-based embeddings. The variable number of mixture components supported by nonparametric methods captures the variable number of vectors supported by attention, and the exchangeability of our nonparametric distributions captures the permutation invariance of attention. This allows NVIB to regularise the number of vectors accessible with attention, as well as the amount of information in individual vectors. By regularising the cross-attention of a Transformer encoder-decoder with NVIB, we propose a nonparametric variational autoencoder (NVAE). Initial experiments on training a NVAE on natural language text show that the induced embedding space has the desired properties of a VAE for Transformers.