PaperlayerView exact revision on arXiv

Source-linked AI summary

The Strong Lottery Ticket Hypothesis for Multi-Head Attention Mechanisms

Hikari Otsuka, Daiki Chijiwa, Yasuyuki Okoshi, Daichi Fujiki, Susumu Takeuchi, Masato Motomura

arXiv:2511.04217v1cs.LGcs.AI

TL;DR

Although SLTH theory covers several neural architectures, its applicability to transformer attention remained an open question. This paper analyzes strong lottery tickets in multi-head attention and extends the result to transformers without normalization layers. It reports that approximation error falls exponentially with source hidden dimension and that its theory-informed weight scaling improves SLTs in practical experiments.

  • Problem

    Existing SLTH theory had not established whether transformer attention, a distinctive component of transformers, contains strong lottery tickets.

    Source: §1 · pp. 1–2

  • Method

    The paper reinterprets the query-key inner product as a neural network to prove an attention approximation result, then extends the theory to transformers without normalization layers.

    Source: §1 · p. 2

  • Results

    Approximation error decays exponentially as source hidden dimension increases, and the authors report improved practical SLTs with a theory-informed weight initialization.

    Source: §1 · p. 2, §6 · p. 8

  • Takeaways & Limitations

    The results extend SLTH theory to multi-head attention and transformers without normalization layers, while indicating a weight-scale strategy for improving practical SLTs.

    Source: §6 · pp. 7–8

    The stated approximation lemma assumes target weights and inputs have norms at most 1 and source weights are uniformly initialized on [−1, 1].

    Source: §2.1 · p. 2

Abstract

from arXiv · show

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.

1 Introduction

The SLTH asks whether randomly initialized networks contain untrained subnetworks that match trained dense networks, but its theory had not established this for transformer attention. This work proves an attention-specific approximation result, extends the theory to transformers without normalization layers, and reports empirical support.

1 Introduction

  • Theorem 1 gives a high-probability approximation of an arbitrary target attention mechanism with hidden dimension n = O(d log(d^(3/2)/ϵ)).The source and target have the same input dimension.Source: §1 · p. 2
  • The query-key inner product can be treated as a neural network, enabling a two-layers-for-one-style approximation without adding layers to the MHA.The query and key projections form the source's two layers, while their merged weight matrix forms the target's single layer.Source: §1 · p. 2
  • The attention result extends the SLTH to randomly initialized transformers without normalization layers that approximate arbitrary transformers of similar structures.The paper states this as Theorem 6.Source: §1 · p. 2
  • Approximation error decays exponentially as source hidden dimension increases and does not diverge as input length increases.The paper also reports that a theoretically motivated weight initialization leads to better SLTs in experiments.Source: §1 · p. 2
  • The authors empirically validate the theory with synthetic experiments designed to test how hidden dimension and input length affect SLT approximation error.They also report that the theory motivates a practical weight initialization for finding better SLTs.Source: §1 · p. 2

2 Preliminaries

The preliminaries review the SLTH and define the multi-head attention setting used in the paper. They contrast the established two-layers-for-one approximation with MHA's per-head computation and output projection.

2 Preliminaries

  • The section introduces the SLTH and reviews its theoretical development from fully connected networks to more complex architectures.The hypothesis concerns subnetworks in randomly initialized networks that match dense trained networks without weight updates.Source: §2 · p. 2, §2.1 · p. 2

2.1 Strong Lottery Ticket Hypothesis

  • Conventional SLTH theory approximates one target weight matrix by pruning two randomly initialized matrices, the two-layers-for-one approximation.The cited passage identifies this approach as a foundation for SLTH theory in more complex architectures.Source: §2.1 · p. 2

2.2 Multi-head Attention Mechanisms

  • MHA computes pairwise relationships among a sequence of embeddings at H heads, and partitioning the output projection makes the result a sum of head outputs.The formulation assumes each embedding attends to at least one other embedding.Source: §2.2 · p. 3, Figure 2 · p. 3

3 Strong Lottery Ticket Hypothesis for Transformers

The section sets up target and randomly initialized source MHAs, then proves that pruning can produce an SLT approximation and extends the result to transformers without normalization layers.

3.1 Setups

  • The setup contrasts an arbitrary-weight target MHA with a randomly initialized, pruned source MHA whose key and value hidden dimensions may differ from the target’s.The source query and key weights are initialized independently from U[−n^1/4, n^1/4], while value and output weights use U[−1, 1].Source: §3.1 · p. 3

3.2 The Existence of SLTs Within an MHA

  • With sufficiently large source hidden dimensions, binary pruning masks make an MHA approximate its target within error ϵ with high probability, with the bound independent of input length T.The proof applies per-head approximations and a union bound across heads.Source: §3.2 · pp. 3, 5
  • Merging adjacent target query-key and value-output projections reduces the approximation task to two merged matrices that pruned source projections approximate.The query-key merge enables a variant of the two-layers-for-one approximation; the value-output pair is treated similarly.Source: §3.2 · pp. 4–5

