Towards Understanding the Shape of Representations in Protein Language Models

2025-09-29

AI Summary

The mechanisms by which protein language models (PLMs) map sequences to structure-relevant hidden representations—and the geometric properties of their representational spaces—remain poorly understood. Method: We introduce square-root velocity (SRV) shape space analysis and graph filtering into PLM representation analysis, integrating Karcher mean computation and effective dimension estimation to conduct multi-scale structural correlation analysis on the SCOP dataset. Contribution/Results: We find that the geometry of PLM representational spaces evolves nonlinearly with network depth; structural fidelity peaks at intermediate layers (e.g., layers 24–28 in ESM-2), while long-range contextual relationships degrade substantially. This reveals a hierarchical, depth-dependent mechanism for structural encoding within PLMs, providing both theoretical foundations and empirical evidence for interpretable modeling and architecture optimization targeting protein folding tasks.

Technology Category

• Machine Learning: Representation Learning
• Natural Language Processing: (Large) Language Models
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

While protein language models (PLMs) are one of the most promising avenues of research for future de novo protein design, the way in which they transform sequences to hidden representations, as well as the information encoded in such representations is yet to be fully understood. Several works have attempted to propose interpretability tools for PLMs, but they have focused on understanding how individual sequences are transformed by such models. Therefore, the way in which PLMs transform the whole space of sequences along with their relations is still unknown. In this work we attempt to understand this transformed space of sequences by identifying protein structure and representation with square-root velocity (SRV) representations and graph filtrations. Both approaches naturally lead to a metric space in which pairs of proteins or protein representations can be compared with each other. We analyze different types of proteins from the SCOP dataset and show that the Karcher mean and effective dimension of the SRV shape space follow a non-linear pattern as a function of the layers in ESM2 models of different sizes. Furthermore, we use graph filtrations as a tool to study the context lengths at which models encode the structural features of proteins. We find that PLMs preferentially encode immediate as well as local relations between residues, but start to degrade for larger context lengths. The most structurally faithful encoding tends to occur close to, but before the last layer of the models, indicating that training a folding model ontop of these layers might lead to improved folding performance.

Problem

• Understanding how protein language models transform sequence representations
• Analyzing structural information encoded in protein representation space
• Investigating context lengths where models encode protein structural features

Innovation

• Using SRV representations for protein shape analysis
• Applying graph filtrations to study structural encoding
• Identifying optimal encoding layers for folding models