Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Uncovering hidden geometry in Transformers via disentangling position and context

Jiajun Song, Yiqiao Zhong

TL;DR

This paper tackles the opacity of transformer embeddings by introducing a mean-based decomposition that splits each hidden state into a global mean, a layer-specific positional basis, a context basis, and a residual term, enabling direct interpretation of how position and context shape representations. The authors show that the positional basis is typically low-rank and smooth, while the context basis forms clear topic-like clusters, and that these two factors are largely incoherent, which supports robust interpretability. They provide both empirical evidence across multiple models and datasets and theoretical justification (via smoothness implying low-rank/low-frequency structure) and further show that the learned weight matrices in self-attention align with the positional basis, enabling a decoupled, kernel-based view of attention. The results suggest that language naturally yields smooth, decoupled geometric structure in hidden states, which has implications for efficient attention, interpretability, and the design of diagnostics beyond attention visualizations. The work also delineates limits by showing nonsmooth patterns in arithmetic tasks, highlighting language-specific geometry and potential boundaries of the decomposition’s applicability.

Abstract

Transformers are widely used to extract semantic meanings from input tokens, yet they usually operate as black-box models. In this paper, we present a simple yet informative decomposition of hidden states (or embeddings) of trained transformers into interpretable components. For any layer, embedding vectors of input sequence samples are represented by a tensor $\boldsymbol{h} \in \mathbb{R}^{C \times T \times d}$. Given embedding vector $\boldsymbol{h}_{c,t} \in \mathbb{R}^d$ at sequence position $t \le T$ in a sequence (or context) $c \le C$, extracting the mean effects yields the decomposition \[ \boldsymbol{h}_{c,t} = \boldsymbolμ + \mathbf{pos}_t + \mathbf{ctx}_c + \mathbf{resid}_{c,t} \] where $\boldsymbolμ$ is the global mean vector, $\mathbf{pos}_t$ and $\mathbf{ctx}_c$ are the mean vectors across contexts and across positions respectively, and $\mathbf{resid}_{c,t}$ is the residual vector. For popular transformer architectures and diverse text datasets, empirically we find pervasive mathematical structure: (1) $(\mathbf{pos}_t)_{t}$ forms a low-dimensional, continuous, and often spiral shape across layers, (2) $(\mathbf{ctx}_c)_c$ shows clear cluster structure that falls into context topics, and (3) $(\mathbf{pos}_t)_{t}$ and $(\mathbf{ctx}_c)_c$ are mutually nearly orthogonal. We argue that smoothness is pervasive and beneficial to transformers trained on languages, and our decomposition leads to improved model interpretability.

Uncovering hidden geometry in Transformers via disentangling position and context

TL;DR

This paper tackles the opacity of transformer embeddings by introducing a mean-based decomposition that splits each hidden state into a global mean, a layer-specific positional basis, a context basis, and a residual term, enabling direct interpretation of how position and context shape representations. The authors show that the positional basis is typically low-rank and smooth, while the context basis forms clear topic-like clusters, and that these two factors are largely incoherent, which supports robust interpretability. They provide both empirical evidence across multiple models and datasets and theoretical justification (via smoothness implying low-rank/low-frequency structure) and further show that the learned weight matrices in self-attention align with the positional basis, enabling a decoupled, kernel-based view of attention. The results suggest that language naturally yields smooth, decoupled geometric structure in hidden states, which has implications for efficient attention, interpretability, and the design of diagnostics beyond attention visualizations. The work also delineates limits by showing nonsmooth patterns in arithmetic tasks, highlighting language-specific geometry and potential boundaries of the decomposition’s applicability.

Abstract

Transformers are widely used to extract semantic meanings from input tokens, yet they usually operate as black-box models. In this paper, we present a simple yet informative decomposition of hidden states (or embeddings) of trained transformers into interpretable components. For any layer, embedding vectors of input sequence samples are represented by a tensor . Given embedding vector at sequence position in a sequence (or context) , extracting the mean effects yields the decomposition where is the global mean vector, and are the mean vectors across contexts and across positions respectively, and is the residual vector. For popular transformer architectures and diverse text datasets, empirically we find pervasive mathematical structure: (1) forms a low-dimensional, continuous, and often spiral shape across layers, (2) shows clear cluster structure that falls into context topics, and (3) and are mutually nearly orthogonal. We argue that smoothness is pervasive and beneficial to transformers trained on languages, and our decomposition leads to improved model interpretability.
Paper Structure (62 sections, 5 theorems, 55 equations, 34 figures, 6 tables)

This paper contains 62 sections, 5 theorems, 55 equations, 34 figures, 6 tables.

