  There is a growing amount of literature on the relationship between wide
neural networks (NNs) and Gaussian processes (GPs), identifying an equivalence
between the two for a variety of NN architectures. This equivalence enables,
for instance, accurate approximation of the behaviour of wide Bayesian NNs
without MCMC or variational approximations, or characterisation of the
distribution of randomly initialised wide NNs optimised by gradient descent
without ever running an optimiser. We provide a rigorous extension of these
results to NNs involving attention layers, showing that unlike single-head
attention, which induces non-Gaussian behaviour, multi-head attention
architectures behave as GPs as the number of heads tends to infinity. We
further discuss the effects of positional encodings and layer normalisation,
and propose modifications of the attention mechanism which lead to improved
results for both finite and infinitely wide NNs. We evaluate attention kernels
empirically, leading to a moderate improvement upon the previous
state-of-the-art on CIFAR-10 for GPs without trainable kernels and advanced
data preprocessing. Finally, we introduce new features to the Neural Tangents
library (Novak et al., 2020) allowing applications of NNGP/NTK models, with and
without attention, to variable-length sequences, with an example on the IMDb
reviews dataset.
