Visualizing LLM Latent Space Geometry Through Dimensionality Reduction

Technical Summary: Visualizing LLM Latent Space Geometry Through Dimensionality Reduction

Problem Statement
While Large Language Models (LLMs) achieve state-of-the-art performance, their internal decision-making processes remain opaque. Mechanistic interpretability seeks to reverse-engineer these internal components to achieve a granular, causal understanding. A promising avenue for this is feature geometry: analyzing the structure and organization of representations within high-dimensional latent spaces. However, these spaces are difficult to interpret directly due to their extreme dimensionality. The challenge lies in extracting, processing, and visualizing these latent states to uncover how abstract features are encoded, transformed, and organized across model layers without losing salient structural information.

Methodology
The authors propose a systematic pipeline to capture and visualize latent state geometries in decoder-only Transformer models (specifically GPT-2 and LLaMa). The methodology consists of three primary stages:

1. Latent Extraction Pipeline:

   • The pipeline processes text data (from the PG-19 dataset) or individual tokens through the Transformer architecture.
   • It captures activations at six distinct points within each Transformer block: the outputs of the pre-attention and pre-MLP normalization layers, the raw outputs of the attention and MLP modules, and the residual stream states after the addition of these components (pre-add and post-add).
   • The authors adopt a "residual stream-centric" view, treating the model as a shared communication channel where components read from and write to a common space.

2. Data Processing and Preprocessing:

   • To manage high dimensionality and memory constraints, latent states are optionally reduced (e.g., from 4096 to 512 dimensions) via PCA.
   • Specific preprocessing steps include converting vectors to unit length to emphasize directional structure over magnitude, and averaging vectors across sample or sequence dimensions to isolate specific sources of variability (e.g., layer-wise evolution vs. token-level noise).
   • The authors explicitly exclude the 0-th sequence position (initial token) in many visualizations due to its anomalously high norm, which can obscure other geometric patterns.

3. Dimensionality Reduction and Visualization:

   • Principal Component Analysis (PCA): Used as a linear baseline to capture directions of maximal variance. It is applied without mean centering to preserve the origin's position.
   • Uniform Manifold Approximation and Projection (UMAP): Used as a non-linear alternative to preserve local neighborhoods and global manifold structure. The authors utilize cosine distance for UMAP to align with the Linear Representation Hypothesis, which posits features as directions in latent space.
   • Visualizations are color-coded by layer (blue for attention, red for MLP) and sequence position (gradient from dark to light).

Key Contributions and Results

• Attention vs. MLP Geometric Separation:
  The most significant finding is a clear, persistent geometric separation between the outputs of attention components and Multi-Layer Perceptron (MLP) components in intermediate layers. In dimensionality-reduced visualizations (both PCA and UMAP), attention outputs consistently occupy distinct regions from MLP outputs. The authors note this pattern has not been documented in prior work.

• High Norm of Initial Sequence Position:
  The study confirms and expands upon observations that the latent state norm at the 0-th sequence position is significantly larger than subsequent positions.

  • In GPT-2 (using learned positional embeddings), this is expected.
  • In LLaMa (using Rotary Positional Encodings/RoPE), the authors find that the high norm is not limited to the <BOS> token but applies to the vast majority of vocabulary tokens when they appear as the initial token. This suggests LLaMa has developed a complex mechanism to detect the sequence start, distinct from the trivial learned embeddings of GPT-2.

• Layerwise Evolution and Norm Dynamics:
  Visualizations of latent state norms across layers reveal that norms generally increase through the middle layers and drop after the final normalization layer. The initial token exhibits a unique "spike" in norm that emerges after a few layers and dissipates in the final layers.

• Positional Embedding Geometries:

  • GPT-2: The visualizations confirm the "helical" structure of GPT-2's learned positional embeddings in high-dimensional space, visible across various PCA dimension pairs.
  • LLaMa: The geometric patterns for RoPE are less distinct. In early PCA dimensions, a "tail" of initial sequence positions leads into a dense cloud of later positions. In later dimensions, this pattern dissolves into a uniform cloud, suggesting the positional geometry in RoPE models may be relatively low-dimensional.

Significance and Claims
The paper claims to provide a robust framework for visualizing the internal geometry of LLMs, making the abstract concept of "feature geometry" accessible. By demonstrating the separation of attention and MLP outputs, the work offers a new perspective on how different Transformer components contribute to the shared residual stream. The authors position this work as a step toward principled interpretability, suggesting that understanding these geometric structures is essential for analyzing how models generalize, reason, and abstract. The authors explicitly state that their findings regarding the attention/MLP separation are novel to their knowledge and that the observed patterns in LLaMa's initial token norms warrant further investigation into the mechanisms of relative positional encodings.

The work does not claim to solve the full problem of interpretability but rather provides a "strong background and exposition" and a set of tools (the visualization pipeline) to facilitate further research into the dynamics of latent states across layers and models.