Gaussian Equivalence for Multi-Head Self-Attention

1 Introduction

• A layer-level comparison is needed because concatenated-head spectral moments mix heads, while value projections may depend on the keys that determine attention.Marginal single-head laws alone do not specify these interactions.Source: §1 · p. 1
• The centered multi-head output and its Gaussian equivalent converge almost surely to the same spectral law, weakly and in every fixed moment.The comparison covers value and output projections that may depend on the keys, under the stated assumptions.Source: §1 · pp. 1–2
• The resulting spectral laws predict frozen-feature readout behavior and show how head count, width allocation, sharing, and inverse temperature relate to estimation error.At fixed parameter count, head count can produce multiple error peaks, and inverse temperature can reverse which Q/K–V width allocation has lower error.Source: §1 · p. 2

1.1 Related Work

• Earlier spectral work studies single-head laws, rank collapse, or attention interpolation; this paper compares spectra across concatenated heads and key-dependent random projections.It also treats proportional growth of token count and projection widths at fixed head count, retaining finite width-to-token ratios.Source: §1.1 · p. 2
• The laws distinguish spectrum-preserving across-head K/V sharing from within-head K/V dependence, which changes the covariance governing the output.The paper connects these spectral distinctions to frozen-feature readout models and their predicted training loss and coefficient estimation error.Source: §1.1 · p. 2

2.1 Multi-Head Self-Attention

• Multi-head self-attention applies each head’s attention matrix to its value features, concatenates the outputs, and applies an output projection.The output projection maps the concatenated head features to the output width.Source: §2.1 · p. 3

2.2 Assumptions

• The analysis uses deterministic orthogonal inputs, Gaussian queries and keys, fixed head count and inverse temperature, and proportionally growing dimensions.Value and output weights may be non-Gaussian, shared or dependent across heads, and dependent on keys; the joint key, value, and output weights are independent of the queries, with operator-norm bounds.Source: §2.2 · p. 3

3 Gaussian Equivalence

• The Gaussian model replaces centered attention scores with linear score terms plus independent noise, with θ1 = e^β² − 1 and θ2 = β².Its projected rows are conditionally independent centered Gaussians with covariance determined by the projected feature covariance.Source: §3 · pp. 3–4
• The centered output and its Gaussian counterpart share an almost-sure limiting spectral law, weakly and in every fixed moment.The theorem assumes the conditional covariance spectrum of the Gaussian output rows converges to a deterministic law.Source: §3 · p. 4
• The bulk limit is the compound free Poisson law π(μVO, γO), while the leading singular value is determined by the projected common-row component.The comparison proceeds through fixed spectral moments after admissible right multiplication.Source: §3 · p. 4, Table 1 · p. 4

4 Proof of Equivalence

• Four transformations—normalizer removal, centering, Gaussian replacement, and covariance identification—preserve fixed spectral moments after admissible right multiplication.The right multiplier is bounded in operator norm and satisfies the stated dependence conditions.Source: §4 · p. 5

4.1 Removing the Row Normalizers

• Replacing row normalizers by their concentrated limits makes the centered softmax features spectrally comparable to exponential feature matrices.The moment comparison holds after right multiplication.Source: §4.1 · p. 5

4.2 Centering the Rows

• Subtracting the conditional row mean preserves fixed spectral moments because the removed component has rank one and vanishing norm.The rows are independent after conditioning on keys and the multiplier.Source: §4.2 · pp. 5–6

4.3 Concentration and Gaussian Comparison

• Gaussian replacement compares the centered feature rows with Gaussian rows of matching conditional covariance, preserving fixed moments almost surely.The argument controls row quadratic forms and applies under the stated multiplier and query conditions.Source: §4.3 · p. 6

4.4 Identifying the Conditional Covariance

• The identified covariance has limiting coordinate variance θ1 = e^β² − 1 and leading off-diagonal coefficient θ2 = β²; independent noise supplies the remaining variance.The linear scores carry the leading correlations in the Gaussian model.Source: §4.4 · p. 7

4.5 Conclusion of the Proof

• The centered MHA output and its Gaussian counterpart have matching limiting moments, and the comparison yields weak convergence for the original output as well.The final step uses the Gaussian model’s compactly supported limit; restoring the common-row component changes the output by rank at most one.Source: §4.5 · p. 7

5.1 Spectra and Readout Predictions

• 0.24%–2.82% mean relative distances show close finite-size spectral agreement between softmax MHA and GE across seven configurations.The configurations include across-head and within-head sharing.Source: §5.1 · p. 8, Table 2 · p. 8
• Increasing head count changes a nonzero training-loss plateau into slow and then rapid decay; zero mass sets the plateau and small positive eigenvalues govern slow decay.Figure 2(a) compares the predicted curves with softmax learning curves.Source: §5.1 · p. 9, Figure 2 · p. 9

