Attention Perturbed: Finite-Width Analysis of Attention-Based Deep Neural Network Architectures at Initialization

Abstract

In this work we extend the methods of an “effective” Theory of Deep Learning to investigate the leading order behavior and implicit bias of large-scale language models relying on “attention” mechanisms, at initialization. Due to the algebraic complexity inherent to these attention-based, block architectures we introduce first-handedly a succinct and universal blueprint for a neural network’s architecture, initialization and training/ testing schedule, which can be repurposed for any modern neural network. Utilizing our modular blueprint as a guide, we then perform a series of perturbative analyses of versions of the Transformer Encoder model. We obtain a promising path to a full recursive power expansions in the model’s embedding dimension 1 d of the expected behavior of the output of the attention layer, at initialization. We also introduce the concept of the network’s Neural Tangent Kernel, at initialization. Our goal is to provide an initial bridge from the recently successful general Theory of Deep Learning to a tangible explanation of the dynamics and high-performance of the specific attention-based residual architecture found in the Transformer, which lays at the root of the recent advancements in Natural Language Processing and Computer Vision. “