🤖 machine learning

Towards Understanding the Shape of Representations in Protein Language Models

This paper investigates the geometric structure of protein language model representations using square-root velocity and graph filtrations, revealing that models encode structural features non-linearly across layers with optimal fidelity just before the final layer, while struggling to maintain long-range contextual relations.

Original authors: Kosio Beshkov, Anders Malthe-Sørenssen

Published 2026-05-29

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Kosio Beshkov, Anders Malthe-Sørenssen

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

Imagine you have a giant library of protein "stories." In the real world, these stories are written in 3D, folded into complex shapes that determine what the protein does. But scientists have recently started using "Language Models" (AI trained on protein sequences) to read these stories. The problem is, we don't fully understand how the AI is rewriting these stories in its own secret language.

This paper is like a detective trying to figure out the shape of the AI's secret language. The authors ask: "When the AI looks at a protein, does it see a flat list of words, or does it see a 3D sculpture? And how does that sculpture change as the AI 'thinks' deeper?"

Here is the breakdown of their investigation using simple analogies:

1. The Two Tools: Measuring Shapes and Filtering Noise

To understand the AI's "mind," the authors used two special tools:

The "Shape Shifter" (SRV Representation):
Imagine you have a piece of clay (a protein). You can stretch it, rotate it, or move it around, but it's still the same piece of clay. The authors used a mathematical trick called Square-Root Velocity (SRV) to turn every protein into a smooth, flexible curve.
- The Analogy: Think of this as turning every protein into a unique, stretchy rubber band. The AI's job is to figure out that two rubber bands are "similar" even if one is twisted or rotated. The authors mapped these rubber bands into a special "Shape Space" where they could measure how far apart different proteins are.
The "Sieve" (Graph Filtration):
Imagine looking at a protein through a series of sieves with different-sized holes. A small sieve only lets you see the immediate neighbors (the amino acids right next to each other). A larger sieve lets you see friends further away.
- The Analogy: The authors used this to test how far the AI can "see" to understand the protein's structure. Does the AI only know its immediate neighbors, or can it see the whole picture?

2. What They Found: The "Expansion and Contraction" Dance

The authors looked at the AI's internal layers (like the different floors of a skyscraper) and found a surprising pattern in how the "Shape Space" changes:

The Early Floors (Expansion): In the beginning, the AI takes the protein and stretches it out into a huge, complex, high-dimensional space. It's like taking a small sketch and blowing it up into a massive, detailed 3D model with infinite possibilities. The "shape" of the data gets very wide and varied here.
The Late Floors (Contraction): As the data moves up the skyscraper to the top layers, the AI starts folding that massive space back down. It squashes the data into a much smaller, tighter space.
- The Result: By the time the data reaches the top, the "shapes" of different proteins are very similar to each other. The AI has stripped away the noise and found a few core "shapes" that describe almost everything.

Key Takeaway: The AI doesn't just get "smarter" at the top; it actually simplifies the geometry of the protein world, compressing it into a few efficient shapes.

3. How Far Can the AI "See"? (Context Length)

Using their "Sieve" tool, the authors tested how many neighbors the AI needs to see to understand the protein's 3D structure. They found a two-step pattern:

The Immediate Neighbor (2 steps): The AI is very good at seeing the amino acids right next to each other. This makes sense; it's like knowing your best friend's name.
The "Sweet Spot" (8 steps): Surprisingly, the AI also gets very good at seeing a group of about 8 neighbors.
- The Analogy: It's as if the AI understands the protein best when it looks at a small cluster of friends (2 people) or a small neighborhood block (8 people).
- The Limit: If the AI tries to look at the whole protein at once (very long distances), its understanding of the 3D shape starts to get fuzzy and degrade.

4. The "Sweet Spot" for Folding

The most practical finding is about where in the AI's brain the 3D structure is best preserved.

The authors found that the top layer (the very last floor of the skyscraper) is actually not the best place to find the 3D shape.
Instead, the layer just before the last one holds the most accurate "map" of the protein's 3D structure.
The Analogy: Imagine a chef cooking a meal. The final dish (the last layer) is the finished product, but the moment the flavors are most perfectly balanced and distinct is just before the final garnish is added. If you want to use the AI to help build new proteins (folding), you should look at that "just-before-the-end" layer, not the final output.

Summary

This paper is a map of how AI "sees" proteins. It shows that the AI starts by stretching protein data into a huge, complex shape, then folds it back down into a simple, efficient form. It understands the 3D structure best when looking at small groups of neighbors, and it holds the clearest picture of the protein's shape just before it finishes its final calculation. This suggests that if we want to use AI to design new proteins, we should listen to what it says in the layer right before it gives its final answer.

Technical Summary: Towards Understanding the Shape of Representations in Protein Language Models

Problem Statement

Protein Language Models (PLMs), such as the Evolutionary Scale Modeling (ESM) family, have demonstrated significant promise in de novo protein design, structure prediction, and functional annotation. While it is empirically understood that PLM representations serve as effective initializations for folding models, the precise nature of the information encoded within these hidden representations remains poorly understood. Existing interpretability efforts have largely focused on how individual sequences are transformed into latent vectors or on identifying specific features (e.g., binding motifs, coevolutionary statistics) within single sequences. However, these approaches neglect the geometric structure of the entire space of protein representations and how different proteins relate to one another within this latent space. Furthermore, standard practices often collapse the sequence dimension (e.g., by averaging), thereby discarding the "shape" information inherent in the full tensor representation. This paper addresses the gap in understanding the global metric geometry of PLM representation spaces and the specific context lengths at which structural features are encoded.