5.2 Width Allocation and Temperature

• At fixed projection parameter count, width allocation produces two coefficient-error peaks near hdV = n and dO = n.With fixed per-head widths and dO = n/2, there is one peak near hdV = dO.Source: §5.2 · p. 9
• Near β = 1, coefficient estimation favors wider V than Q/K at a fixed projection budget, while the preferred allocation depends on temperature.The ranking of (64, 192) and (128, 128) reverses as β varies.Source: §5.2 · p. 9

5.3 Parameter Sharing

• Across-head K/V sharing preserves the limiting spectrum and readout predictions at fixed head count and widths, while reinvesting saved weights in wider heads can shorten training time.In the reported example, reallocating the savings reduces the time to relative loss 0.1.Source: §5.3 · p. 9
• At equal widths and parameter count, GQA has lower coefficient error for h = 2, 4, while complete within-head K/V tying is lower for h = 8, 16, 32.The comparison concerns a fixed total width divided among heads.Source: §5.3 · p. 10, Figure 3 · p. 10

6 Discussion

• The theory predicts training loss and coefficient error from architectural parameters, helping identify uses of a fixed parameter budget in frozen-feature readouts.The proposed recommendations concern width allocation and parameter sharing.Source: §6 · p. 10

6.1 Implications for Learning

• For coefficient estimation near β = 1, use wider V than Q/K at a fixed budget; at low β, preserving Q/K width avoids weak feature directions.Increasing head count cannot eliminate a loss plateau caused by insufficient output width.Source: §6.1 · p. 10
• Across-head sharing saves parameters while maintaining performance at fixed widths, whereas narrow-head coefficient estimation near β = 1 favors within-head K/V tying at the same width and budget.Reinvesting across-head sharing savings in wider heads can accelerate readout training.Source: §6.1 · p. 10

6.2 Extensions

• For general inputs, input-Gram spectral interactions with width constraints and K/V dependence remain to be studied to determine their effects on preferred allocations and sharing schemes.Source: §6.2 · p. 10
• The admissible growth rate of β with n that preserves Gaussian equivalence remains an open question.A cited candidate scale is β = Θ(√log n), where bulk growth may change separation from the dominant common mode.Source: §6.2 · p. 10

• Gaussian equivalence for multi-head self-attention includes value and output projections that depend on the keys.Source: §7 · p. 11
• The resulting random-matrix model predicts readout training dynamics and coefficient estimation error for frozen-feature learning.Source: §7 · p. 11
• The equivalence provides an analytical foundation for understanding how multi-head attention shapes representations used in subsequent learning and computation.Source: §7 · p. 11

• Tables A.1–A.5 organize general notation, dimensions and parameters, matrices, functions, and probability laws.Source: A Notation · p. 14
• A subscript ⊕ denotes a direct sum over heads, and scalar functions of score matrices act entrywise.Source: A Notation · p. 14
• Matrix definitions use the conditioning sigma-field Fn = σ(K1,n, . . . , Kh,n, Tn).Source: A Notation · p. 15
• The transforms use a compactly supported law ν on [0, ∞); χν and Sν require positive mean, with specified local domains for χν⁻¹.Source: A Notation · p. 16
• Table A.5’s value and output covariance limits use the assumptions of Section C.1.4 and Theorem 3.1, respectively.Source: A Notation · p. 17, Table A.5 · p. 17

B.1 Input Scaling

• A corresponding rescaling of inputs and weights preserves the projected queries, keys, attention spectra, values, and MHA output; it is separate from head-comparison and attention-bulk scalings.Source: §B.1 · p. 18

B.2 Weak Convergence and Moment Convergence; B.3 Moment Perturbation Bounds

• The appendix relates weak and moment convergence under compact-support conditions and gives singular-value and Gram-rank perturbation controls for comparing spectral moments.For random measures, the convergence implications apply pathwise when the respective support conditions hold.Source: §B.2 · p. 18, §B.3 · p. 18

B.4 Block Representation

• The block representation preserves mixed head products needed for the concatenated spectrum, while independent query heads yield block-diagonal conditional covariances even when keys depend across heads.The proof also uses per-head key Gram matrices and distinguishes the key-Gram and feature-matrix dimension ratios.Source: §B.4 · p. 19

B.5 Conditional Means and Covariances

• Conditional covariance bounds follow from regular Gaussian key Gram matrices and the block-diagonal structure induced by independent query heads.The key regularity result allows arbitrary dependence among key weights across heads and across sizes.Source: §B.5 · p. 20

