Small differences between internal representations can shrink or grow as they pass through a Transformer layer, affecting how information is preserved and how sensitive the model is to perturbations. Yet it remains unclear when changes in attention weights can turn a representation difference that would otherwise shrink into one that grows. We develop a Mixture-Gated Theory of Anisotropic Representation Dynamics that separates self-attention into value transport, query gating, and key gating. Our analysis reveals how the geometry of attended representations constrains gating and how interactions between these mechanisms can reverse contraction induced by value transport. We formalize this phenomenon as gating-induced regime switching, where a representation difference contracts under fixed attention routing but expands under full attention. We derive a necessary and sufficient spectral criterion that certifies the existence or absence of switching and identifies switching directions without exhaustive search. Beyond existence, our theory characterizes maximum relative switching strength, certifies the severity of contraction that gating can overcome, and determines the largest dimension of a subspace in which every nonzero direction switches. Experiments across four Transformer architectures validate these predictions through analytical and finite-forward comparisons, revealing distinct and joint query--key gating mechanisms and nontrivial switching subspaces. Our results provide a mechanistic and spectral understanding of how attention routing can alter local representation dynamics.
Figures & tables
Figure 1: Conceptual illustration of gating-induced regime switching. Two nearby representations initially differ by a small perturbation δzℓ . Under fixed attention weights, value transport contracts this difference. However, a small change in the query can redistribute attention among token representations, introducing an additional gating contribution. When this contribution is sufficiently strong, the full attention computation amplifies the same representation difference rather than suppressing it. The illustration motivates our analysis of how attention routing can reverse transport contraction.
Figure 2: Geometric support and null directions of query gating. Centered attended representations span a subspace Ur , whose dimension may be substantially smaller than the ambient representation dimension. The transformed subspace WQKUr determines which query perturbation directions can influence the attention output through query gating. Perturbations orthogonal to this subspace produce no first-order query-gating change. Consequently, query gating is insensitive to at least d−dim(Ur) independent input directions. The three-dimensional projections provide a schematic visualization of the higher-dimensional geometry.
Figure 3: From self-attention decomposition to gating-induced regime switching. Self-attention induces three local mechanisms: value transport, query gating, and key gating, represented by their corresponding Jacobian components. These components jointly determine directional representation growth through the induced growth operator GM=J⊤MJ−M . The three-dimensional illustration shows how different perturbation directions can exhibit contraction or expansion. Interactions between value transport and attention gating can reverse transport contraction, causing a representation difference to grow under full attention.
Figure 4: Geometric support of query gating. (a) Schematic three-dimensional slice illustrating directions orthogonal to the query-gating support. (b) Analytical response versus support overlap, including near-null directions and the theoretical upper bound. (c) Finite-forward changes in attention weights under fixed-key query perturbations.
Figure 5: Theory-predicted regime geometry and gating-pathway ablations. (a,b) Analytical and finite-forward regime maps on the same selected directional slice, with exact switching and mechanism-label agreement across the tested directions. (c) Growth-sign boundaries separating the observed regimes. (d) Directional growth under transport-only, single-gating-pathway, and full-attention counterfactuals for representative directions.
Figure 6: Switching across architectures and depths. (a) Median normalized effective switching dimension across seven normalized depths and four Transformer architectures. Normalization uses the projected transport-contractive subspace. (b) Mechanism composition among sampled switching directions. Measurements use two WikiText-2 validation prompts per model and 64-dimensional random orthonormal probe subspaces.
Figure 7: Cross-architecture regime geometry and spatial adjacency. (a) Analytic regime spheres: Two orthographic views (View A and B) of selected spherical slices across four Transformer architectures, colored by their predicted growth regimes. (b) Spatial adjacency graphs: Node sizes indicate regime proportions, and edges represent shared regime boundaries. NC: transport non-contractive; NS: contraction without switching; Q/K: query-/key-driven switching; E: either pathway sufficient; J: joint query–key required. Percentages are computed over all sampled sphere directions, rather than only switching directions as in Figure 6 (b). Results characterize the selected slices, not the full representation space.
Figure 9: Search quality and amortized cost of spectral guidance. (a) Blind random search finds no switch within the tested budget, while oracle-assisted sampling in the transport-contractive subspace finds ordinary switches but misses severe reversals. (b) Spectral preprocessing incurs a one-time setup cost but becomes advantageous for sufficiently many directional evaluations. Total-cost curves are projected from measured setup and throughput timings.
Notation
Meaning
Local representation dynamics
ℓ,zℓ,δzℓ,Tℓ
Layer index, incoming representation, local perturbation, and Transformer block map.
Jℓ,v,M
Block Jacobian, nonzero perturbation direction, and symmetric positive-definite representation metric.
GℓM
Local growth operator Jℓ⊤MJℓ−M , governing directional expansion and contraction.
Self-attention and query-gating geometry
xi,r,i,m
Token representation; query-token, attended-token, and perturbed-token indices.
Table 1: Summary of the main theoretical notation.
Figure 10: Attended-representation geometry. Three-dimensional projection of the centered representations xi−μr ; marker size and color indicate attention weights αri . The plot visualizes a selected head, not the full attended-variation subspace.
Figure 11: Progressive classification of local switching regimes. Five columns reveal the conditions defining transport non-contraction (blue), persistent contraction (red), joint query–key necessity (cyan), query-only or key-only sufficiency (green or purple), and either-pathway sufficiency (orange). The top and bottom rows apply analytic and finite-forward tests to the same selected directional slice. The classifications concern counterfactual growth signs, not an additive partition of total growth.
Figure 12: Growth margins under pathway ablations. (a) Query gating alone crosses the zero-growth boundary. (b) Neither query nor key gating alone produces expansion, whereas their joint inclusion does. Bars in each panel evaluate the same perturbation direction under different counterfactual dynamics; dashed lines mark zero growth.
Figure 13: Switching spectrum and severity-restricted certification. (a) Ordered eigenvalues of W=C−1/2RC−1/2 , shown for comparison with the severity test; the switching boundary is λ=1 . (b) Leading eigenvalue of Wτ0 on subspaces constrained to have transport severity at least τ0 . The green reference line indicates a separately computed τcert≈0.967 ; the displayed orange sweep does not reach that threshold and is not, on its own, the certificate.
Figure 14: Relative strength and absolute contraction severity. (a) Relative switching ratios for the relative optimum and a numerically identified severe switch (logarithmic axis; dashed line at Γ=1 ). (b) Fractions of incoming squared magnitude removed by transport for the same directions. The green reference marks τcert≈0.967 . The two panels compare different objectives, not two estimates of the same quantity.
Direction
gTM
gT+QM
gT+KM
gfullM
τ(v)
Γ(v)
vrel
−0.003361
5.7437
−0.002801
5.8107
0.003361
1729.98
Severe switch
−0.967349
−0.1842
−0.952478
0.0722
0.967349
1.07465
Table 2: Relative-optimal and severe-switch counterfactuals. Directional growth, transport severity, and relative strength at the selected GPT-2 Medium operating point. Values are rounded.
Figure 15: Projected spectra across architecture and depth. (a) Median leading switching eigenvalue and the λ=1 threshold. (b) Median normalized effective switching dimension dsw/dim(S) within the retained probe subspace. Shaded interquartile bands summarize variation across the sampled prompts and are not confidence intervals.
Figure 16: Mechanism composition among projected switching directions. Fractions are pooled over the sampled switching directions within each architecture. Query-only, key-only, joint query–key, and either-pathway-sufficient classes are disjoint under the specified counterfactual tests. Fractions are conditional on full-dynamics switching and are not ambient-space probabilities.
Figure 17: Probe-dimension sensitivity of projected switching statistics. Nested orthonormal probe spaces of dimensions 32 , 64 , 128 , and 256 are evaluated at matched model–prompt–depth operating points. (a) Median projected λmax(W) ; (b) median normalized effective switching dimension; (c) median query-sufficient fraction among sampled switches. Shaded regions indicate interquartile ranges across operating points, not confidence intervals or convergence guarantees.
The projection of queries and keys are central to the attention mechanism in Transformer architectures. While they are mathematically symmetric, they play different roles in attention mechanisms. The question of whether there is an effect from their functional distinction on their geometric development in training remains unanswered. We investigate the problem through the training of small GPT-like Transformers on character-level WikiText-103 for three different depths (4, 6, and 8 layers), three types of initialization for queries and keys, and four random seeds, resulting in 36 runs and 54 trajectories of average layers across seeds. We track the effective dimensionality of those layers using participation ratios and discover that effective dimension of queries expand while keys shrink, and that PRQ−PRK is positive in all trajectories studied. In connection to attention, the shrinking of keys leads to a narrower spectrum of QK⊤ and more peaked attention weights. In order to determine if this connection is causal or coincidental, we directly control the spectrum of keys during training across five seeds: restricting it to make it shrink sharpens the attention with high directional confidence, while keeping it constant to the level of initial dispersion makes attention softer. Additional token-level checkpoint analyses show that the monotonic paired-contrast trend is not universal across pretrained families, but survives as an early-training regime that later decays over a full pretraining run, and the link between interaction-rank geometry and attention entropy remains visible in several models.
Transformer representations evolve through learned additive transformations that either preserve their current direction or redirect it. We study this evolution as a functional geometry, decomposing learned updates into parallel and perpendicular components. Across pretrained models, we find substantial parallel components beyond the residual identity path. We then apply the decomposition in two spaces: to attention and MLP updates relative to the hidden state, and to attention value aggregation relative to the current token's value. Targeted edits reveal a strongly space-dependent asymmetry: exclude-self value-space parallel manipulation is markedly more robust than residual-space and perpendicular counterparts, preserving the direct self message while scaling only the non-self aggregate. The same decomposition gives a component-resolved description of compression-induced update error: perpendicular error separates compression methods more clearly than parallel error. Extensive experiments further demonstrate that full-aggregate parallel suppression during from-scratch pretraining lowers validation-loss trajectories and improves downstream averages, with the value-space variant strongest. Together, these results connect representation geometry to editing robustness, compression diagnosis, and training-time intervention. Code is available in the project repository.
Shwai He, Haichao Zhang, Shen Yan
University of Maryland, College Park · Northeastern University · ByteDance
Standard transformer attention computes pairwise similarity between queries and keys, treating all tokens as equally salient regardless of their intrinsic informational content. In turbulent fluid dynamics, coherent structures -- the energetically dominant, spatially organized patterns that persist amid background chaos -- carry a disproportionate fraction of total energy and govern all transport. We propose that tokens play an analogous role in transformer attention: informationally dense positions (morphological boundaries, syntactic heads, discourse markers) concentrate spectral energy and should attract proportionally more attention than background tokens (function words, repeated patterns, low-information filler). We propose Energy-Gated Attention (EGA): a simple modification that gates value aggregation by the spectral energy of key token embeddings, computed by a single learned linear projection that discovers the dominant spectral mode of the embedding field. On TinyShakespeare, EGA achieves +0.103 validation loss improvement with only 12,480 additional parameters (<0.26% overhead) and no measurable computational cost. The result is consistent on Penn Treebank (+0.101), demonstrating dataset independence. A systematic ablation across three wavelet families (fixed Morlet, Daubechies db2/db4, and a parametric Morlet) establishes that fixed structured bases are suboptimal -- the optimal energy direction is data-adaptive and non-sinusoidal -- while identifying learned wavelet packets as a promising open direction. The learned energy threshold converges to tau ~= 0.35 independently of initialization, corresponding to the fraction (~36%) of tokens carrying above-average spectral energy in English text, a stable linguistic property consistent with the fraction of content words in running English text.