Table of Contents

  1. Introduction
  2. A mean-based decomposition
  3. Positional basis vs. positional embeddings.
  4. Connections to ANOVA
  5. On terminology.
  6. Accessible reproducibility
  7. Notations
  8. Pervasive geometrical structure
  9. Sinusoidal patterns are learned, not predefined.
  10. Consistent pattern across layers and models is not solely explained by residual connections.
  11. Some exceptions.
  12. Key properties of positional basis: low rank and low frequency
  13. Measuring smoothness.
  14. Low rank via spectral analysis
  15. Low rank.
  16. ...and 47 more sections

Key Result

Theorem 1

Fix positive integers $k \le T$ and $m$. Define the low-frequency vector ${\bm{f}}_s = (1, \cos((s-0.5)\pi/T), \ldots, \cos ((s-0.5)(T-1)\pi / T))^\top \in \mathbb{R}^T$ where $s=1,\ldots,k$, and denote ${\bm{F}}_{\le k} = [{\bm{f}}_1,\ldots,{\bm{f}}_k] \in \mathbb{R}^{T \times k}$. Then there exist

Figures (34)

  • Figure 1: PCA visualization of positional basis (blue) and cvecs ( red) from GPT-2 on OpenWebText. For every layer $\ell$, each ${\mathbf{pos}}_t^{(\ell)}$ and randomly selected ${\mathbf{cvec}}_{c,t}^{(\ell)}$ are projected using top-$2$ principal directions of $({\mathbf{pos}}_t^{(\ell)})_{t \le T}$. Darker blue/red colors correspond to larger $t$. Principal components have dramatically increasing scales across layers, but for aesthetic purposes we rescaled all plots.
  • Figure 2: Normalized Gram matrix$[\bar{{\bm{P}}}, \bar{{\bm{C}}}]^\top [\bar{{\bm{P}}}, \bar{{\bm{C}}}]$ where $\bar{{\bm{P}}}=[\frac{{\mathbf{pos}}_1}{\lVert{\mathbf{pos}}_1\rVert},\ldots,\frac{{\mathbf{pos}}_T}{\lVert{\mathbf{pos}}_{T}\rVert}]$ and $\bar{{\bm{C}}} = [\frac{{\mathbf{ctx}}_1}{\lVert{\mathbf{ctx}}_1\rVert},\ldots,\frac{{\mathbf{ctx}}_{C}}{\lVert{\mathbf{ctx}}_C\rVert}]$ based on GPT-2. Here, $T=128$, and ${\mathbf{ctx}}_c$ is sampled from $4$ documents with sample size $32$ in OpenWebText. We find (i) Smoothness, ${\mathbf{pos}}$-${\mathbf{pos}}$ part (top left) of Gram matrix is smooth; (ii) Incoherence, ${\mathbf{pos}}$-${\mathbf{ctx}}$ part (top right/bottom left) has values close to $0$; (iii) Clustering, ${\mathbf{ctx}}$-${\mathbf{ctx}}$ part (bottom right) shows strong cluster patterns.
  • Figure 3: Spectral and Fourier analysis based on GPT-2 model and OpenWebText. Left: Top-60 singular values of ${\bm{P}}$. Right: Applying 2D discrete cosine transform to $\bar{{\bm{P}}}^\top \bar{{\bm{P}}}$, we show first $10$ frequency coefficients.
  • Figure 4: Decoupling trained weight matrices. For $12$ attention heads (layer $L=6$ shown here) in GPT-2, we study the matrix ${\bm{W}} = {\bm{W}}^q ({\bm{W}}^k)^\top / \sqrt{d_{{\mathrm{head}}}} \in \mathbb{R}^{d \times d}$. Red: diagonal entries ${\color{red}{\bm{D}}}:={\mathrm{diagg}}({\bm{W}})$. Blue: take off-diagonal matrix ${\bm{W}} - {\mathrm{diagg}}({\bm{W}})$, rotate it by the right singular vectors of positional basis ${\bm{V}}$, then apply denoising. Large absolute values concentrate in small top-left part ${\color{blue}{\bm{L}}}$.
  • Figure 5: Addition task trained on NanoGPT exhibits nonsmooth patterns: discontinuity as a consequence of non-language data training. Left: Gram matrix of normalized positional basis. Compare with top-left of plots in Figure \ref{['fig:Gram']}. Right: QK matrix.
  • ...and 29 more figures

Theorems & Definitions (9)

  • Theorem 1
  • Theorem 2
  • Remark 1
  • Lemma 1
  • proof
  • Lemma 2
  • proof
  • Lemma 3
  • proof : Proof of Lemma \ref{['lem:A-lowfreq']}