🤖 machine learning

Analogies between Transformer Layers and Power Method

This paper establishes an analogy between transformer layers and the power method, demonstrating that tokens align with the principal eigenvector of the product of value and output weight matrices, a phenomenon that is analytically provable in shared-weight models and can be leveraged to steer transformer outputs toward arbitrary directions.

Original authors: Chenglong Li, Claudio Altafini

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Chenglong Li, Claudio Altafini

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Transformers as a "Power Method" Machine

Imagine a Transformer (the brain behind AI chatbots) not as a complex neural network, but as a giant, multi-story building. Each floor of the building is a "layer." When a piece of information (a "token," like a word in a sentence) enters the building, it travels from the bottom floor to the top, getting processed at every stop.

The authors of this paper discovered a surprising secret: The way a token moves through these layers is mathematically almost identical to a classic math trick called the "Power Method."

The Analogy: The Magnetic Compass

Think of the "Power Method" like a compass that is trying to find North.

The Setup: You have a compass needle (the token) that is pointing in a random direction.
The Process: You place the compass on a map that has a giant, invisible magnet hidden underneath it. The magnet pulls the needle slightly toward "North."
The Normalization: After the needle moves, you force it to stay the same length (you don't let it grow or shrink, you just rotate it).
The Result: If you do this over and over again, the needle eventually stops wobbling and points almost perfectly at North.

The paper argues that each layer of a Transformer does exactly this.

The "Magnet" is a specific mathematical matrix (a grid of numbers) created by the layer's weights.
The "North" is the Principal Eigenvector. This is a special, dominant direction in the data space.
The "Normalization" is the Layer Normalization step, which keeps the token's size constant.

So, as a token travels up the Transformer, it is constantly being "tilted" or pulled toward this dominant direction, just like the compass needle.

Two Types of Buildings

The paper looks at two different kinds of Transformer buildings, and they behave differently.

1. The "Universal" Building (Shared Weights)

Some Transformers, like ALBERT, use the exact same blueprints for every single floor. The "magnet" is identical on every level.

What happens: Because the magnet is the same everywhere, the token gets pulled in the same direction at every step. It's like walking up a staircase where every step is tilted slightly toward the same wall.
The Result: By the time the token reaches the top, it is almost perfectly aligned with that one dominant direction. The paper proves mathematically that in this case, all tokens eventually converge to point in the same direction (a state called "consensus").

2. The "Custom" Building (Layer-Varying Weights)

Most popular Transformers, like GPT-2, BERT, and GPT-Neo, have different blueprints for every floor. The "magnet" changes from layer to layer.

What happens: On Floor 1, the magnet pulls the token toward "North." On Floor 2, the magnet changes and pulls it toward "North-East." On Floor 3, it pulls toward "South."
The Result: Because the target keeps moving, the token never gets perfectly aligned with a single direction. However, the paper shows that at every single step, the token still tries to align with the current floor's dominant direction. It's like a hiker trying to follow a guide who keeps changing their mind about which way is "best." The hiker doesn't reach a perfect destination, but they are still being pushed in a specific direction at every moment.

The "Steering" Trick

The most exciting part of the paper is a practical application of this discovery. The authors realized that because the "Power Method" relies on the strength of the magnet (the spectral gap), we can hack the system.

The Analogy:
Imagine you are in the "Custom Building" (like GPT-2). The tokens are wandering around because the magnets on each floor are weak and conflicting. But, you are allowed to change the magnet on the very last floor before the token exits.

The Move: You replace the last floor's magnet with a super-strong, custom-made magnet that points exactly where you want the token to go (e.g., "Be helpful" or "Be creative").
The Effect: Because this new magnet is so much stronger than the others (creating a huge "spectral gap"), it overrides all the previous confusion. The token snaps into alignment with your chosen direction instantly.

The paper claims this allows us to steer the output of a large language model toward any specific direction we want, simply by tweaking the math of the final layer.

Summary of Findings

The Mechanism: Transformers work like a repeated math exercise (the Power Method) that tries to align data with a "main direction."
Shared Weights (ALBERT): If the layers are identical, the tokens converge perfectly to that main direction.
Different Weights (GPT/BERT): If the layers are different, the tokens don't converge perfectly, but they still show a tendency to align with the current layer's main direction at every step.
The Hack: By adding a strong, custom "kick" to the final layer, we can force the tokens to align with any direction we choose, effectively steering the AI's output.

Important Note: The paper specifically ignores the "Feed-Forward Network" (the part of the layer that does complex non-linear thinking) to make this mathematical analogy work. They focus only on the attention and normalization parts. They also note that in real-world models, this alignment is often weak because the "magnets" aren't strong enough, but the principle holds true.

Technical Summary: Analogies between Transformer Layers and Power Method

Problem Statement
The paper addresses the lack of understanding regarding the specific vector direction to which token representations converge in deep transformers, a phenomenon known as "oversmoothing," "rank collapse," or "token-uniformity." While it is established that deep transformers tend toward a consensus state where all tokens become identical, the nature of this consensus vector and the mechanism driving it remain unclear. Existing literature often models self-attention as a continuous-time system or focuses on single-head scenarios, leaving the discrete-time, multi-head behavior of practical models (like GPT2, BERT, and ALBERT) less explored. Specifically, the paper seeks to determine if the iterative process of passing tokens through transformer layers resembles a known numerical method for finding principal eigenvectors.

Methodology
The authors model a transformer as a discrete-time dynamical system where the layer index $t$ represents time. They focus on the update law of token vectors $x_i(t)$ through a single layer, explicitly excluding the Feed-Forward Network (FFN) and utilizing $L_2$ normalization (replacing standard LayerNorm) to isolate the geometric effects of the attention and projection mechanisms.

The core of the methodology involves defining a "layer matrix" $F_t$ for each layer $t$ , constructed as the product of the output weight matrix ( $O$ ) and the concatenated value weight matrices ( $V$ ) of the attention heads:
$F_t = O_t V_t$
The update rule is analyzed as a step in a dynamical system:
$x_i(t+1) = \text{LayerNorm}\left(x_i(t) + \text{AttentionOutput}\right)$
The authors draw a direct analogy between this update and the Power Method, an iterative algorithm used to compute the principal eigenvector of a matrix. In the Power Method, a vector is repeatedly multiplied by a matrix and normalized, converging to the dominant eigenvector if a spectral gap exists.

The analysis is divided into two main scenarios:

Shared Weights (Time-Invariant): Analyzed using ALBERT models where weights are identical across all layers ( $F_t = F$ ). The authors provide rigorous analytical characterizations for cases where the principal eigenvalue is real and strictly dominant, complex, or non-dominant.
Layer-Varying Weights (Time-Varying): Analyzed using BERT, GPT2, and GPTNeo. Here, $F_t$ changes at each layer. The authors investigate whether tokens still exhibit a tendency to align with the current layer's principal eigenvector $v_{1,t}$ despite the changing target.

To validate the theory, the authors conduct numerical experiments on pre-trained models, measuring the correlation between token vectors and the principal eigenvectors of the layer matrices. They also propose a method for "steering" the output by modifying the final layer.

Key Contributions

Analogy to the Power Method: The paper establishes that the operation of a transformer layer (disregarding FFN) acts analogously to a single step of the Power Method. Tokens are projected and normalized, tending to tilt toward the principal eigenvector of the layer matrix $F_t$ .
Analytical Characterization of Convergence:
- For shared-weight transformers (e.g., ALBERT), the authors prove that if the layer matrix $F$ has a strictly dominant real eigenvalue, the system converges to a consensus state where all tokens align with the principal eigenvector $v_1$ .
- They identify conditions for stability, noting that if the principal eigenvalue is complex, tokens converge to a limit cycle in the 2D subspace spanned by the real and imaginary parts of the eigenvectors.
- They show that "bipartite consensus" (tokens splitting into two groups aligned with $\pm v_k$ ) is theoretically possible but empirically rare and unstable in multi-head settings.
Empirical Validation on Real Models:
- Experiments on ALBERT confirm that tokens converge to the principal eigenvector of the shared layer matrix. However, the convergence speed is often limited by small spectral gaps (where $\lambda_1 \approx \lambda_2$ ) or complex eigenvalues.
- Experiments on layer-varying models (BERT, GPT2, GPTNeo) show that while tokens do not accumulate alignment toward a single global eigenvector (due to changing $F_t$ ), they still exhibit a statistically significant increase in correlation with the current layer's principal eigenvector $v_{1,t}$ at each step. This confirms the "Power Method" behavior persists locally even in time-varying systems.
Token Steering via Rank-1 Perturbation: The paper proposes a constructive method to steer tokens toward an arbitrary desired direction. By adding a sufficiently large rank-1 perturbation to the output matrix of the final layer, the authors can create a large spectral gap. This forces the final layer matrix to have a new, user-defined principal eigenvector, causing all tokens to align with this specific direction.

Results

ALBERT: In models with shared weights, tokens rapidly converge to consensus. For ALBERT-base and large, the principal eigenvalue is real but close to the second, leading to convergence within the span of the first two eigenvectors. For xlarge and xxlarge, complex eigenvalues cause tokens to rotate in a 2D subspace.
Layer-Varying Models: In BERT and GPT variants, the alignment with the principal eigenvector of the current layer ( $\gamma_t$ ) is positive but generally weak. This weakness is attributed to the small spectral gaps ( $\lambda_{1,t} - \lambda_{2,t}$ ) in real-world transformers. However, the probability of alignment increasing at each step ( $\rho_t > 0.5$ ) is consistently observed, validating the analogy.
Steering: When the final layer is modified with a rank-1 perturbation, the spectral gap becomes large, and the tokens achieve near-perfect alignment with the newly induced principal eigenvector, regardless of the model architecture.

Significance and Claims
The paper claims that viewing transformers through the lens of the Power Method provides a principled explanation for token oversmoothing and consensus. It suggests that the "direction" of the consensus is not arbitrary but is determined by the spectral properties of the product of value and output weight matrices.

The authors emphasize that while the natural alignment in standard transformers is weak due to small spectral gaps, the mechanism is robust enough to be exploited. The proposed steering method offers a principled alternative to existing techniques (like prompt engineering or activation editing) for controlling model output, as it directly manipulates the geometric projection properties of the final layer. The paper concludes that this analogy holds even when disregarding the FFN, suggesting that the attention mechanism and layer normalization are the primary drivers of this eigenvector alignment behavior.