Methodology

The authors propose a framework to analyze the "shape" of protein representations by mapping them to metric spaces using two primary mathematical approaches:

1. Square-Root Velocity (SRV) Representations

To analyze the global geometry of representation spaces, the authors treat protein sequences and their PLM embeddings as ordered point clouds.

Curve Interpolation: Point clouds (representing either 3D structures or PLM embeddings) are interpolated using quadratic splines to form continuous curves $\gamma(t)$ in $\mathbb{R}^m$ .
SRV Transformation: These curves are transformed into Square-Root Velocity (SRV) representations, $q(t) = \dot{\gamma}(t)/\sqrt{\|\dot{\gamma}(t)\|^2}$ , which projects them onto an infinite-dimensional sphere $S^\infty$ .
Quotienting: To enforce invariance to translation and rotation (isometries of $SE(m)$), the authors quotient out the rotation group $SO(m)$ using Singular Value Decomposition (SVD) to find the optimal rotation aligning two curves.
Metric Space: This process creates a Riemannian shape space where distances between proteins can be computed. The authors utilize this space to calculate:
- Fréchet Radius: The mean distance of all representations from their Fréchet mean, serving as a measure of the spread of shapes.
- Effective Dimension: Calculated via Tangent Principal Component Analysis (Tangent PCA) on the shape space manifold, measuring the intrinsic dimensionality of the variations in shape.

2. Graph Filtrations

To investigate the context length sensitivity of PLMs, the authors employ a graph filtration approach from Topological Data Analysis.

Graph Construction: For a given protein and its PLM representation, $k$ -nearest neighbor graphs are constructed based on Euclidean distances between residues.
Filtration: By varying the number of neighbors $k$ (or a distance threshold $t$ ), a filtration of graphs is generated.
Metric Comparison: The authors compare the adjacency matrices of the true 3D protein structure graphs against the adjacency matrices of the PLM representation graphs.
Graph Filtration Moment: A normalized metric is defined as the ratio of the distance between the protein structure and its PLM embedding to the distance between the protein structure and random point clouds. This quantifies how well the PLM preserves structural proximity at different scales (context lengths).

The analysis was conducted on 1,377 proteins sampled from the SCOPe dataset, processed through various ESM2 models of different sizes and the Ankh model.

Key Results

1. Geometry of PLM Shape Spaces

Dimension Expansion and Contraction: The effective dimension of the shape space follows a non-linear pattern across layers. Initial layers exhibit a significant expansion in dimensionality, particularly in larger models, followed by a contraction in deeper layers.
Low-Dimensional Subspace: In the final layers, the effective dimension becomes very low, significantly lower than the ambient embedding dimension. This suggests that while PLMs span a high-dimensional space, the variations in protein shapes can be described by a few specific shape descriptors.
Fréchet Radius: The Fréchet radius generally decreases with depth, indicating that representations become more clustered in shape space. Notably, the variability among shapes is low across all models, and the radius does not vary significantly with model size.
Generalizability: This expansion-contraction pattern was observed in both ESM2 models and the Ankh model, suggesting it is a fundamental property of PLM architecture.

2. Context Length Sensitivity

Bimodal Encoding: The graph filtration analysis reveals a bimodal pattern in structural encoding. PLMs optimally encode 3D structural features at two distinct scales:
1. Immediate Neighbors: At very short context lengths (approximately 2 residues).
2. Local Neighborhoods: At slightly longer context lengths (approximately 8 residues).
Layer Dependence: Structural encoding is present throughout the network but is not optimal in the final layer. The best structural fidelity is found in layers immediately preceding the last one.
Class Variations: Different protein classes exhibit varying degrees of structural faithfulness. Alpha/Beta proteins show the highest similarity between their PLM representations and 3D structures, while small and designed proteins show more distinct shapes.
Independence from Length: In deeper layers, the encoding of local context structure appears independent of the total protein length.

Significance and Claims

The paper claims to provide a novel geometric perspective on PLM representations, moving beyond individual sequence analysis to understand the collective structure of the representation space.

Interpretability: The work demonstrates that PLMs encode protein structure implicitly, despite being trained solely on masked language modeling tasks without explicit structural supervision. The finding that structure is encoded at specific context lengths (2 and 8 residues) suggests that the model learns local and semi-local geometric constraints as a byproduct of sequence prediction.
Model Architecture Insights: The observed expansion-contraction pattern in effective dimension aligns with theories of high-abstraction regimes in early layers and specific, semantically rich regimes in later layers. The low effective dimension in deep layers implies that the complex task of protein folding may be navigable via a small number of shape transformations.
Practical Implications for Folding: The authors suggest that the suboptimal structural encoding in the final layer, contrasted with the high fidelity in preceding layers, offers a potential avenue for improving protein folding models. Specifically, initializing folding models with representations from layers just before the final one (rather than the final layer itself) could yield better performance. However, the authors note that initial attempts to verify this with linear models or small networks did not generalize, indicating that further research with larger models is required.

In summary, the paper establishes that PLM representation spaces possess a distinct geometric structure characterized by dimension expansion followed by contraction, and that these representations capture protein structural features most effectively at specific, short-to-moderate context lengths, particularly in layers preceding the final output.