The Strong Lottery Ticket Hypothesis for Multi-Head Attention Mechanisms

TL;DR

This paper provides a theoretical foundation for the strong lottery ticket hypothesis in multi-head attention mechanisms of transformers, showing that high-performing subnetworks exist and can approximate arbitrary MHAs with high probability.

Contribution

It introduces the first theoretical analysis of SLTs in MHAs and extends the SLTH to transformers without normalization layers.

Findings

SLTs exist in randomly initialized MHAs with high probability

Approximation error decreases exponentially with hidden dimension

Theory validated through empirical experiments

Abstract

The strong lottery ticket hypothesis (SLTH) conjectures that high-performing subnetworks, called strong lottery tickets (SLTs), are hidden in randomly initialized neural networks. Although recent theoretical studies have established the SLTH across various neural architectures, the SLTH for transformer architectures still lacks theoretical understanding. In particular, the current theory of the SLTH does not yet account for the multi-head attention (MHA) mechanism, a core component of transformers. To address this gap, we introduce a theoretical analysis of the existence of SLTs within MHAs. We prove that, if a randomly initialized MHA of $H$ heads and input dimension $d$ has the hidden dimension $O(d\log(Hd^{3/2}))$ for the key and value, it contains an SLT that approximates an arbitrary MHA with the same input dimension with high probability. Furthermore, by leveraging this theory for MHAs, we extend the SLTH to transformers without normalization layers. We empirically validate our theoretical findings, demonstrating that the approximation error between the SLT within a source model (MHA and transformer) and an approximate target counterpart decreases exponentially by increasing the hidden dimension of the source model.