Rethinking Self-Attention Parameterization: Decoupling Self- and Cross-Token Interactions

Work in Progress — This article is an ongoing exploration of the topic. Further analysis, experiments, and discussions will be added.

1. Starting Point: Two Types of Token Interactions

The Query-Key interaction is the central mechanism of self-attention. Given an input sequence \(X\in\mathbb{R}^{n\times d}\), a standard Transformer first projects it into queries and keys

\[ Q=XW_Q,\quad K=XW_K, \]

where \(W_Q,W_K\in\mathbb{R}^{d\times d}\) are the projection matrices for query and key, respectively. The attention score matrix is then computed as

\[ S=QK^\top=XW_QW_K^\top X^\top. \]

The bilinear form \(x_iW_QW_K^\top x_j^\top\) is the classic and natural parameterization of pairwise interactions, and one rarely questions whether it is equally suitable for every entry of \(S\).

A closer examination, however, reveals an important distinction. The off-diagonal entries of \(S\) capture interactions between distinct tokens, whereas the diagonal entries measure the interaction of each token with itself. Although both are computed using the same unified formula, their underlying mathematical structures are fundamentally different.

For interactions between distinct tokens (\(i\neq j\)), the score is a general bilinear form

\[ s_{ij}=x_iW_QW_K^\top x_j^\top, \]

which is in general asymmetric — the score from token \(i\) to token \(j\) is typically not equal to the score in the reverse direction, i.e.,

\[ x_iW_QW_K^\top x_j^\top \;\ne\; x_jW_QW_K^\top x_i^\top. \]

In fact, prior studies have shown that such asymmetric or directional attention pattern naturally emerges in autoregressive Transformers and plays an important role in their information processing[1]. Mechanistic circuit analysis has also clarified this point, showing that the query and key subspaces serve distinct, asymmetric roles in in-context learning[2],[3].

However, when the query and key correspond to the same token (\(i=j\)), the score becomes a quadratic form \[ s_{ii}=x_iW_QW_K^\top x_i^\top, \] where the same input token \(x_i\) appears on both sides. Intuitively, there should be no asymmetry in a token's interaction with itself.

2. Self-Interaction Scores Depend Only on the Symmetric Component

To better understand the distinction between self-interactions and interactions between distinct tokens, we decompose \(W_QW_K^\top\) into its symmetric and skew-symmetric components. For notational convenience, let

\[ A = W_QW_K^\top. \]

We decompose \(A\) into its symmetric and skew-symmetric components as

\[ A = A_{\mathrm{sym}} + A_{\mathrm{skew}}, \]

where

\[ A_{\mathrm{sym}} = \frac{A+A^\top}{2}, \quad A_{\mathrm{skew}} = \frac{A-A^\top}{2}. \]

By construction, \(A_{\mathrm{sym}}\) is symmetric, satisfying \(A_{\mathrm{sym}}^\top=A_{\mathrm{sym}}\), while \(A_{\mathrm{skew}}\) is skew-symmetric, satisfying \(A_{\mathrm{skew}}^\top=-A_{\mathrm{skew}}\).

A particularly interesting property emerges when the query and key come from the same token. For any vector \(x\in\mathbb{R}^{1\times d}\), the quadratic form associated with a skew-symmetric matrix is always zero

\[ xA_{\mathrm{skew}}x^\top = 0. \]

To see this, note that the scalar \(xA_{\mathrm{skew}}x^\top\) is equal to its transpose

\[ xA_{\mathrm{skew}}x^\top = \left(xA_{\mathrm{skew}}x^\top\right)^\top = xA_{\mathrm{skew}}^\top x^\top = -xA_{\mathrm{skew}}x^\top. \]

Therefore, it must be zero. Consequently, the self-interaction score can be written entirely in terms of the symmetric component

\[ s_{ii} =x_iAx_i^\top=x_i(A_{\mathrm{sym}}+A_{\mathrm{skew}})x_i^\top=x_iA_{\mathrm{sym}}x_i^\top. \]

This leads to the following important observation.

For self-interactions, the skew-symmetric component of \(W_QW_K^\top\) is functionally redundant: although it is important for parameterizing interactions between distinct tokens, it contributes nothing to the self-interaction attention score.

3. Optimization Friction: Redundancy and Implicit Gradient Conflicts

The analysis above reveals that standard attention parameterization is not as flawless or seamless as traditionally assumed. Mathematically, the parameterization is very sound: while it contains redundant parameters, their contribution is identically zero. However, when viewed through the lens of optimization, this setup exposes deeper structural friction. Specifically, while \(W_Q W_K^\top\) unifies both self- and cross-token interactions under a single formulation, this shared parameterization introduces optimization inefficiencies during gradient-based learning.

  1. Parameter Redundancy in Self-Interactions. As established, the skew-symmetric component \(A_{\mathrm{skew}}\) is completely inert on the diagonal (\(xA_{\mathrm{skew}}x^\top = 0\)). This creates parameter redundancy where a subset of the model's capacity contributes zero expressive value to self-interactions, yet still remains active in cross-token interactions.
  2. Implicit Multi-Objective Gradient Conflict. The symmetric component \(A_{\mathrm{sym}}\) is forced to simultaneously accommodate two distinct optimization objectives: modeling self-interactions (\(s_{ii}\)) and capturing directional cross-token relations (\(s_{ij}\)). When these two objective demand opposing update directions, they create implicit multi-objective gradient conflicts. This shared parameterization leads to gradient interference and therefore impairs optimization efficiency.

These issues demonstrate that the current attention parameterization—though seemingly unproblematic—creates structural friction during training. This motivates us to rethink and redesign a more effective attention parameterization.

4. A Decoupled Parameterization Strategy

Motivated by the preceding observation and analysis, we propose a simple modification to the attention score parameterization. The key idea is to assign different parameterizations to the two types of interactions:

Specifically, we introduce an additional learnable vector \(v\in\mathbb{R}^d\) and replace the self-interaction score with a feature-wise quadratic form

\[ s_{ii}=\sum_{k=1}^{d}v_k\,x_{ik}^2, \] or equivalently, \[ s_{ii}=x_i\operatorname{diag}(v)\,x_i^\top. \]

The complete parameterization is therefore

\[ s_{ij}= \begin{cases} x_iW_QW_K^\top x_j^\top, & i\neq j,\\[6pt] x_i\operatorname{diag}(v)\,x_i^\top, & i=j. \end{cases} \]

In this formulation, each component of \(v\) acts as an independent gate on the corresponding feature dimension, fully disentangling the self-interaction score from \(W_Q\) and \(W_K\). This decoupling eliminates the gradient conflicts identified in the previous section, as the two families of scores now receive gradients through entirely separate parameters.

References

  1. Matteo Saponati, Pascal Sager, Pau Vilimelis Aceituno, Thilo Stadelmann, Benjamin Grewe. The underlying structures of self-attention: symmetry, directionality, and emergent dynamics in Transformer training. In International Conference on Machine Learning (ICML), 2025. [link]
  2. Hyunmin Cho, Woo Kyoung Han, Kyong Hwan Jin. Balancing fidelity and diversity in diffusion models via symmetric attention decomposition: Hopfield perspective. In International Conference on Machine Learning (ICML), 2026. [link]
  3. Catherine Olsson, Nelson Elhage, Neel Nanda, et al. In-context learning and induction heads. Transformer Circuits Thread, 2022. [link]
  4. Salman Quazi. Induction heads: The circuit behind in-context learning. Blog post, 2025. [link]
Previous ← Back to Home → Next