B.6.1 Gaussian Estimates

• For Gaussian coordinates with bounded variances, the maximum has exponential moments controlled by a term of order sqrt(log N), including under conditional Gaussian laws.For N = O(n^k) with fixed k, the cited bound gives E[e^(aL)] = n^o(1) for fixed a and variance bound.Source: §B.6.1 · p. 21

B.6.2 Truncation and Control of the Intermediate Matrices

• Rowwise truncation and matrix Chernoff bounds control intermediate covariance operators, while fixed trace moments concentrate for identically distributed feature or Gaussian rows.The resulting operator norm bound is almost surely subpolynomial in n.Source: §B.6.2 · pp. 21–22

B.7 Concentration of the Row Normalizers; B.7.1 A Finite-Size Normalizer Identity

• The row normalizers concentrate almost surely near one, and the finite-size analysis derives exact mean and variance formulas with a variance expansion when dK > 4β^2.The normalizer argument conditions on a query row and applies concentration to its scores across keys.Source: §B.7 · pp. 22–23, §B.7.1 · p. 23

B.8 Gaussian Replacement from Concentration

• For every fixed integer q ≥ 1, m_q(ν˚Y_fT) and m_q(ν_YG[Σ_f]T) differ by a quantity tending to zero almost surely, uniformly over deterministic symmetric H on Ω_n(M), for fixed M ≥ 1.For the one-row replacement, |E m_q(ν_X) − E m_q(ν_X′)| ≤ n^-3/2+o(1), uniformly on Ω_n(M), over the replaced row and other-row types, for each fixed M.Source: §B.8 · pp. 24–25

B.9 Completing the Gaussian Comparison

• Entrywise covariance matching yields almost-sure equality of fixed spectral moments between the Gaussian models, and hence vanishing differences for bounded continuous tests and resolvent traces.The covariance entries must match the specified diagonal and off-diagonal forms up to errors controlled by δn.Source: §B.9 · pp. 25–26

B.10 Leading Singular Value of the Output

• The output decomposition separates a rank-at-most-one term from the perpendicular attention component.The term is B = [PV1, . . . , PVh]WO in the displayed decomposition.Source: §B.10 · p. 26

C.1.1 Transforms and Gaussian Covariance Limits

• Under bounded spectral support, Gaussian sample covariance matrices converge weakly and in every fixed moment to the compound free-Poisson law π(ν, t).The result assumes Dn has limiting law ν, bounded operator norm, and pn/n → t; the corresponding rectangular covariance limit is also stated for m/n → t.Source: §C.1.1 · p. 27

C.1.2 The Attention Law

• The attention law is identified through the transform of ργK ⊠ MP1/h, with weak convergence and convergence of every fixed moment.The derivation uses only the common marginal law of the h key blocks.Source: §C.1.2 · pp. 27–28

C.1.3 The Attention-Matrix Corollary

• The attention-matrix corollary gives the same weak limit for the normalized leading singular value and the Gaussian comparison, while the second singular value vanishes.The rank-one difference between the attention matrix and its perpendicular component preserves the Gram-distribution weak limit.Source: §C.1.3 · p. 28

C.1.4 Independent Value and Output Projections

• With independent values, the stacked head output has limiting squared-singular-value law ν∞,h ⊠ πγV, with convergence in every fixed moment.The independent-value and output-weight setup permits value weights to be dependent or shared across heads.Source: §C.1.4 · p. 29

C.1.5 Head Aggregation

• For independent keys, the head-sum law is ργK ⊠ π1; sharing one key matrix changes the ratio to γK/h.The shared-key case reduces to a single-head model with width-to-token ratio γK/h.Source: §C.1.5 · pp. 29–30

C.2 Key-Dependent Value Projections

• The key-dependent value-projection results are stated for fixed head count and inverse temperature, with hdK/n converging to γK.These conditions specify the asymptotic regime for the key-dependent projection analysis.Source: §C.2 · p. 30

C.2.1 The Spectral Law for Shared Key and Value Weights

• The shared key–value covariance law is built by mapping the key-spectrum law through the limiting covariance polynomial.The construction defines µKV as the pushforward of πγK/h under φKV.Source: §C.2.1 · p. 30
• Across-head dependence may be arbitrary, including a common key and value matrix, while the proof uses each key matrix’s marginal law.The output law’s free cumulants characterize the resulting spectrum.Source: §C.2.1 · p. 30

C.3 Projection Summary

• The projection summary specifies that the head-sum row assumes independent key matrices; with a common key matrix, ργK is replaced by ργK/h.The tied-value law is the one defined for Theorem C.7.Source: §C.3 · p. 31