3.3 The Existence of SLTs Within a Transformer

  • For transformers without normalization layers and with B ≥ 2 blocks, sufficiently large source MHA and fully connected hidden dimensions yield a pruned model approximating the target with high probability.The proof combines module-level SLT approximations, error-propagation bounds, and a union bound.Source: §3.3 · p. 5

4 Experimental Results

Experiments support the predicted exponential decrease in approximation error as hidden dimensions grow, and show that error remains bounded as sequence length increases. In practical GPT-2 settings, query/key weight scaling near the theoretical value is associated with lower loss.

4.1 Experimental Settings

  • The synthetic experiments use angular-velocity sequences and theoretical weight initialization, while practical experiments train GPT-2-family models on WikiText-103.The synthetic setup approximates targets with 100 randomly initialized source MHAs; practical SLTs are found with edge-popup.Source: §4.1 · pp. 5–6

4.2 Verification of Main Theorems

  • ϵ = 0.8 exp(−0.06nK) fits the MHA approximation errors, supporting exponential decay as hidden dimensions increase.The result supports the requirement that each source hidden dimension scale as O(log(1/ϵ)).Source: Figure 4 · p. 6, §4.2 · p. 6
  • As sequence length T increases, MHA approximation error remains bounded, with larger hidden dimensions lowering the bound.This is consistent with the theoretical bound being independent of T.Source: Figure 5 · p. 6, §4.2 · p. 6
  • Across transformer block counts, approximation error decreases rapidly with hidden dimensions; a shared decay rate fits the blocks while their scale factors vary.The experiment uses equal MHA and feed-forward hidden dimensions and an untrained target model.Source: Figure 6 · p. 6, §4.2 · pp. 6–7

4.3 Behavior of SLTs in Practical Settings

  • In GPT-2 models, SLTs with query/key scaling tend to have lower validation loss, and loss is lowest near nK^1/4 across models.Increasing the scale from 1 lowers loss up to a point, while scaling beyond nK^1/4 raises it.Source: Figure 7 · p. 7, §4.3 · p. 7

5 Related Work

Prior SLTH work developed theory for several architectures, while empirical studies explored randomly weighted transformers and SLTs. This paper contributes theoretical support for SLT existence in attention mechanisms and transformers.

  • SLTH research has expanded through methods that relax source-network overparameterization, impose additional constraints, or extend the theory to new architectures.This work belongs to the architectural-extension direction.Source: §5 · p. 7
  • Empirical studies report capabilities in randomly weighted transformers and SLTs, including toy tasks and language tasks.The cited studies include comparable accuracy on translation and language-understanding tasks, and high accuracy on toy tasks as hidden dimension grows.Source: §5 · p. 7

6 Conclusion

The paper proves that randomly initialized MHAs with logarithmically large hidden dimensions contain SLTs approximating arbitrary MHAs with high probability. It also extends the theory to transformers without normalization layers and reports empirical support, including improved practical SLTs with theoretically motivated query/key scaling.

  • A randomly initialized MHA with logarithmically large hidden dimensions contains an SLT approximating an arbitrary MHA with high probability.The conclusion describes this as extending SLTH theory to MHAs.Source: §6 · p. 7
  • MHA approximation requires no additional layers, unlike existing theories that approximate a single-layer structure with a two-layer one.The MHA result is also used to establish SLTH for transformers without normalization layers.Source: §6 · p. 8
  • Empirical results are consistent with the theory, and the theoretically motivated query/key weight scale improves SLT performance in practical settings.Source: §6 · p. 8

A Proofs of Main Theorems

The proof combines masked random-weight approximation with a bound on softmax perturbations. These ingredients establish SLT existence for attention mechanisms and, by extension, transformer blocks and transformers without normalization layers.

A Proofs of Main Theorems

  • The main proofs use a lemma for approximating target weights by pruning random matrices and a lemma bounding softmax perturbations.Together, these establish SLTH for attention mechanisms and extend it to transformer blocks and transformers without normalization layers.Source: A Proofs of Main Theorems · p. 10

A.1 Weight Approximation

  • For a target matrix with entries in [−1, 1], pruning two random matrices yields a masked matrix product approximation under the lemma’s width condition, with probability at least 1 − ϵ.Source: §A.1 · p. 10
  • Binary masks select products in a subset sum to approximate target-matrix entries; a union bound gives simultaneous approximation of all entries with probability at least 1 − ϵ.The construction applies to entries formed from random matrices with i.i.d. U[−1, 1] entries.Source: §A.1 · p. 10

