Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers

TL;DR

This paper develops a statistical mechanics theory for a tractable deep attention model, revealing how attention paths and kernel combinations influence generalization, with experimental validation and implications for network pruning.

Contribution

It introduces an analytically solvable model linking attention path interactions to kernel-based learning and generalization in Transformers.

Findings

Predictor statistics as a sum of attention path kernels

Kernel combination aligns with task labels, improving generalization

Pruning less relevant attention heads based on theory

Abstract

Despite the remarkable empirical performance of Transformers, their theoretical understanding remains elusive. Here, we consider a deep multi-head self-attention network, that is closely related to Transformers yet analytically tractable. We develop a statistical mechanics theory of Bayesian learning in this model, deriving exact equations for the network's predictor statistics under the finite-width thermodynamic limit, i.e., $N,P\rightarrow\infty$, $P/N=\mathcal{O}(1)$, where $N$ is the network width and $P$ is the number of training examples. Our theory shows that the predictor statistics are expressed as a sum of independent kernels, each one pairing different 'attention paths', defined as information pathways through different attention heads across layers. The kernels are weighted according to a 'task-relevant kernel combination' mechanism that aligns the total kernel with the task labels. As a consequence, this interplay between attention paths enhances generalization performance. Experiments confirm our findings on both synthetic and real-world sequence classification tasks. Finally, our theory explicitly relates the kernel combination mechanism to properties of the learned weights, allowing for a qualitative transfer of its insights to models trained via gradient descent. As an illustration, we demonstrate an efficient size reduction of the network, by pruning those attention heads that are deemed less relevant by our theory.