C.4 Effective Degrees of Freedom per Token

• Across-head sharing preserves the entire limiting law and effective degrees of freedom at fixed positive penalty, whereas within-head tying yields fewer effective degrees of freedom under sufficiently strong regularization after matching means.The tied model has a larger first moment and normalized second moment for every fixed head count.Source: §C.4 · p. 32

C.5 Training Loss and Estimation Error from the Bulk Law

• Centered features have limiting average squared row norm one, while the common component remains in the features used for learning.The uncentered attention features share the centered features’ weak limit because their difference has rank at most one.Source: §C.5 · p. 32

C.5.1 Fixed-Target Gradient Flow

• For fixed-target gradient flow, weak spectral convergence yields the limiting loss prediction at every fixed time.The target is deterministic, fixed across feature initializations, and orthogonal to the uniform component.Source: §C.5.1 · p. 33

C.5.2 Ridge Estimation with an Isotropic Teacher

• The ridge coefficient error is expressed through a spectral integral, and the full MHA output uses the output-width ratio γO.The integral limit applies for fixed λ > 0.Source: §C.5.2 · p. 33

C.6.1 The Normalized Cauchy Equation

• The normalized-law equation incorporates head, value, and output width factors for MHA, omitting projection factors for attention-only or value-only laws.The normalized laws are scaled to have mean one before deriving the equation.Source: §C.6.1 · p. 34

D Experimental Details and Additional Comparisons

• All experimental inputs, random weights, and targets are synthetic; public models provide architecture dimensions, with no pretrained weights or external datasets used.Numerical computations use float64, and the reported intervals are nominal pointwise 95% confidence intervals over independent initializations.Source: D Experimental Details and Additional… · p. 35

D.1 Finite-Size GE at Dimensions from Publicly Available Models

• Finite-size GE tests retain model-derived widths, head counts, and K/V sharing, using n ≤ d, fixed inputs, fan-in Gaussian weights, and β = 1.Each configuration uses ten independent initializations.Source: §D.1 · p. 36
• The finite-size comparison measures softmax–GE spectral agreement across all n eigenvalues, including the largest, using relative W1 and KS distances.The main figures separately test limiting predictions for loss and coefficient error.Source: §D.1 · p. 36, Table D.2 · p. 36

D.2 Head-Count Training Curves

• For the plotted times and three head-count conditions, the largest absolute difference between measured mean relative loss and the limiting prediction is below 0.003.Implementation checks compare the spectral formula with direct matrix-exponential calculations.Source: §D.2 · p. 37
• For h = 32, the limiting law has mass 1/2 at zero, yielding a loss plateau; the other two head counts have no zero atom and show different decay rates.The measurements retain alignment between a fixed target and the empirical eigenvectors.Source: §D.2 · p. 37

D.3 Head Count and Projection Budget

• The projection-budget map encodes continuous estimation error by color, with boundaries marking changes in the smallest width.The fixed-budget path keeps projection parameter count at 6n2.Source: §D.3 · p. 37
• Across 19 configurations, the largest relative difference between measured mean coefficient error and GE is 1.72%.The comparisons vary head count and output width along fixed-output and fixed-parameter-budget paths.Source: §D.3 · p. 37

D.4 Temperature and Q/K–V Width Allocation

• βc ≈ 0.559 marks the crossover between equal-parameter width allocations, while GE’s softmax-ratio estimates differ by at most 1.52% across six temperatures.The ratio compares the (64, 192) allocation with (128, 128); below one favors wider V.Source: §D.4 · p. 38

D.5 Across-Head Sharing and Gradient-Flow Time

• At fixed widths, the predicted arrival time is 41.22; at fixed parameter count it falls to 7.80 (7.76 at rounded width) for hKV = 1, with a maximum measured–predicted gap of 5.14%.For full sharing, measured times are 39.11 at fixed widths and 7.57 at fixed parameter count.Source: §D.5 · p. 39

D.6 Head Count and Sharing across or within Heads

• At h = 4, tied-minus-GQA coefficient error is 0.04600 ± 0.00379, and the softmax mean differs from GE by at most 0.75% across displayed points.Difference intervals exclude zero at all five head counts, supporting the reported ordering.Source: §D.6 · p. 39

D.7 MHA Equivalence and Architectural Comparisons

• Figure D.1 shows head splitting concentrates the attention bulk, output width changes MHA scale and shape, and within-head tying changes the limit while across-head sharing preserves it.The theoretical curves depend on β and dimension ratios; histograms show original-model spectra.Source: Figure D.1 · p. 40, §D.7 · p. 40