A.2 Spectral Norm of Softmax Difference

  • The softmax perturbation bound depends on the maximum input perturbation and is independent of the input index.The stated upper bound is d1 · α · (exp(2ϵmax) − 1).Source: §A.2 · p. 11

A.3 SLT Existence within Attention Mechanisms

The section proves that pruning a randomly initialized multi-head attention mechanism can approximate an arbitrary target MHA with high probability. Its proof merges target projections, approximates the resulting weights, and bounds how errors propagate through softmax and across heads.

A.3 SLT Existence within Attention Mechanisms

  • Theorem 9 establishes high-probability approximation of an arbitrary target MHA by a suitably pruned source MHA under the stated input and dimension conditions.The proof combines per-head approximation bounds and a union bound.Source: §A.3 · pp. 11, 13
  • The proof first merges each target head’s query/key and value/output projections, then uses weight approximation to construct corresponding pruned source weights.The merging represents each target head using two combined matrices.Source: §A.3 · pp. 12–13
  • Query/key weight errors are bounded through softmax, while value/output weight approximation and total-error analysis bound the resulting MHA output discrepancy.The bounds are combined across heads before the final probability guarantee.Source: §A.3 · pp. 12–13

A.4 SLT Existence Within Transformer Blocks

The section extends SLT existence from attention mechanisms to transformer blocks by combining the attention result with a fully connected network approximation theorem. Under the theorem’s conditions, a pruned random source block approximates a target block with high probability.

A.4 SLT Existence Within Transformer Blocks

  • Theorem 11 establishes that a pruned random transformer block approximates the target block within ε with probability at least 1 − ε under the stated hidden-dimension conditions.The source block pairs a pruned MHA with a pruned fully connected network.Source: §A.4 · pp. 14–15
  • The proof first applies the attention SLT theorem, then applies the fully connected SLT theorem to the source network receiving the attention output.The target fully connected network has L layers, while the source construction uses 2L layers.Source: §A.4 · pp. 14–15
  • The combined guarantee relies on the attention approximation controlling the input passed to the fully connected stage, after which the two approximation errors are bounded together.A union bound combines the success probabilities for the two stages.Source: §A.4 · pp. 14–15

A.5 SLT Existence Within Transformers Without Normalization Layers

The section extends the transformer-block result to transformers without normalization layers by controlling approximation error across successive blocks. It also notes that the experiments can be verified using four NVIDIA H100 SXM5 GPUs.

A.5 SLT Existence Within Transformers Without Normalization Layers

  • The theorem guarantees a pruned random transformer approximates the target transformer within ε with probability at least 1 − ε, given the stated source hidden-dimension conditions.The proof propagates approximation error through the blocks and bounds the final discrepancy.Source: §A.5 · pp. 16, 21
  • The proof analyzes block errors sequentially, using each successful block approximation to bound the input error passed to the next block.It carries these bounds through the final block and combines blockwise success probabilities with a union bound.Source: §A.5 · pp. 16–21

B Experimental Details

B.1 Synthetic Data Experiment

The experiment uses synthetic angular-velocity estimation to evaluate MHA and transformer approximations. It varies source hidden dimension and reports approximation error across randomly initialized source networks.

B.1 Synthetic Data Experiment

  • The synthetic task estimates angular velocity from sequences of two-dimensional vectors generated by an angular velocity and initial phase.Training, validation, and test sets each contain 10,000 samples, with sequence lengths from 4 to 256 during training.Source: §B.1 · p. 22
  • The MHA experiment uses single-head source and target modules with input dimension 2 and output dimension 1; the transformer experiment uses single-head models with input and output dimensions 2.The transformer target query and key dimensions are 8, and source and target share the same regression token.Source: §B.1 · p. 22
  • For each target, the study tests 100 randomly initialized source networks, solves the pruning subset-sum problem with Gurobi, and reports mean and standard deviation of approximation error.The hidden-dimension experiments fix input length at 4 and fit exponential decay curves to the error.Source: §B.1 · p. 22

B.2 Language Modeling Experiment

The language-modeling experiment searches for strong lottery tickets in GPT-2-family models trained on WikiText-103. It uses edge-popup at 30% retained connections and reports the mean and standard deviation across three seeds for each model size.

B.2 Language Modeling Experiment

  • Three random seeds per model size are used, with the mean and standard deviation of the best performance reported.Source: §B.2 · p. 22
  • GPT-2-family models are trained on WikiText-103 with a maximum sequence length of 1024.Source: §B.2 · p. 22
  • Edge-popup searches for subnetworks by scoring connections and retaining the top 30% during training.Source: §B.2 · p. 22
  • Training lasts 50 epochs of 227 steps, using AdamW with cosine-decayed learning rates and a one-epoch linear warm-up.The learning rate starts at 0.0001 and decays to 0.00001.Source: §B.2 · p. 22
Loading 2511.04217v1…