The geometry of hidden representations of large transformer models

🎯 Pitch

Hidden representations in large transformers spontaneously expand into high-dimensional manifolds before sharply compressing into low-dimensional valleys, and these valleys consistently house the most semantically rich features—revealing a universal, unsupervised way to find the best layers for downstream tasks without any labels.

1. Executive Summary

This paper studies how the geometric and statistical properties of hidden representations evolve across layers in large transformer models trained with self-supervised objectives, analyzing both protein language models (ESM-2 family, 35M–3B parameters) and image transformers (iGPT family, 76M–1.4B parameters) on the ProteinNet, SCOPe, and ImageNet datasets. The core contribution is the discovery of a consistent intrinsic dimension (ID) profile that reveals distinct phases of representation processing—a rapid expansion to a high-dimensional peak in early layers (reaching ID values of ~20–32), followed by compression to a low-dimensional plateau or local minimum (ID values of ~5–7 for proteins, ~22 for images), and then a final ascent or second shallow peak in the decoding layers—establishing that the intermediate representations at the ID minimum encode the most semantically meaningful features of the data. The paper demonstrates that the neighborhood overlap with ground-truth labels—remote homology superfamilies for proteins and ImageNet class labels for images—peaks precisely at these low-ID regions (reaching ~0.8 for remote homology detection in ESM-2 and ~0.35 for class-label overlap in iGPT-L), providing an explicit unsupervised strategy to identify the layers most suitable for downstream learning tasks without requiring annotation. The finding holds across model sizes, training objectives (masked language modeling and next-pixel prediction), and data modalities, establishing that the encoding–compression–decoding computational strategy emerges spontaneously from self-supervised training but with the critical boundary condition that semantic abstraction is maximized only after the manifold has been expanded and then compressed into a low-dimensional representation.

2. Context and Motivation

The Core Problem: We Don't Understand How Representations Evolve in Self-Supervised Transformers

The fundamental question this paper tackles is deceptively simple: what happens geometrically to data as it passes through the layers of a large transformer trained with self-supervision? Modern transformers—whether processing protein sequences, images, or text—are stacks of architecturally identical self-attention blocks that transform input tokens through a sequence of representations, all residing in vector spaces of identical dimensionality. Yet despite this architectural uniformity, the representations at different layers encode radically different information. The paper asks: is there a systematic, predictable pattern to how the geometry of these representations changes across layers, and can we use that pattern to identify where abstract semantic knowledge emerges?

This gap matters for several reasons the authors develop across Sections 1 and 2:

Scientific understanding of self-supervised learning: Self-supervised objectives like masked language modeling (used in ESM-2) and next-token prediction (used in iGPT) have proven remarkably effective at producing representations that transfer to diverse downstream tasks—homology detection in proteins, image classification, syntactic parsing in language. Yet the mechanism by which these objectives produce abstract understanding remains opaque. The training signal comes from reconstructing low-level input features (amino acid identities, pixel colors), not from any explicit semantic supervision. Understanding where and how high-level concepts emerge from this low-level training signal is a fundamental scientific question about the nature of representation learning.
Practical deployment of intermediate representations: The paper explicitly notes (Section 2) that "in models trained by self-supervision, the last representation generally has the role of allowing reconstruction of the input representation. Therefore, the most semantically rich representation is likely to arise within the intermediate hidden layers of the network." If this is true, then using the final-layer representation for downstream tasks—as is standard practice in much of the protein language modeling literature (e.g., Rives et al., 2021)—is suboptimal. The paper demonstrates precisely this for remote homology detection: the last hidden layer achieves a neighborhood overlap with superfamily labels of approximately 0.4, while plateau layers achieve approximately 0.8 (Figure 4, left). This near-doubling of semantic signal has direct practical implications for every application that builds on top of pretrained transformer embeddings.
Unsupervised identification of optimal layers: In many domains where self-supervised transformers are most valuable—protein function prediction, metagenomic analysis, rare disease diagnosis—labeled data is scarce or expensive. The ability to identify which layers encode the most semantically useful representations without any labels would be transformative. The paper's central proposition is that the ID profile itself provides this unsupervised signal: the layers at or near the ID minimum are consistently the most semantically rich, regardless of domain or architecture.

Conflicting Evidence and Open Questions in Prior Work

The paper is motivated by a landscape of intriguing but fragmented observations about transformer representations, none of which had been synthesized into a coherent framework.

On the geometry side, prior work had established that representations in deep neural networks undergo systematic geometric changes. Ansuini et al. (2013) showed that in convolutional networks trained with supervision on ImageNet, the intrinsic dimension of representations follows a characteristic "hunchback" shape—rising in early layers and falling in later ones—with the ID at the final layer correlating with classification accuracy. Doimo et al. (2020) introduced the neighborhood overlap metric to track how representations rearrange across layers, finding in convolutional classifiers that significant rearrangement occurs primarily in the final layers where the network must comply with the supervised classification signal. These studies established the tools (TwoNN estimator, neighborhood overlap) and demonstrated that geometry matters, but they were limited to supervised convolutional architectures—a fundamentally different setting from the self-supervised transformers that now dominate the field.

On the representation-content side, multiple studies had probed where different types of information reside in transformer layers, with complex and sometimes contradictory results:

In protein language models, Rives et al. (2021) observed that Euclidean distances among representations in the last hidden layer encode remote homology information—proteins with similar structures but dissimilar sequences cluster together. This established that self-supervised protein models learn biologically meaningful features, but it left open whether the last layer is optimal. The paper shows it emphatically is not: homology information peaks in the plateau region and then sharply degrades in the final layers (Figure 4, left).
In image transformers, Chen et al. (2010) demonstrated with linear probes that features extracted from intermediate layers of iGPT encode semantic class information, with the best-performing layer varying by model size. But the probing approach requires labeled data and provides no geometric explanation for why certain layers work better than others. The paper connects this finding to geometry, showing that class-label overlap peaks precisely where the ID reaches its local minimum after the first expansion (Figure 4, right).
In natural language processing, the localization of linguistic information in transformer layers had been actively debated. Hewitt and Manning (2019) found that syntactic information (tree depth) is most prominent in middle layers of BERT. However, "the localization of semantic information has yielded contrasting results" (the paper notes in Section 4, citing Belinkov et al., 2017; Blevins et al., 2018; Liu et al., 2019; Tenney et al., 2019; Rogers et al., 2020). The paper's preliminary NLP results in the Appendix (Section E, Figure S6) on Llama-2-70B suggest their geometric framework may help resolve this: "the highest overlap with class partition, determined by sentence sentiment, occurs in correspondence with the first local minimum" of the ID profile.

The critical missing piece across all this prior work was a unifying framework that: (1) characterizes the geometric evolution of representations across the full depth of the network, not just at the output; (2) systematically relates geometric properties (intrinsic dimension) to semantic content (overlap with ground-truth labels); and (3) demonstrates that these relationships are consistent across architectures, training objectives, data modalities, and model sizes. The paper explicitly positions itself to fill this gap.

Where Existing Approaches Fall Short

The paper identifies specific limitations in prior work along several axes:

Final-layer myopia. The standard approach in applied transformer work—extract the last hidden layer and use it for downstream tasks—rests on an implicit assumption that the network progressively refines representations, culminating in the most informative embedding at the output. This paper shows that for self-supervised models, this assumption is wrong. The task of reconstructing the input (predicting masked amino acids or next pixels) forces the final layers to recover low-level detail that is irrelevant or even antagonistic to abstract semantics. The paper's three-phase model (encoding expansion → semantic compression → decoding reconstruction) explains why: the final layers function as a decoder, not a semantic encoder. Concretely, for protein remote homology (Figure 4, left), the neighborhood overlap with superfamilies drops from ~0.8 in the plateau to ~0.4 in the final layer—a catastrophic loss of biologically relevant signal.

Domain-specific, not cross-domain analysis. Prior work typically studied one model type in isolation: Ansuini et al. (2013) on supervised CNNs for images, Rives et al. (2021) on protein language models, Hewitt and Manning (2019) on NLP transformers. Each study provided valuable domain-specific insights but could not distinguish universal computational principles from domain-specific artifacts. The paper's deliberate comparison of protein models (ESM-2 family, trained with MLM on amino acid sequences) and image models (iGPT family, trained with next-pixel prediction on ImageNet) is designed to identify what is invariant across domains. The finding that both exhibit an ID expansion followed by compression—despite completely different data modalities, token vocabularies (~20 amino acids vs. 512 color clusters), and training objectives—is strong evidence for a general computational strategy.

Lack of geometric grounding for semantic probes. Linear probing (Chen et al., 2020) and nearest-neighbor-based evaluation (Rives et al., 2021) measure whether intermediate representations contain semantic information, but they don't explain why the information appears where it does or provide a label-free method for finding it. The paper's neighborhood-overlap-with-ground-truth metric (χ^{l,gt}) serves a dual role: it validates that geometry predicts semantics (the ID minimum coincides with the overlap maximum), and it provides the ground-truth benchmark against which the unsupervised ID-based method can be evaluated.

No framework for test-time layer selection. Even if one knows that intermediate layers are better than final layers, which intermediate layer should be used? Prior work provided no principled answer—Chen et al. (2010) had to sweep all layers with linear probes (requiring labels) to find the best one. The paper's ID profile provides an immediate, unsupervised answer: use the layer at the ID minimum after the first peak. Figure 4 shows this is remarkably consistent: for ESM-2 models, the entire plateau region works well; for iGPT models, the precise local minimum at ~0.4 relative depth maximizes class-label overlap.

How This Paper Positions Itself

The paper frames its contribution not as proposing a new model or training method, but as characterizing a universal geometric phenomenon that emerges spontaneously from self-supervised transformer training. The key conceptual move is to repurpose the intrinsic dimension—a geometric property measurable without any labels—as an unsupervised proxy for semantic content.

This positioning is made explicit in the Introduction and Section 4: "we point out an explicit strategy to identify, without supervision, the layers that maximize semantic content: representations at intermediate layers corresponding to a relative minimum of the ID profile are more suitable for downstream learning tasks."

The paper's interpretative framework—the autoencoder analogy—provides the conceptual scaffolding. In a traditional autoencoder, an encoder compresses input to a low-dimensional bottleneck (maximally abstract representation), and a decoder reconstructs from it. The paper argues that large self-supervised transformers learn a similar strategy spontaneously: early layers expand the data manifold into a high-dimensional space (analogous to the kernel trick in kernel methods, where non-linear feature expansion makes patterns more separable), intermediate layers compress it into a low-dimensional, semantically rich representation, and final layers decode back toward the input distribution. Section 4 makes this explicit: "large transformers behave essentially like sophisticated autoencoders, in which the data are first encoded into a low-dimensional and abstract representation and are successively decoded from it."

This differs from the supervised CNN picture of Ansuini et al. (2013), where the ID monotonically decreases after an early peak as the network progressively discards task-irrelevant variation. In self-supervised transformers, the ID increases again in the final layers (the final ascent in pLMs, the second peak in iGPTs) because the reconstruction objective forces the model to recover input-level detail. This is a crucial distinction: the compression is not terminal but temporary, and the semantic content is maximized at the compression minimum, not at the network output.

The paper also positions itself relative to the kernel methods literature in machine learning. The observation that the data manifold expands in early layers—reaching ID values of 20–32 from much lower input dimensionality—is explicitly analogized to the kernel trick (Section 1): "This expansion resembles the strategy followed in kernel methods, where one implicitly expands the feature space by introducing non-linear functions of the input features." The idea is that high-dimensional expansion makes patterns more linearly separable (per Cover's theorem), after which the network can learn a compression that captures the essential structure. This connection to classical ML theory provides intellectual grounding for why the expansion-then-compression pattern might be a broadly optimal computational strategy rather than an architectural accident.

The paper's domain positioning is also strategic. By choosing protein language models and image transformers—two domains where (a) large pretrained models are publicly available, (b) benchmark datasets with ground-truth semantic labels exist (SCOPe superfamilies for proteins, ImageNet classes for images), and (c) the data modalities and training objectives are maximally different—the paper makes the strongest possible case for universality. If the same geometric pattern appears in both, it's unlikely to be a domain-specific artifact. The preliminary NLP results in Appendix E (Llama-2-70B on sentiment analysis) extend this argument further, suggesting the framework may generalize to language as well.

Finally, the paper addresses a practical bottleneck in protein biology explicitly. The Appendix (Figure S2) shows that performing nearest-neighbor remote homology search in a plateau layer instead of the final layer improves accuracy by approximately 6% for the ProtT5-XL-U50 model—without any additional training. The paper notes that this "performance gain of ~6% is obtained without any further training," highlighting that the geometric insight translates immediately to improved downstream performance in a domain where state-of-the-art methods (Schütze et al., 2022) already use transformer embeddings for homology detection. This grounds the theoretical contribution in a concrete, quantifiable application.

3. Technical Approach

3.1 Reader Orientation

This paper is a geometric analysis study that characterizes how data representations—viewed as point clouds in high-dimensional vector spaces—change their shape and semantic content as they pass through the layers of large self-supervised transformer models. The "system" being analyzed is not something the authors built, but rather the internal representation spaces of existing pretrained models (ESM-2 for proteins, iGPT for images) probed at each layer. The problem the paper solves is: given a pretrained transformer with no access to labeled data, how can we identify which layers encode the most useful abstract features for downstream tasks? The solution has the shape of a geometric diagnostic: measure the intrinsic dimension of the data manifold at each layer, find the layer where the ID reaches a relative minimum after an early expansion, and that layer will contain the most semantically rich representations—without ever consulting a label.

3.2 Big-Picture Architecture (Diagram in Words)

The analytical pipeline has five major components, though none are trained by the authors—they are measurement and evaluation procedures applied to frozen pretrained models:

Pretrained transformer models (fixed, off-the-shelf): The ESM-2 protein language model family (35M, 650M, 3B parameters) and the iGPT image transformer family (76M, 455M, 1.4B parameters). These process input tokens through a stack of identical self-attention blocks, producing a representation at each layer. All representations have the same vector-space dimensionality $d$ , but their geometric properties differ radically.
Average pooling across the sequence dimension: Since protein sequences have variable lengths and image patches have fixed spatial extent, the representations at each layer are matrices of shape $\mathbb{R}^{l \times d}$ (tokens × embedding dimension). To compute distances between data points for geometric analysis, the authors collapse each sequence into a single vector in $\mathbb{R}^{d}$ by averaging over the token positions. This yields $B+1$ collections of $N$ vectors—one collection per layer—where $N$ is the number of data points in the analysis dataset.
Intrinsic dimension estimation (TwoNN): For each layer's representation, the authors estimate the dimensionality of the manifold on which the data approximately lies using the TwoNN algorithm (Facco et al., 2017). This requires only the distances from each point to its first two nearest neighbors, making it computationally tractable for large datasets and high embedding dimensions.
Neighborhood overlap measurement: To quantify how representations change between layers and how they relate to semantic labels, the authors compute the fraction of shared $k$ -nearest neighbors between two representations of the same data point. When comparing consecutive layers $l$ and $l+1$ , this measures the rate of representational rearrangement. When comparing layer $l$ against ground-truth class labels, this measures semantic content without requiring linear classifiers or probes.
Semantic validation via label overlap: Using benchmark datasets with known ground-truth categories—SCOPe protein superfamilies for remote homology, ImageNet class labels for image semantics—the authors compute the neighborhood overlap between each layer's representation and the label partition. This validates that the ID minimum corresponds to maximal semantic content and provides the quantitative evidence for the paper's central claim.

Information flows as follows: a batch of sequences or images enters the frozen pretrained model → the model computes representations at every layer → average pooling collapses the sequence dimension → TwoNN estimates the ID for each layer → neighborhood overlap computes similarity between consecutive layers and between each layer and ground-truth labels → the ID profile and overlap profiles are plotted against relative depth (layer index divided by total layers) → the patterns reveal three to four phases of geometric evolution, with semantic content maximized at the ID minimum.

3.3 Roadmap for the Deep Dive

First, the intrinsic dimension estimator (TwoNN) and its underlying assumptions, because the ID profile is the central unsupervised diagnostic and all subsequent claims depend on its validity.
Second, the neighborhood overlap metric ( $\chi^{l,m}_k$ ) and its label-conditioned variant ( $\chi^{l,gt}_k$ ), because this is the quantitative bridge between geometry and semantics—it measures both how representations rearrange across layers and how well they encode abstract concepts.
Third, the data extraction and preprocessing pipeline—model checkpoints, layer extraction points, average pooling, and dataset choices—because these implementation details determine what is actually being measured and whether the measurements are comparable across models and domains.
Fourth, the models and datasets (ESM-2 family, iGPT family, ProteinNet, SCOPe, ImageNet) with their training objectives and architectural parameters, because the paper's universality claim depends on understanding exactly what differs between the domains being compared.
Fifth, the validation methodology—how the overlap with ground-truth labels is computed, why remote homology requires excluding same-family neighbors, and how the ID profile is validated against training dynamics—because this connects the geometric measurements to practical utility.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an observational and analytical paper whose core idea is that the intrinsic dimension of hidden representations reveals a universal encoding-compression-decoding strategy in self-supervised transformers, and that the ID minimum provides an unsupervised signal for identifying semantically rich layers.

Intrinsic Dimension Estimation via TwoNN

The paper needs to measure the dimensionality of the data manifold at each layer without knowing anything about the manifold's shape, curvature, or parametric form. This is a hard problem in general, but the TwoNN estimator (Facco et al., 2017) provides an elegant solution that requires only local distance information and makes minimal assumptions.

The conceptual foundation: Real-world high-dimensional datasets—images, protein sequences, text embeddings—do not fill their ambient space uniformly. Images of natural scenes vary primarily along a modest number of meaningful axes (object identity, pose, lighting, background) despite having millions of pixel intensities. Protein sequences that fold into similar structures vary along evolutionary trajectories constrained by biophysics. The intrinsic dimension (ID) is the number of independent coordinates needed to describe the data approximately without losing essential structure—it is the dimensionality of the manifold that the data points lie on or near. The paper defines it as "the minimum number of coordinates that allow specifying a data point approximately without information loss" (Section 2.1).

Why TwoNN rather than alternatives: Global dimensionality reduction methods like PCA assume linear structure and fail on highly nonlinear manifolds. Local PCA requires choosing a neighborhood radius and can be sensitive to curvature. Graph-based methods (Isomap, diffusion maps) require building and analyzing large similarity matrices, which is computationally prohibitive for the $N = 25,299$ sequences in ProteinNet or the $N = 90,000$ images analyzed from ImageNet, especially when repeated across dozens of layers. TwoNN requires only the distances to the first two nearest neighbors of each point—a quantity that can be computed efficiently with approximate nearest-neighbor libraries—and makes no assumptions about global manifold structure.

The core insight of TwoNN: Under the assumption that data points are sampled from a probability distribution that is locally constant in density on a manifold of dimension $d$ , the ratio $\mu_i = r_{i2} / r_{i1}$ —where $r_{i1}$ is the distance from point $i$ to its nearest neighbor and $r_{i2}$ is the distance to its second-nearest neighbor—follows a Pareto distribution whose shape parameter is exactly the intrinsic dimension:

$p(\mu_i \mid d) = d \cdot \mu_i^{-d-1}$

where $\mu_i \geq 1$ is the ratio of second-to-first neighbor distances, and $d > 0$ is the intrinsic dimension.

What this equation means: For a given point $i$ , you find its two closest neighbors, measure how much further the second neighbor is compared to the first (the ratio $\mu_i$ ), and this ratio—across all points—follows a distribution that depends only on the manifold dimension. On a low-dimensional manifold (small $d$ ), the second neighbor is typically much further than the first, producing large ratios; on a high-dimensional manifold, the distances are more uniform, producing ratios closer to 1. The Pareto distribution captures this precisely: for large $d$ , the density $p(\mu_i \mid d)$ is heavily concentrated near $\mu_i = 1$ .

The estimation procedure, step by step:

For a given layer's representation (a set of $N$ vectors in $\mathbb{R}^d$ ), compute the Euclidean distance from each point to every other point, or efficiently find only the first two nearest neighbors using approximate methods.
For each point $i$ , compute $\mu_i = r_{i2} / r_{i1}$ where $r_{i1}$ is the distance to the nearest neighbor and $r_{i2}$ is the distance to the second-nearest neighbor.
Sort the $N$ ratios in ascending order, yielding $\mu_{(1)} \leq \mu_{(2)} \leq \cdots \leq \mu_{(N)}$ .
The cumulative distribution function (CDF) of the Pareto distribution is $F(\mu) = 1 - \mu^{-d}$ . Taking logs: $\log(\mu) = -\frac{1}{d} \log(1 - F(\mu))$ .
The empirical CDF at the $i$ -th sorted point is $F^{\text{emp}}(\mu_{(i)}) = i/N$ , so the authors regress $\log(\mu_{(i)})$ against $-\log(1 - i/N)$ and extract the ID as the slope of the best-fit line through the origin. The paper states (Section 2.1): "The intrinsic dimension can be inferred by maximum likelihood or linear regression from the cumulative distribution function."

Why this form works: The linear relationship $\log(\mu) = -\frac{1}{d} \log(1 - F(\mu))$ converts the problem of estimating a Pareto shape parameter into the simpler problem of estimating a slope. Under the locally constant density assumption, this relationship is exact—not an approximation—which means the estimated slope has a direct geometric interpretation as the inverse of the manifold dimension.

The locally constant density assumption and its validation: The TwoNN derivation assumes that within the neighborhood containing the first two nearest neighbors, the probability density is approximately constant. This is a much weaker assumption than requiring constant density globally, but it still needs checking. The paper explicitly validates it using the Point Adaptive kNN (PAk) method (Rodriguez et al., 2018), which "determines the extent of the neighborhood over which the probability density can be considered constant for each data point, subject to a specified confidence level" (Appendix B). The validation showed that "on average, the density can be considered constant within the first 6 neighboring data points," which comfortably covers the first two neighbors used in the estimator. The paper states: "This analysis allows the conclusion that, at this scale, the assumption of local density holds."

The noise-robustness procedure: In practice, measurement noise and violations of the local-density assumption (particularly at very local scales where data is discrete or quantized) can cause the TwoNN estimate to overestimate the ID. The paper follows the multiscale protocol from Facco et al. (2017): apply the TwoNN estimator to the full dataset of $N$ points, then to random subsets of decreasing size $n = [N, N/2, N/4, \ldots]$ , and select the ID value where the curve $\hat{d}(n)$ plateaus—meaning the estimate stabilizes and is no longer sensitive to the specific subset size. The paper reports (Appendix B): "the ID estimated with the TwoNN was less scale-dependent across dataset sizes between $[N/2, N/8]$ . Thus, in the main text, we report the TwoNN estimates on the dataset decimated by a factor of 4." This means the ID values plotted in Figure 1 are computed on 25% random subsets of the full analysis dataset, balancing statistical stability against computational cost.

Alternative considered but not used: The paper notes (Appendix B) that a more recent multiscale approach by Denti et al. (2022) "considers nearest neighbors of higher order instead of decimating the dataset" and "reduces the variance of the likelihood, giving more stable estimates." The authors state this method "gave similar results in the experiments of this work," indicating that the simpler decimation approach is adequate for their purposes, though the higher-order method might be preferable in more challenging settings.

Why this matters for the paper's claims: The entire unsupervised layer-selection strategy—identify the ID minimum after the first peak and use those layers for downstream tasks—rests on the TwoNN estimator producing reliable, comparable ID values across layers with vastly different geometric properties. If the estimator were biased by changes in data density or curvature (which vary systematically across layers as the paper shows), the ID profile might be an artifact of the estimator rather than a genuine geometric phenomenon. The paper's validation of the local density assumption and the multiscale robustness check are therefore essential—not peripheral—to the central argument.

Neighborhood Overlap as a Measure of Representational Similarity

Knowing the ID at each layer tells us about the global dimensionality of the manifold, but it does not tell us whether the neighborhood structure—which points are close to which others—remains stable or changes across layers. Two representations could have the same ID but organize the data completely differently. The paper needs a metric that captures local rearrangements: do the same data points remain close to the same neighbors, or does the similarity structure get rewritten?

The neighborhood overlap $\chi^{l,m}_k$ : Introduced by Doimo et al. (2020), this metric measures the average fraction of shared $k$ -nearest neighbors between two representations of the same dataset. Formally:

$\chi^{l,m}_k = \frac{1}{N} \sum_{i=1}^N \frac{1}{k} \sum_{j=1}^N A^l_{ij} A^m_{ij}$

where $N$ is the number of data points, $k$ is the number of nearest neighbors considered, and $A^l_{ij} = 1$ if point $j$ is among the $k$ nearest neighbors of point $i$ in the representation at layer $l$ (based on Euclidean distance), and $0$ otherwise. $A^m_{ij}$ is the analogous adjacency matrix for layer $m$ .

What this equation computes: For each data point $i$ , look at its $k$ nearest neighbors in representation $l$ and its $k$ nearest neighbors in representation $m$ . Count how many neighbors appear in both sets. Average this count across all $N$ data points and normalize by $k$ . The result $\chi^{l,m}_k$ lies in $[0, 1]$ , where $1$ means the two representations have identical local neighborhood structure (every point has exactly the same $k$ nearest neighbors in both) and $0$ means the neighborhood structure is completely disjoint. A value of $0.5$ means that, on average, half of each point's $k$ nearest neighbors in layer $l$ are also among its $k$ nearest neighbors in layer $m$ .

Why this form: Unlike correlation-based metrics (e.g., CCA, SVCCA) that measure linear alignment between representation spaces, neighborhood overlap measures similarity of relative positions—it is invariant to rotations, translations, and global scaling of the representation. If two representations differ only by an orthogonal transformation (a global rotation in $\mathbb{R}^d$ ), all pairwise Euclidean distances are preserved, so all nearest-neighbor sets are identical, and $\chi^{l,m}_k = 1$ . This is the correct behavior: a rotation doesn't change the information content or the local geometry, and the metric should reflect that invariance. In contrast, if the representation undergoes a nonlinear transformation that preserves global structure but scrambles local neighborhoods—say, mapping all cats close to all dogs while keeping cats clustered together—the overlap would decrease, correctly signaling that the local semantics have changed.

Two distinct uses of the overlap metric:

Use 1: Inter-layer overlap $\chi^{l,l+1}_k$ measures how much the representation changes between consecutive layers. The paper uses this to identify which parts of the network are in a phase of rapid reorganization (low overlap) versus stable refinement (high overlap). The choice of consecutive layers $(l, l+1)$ rather than comparing arbitrary layer pairs is deliberate: it measures the rate of change rather than the accumulated change. A low $\chi^{l,l+1}_k$ means the representation is being substantially rewritten at that specific block. A high $\chi^{l,l+1}_k$ means the block is making only minor adjustments to the local neighborhood structure.

Use 2: Ground-truth overlap $\chi^{l,gt}_k$ measures how well the neighborhood structure at layer $l$ aligns with a semantic partition of the data defined by external labels. Here, the "adjacency matrix" for the ground truth is $A^{gt}_{ij} = 1$ if data points $i$ and $j$ share the same label (e.g., same protein superfamily, same ImageNet class) and $0$ otherwise. Then $\chi^{l,gt}_k$ is computed identically: for each point, what fraction of its $k$ nearest neighbors in the representation share its label? The paper states (Section 2.2): "Both $\chi^{l,m}_k$ and $\chi^{l,gt}_k$ lie in $[0, 1]$ and depend on the choice of neighborhood size $k$ ."

The label overlap as a semantic probe: The ground-truth overlap is effectively a $k$ -nearest-neighbor classifier accuracy measured without training any classifier. If proteins from the same superfamily cluster tightly together in the representation space, then a random protein's $k$ nearest neighbors will mostly share its superfamily label, and $\chi^{l,gt}_k$ will be high. This is a geometric measure of semantic content: it captures whether the representation has organized the data such that semantically similar items are close in Euclidean distance. The paper emphasizes that this is not a linear probe—it does not test whether semantic information is linearly decodable from the representation, only whether it is reflected in the local similarity structure.

Choice of $k$ and robustness: The paper uses $k = 10$ for the SCOPe protein superfamily analysis and $k = 30$ for the ImageNet analysis (Section 2.2). The paper states that "the qualitative trend of the overlap is robust to changes in $k$ " and demonstrates this in Appendix Figure S5, which shows $\chi^{l,l+1}_k$ and $\chi^{l,gt}_k$ for $k \in \{1, 2, 5, 10, 20, 50\}$ on ESM-2 (650M) and iGPT-L. The curves are qualitatively unchanged—the peaks and troughs appear at the same relative depths—confirming that the choice of $k$ does not drive the observed patterns.

A subtlety for protein remote homology: When computing $\chi^{l,gt}_k$ for the SCOPe dataset, the paper specifically excludes neighbors that belong to the same family as the query protein, even if they share the same superfamily label. The paper states (Section 3.2): "Here, 'gt' represents the classification by superfamily, with the exclusion of neighbors in the same family, allowing us to focus on remote homology specifically." This is crucial: proteins in the same family have highly similar sequences (often >30% sequence identity) and are trivially detectable by sequence alignment methods. Remote homology—proteins that share a common evolutionary ancestor and thus similar 3D structure, but whose sequences have diverged beyond recognition by alignment—is the challenging biological problem. By excluding same-family neighbors from the overlap computation, the paper ensures that $\chi^{l,gt}_k$ measures the representation's ability to capture deep structural relationships, not just sequence similarity. The SCOPe dataset filtering (Section 2.3) ensures this: "we select proteins that belong to superfamilies with at least ten members and ensure that each superfamily consists of at least two families," producing "a dataset of 10,256 sequences and 288 superfamilies."

Why this metric is better than linear probes for this paper's goals: Linear probing (training a logistic regression or linear SVM on frozen representations) measures whether semantic information is linearly separable—a property that depends on both the representation quality and the probe's capacity. Two representations might have the same linear probe accuracy despite one having a much cleaner cluster structure (e.g., if the good representation has nonlinear but simple decision boundaries), or vice versa. The ground-truth overlap $\chi^{l,gt}_k$ measures something different and more fundamental: whether semantic neighbors are also geometric neighbors. This directly tests the paper's geometric hypothesis—that semantic structure emerges through changes in the manifold geometry—and provides a label-based validation that is philosophically consistent with the label-free ID measurement.

Data Extraction and Preprocessing Pipeline

The paper analyzes representations from pretrained models—it does not train any models itself (except for the preliminary training-dynamics analysis using provided checkpoints). The extraction pipeline must therefore be consistent across all layers and all models to enable fair comparison of geometric properties.

Layer extraction point: The authors extract representations "after the first normalization layer of each block" (Section 2.3). In a standard transformer block (Vaswani et al., 2017), the sequence of operations is: input → LayerNorm → Multi-Head Self-Attention → residual addition → LayerNorm → MLP → residual addition → output. Extracting after the first LayerNorm means the representation has been normalized but has not yet been modified by the attention mechanism of that block. This choice is somewhat arbitrary—one could extract after attention or after the MLP—and the paper briefly validates robustness: Appendix Figure S4 compares ID and $\chi^{l,gt}$ for representations extracted "after the first normalization layer" versus "after the attention maps of each self-attention block" for iGPT-L. The profiles are "consistent," indicating "the robustness of our analysis with respect to the layer choice."

Average pooling rationale and procedure: Each input (protein sequence or image) is a sequence of $l$ tokens, and the model produces a representation of shape $\mathbb{R}^{l \times d}$ at each layer, where $l$ varies across data points (proteins have different lengths; images have the same number of pixels per model size but different model sizes use different resolutions). To compare representations across data points with different sequence lengths, the paper collapses the sequence dimension by averaging:

$f_i(x) \rightarrow \frac{1}{l} \sum_{j=1}^{l} (f_i(x))_j$

where $f_i(x)$ is the $\mathbb{R}^{l \times d}$ representation at layer $i$ , and the result is a single vector in $\mathbb{R}^d$ . The paper states this choice is motivated by practicality: "this reduction retrieves sufficient biological information to solve, directly or possibly after finetuning, homology, structural and evolutionary tasks" (Appendix A). For the image case, average pooling also dramatically reduces computational cost: "calculating distance matrices for many images (e.g., 90,000) in the entire feature space ( $l \times d = 524,288$ ) would have been unfeasible given our current computational resources" (Section 4).

Does average pooling distort the ID profile? This is a critical question because averaging over $l$ positions is a many-to-one mapping that could artificially reduce the ID or smooth out geometric structure. The paper addresses this in Appendix D with a control experiment on CIFAR-10 using a WideResNet-28-8, where they can compute the ID with and without average pooling over the spatial dimensions. Appendix Figure S3 shows that "we observe a qualitative consistency in the shape of the ID profiles in both cases, as they conform to the typical bell-shaped curve characteristic of CNN architectures." Quantitatively, the ID after pooling is lower—"a downward shift as a consequence of the averaging procedure"—but the relative pattern (which layers have higher or lower ID) is preserved. The paper notes this shift is "particularly pronounced at the initial stages of the architecture" likely due to "the low number of channels after the first block (16) compared to the later blocks (128, 256, and 512, respectively)," a concern that does not apply to transformers where the embedding dimension is constant across all layers and is "substantially higher than in early CNN layers (ranging from 512 to several thousand in modern large language models)." The paper concludes: "we can be confident about the qualitative robustness of the profiles as long as the number of 'channels'... significantly surpasses the ID as in large transformer models."

Distance computation: After average pooling, each dataset becomes a collection of $N$ vectors in $\mathbb{R}^d$ per layer. The paper computes pairwise Euclidean distances between all vectors within each layer's collection. These distances are the sole input to both the TwoNN ID estimator and the neighborhood overlap computation. The choice of Euclidean distance is standard but non-trivial: it assumes the representation space is approximately Euclidean at the local scale, and it is sensitive to the overall scale (norm) of the representations. Since LayerNorm is applied at the start of each transformer block, the representations are approximately normalized, making Euclidean distance a reasonable choice. However, the paper does not explore alternative distance metrics (cosine similarity, Mahalanobis distance) or their effect on the geometric measurements.

Computation of $\chi^{l,l+1}_k$ : For consecutive layers $l$ and $l+1$ , the Euclidean distances within each layer's representation are computed independently. For each point $i$ in layer $l$ , its $k$ nearest neighbors are identified; the same is done in layer $l+1$ . The overlap is the average over points of the fraction of shared neighbors. This is repeated for all consecutive layer pairs to produce the $\chi^{l,l+1}_k$ profile plotted in Figure 2.

Computation of $\chi^{l,gt}_k$ : For each layer $l$ and each point $i$ , the $k$ nearest neighbors in the representation at layer $l$ are identified. The label adjacency matrix $A^{gt}$ is constructed from the ground-truth annotations (SCOPe superfamilies for proteins, ImageNet classes for images). The overlap is the average fraction of neighbors that share the point's label, with the same-family exclusion applied for the SCOPe remote homology task.

Models and Datasets: The Two Domains Under Study

The paper's universality claim depends on comparing models from maximally different domains. Understanding exactly what differs between the ESM-2 and iGPT settings—and what is held constant—is essential for interpreting the results.

Protein language models: ESM-2 family. The ESM-2 models (Lin et al., 2023) are trained with a masked language modeling (MLM) objective: random amino acids in a protein sequence are replaced with a [MASK] token, and the model must predict the original amino acid from the surrounding context. This is a bidirectional objective—the model can attend to tokens both before and after the masked position. The training data is UniRef50, a clustered version of the UniProt database where sequences with >50% identity are collapsed into a single representative, comprising millions of protein sequences. The three model sizes studied are:

Model	Blocks	Embedding dim. $d$	Attention heads	Parameters
ESM-2 (35M)	12	480	20	35M
ESM-2 (650M)	33	1280	20	650M
ESM-2 (3B)	36	2560	40	3B

The token vocabulary comprises "an alphabet of ≃20 amino acids" (Section 2.3). Each token is embedded into a vector of size $d$ , and a learned positional embedding (common to all sequences) is added. The MLM head $f_{\text{out}}$ at the output projects the final representation to a distribution over the 20 amino acids at each masked position.

Image transformers: iGPT family. The iGPT models (Chen et al., 2010) are trained with an autoregressive next-token prediction objective: given the sequence of pixels up to position $i-1$ , predict the pixel at position $i$ . This is a unidirectional (causal) objective—the model can only attend to previous tokens. To make the autoregressive objective computationally feasible on images, the authors quantize the RGB color space: "the $\mathbb{R}^3$ color space in which each pixel is represented by a triplet of real numbers $(R, G, B)$ is quantized with k-means clustering ( $k = 512$ ), and each pixel is described by the discrete 'code' of the cluster where it belongs" (Appendix A). The three model sizes studied are:

Model	Blocks	Embedding dim. $d$	Attention heads	Parameters	Resolution
iGPT-S (76M)	24	512	8	76M	32×32
iGPT-M (455M)	36	1024	8	455M	48×48
iGPT-L (1.4B)	48	1536	16	1.4B	64×64

The larger models use higher-resolution inputs, which means more tokens per image and thus a larger sequence length $l$ . The training data is the ImageNet training set (1.28 million images across 1000 classes), though the models are trained purely on pixel prediction without class labels. The input resolution is deliberately reduced from the standard 224×224: "Due to the high memory footprint required by the attention layers, Chen et al. [10] reduce the image from the standard ImageNet size."

Key differences between the domains that make the comparison powerful:

Training objective: MLM (bidirectional, predicting masked tokens from context) vs. autoregressive (unidirectional, predicting next token from history). If the same geometric pattern emerges, it is not an artifact of the specific self-supervision objective.
Vocabulary size: ~20 amino acids vs. 512 color clusters. The token space in iGPT is much richer, which might affect the complexity of the learned representations.
Data structure: Protein sequences are one-dimensional strings with long-range dependencies (residues far apart in the sequence may be close in 3D structure). Images are two-dimensional with strong local correlations and spatial structure, rearranged into 1D raster order for the transformer. The fact that the same geometric pattern appears despite this dramatic difference in data structure is strong evidence for universality.
What's held constant: Both model families use the same core architecture—a stack of identical self-attention blocks with residual connections and layer normalization. Both produce representations of constant dimensionality $d$ across all layers. Both are trained on large-scale datasets with self-supervision and have publicly available pretrained weights.

Analysis datasets: The paper uses three datasets for extracting and analyzing representations:

ProteinNet (training set, $N = 25,299$ sequences): Used for the main ID curves (Figure 1, left) and the inter-layer neighborhood overlap (Figure 2, left). ProteinNet is a standard benchmark "for evaluating the relationships between protein sequences and their structures" (Section 2.3). It provides diverse protein sequences with known 3D structures, making it suitable for studying how structural information is encoded in the representations.
SCOPe v2.08 (filtered, $N = 10,256$ sequences, 288 superfamilies): Used for the remote homology analysis (Figure 4, left). The filtering procedure recommended by Söding and Remmert (2011) and the additional constraint that "each superfamily consists of at least two families" ensures the dataset tests deep evolutionary relationships rather than trivial sequence similarity. The paper notes this produces "a dataset of 10,256 sequences and 288 superfamilies."
ImageNet training set (subset, $N = 90,000$ images, 300 classes): Used for the image ID curves (Figure 1, right), inter-layer overlap (Figure 2, right), and class-label overlap (Figure 4, right). The paper "randomly select[s] 300 classes and keep[s] 300 images per class." This sampling is for computational tractability—computing pairwise distances for 90,000 images in 1536-dimensional space is already heavy, and the full ImageNet training set (1.28 million images) would be infeasible.

The ID Profile and the Three-Phase Decomposition

The core empirical observation—the ID profile has a characteristic shape with distinct phases—is established through Figures 1 and 2. The paper does not train models to produce this profile; it observes it in fully trained, off-the-shelf models.

The typical ID profile in protein language models (Figure 1, left): The ID is plotted against "relative depth"—the block number divided by the total number of blocks, normalizing the x-axis to $[0, 1]$ so that models with different numbers of layers can be compared. Three phases are consistently observed across all three ESM-2 model sizes:

Phase 1: Expansion (relative depth ~0–0.3). The ID rises rapidly from the embedding layer's value to a prominent peak. For ESM-2 35M, the peak ID is approximately 20; for 650M, approximately 25; for 3B, approximately 32. The paper notes that "the ID at the peak grows with the model size"—larger models expand the data into higher-dimensional manifolds.
Phase 2: Compression to a plateau (relative depth ~0.3–0.85). The ID contracts from the peak and stabilizes at "notably low values, between 5 and 7." The paper emphasizes a surprising quantitative consensus: "During the plateau phase, we observe a quantitative solid consensus in the ID, especially in the layer at the elbow before the final ascent, independent from the embedding dimension that varies from 480 in ESM-2 35B to 2560 in ESM-2 3B." This means the ID of the compressed representation is not an artifact of the embedding dimension—the 3B model's 2560-dimensional space supports a manifold of ID ~6, just like the 35M model's 480-dimensional space. The paper connects this to an independent measurement by Facco et al. (2019), who found ID values "between 6 and 12" using "different metrics for computing distances directly applied to pairwise alignments of protein sequences." The fact that a self-supervised transformer spontaneously discovers the same intrinsic dimensionality as direct sequence-alignment-based methods suggests the model is recovering something fundamental about the underlying biology.
Phase 3: Final ascent (relative depth ~0.85–1.0). The ID "grows again, returning progressively to values close to the ID computed on the input representation after the positional embedding." The paper attributes this to the MLM objective: "This is a consequence of the masked language modeling objective, which focuses on using the context information to recover the missing tokens." The final layers must reconstruct the input-level detail needed for token prediction, which requires expanding back toward the input's dimensionality.

The typical ID profile in image transformers (Figure 1, right): The iGPT models show a related but distinct pattern:

Phase 1: Expansion (relative depth ~0–0.4). Similar to pLMs, the ID rises to a peak. For iGPT-S, the peak is approximately 25; for iGPT-M, approximately 28; for iGPT-L, approximately 32. The paper notes the peak is "significantly smaller" than in convolutional models (which reach ID values "in the range of $\mathcal{O}(10^4)$ to $\mathcal{O}(10^6)$ "), attributing this partly to "the different embedding dimensions: equal to 512 in the iGPT-S transformer, whereas in convolutional architectures, it fluctuates across layers and is considerably larger."
Phase 2: Compression to a local minimum (relative depth ~0.3–0.4). The ID decreases from the peak to a minimum of approximately 22, common across all three model sizes. This minimum is "compatible with the values observed by Ansuini et al. [13] at the output of a broad range of convolutional classifiers trained on the ImageNet dataset," where "the ID varied from 13 to 24 and was found to be correlated with the classifier's accuracy." The paper emphasizes this correspondence: the self-supervised transformer's intermediate representation has the same ID as a supervised CNN's final representation, suggesting similar semantic content.
Phase 3: Second peak (relative depth ~0.4–0.85). Unlike pLMs, iGPT models show "another, more gradual increase, forming a second, shallower peak near the end of the network." The iGPT-S model (76M parameters) does not exhibit this second peak, while iGPT-M and iGPT-L do. The paper interprets this as evidence that larger models develop a more explicit decoder phase, making the autoencoder analogy more visually apparent.
Phase 4: Final compression (relative depth ~0.85–1.0). The ID decreases again toward the output, approaching the input ID. This contrasts with pLMs, where the ID increases in the final layers.

The connection between ID and neighborhood rearrangement (Figure 2): The inter-layer neighborhood overlap $\chi^{l,l+1}_k$ reveals that the ID phases correspond to distinct dynamical regimes. For pLMs (Figure 2, left), during the expansion phase (relative depth < 0.4), " $\chi^{l,l+1}$ remains approximately at 0.5, meaning that the neighborhood composition undergoes substantial changes in each layer, significantly altering representations in the initial part of the network." In the plateau phase, "the rate of representation evolution is much slower, with over 90% of the neighbors shared between consecutive layers in the larger models." The paper notes an interesting size effect: "In the smallest model, characterized by a higher perplexity, the neighborhood composition in the plateau phase is less consistent, with some rearrangements also occurring in successive layers." This means the 35M model, which makes more errors (higher perplexity), has to continue adjusting its representations even in the plateau, while the larger models achieve a stable compressed representation. For iGPT models (Figure 2, right), the pattern is "qualitatively similar" but with "a more gradual" transition between the initial rapid-rearrangement phase ( $\chi^{l,l+1} \sim 0.7$ ) and the stable phase ( $\chi^{l,l+1} \sim 0.9$ ). A significant rearrangement is "always observed in the last layers, where the reconstruction task is carried on."

Training dynamics of the ID profile (Figure 3): The paper analyzes intermediate training checkpoints to understand how the ID profile emerges over the course of training. For ESM-2 650M (Figure 3, left), checkpoints are at $[0, 10^4, 3 \times 10^4, 10^5, 5 \times 10^5]$ training steps. The emergence follows a two-stage process: "In the initial stages of training, we observe a rapid formation of a peak in the early layers, while the ID curve in the remaining part of the network closely resembles that of an untrained model." Then, "between $10^4$ and $3 \times 10^4$ training steps, the ID of the plateau layers substantially decreases," while "the ID measured in the last layers progressively increases towards the input ID." From $3 \times 10^4$ steps to convergence, the curve takes its final shape with minor adjustments. The paper summarizes: "The final ID curve is essentially achieved in two stages: first, the initial ID peak emerges, and only in a second phase during training the data representation in the plateau layers is compressed to a lower dimension." For iGPT-L (Figure 3, right), checkpoints are at $[0, 1.3 \times 10^5, 2.6 \times 10^5, 5.2 \times 10^5, 10^6]$ training steps. A similar two-stage pattern appears, with the additional nuance that "a second peak emerges in the last third of the model's hidden layers" in later training stages. This temporal separation—expansion learned first, compression learned second—suggests these are genuinely distinct computational phases rather than an artifact of the architecture.

Semantic Validation: Connecting Geometry to Content

The paper's central claim—that the ID minimum identifies semantically rich layers—requires demonstrating that the geometric measurement actually predicts semantic content. This is done through the ground-truth overlap $\chi^{l,gt}_k$ .

Protein remote homology (Figure 4, left): For all three ESM-2 models, the overlap with SCOPe superfamilies (excluding same-family neighbors) shows a consistent pattern: "structural homology information is absent in the positional embedding layer; it grows smoothly in the peak phase, reaching a stationary maximum $\chi^{l,gt}_k \sim 0.8$ in the plateau phase, and it suddenly decreases in the final ascent phase." The key numbers: the last hidden layer scores approximately 0.4, while the plateau layers score approximately 0.8—a factor-of-two improvement. The paper states this explicitly: "the predictive power of remote homology by nearest neighbor search is considerably lower in the last hidden layer, scoring at 0.4 instead of 0.8." This has direct practical implications: "searching for the closest homologs in the plateau layers where protein relationships are better expressed can improve state-of-the-art methods based on representations from the last hidden layer."

Image class semantics (Figure 4, right): For the iGPT models, the overlap with ImageNet class labels shows a peak "around a relative depth of ~0.4 where the ID in the larger models iGPT-M and iGPT-L is lower (see Fig. 1)." The peak becomes sharper and the absolute overlap increases with model size: "the peak value of $\chi^{l,gt}$ is 0.15 in iGPT-S, 0.27 in iGPT-M, and 0.35 in iGPT-L." The paper notes that larger models not only achieve higher peak semantic content but also localize it more sharply to the ID minimum—iGPT-S has a broad, low peak, while iGPT-L has a narrow, high peak precisely at the ID trough.

The consistency of the geometry-semantics relationship: The paper emphasizes that this pattern holds across both domains despite their fundamental differences: "Similar to our observations in pLMs, the representations where the semantic abstractions are better encoded are also those where the ID is low." This is the empirical foundation for the paper's proposed unsupervised strategy: if you can measure the ID profile (which requires no labels), you can identify the layers where semantic content is maximized.

Quantifying the practical benefit: The Appendix (Figure S2) demonstrates a concrete downstream improvement. For the ProtT5-XL-U50 pLM on the SCOPe remote homology task (using $k=1$ nearest neighbor, mirroring the experimental setup of Schütze et al., 2022), "performing the search on a plateau layer instead of the last layer before the output" improves accuracy by "~6%." The paper emphasizes that "the performance gain of ~6% is obtained without any further training"—it requires only using a different layer's representation for the nearest-neighbor lookup.

Why does semantic content peak at the ID minimum? The paper offers an interpretation grounded in the autoencoder analogy. The early expansion phase (high ID) is where the model enriches the representation with nonlinear features, analogous to the kernel trick. The compression phase (low ID) is where the model distills this rich representation into a compact, abstract form that captures the essential structure of the data—protein fold relationships, image class structure—while discarding nuisance variation (sequence details, pixel-level texture). The final ascent or second peak is the decoder phase, where the model must recover the input-level detail needed for the reconstruction objective, necessarily reintroducing variation that dilutes the semantic signal. The ID minimum is the "bottleneck" in this information-theoretic sense: it is the layer where the representation is most compressed and therefore most abstract.

A caveat the paper acknowledges: "The ID of the compressed representations can be larger or smaller than the one measured in the early layers and close to the output, where the network reconstructs the inputs." This means the ID minimum is defined relationally—it is a local minimum after the first peak—not absolutely. In pLMs, the compressed ID (~6) is lower than the input ID; in iGPT, the compressed ID (~22) may be higher than the input ID. The key is not the absolute value but the dip in the profile. The paper attributes the absolute differences to domain-specific factors, including "variations of the pixel brightness and details of data preprocessing, including the substantial reduction of the input resolution" in iGPT, and the "distinct training objectives (MLM in protein language models and next pixel prediction in iGPTs)." The robust cross-domain finding is the shape of the profile, not the absolute ID values.

Design Choices and Their Justifications

Choice 1: Average pooling over sequence dimension rather than using per-token representations. Justification: The paper needs a fixed-dimensional vector per data point to compute Euclidean distances across a dataset of variable-length sequences. Alternatives like treating each token as a separate data point would multiply the dataset size, making distance computations infeasible, and would mix positional and semantic information. The Appendix D validation on CIFAR-10 shows that pooling preserves the qualitative ID profile shape. The paper also notes that pooled representations "retrieve sufficient biological information to solve, directly or possibly after finetuning, homology, structural and evolutionary tasks" (Appendix A), citing the established practice in the protein representation learning literature (Rives et al., 2021).

Choice 2: TwoNN ID estimator rather than PCA variance explained, local PCA, or graph-based methods. Justification: TwoNN requires only nearest-neighbor distances (computable efficiently with approximate nearest-neighbor algorithms), makes minimal assumptions (local density constancy, validated with PAk), and provides a single scalar ID per layer that is directly interpretable as the manifold dimensionality. PCA would miss nonlinear structure; local PCA requires choosing a neighborhood radius; graph-based methods are computationally expensive for $N \sim 10^4$ – $10^5$ data points and are sensitive to the graph construction hyperparameters.

Choice 3: Neighborhood overlap rather than linear probing, CCA, or centered kernel alignment. Justification: The paper's research question is about geometric structure—do semantically similar items become geometric neighbors?—not about linear decodability. The ground-truth overlap $\chi^{l,gt}_k$ directly measures this geometric-semantic alignment without training any classifier. Moreover, the method is philosophically consistent with the ID measurement: both rely solely on Euclidean distances and nearest-neighbor relationships, so the connection between ID and $\chi^{l,gt}$ is a clean relationship between two geometric quantities, not confounded by the capacity of a learned probe.

Choice 4: Relative depth (layer index / total layers) as the x-axis. Justification: Normalizing by total layers enables visual comparison across models with different depths (12 layers for ESM-2 35M, 36 for ESM-2 3B, 48 for iGPT-L). This assumes that the fraction of the network traversed, not the absolute layer count, determines which computational phase is active—an assumption the paper validates by showing that the phases align across model sizes when plotted against relative depth.

Choice 5: Five difficulty quintiles versus continuous difficulty. In this paper, difficulty is not a variable being binned, so the equivalent design decision is the choice to compare three model sizes (small, medium, large) within each architecture family rather than a continuous scaling analysis. This discrete sampling of the model-size axis is sufficient to establish the key trends: the ID at the peak grows with model size, the ID at the plateau/minimum is remarkably consistent across sizes, and the semantic content at the ID minimum increases with model size. A continuous scaling analysis (e.g., training models at many intermediate sizes) would be interesting but would require training models from scratch—beyond the scope of an observational study using publicly available checkpoints.

Choice 6: Using publicly available pretrained checkpoints rather than training models specifically for the analysis. Justification: This ensures the results reflect real-world model behavior, not artifacts of a particular training setup. It also makes the study reproducible and low-cost. The cost is that the authors cannot control for confounds like differences in training data, optimization hyperparameters, or architecture details between model families. However, the fact that the same geometric pattern appears across independently developed model families (ESM-2 by Meta, iGPT by OpenAI, ProtTrans by Rost lab—shown in Appendix Figure S1) with different training recipes strengthens the universality claim precisely because these confounds exist and yet the pattern persists.

Choice 7: SCOPe superfamily labels with same-family exclusion for remote homology. Justification: This ensures the semantic content being measured is genuinely abstract—detecting deep evolutionary relationships that are not trivially recoverable from sequence similarity. The Appendix (Figure S2) demonstrates the practical relevance by connecting to the benchmark of Schütze et al. (2022).

Choice 8: Multi-scale ID estimation with dataset decimation. Justification: The paper reports ID estimates on a 25% random subset (decimation factor 4) because the ID estimate stabilizes at this scale. This addresses the potential confound that measurement noise at the full dataset scale might inflate or deflate the estimated ID. The Appendix notes that the more recent Denti et al. (2022) higher-order neighbor method gives similar results, so the simpler decimation approach is adequate.

4. Key Insights and Innovations

Innovation 1: The Intrinsic Dimension Profile as an Unsupervised Diagnostic for Semantic Content

The paper's most fundamental conceptual contribution is the discovery that the intrinsic dimension of hidden representations—a purely geometric quantity measurable without any labels—serves as a reliable proxy for where abstract semantic knowledge emerges in self-supervised transformers. This is not an incremental improvement to an existing probing technique; it is a fundamentally new diagnostic concept. Before this work, identifying which layers of a pretrained transformer encode the most useful features required either (a) training linear probes on labeled data at every layer (as in Chen et al., 2020 for iGPT), which defeats the purpose of unsupervised pretraining, or (b) adopting the default assumption that the final layer is best (as in Rives et al., 2021 for protein language models), which the paper shows is wrong by a factor of roughly 2× on remote homology detection (Figure 4, left: χ^{l,gt} ~0.8 in plateau layers vs. ~0.4 in the final layer).

The intellectual shift is from treating the ID as a descriptive statistic—something you measure to characterize a representation after you've already identified which layer to use—to treating it as a decision-making signal. The paper's explicit strategy statement captures this: "representations at intermediate layers corresponding to a relative minimum of the ID profile are more suitable for downstream learning tasks" (Section 1). The ID profile becomes a map, and the local minimum after the first expansion marks the location of maximal semantic compression. This is conceptually analogous to how a validation loss curve guides early stopping, but it operates without labels: the ID minimum tells you where to extract representations simply by measuring the geometry of the data manifold at each layer.

What makes this discovery surprising—and not just a restatement of the fact that intermediate layers often work well—is that the relationship between ID and semantic content is non-monotonic and precisely localized. The ID rises to a peak, falls to a minimum, and then rises again (in pLMs) or forms a second peak (in larger iGPT models). The semantic content follows the inverse trajectory: it rises through the expansion phase, peaks precisely at or near the ID minimum, and then degrades in the decoding layers. This is not "intermediate layers are better than final layers," a claim that could be explained by many mechanisms. It is a specific geometric-semantic coupling: the layer where the data manifold is most compressed is the layer where abstract concepts are most explicitly encoded. The fact that this holds across protein models (MLM objective, amino acid sequences, 1D data structure) and image models (autoregressive objective, quantized pixels, 2D data rasterized to 1D) argues that the coupling is not an accident of architecture or training objective but reflects something fundamental about how self-supervised transformers organize computation.

The practical consequence is an unsupervised layer-selection strategy that requires no labels, no probing, and no downstream task definition—only the ability to compute pairwise distances in the representation space and run the TwoNN estimator. This is significant beyond performance gains because it means the diagnostic works in domains where labels are scarce or absent, which are precisely the domains where self-supervised pretraining is most valuable (the paper explicitly mentions rare disease diagnosis in the Acknowledgments). The 6% improvement in remote homology detection from simply switching to a plateau layer (Appendix, Figure S2) demonstrates that the unsupervised diagnostic translates directly to practical accuracy gains, without model modification or fine-tuning.

Innovation 2: The Encoding–Expansion–Compression–Decoding Model of Self-Supervised Transformer Computation

Prior to this work, the dominant conceptual model for how transformers process information was implicit and underdeveloped. The field knew that representations changed across layers—various probing studies had shown that syntactic information peaks in middle layers (Hewitt and Manning, 2019) while semantic information had yielded "contrasting results" (Section 4, citing the NLP probing literature). Architectural analyses had noted that self-attention and MLP blocks alternate, but what this alternation accomplishes geometrically was unknown. The paper proposes and empirically substantiates a specific four-phase computational model that reframes the transformer not as a monolithic sequence processor but as an autoencoder-like system that spontaneously learns to encode, compress, and decode:

Expansion (early layers): The data manifold is expanded into a high-dimensional space, reaching ID values of ~20–32 depending on model size. The paper explicitly connects this to kernel methods: "This expansion resembles the strategy followed in kernel methods, where one implicitly expands the feature space by introducing non-linear functions of the input features" (Section 1). The neighborhood overlap in this phase is low (~0.5 for pLMs, ~0.7 for iGPTs), indicating rapid local reorganization—the representation is being actively rewritten at every layer.
Compression (intermediate layers): The manifold contracts to a low-dimensional representation, with ID values of ~5–7 for proteins and ~22 for images. The neighborhood overlap rises to ~0.9, meaning the representation stabilizes—fewer than 10% of neighbors change between consecutive layers. This is the phase where abstract semantic structure (remote homology, class identity) emerges and peaks.
Decoding (final layers): The manifold either expands back toward the input dimensionality (the final ascent in pLMs) or undergoes a second compression-expansion cycle (the second peak in larger iGPT models). The neighborhood overlap decreases again in these layers, and semantic content degrades sharply, because the reconstruction objective forces the model to recover input-level detail.

The intellectual contribution here is a phase decomposition that was not previously articulated. Prior work had measured ID profiles in supervised CNNs (Ansuini et al., 2013), finding a "hunchback" shape where the ID rises and then monotonically falls—consistent with a network that progressively discards task-irrelevant variation to converge on the classification output. The self-supervised case is fundamentally different: the ID rises again in the final layers because the task is reconstruction, not classification. The network cannot simply compress and stop; it must decompress to produce output tokens. This means the semantic bottleneck is not at the network output but at an intermediate compression minimum, and the final layers actively degrade the semantic signal. This reframes the standard practice of using final-layer representations for downstream tasks as a category mistake—using the decoder output when the encoder bottleneck is what carries the abstract content.

The training dynamics analysis (Figure 3) adds temporal evidence for the phase model: the expansion peak emerges first during training, and only later does the compression to low ID develop. The paper states: "The final ID curve is essentially achieved in two stages: first, the initial ID peak emerges, and only in a second phase during training the data representation in the plateau layers is compressed to a lower dimension." This temporal separation—expansion learned before compression—suggests the phases are genuinely distinct computational regimes, not just a continuous deformation of a single representation. The autoencoder structure is not baked into the architecture (there is no bottleneck layer, no asymmetric encoder-decoder design); it emerges from the training dynamics of a homogeneous stack of identical blocks.

This model is significant beyond description because it makes predictions: if you change the training objective, the final layers should change behavior (they do—MLM produces a final ascent, next-pixel prediction produces a second peak), but the encoding-expansion-compression sequence should remain (it does, across both objectives). If you increase model size, the expansion should become more pronounced (it does—the ID at the peak grows from ~20 in ESM-2 35M to ~32 in ESM-2 3B), because larger models can support richer feature expansions. The model also explains previously puzzling findings: why linear probes on intermediate iGPT layers outperform probes on final layers (because the final layers are decoding, not encoding semantics), and why the last hidden layer of BERT performs worse than intermediate layers on certain semantic tasks (because the masked token prediction objective forces final-layer decompression).

Innovation 3: The Autoencoder as an Emergent Computational Strategy in Architecturally Symmetric Networks

This innovation is more subtle than the phase decomposition itself—it is the claim that the encoding–compression–decoding structure emerges spontaneously in networks that were not designed to have it. Standard autoencoders enforce the bottleneck through architecture: an encoder that progressively reduces dimensionality, a low-dimensional latent layer, and a decoder that expands back. Transformers have no such architectural asymmetry—every block is identical, every layer has the same dimensionality, and the computation is a sequence of mappings from R^{l×d} to R^{l×d}. There is no architectural reason for the middle layers to form a compressed representation; the network could, in principle, maintain the same manifold geometry throughout and simply refine the token-level predictions.

The paper's evidence that an autoencoder-like structure emerges nonetheless is strongest in the iGPT-L model, where "a second peak in the ID profile is present, approximately mirroring the first, and the overlap with the ground truth labels also varies in an almost symmetric manner in the first and the second half of the network (see Fig. 4 right)." The paper explicitly draws the conceptual conclusion: "This makes this model akin to a symmetric autoencoder, in which the decoder performs operations on the representations that are dual to those performed by the encoder" (Section 4). The symmetry is approximate—the second peak is shallower and broader than the first—but the qualitative pattern is unmistakable in the larger models (iGPT-M and iGPT-L) while absent in the smallest (iGPT-S, 76M parameters), suggesting this structured computation requires sufficient capacity to emerge.

The intellectual significance here is connecting to a broader principle in representation learning: that compressing data through a low-dimensional bottleneck forces the representation to capture the essential structure and discard noise. This principle has been understood theoretically for autoencoders and information-bottleneck methods (Tishby et al.), but those frameworks assumed an architectural bottleneck or an explicit information-theoretic objective. The paper's contribution is demonstrating that self-supervised transformers discover this strategy on their own, without architectural constraints or explicit compression objectives, as a consequence of optimizing the reconstruction loss. The training dynamics (Figure 3) support this: compression is not present at initialization (the untrained model has a flat ID profile) and develops during training after the expansion has formed.

Why does this matter? It suggests that the encoding–compression–decoding decomposition is not just an artifact of the models studied, but a convergent computational strategy that gradient-based optimization discovers when tasked with self-supervised reconstruction in high-capacity architectures. If true, this has implications for architecture design (you may not need to build bottlenecks; the network will find them), for interpretability (the compression minimum is a natural point to extract representations, regardless of architecture), and for transfer learning (the compressed representations may be more universal across tasks than final-layer representations, which are task-specifically decompressed).

The paper is appropriately cautious about this claim. It notes alternative explanations for why the final layers might differ between pLMs and iGPTs: "variations of the pixel brightness and details of data preprocessing" in iGPT, including the aggressive resolution reduction, "can artificially alter the ID significantly" (Section 4). The second peak in iGPT might partly reflect the specific decoding demands of next-pixel prediction on quantized images rather than a universal decoder structure. But the core observation—that a compressed intermediate representation emerges spontaneously and carries maximal semantic content—holds across all models studied, regardless of whether the subsequent decoding manifests as a final ascent or a second peak.

Innovation 4: Reconciling Contradictory Findings Across the Transformer Probing Literature Through a Unified Geometric Lens

The paper does not frame this as its primary contribution, but it represents a significant intellectual move: using the geometric framework to explain why prior work on transformer representations produced apparently contradictory results. The NLP probing literature had found that syntactic information peaks in middle layers (Hewitt and Manning, 2019), but "the localization of semantic information has yielded contrasting results" (Section 4, citing six studies from 2017–2020). Some studies found semantic features in middle layers, others in final layers, others with no clear peak. The protein language modeling literature had settled on final-layer representations for homology detection (Rives et al., 2021), even though the present paper shows this is suboptimal by a factor of ~2× on χ^{l,gt}. The image transformer literature (Chen et al., 2010) had found that the best intermediate layer for linear probing varied with model size, with no explanation for why certain layers worked better.

The geometric framework provides a unified explanation: the layer that maximizes semantic content is the layer at or near the ID minimum after the first expansion. The reason prior studies reached different conclusions is that the location of this minimum varies—across domains (it's a broad plateau in pLMs, a sharp dip in iGPT), across model sizes (it sharpens with scale, as shown by the χ^{l,gt} peaks in Figure 4 right), and possibly across tasks (syntactic vs. semantic information may have different optimal layers, as the NLP probing literature suggests). A study that probed only a few layers, or that used a model size where the ID minimum was less pronounced, or that used a task whose semantic content emerged at a different depth, would naturally reach different conclusions. The geometric framework doesn't invalidate these prior findings; it contextualizes them by providing the missing independent variable—the ID minimum's location—that predicts where each type of information will be found.

The paper's preliminary NLP result in Appendix E (Figure S6) on Llama-2-70B is particularly revealing: the ID profile shows "tree peaks and two local minima across the hidden layers," and "the highest overlap with class partition, determined by sentence sentiment, occurs in correspondence with the first local minimum." This is consistent with the framework while acknowledging that language data may be more complex: "After the second peak, the ID profile shows a more complex behavior, which we are currently investigating" (Section 4). The complexity of the Llama-2 ID profile (three peaks rather than the one or two seen in ESM-2 and iGPT) suggests that language models may go through multiple expansion–compression cycles, perhaps corresponding to different levels of linguistic abstraction (morphology, syntax, semantics, pragmatics), but the key finding—that semantic content peaks at an ID minimum—remains consistent.

This reconciliation is significant because it converts a confusing empirical landscape into a coherent framework with a clear diagnostic. Instead of asking "which layer should I use for my task?" and searching across layers with probes, one can measure the ID profile and focus on the layers near the ID minimum. The paper demonstrates this for two domains; the preliminary NLP result suggests broader applicability. The framework also generates specific, testable predictions: for any self-supervised transformer, (a) the ID profile will show at least one expansion–compression cycle, (b) the semantic content of the data (as defined by some coherent label partition) will peak at or near the ID minimum within each cycle, and (c) the sharpness of the ID minimum and the corresponding semantic peak will increase with model scale. These predictions are falsifiable and span domains, making the framework more than a post-hoc description of observed patterns.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper uses three analysis datasets: (1) ProteinNet training set (25,299 protein sequences) for the main ID curves (Figure 1, left) and inter-layer neighborhood overlap (Figure 2, left); (2) SCOPe v2.08 filtered to 10,256 protein domain sequences across 288 superfamilies (each containing at least two families, with each superfamily having at least ten members), used for the remote homology overlap analysis (Figure 4, left); (3) ImageNet training set subset of 90,000 images (300 randomly selected classes, 300 images per class), used for image transformer ID profiles (Figure 1, right), inter-layer overlap (Figure 2, right), and class-label overlap (Figure 4, right). The ProteinNet and SCOPe datasets are standard benchmarks in protein structure prediction; the ImageNet subset is sampled for computational tractability given the cost of computing pairwise distance matrices for 90,000 points in 1536-dimensional space.
Base model(s). The paper analyzes two pretrained transformer families: ESM-2 (Lin et al., 2023) protein language models at three scales—35M parameters (12 blocks, embedding dimension 480), 650M parameters (33 blocks, dimension 1280), and 3B parameters (36 blocks, dimension 2560)—trained with masked language modeling on UniRef50; and iGPT (Chen et al., 2010) image transformers at three scales—76M parameters (24 blocks, dimension 512), 455M parameters (36 blocks, dimension 1024), and 1.4B parameters (48 blocks, dimension 1536)—trained with autoregressive next-pixel prediction on ImageNet. The paper also includes supplementary results on ProtBert, ProtT5-XL-U50, ESM-1v, and ESM-1b (Appendix Figure S1), and preliminary NLP results on Llama-2-70B (Appendix Figure S6). All models are used off-the-shelf with publicly available pretrained weights; no models are trained by the authors except for the training-dynamics analysis which uses provided intermediate checkpoints.
Metrics. The paper measures three quantities, all derived from Euclidean distances in the representation space after average pooling:

(1) Intrinsic dimension (ID): Estimated via the TwoNN method (Facco et al., 2017), which fits a Pareto distribution to the ratios $\mu_i = r_{i2} / r_{i1}$ of second-to-first nearest-neighbor distances. The ID is the slope extracted from regressing $\log(\mu_{(i)})$ against $-\log(1 - i/N)$ for the sorted empirical cumulative distribution. To ensure robustness against measurement noise, the ID is computed on random subsets decimated by a factor of 4 relative to the full dataset, as the paper reports the estimate stabilizes between $N/2$ and $N/8$ subset sizes.

(2) Inter-layer neighborhood overlap $\chi^{l,l+1}_k$ : The average fraction of shared $k$ -nearest neighbors between consecutive layer representations of the same data point, computed as $\chi^{l,l+1}_k = \frac{1}{N} \sum_i \frac{1}{k} \sum_j A^l_{ij} A^{l+1}_{ij}$ , where $A^l_{ij} = 1$ if point $j$ is among the $k$ nearest neighbors of $i$ at layer $l$ . The paper states that the qualitative trends are robust to the choice of $k$ (Appendix Figure S5, tested for $k \in \{1, 2, 5, 10, 20, 50\}$ ).

(3) Ground-truth overlap $\chi^{l,gt}_k$ : The average fraction of $k$ -nearest neighbors in the representation at layer $l$ that share the same ground-truth label as the query point. For SCOPe, labels are superfamily classifications, and neighbors from the same family are excluded to isolate remote homology detection specifically ( $k = 10$ ). For ImageNet, labels are the 300 class identities ( $k = 30$ ). The paper uses $k = 10$ for proteins and $k = 30$ for images, consistent with the value used in Doimo et al. (2020).
Baselines. The paper's "baselines" are not competing methods in the traditional sense, as this is an observational study rather than a methods comparison. The key comparisons are (1) final-layer representations vs. intermediate representations, motivated by the standard practice in the protein language modeling literature of using the last hidden layer for downstream tasks (Rives et al., 2021); (2) the ID profile's prediction of semantic content vs. the ground-truth label overlap, which validates that the ID minimum correctly identifies semantically rich layers; and (3) cross-domain and cross-model-size consistency, where the baseline expectation would be domain-specific or size-specific patterns. For the nearest-neighbor remote homology search (Appendix Figure S2), the baseline is the accuracy achieved using the last hidden layer of ProtT5-XL-U50, following the experimental protocol of Schütze et al. (2022).
Generation budget / compute accounting. This paper does not use "compute budget" in the sense of generation counts, since it analyzes frozen representations rather than sampling from models. The relevant computational constraint is the feasibility of computing pairwise distance matrices: the paper samples 90,000 images from ImageNet (rather than the full 1.28 million training images) because "calculating distance matrices for many images (e.g., 90,000) in the entire feature space ( $l \times d = 524,288$ ) would have been unfeasible given our current computational resources" (Section 4). The ID estimates are computed on subsets decimated by a factor of 4 to stabilize the TwoNN estimator, as reported in Appendix B.
Cross-validation / statistical protocol. The paper does not employ cross-validation in the conventional machine learning sense because no models are trained and no hyperparameters are tuned. The key validation procedures are: (1) Multiscale ID estimation: TwoNN is applied to the full dataset and to random subsets of size $N, N/2, N/4, \ldots$ , and the reported ID is the value where the curve stabilizes—specifically on the subset decimated by a factor of 4 (Appendix B). (2) Local density assumption validation: The Point Adaptive kNN (PAk) method (Rodriguez et al., 2018) is used to verify that "on average, the density can be considered constant within the first 6 neighboring data points" (Appendix B), which validates the TwoNN assumption at the scale of the first two neighbors. (3) Robustness to $k$ in neighborhood overlap: Appendix Figure S5 shows that the qualitative shape of $\chi^{l,l+1}_k$ and $\chi^{l,gt}_k$ is unchanged for neighborhood sizes $k \in \{1, 2, 5, 10, 20, 50\}$ , confirming the patterns are not artifacts of a particular $k$ choice. (4) Robustness to layer extraction point: Appendix Figure S4 shows that ID and $\chi^{l,gt}$ are consistent whether representations are extracted after the first LayerNorm or after the self-attention maps in iGPT-L.

Main Quantitative Results

Intrinsic Dimension Profiles Across Domains and Model Sizes (Figure 1)

Protein language models (Figure 1, left): The ID of ESM-2 representations on the ProteinNet dataset follows a consistent three-phase profile across all three model sizes:

The ID at the expansion peak grows systematically with model size: approximately 20 for ESM-2 35M, approximately 25 for ESM-2 650M, and approximately 32 for ESM-2 3B. All three peaks occur in roughly the first third of the network (relative depth ~0–0.3).
During the plateau phase, the ID compresses to remarkably consistent values across model sizes: between 5 and 7 for all three ESM-2 variants. The paper emphasizes: "we observe a quantitative solid consensus in the ID, especially in the layer at the elbow before the final ascent, independent from the embedding dimension that varies from 480 in ESM-2 35B to 2560 in ESM-2 3B." This convergence suggests the compressed ID reflects an intrinsic property of the protein sequence data manifold rather than an artifact of model architecture.
In the final ascent (relative depth ~0.85–1.0), the ID rises toward values comparable to the input representation after positional embedding.

Image transformers (Figure 1, right): The iGPT ID profiles on the ImageNet subset show a related but distinct pattern:

The first expansion peak reaches approximately 25 for iGPT-S (76M), approximately 28 for iGPT-M (455M), and approximately 32 for iGPT-L (1.4B). The peak values are comparable to those in pLMs at similar parameter scales, and they also grow with model size.
After the first peak, the ID decreases to a local minimum at approximately 0.3–0.4 relative depth, reaching values around 22—common across all three model sizes. The paper notes this is "compatible with the values observed by Ansuini et al. [13] at the output of a broad range of convolutional classifiers trained on the ImageNet dataset," where ID values ranged from 13 to 24 and correlated with classifier accuracy.
A distinguishing feature: a second, shallower peak emerges in the larger models. For iGPT-M and iGPT-L (but not iGPT-S, which lacks this feature), the ID rises after the local minimum, forming "a second, shallower peak near the end of the network." The paper interprets this as a decoder phase mirroring the encoder's expansion.
In the final layers (relative depth ~0.85–1.0), the ID contracts again toward the output, approaching the input ID value. This contrasts with pLMs, where the final layers expand the ID.

Supplementary model families (Appendix Figure S1): The three-phase ID profile is replicated across additional pLMs—ProtBert (trained on UniRef100), ProtT5-XL-U50 (UniRef50), ESM-1v (UniRef90), and ESM-1b (UniRef50)—confirming that "despite the significant differences of the pLMs considered in the analysis, the consistency of the three-phased behavior of the ID remains."

Neighborhood Rearrangement Mirrors the ID Phases (Figure 2)

Inter-layer overlap in pLMs (Figure 2, left): The neighborhood overlap between consecutive layers $\chi^{l,l+1}_k$ on ProteinNet reveals a direct correspondence with the ID profile:

In the expansion phase (relative depth < ~0.4), $\chi^{l,l+1}$ remains approximately at 0.5 for all three ESM-2 models, meaning roughly half of each point's neighbors change between consecutive layers—rapid and substantial representational reorganization.
In the plateau phase, the overlap rises sharply: over 90% of neighbors are shared between consecutive layers in the larger models (ESM-2 650M and 3B). The paper observes a size-dependent effect: "In the smallest model, characterized by a higher perplexity, the neighborhood composition in the plateau phase is less consistent, with some rearrangements also occurring in successive layers."
The final ascent phase (relative depth > ~0.85) shows increased rearrangement, with the overlap decreasing from the plateau's high values.

Inter-layer overlap in iGPTs (Figure 2, right): The pattern is "qualitatively similar" but with a more gradual transition between phases:

In the initial stage coinciding with the first ID peak (relative depth < ~0.4), $\chi^{l,l+1} \sim 0.7$ —lower overlap than the subsequent stable phase, but higher than the ~0.5 observed in pLMs, suggesting less dramatic per-block reorganization.
In the intermediate region where the ID is at its local minimum, $\chi^{l,l+1} \sim 0.9$ —comparable to the pLM plateau, indicating a stable representation.
"A significant neighborhood rearrangement is always observed in the last layers, where the reconstruction task is carried on," with the overlap dropping again.
As with pLMs, " $\chi^{l,l+1}$ is lower in shallower models that need faster rearrangements to achieve consistent results in fewer blocks"—iGPT-S (24 blocks) shows lower overlap throughout than iGPT-L (48 blocks).

Training Dynamics of the ID Profile (Figure 3)

ESM-2 650M training progression (Figure 3, left): Analysis of intermediate checkpoints at $[0, 10^4, 3 \times 10^4, 10^5, 5 \times 10^5]$ training steps reveals a two-stage emergence:

In the earliest training stages (≤10^4 steps), "a rapid formation of a peak in the early layers" occurs, while the rest of the ID curve "closely resembles that of an untrained model"—a relatively flat profile.
Between 1×10^4 and 3×10^4 steps, "the ID of the plateau layers substantially decreases," while "the ID measured in the last layers progressively increases towards the input ID."
From 3×10^4 steps to convergence (5×10^5), the curve adopts its final shape with minor quantitative adjustments: "a slight increase of the ID measured at the peak and a minor compression at the plateau."
The paper summarizes: "The final ID curve is essentially achieved in two stages: first, the initial ID peak emerges, and only in a second phase during training the data representation in the plateau layers is compressed to a lower dimension."

iGPT-L training progression (Figure 3, right): Checkpoints at $[0, 1.3 \times 10^5, 2.6 \times 10^5, 5.2 \times 10^5, 10^6]$ training steps show a similar but more complex trajectory:

At initialization (0 steps), the ID curve is flat, as in the untrained pLM.
In early training (<1.3×10^5 steps), "the first ID peak emerges in the initial layers of the network."
In later stages, "the ID slightly increases at the first peak and decreases at 0.4 of relative depth, forming the local minimum."
Distinct from pLMs: "a second peak emerges in the last third of the model's hidden layers" during the later training stages, suggesting the decoder structure develops after the encoder bottleneck is established.

Semantic Content Peaks at the ID Minimum (Figure 4)

Remote homology in protein language models (Figure 4, left): The ground-truth overlap $\chi^{l,gt}_{10}$ with SCOPe superfamilies (same-family neighbors excluded) shows:

Homology information is absent in the positional embedding layer (overlap near zero).
The overlap "grows smoothly in the peak phase" as the ID expands, rising from near-zero to high values.
It reaches a stationary maximum of $\chi^{l,gt}_{10} \sim 0.8$ in the plateau phase for all three model sizes. This means that, on average, 8 out of 10 nearest neighbors of a protein in the plateau representation belong to the same superfamily (excluding same-family neighbors), indicating strong remote homology signal.
The overlap "suddenly decreases in the final ascent phase," dropping sharply in the layers where the ID rises again. The last hidden layer scores approximately 0.4—half the plateau value. The paper states: "the predictive power of remote homology by nearest neighbor search is considerably lower in the last hidden layer, scoring at 0.4 instead of 0.8."
The three model sizes show nearly identical $\chi^{l,gt}$ curves in the plateau, suggesting that remote homology encoding is robust to model scale once the compression phase is reached.

Image class semantics in iGPT (Figure 4, right): The ground-truth overlap $\chi^{l,gt}_{30}$ with ImageNet labels shows:

A clear peak of $\chi^{l,gt}$ "around a relative depth of ~0.4 where the ID in the larger models iGPT-M and iGPT-L is lower."
The peak value increases substantially with model size: $\chi^{l,gt} \sim 0.15$ in iGPT-S, ~0.27 in iGPT-M, and ~0.35 in iGPT-L. Unlike the pLM case where all models reach the same plateau value, larger iGPT models achieve higher peak semantic content.
The peak becomes sharper as model size increases: iGPT-S shows a broad, low peak spanning much of the network; iGPT-L shows a narrow, high peak precisely at the ID local minimum (~0.4 relative depth).
The overlap profile is approximately symmetric: it rises through the encoding expansion, peaks at the compression minimum, and falls through the decoding second peak. The paper notes that in iGPT-L, "the overlap with the ground truth labels also varies in an almost symmetric manner in the first and the second half of the network."

Practical downstream improvement (Appendix Figure S2): Using ProtT5-XL-U50 on SCOPe with $k=1$ nearest neighbor search (matching the protocol of Schütze et al., 2022):

Performing the homology search using a plateau layer representation instead of the last hidden layer improves accuracy by approximately 6% (absolute improvement, based on visual inspection of Appendix Figure S2).
The paper emphasizes: "the performance gain of ~6% is obtained without any further training." The improvement comes solely from extracting representations from a different layer of the already-trained model.

Preliminary NLP Analysis (Appendix Figure S6)

The paper includes early results on Llama-2-70B applied to the Stanford Sentiment Treebank (SST) dataset:

The ID profile "is more complex, showing tree peaks and two local minima across the hidden layers." This is qualitatively different from the one-peak (pLM) or two-peak (iGPT) profiles, suggesting additional phases of expansion and compression in language models.
Critically, the central finding persists: "the highest overlap with class partition, determined by sentence sentiment, occurs in correspondence with the first local minimum" of the ID profile.
The paper is cautious about interpretation beyond this: "After the second peak, the ID profile shows a more complex behavior, which we are currently investigating."

Cross-Domain Quantitative Summary

The following table synthesizes the key ID values across domains and model sizes, extracted from Figures 1 and 4:

Model	Peak ID	ID at Minimum/Plateau	Peak $\chi^{l,gt}$
ESM-2 35M	~20	~5–7	~0.8
ESM-2 650M	~25	~5–7	~0.8
ESM-2 3B	~32	~5–7	~0.8
iGPT-S (76M)	~25	~22	~0.15
iGPT-M (455M)	~28	~22	~0.27
iGPT-L (1.4B)	~32	~22	~0.35

The pattern is consistent: the ID at the peak grows with model size in both domains, the ID at the minimum/plateau is remarkably stable across model sizes within each domain (but differs between domains—~6 for proteins vs. ~22 for images), and the peak semantic content increases with model size for images but plateaus for proteins.

Ablation Studies and Robustness Checks

Robustness to neighborhood size $k$ (Appendix Figure S5): The paper tests $\chi^{l,l+1}_k$ and $\chi^{l,gt}_k$ for $k \in \{1, 2, 5, 10, 20, 50\}$ on ESM-2 650M and iGPT-L. The qualitative shape of both overlap profiles is preserved across all $k$ values—peaks and troughs appear at the same relative depths for both the inter-layer and ground-truth overlaps. The paper notes one quantitative effect: when $k = 50$ in ESM-2, the absolute overlap values decrease because "certain superfamilies have fewer than 50 elements" in the SCOPe dataset, meaning some points cannot possibly have 50 same-superfamily neighbors.

Robustness to layer extraction point (Appendix Figure S4): For iGPT-L, the paper compares representations extracted after the first LayerNorm ("post_ln") with those extracted after the self-attention maps ("post_attn"). Both the ID profile and the ground-truth overlap $\chi^{l,gt}$ are "consistent" between the two extraction points, "indicating the robustness of our analysis with respect to the layer choice."

Robustness to average pooling (Appendix Figure S3): The paper tests the effect of average pooling on a CIFAR-10 WideResNet-28-8 model where the ID can be computed both in the full feature space and after global average pooling over spatial dimensions. The ID profile shows "qualitative consistency in the shape" between the two settings, conforming to the typical CNN hunchback shape in both cases. Quantitatively, average pooling produces "a downward shift," which is "particularly pronounced at the initial stages of the architecture" where the number of channels is low (16) compared to later blocks (128–512). Since transformer embedding dimensions are substantially larger (512–2560) and constant across layers, the paper argues this concern is minimal for the main experiments.

Validation of the TwoNN local density assumption (Appendix B): Using the Point Adaptive kNN (PAk) method, the paper verifies that "on average, the density can be considered constant within the first 6 neighboring data points." Since the TwoNN estimator only requires the first two neighbors, the locally constant density assumption holds at the relevant scale.

Multiscale ID stability (Appendix B): The paper reports that ID estimates are "less scale-dependent across dataset sizes between $[N/2, N/8]$ ," leading to the choice of reporting estimates on subsets decimated by a factor of 4. The more recent Denti et al. (2022) higher-order neighbor estimator "gave similar results in the experiments of this work."

Cross-architecture consistency in pLMs (Appendix Figure S1): The three-phase ID profile (peak, plateau, final ascent) is replicated on four additional protein language models beyond the ESM-2 family: ProtBert, ProtT5-XL-U50, ESM-1v, and ESM-1b. These models differ in training data (UniRef100, UniRef50, UniRef90), architecture (BERT-style vs. T5-style), and parameter counts, yet "the consistency of the three-phased behavior of the ID remains." This is a strong robustness check against the possibility that the observed pattern is specific to the ESM-2 training recipe.

Negative result: iGPT-S lacks the second peak. The smallest iGPT model (76M parameters, 24 blocks) does not exhibit the second ID peak observed in iGPT-M and iGPT-L. The paper interprets this as evidence that the full encoding–compression–decoding structure requires sufficient model capacity to emerge: "a second shallow peak in the ID is observed, with the exception of the iGPT-S model, which does not have one." This is a meaningful negative result because it establishes a boundary condition on the universality claim—the four-phase profile may require models above a certain depth or parameter threshold.

Negative result: pLM plateau compression degrades in the smallest model. While ESM-2 650M and 3B show $\chi^{l,l+1} \sim 0.9$ in the plateau (meaning 90% of neighbors preserved between consecutive layers), the 35M model shows lower overlap in the plateau, with "some rearrangements also occurring in successive layers." The paper attributes this to "a higher perplexity" in the smallest model, suggesting that models with insufficient capacity cannot fully stabilize the compressed representation.

Critical Assessment

Does the ID Profile Genuinely Predict Semantic Content Without Labels?

The paper's central claim is that the ID minimum provides an unsupervised signal for identifying the layers that maximize semantic content. The evidence for this claim is correlational rather than causal: the ground-truth overlap $\chi^{l,gt}_k$ peaks at or near the ID minimum in every model-domain combination tested (Figure 4). The correlation is striking and consistent, but correlation does not establish that the ID measurement alone would reliably guide layer selection in a genuinely unsupervised setting where one cannot validate against labels.

Specifically: the paper demonstrates that for ESM-2 models, the entire plateau region (which spans over half the network's depth) achieves high semantic content, meaning the ID profile's "minimum" is not a sharp pinpoint but a broad basin. Selecting any layer within the plateau would work well—but the ID profile does not distinguish between the beginning, middle, or end of the plateau, and the paper does not test whether the absolute ID value correlates with fine-grained differences in semantic content within the plateau. For iGPT models, the minimum is sharper, but even there, the iGPT-S model has no clear minimum and the semantic peak is broad and low. The "unsupervised strategy" claim is better supported for larger models (iGPT-L, ESM-2 3B) than for smaller ones, and the paper does not characterize this size-dependence systematically.

A stronger test would be: given only unlabeled data and a pretrained transformer, can the ID profile correctly rank layers by downstream task performance across a diverse set of tasks, without any task-specific tuning? The paper tests only one semantic property per domain (remote homology for proteins, class identity for images) and only one task format (nearest-neighbor overlap). The preliminary NLP result (Appendix Figure S6) adds sentiment analysis on Llama-2-70B, which is encouraging but is a single data point. The universality claim across "many models trained on diverse datasets" (as stated in the abstract) is supported for the specific models and tasks tested but is not yet demonstrated to generalize to, for example, diverse NLP tasks (named entity recognition, question answering, natural language inference) or to different semantic properties in proteins (enzyme function, binding affinity, subcellular localization).

The Autoencoder Analogy Is Compelling but Incompletely Validated

The paper's interpretive framework—that transformers spontaneously develop an encoder-bottleneck-decoder structure—is supported by the symmetry of the ID profile in iGPT-L (Figure 1, right, where the second peak approximately mirrors the first) and the symmetry of the semantic overlap profile (Figure 4, right). However, several aspects of this analogy are asserted rather than demonstrated:

Causal direction is untested. The paper observes that compression follows expansion in the layer sequence and during training, but it does not test whether the expansion enables the compression. Alternative hypotheses are consistent with the data: expansion and compression could be independent processes that happen to occur in sequence; compression might be driven by the increasing receptive field of deeper layers rather than by the preceding expansion; or the expansion might be a side effect of the self-attention mechanism's tendency to mix information rather than a deliberate computational strategy. The kernel-method analogy (Section 1: "this expansion resembles the strategy followed in kernel methods, where one implicitly expands the feature space by introducing non-linear functions of the input features") is interesting but speculative—no experiment tests whether the expanded representation actually makes the data more linearly separable, as the kernel trick would predict.

The decoder phase is not functionally characterized. The paper states that the final layers "address the minute decision-making needed to reconstruct the data" and that "the reconstruction task is carried on" in these layers (Section 1, Figure 2), but it does not probe what these layers are actually computing. Are they recovering positional information? Fine-grained token-level detail? Sequence-level features? The neighborhood rearrangement in the final layers (decreasing $\chi^{l,l+1}$ ) could reflect many types of computation, not necessarily a systematic decoding process. The paper's interpretation that iGPT's second peak represents a decoder mirroring the encoder would be strengthened by showing that the representations at corresponding points in the first and second peaks are geometrically similar (e.g., have high cross-layer neighborhood overlap at symmetric depths) or that they encode complementary information (e.g., the encoder layers specialize in semantic abstraction while the decoder layers specialize in spatial or textural detail).

The compression is not shown to be information-theoretically optimal. The autoencoder analogy implies that the compressed representation captures the essential information needed for reconstruction while discarding noise. The paper does not measure whether the plateau representations can actually be decoded back to the input (a standard test for autoencoder bottlenecks), nor does it characterize what information is lost in the compression. The fact that the compressed ID (~5–7 for proteins) is consistent with independent estimates from sequence-alignment methods (Facco et al., 2019) is suggestive but does not establish that the model's compression is information-preserving in any formal sense.

Model Scale Is Underexplored

The paper compares three model sizes within each architecture family—a small, medium, and large variant. This is sufficient to establish that the peak ID grows with scale and that the compressed ID is stable across scales within a domain. However, several scale-dependent claims would benefit from a more systematic treatment:

The emergence of the second peak in iGPT is attributed to model size but not rigorously isolated. The paper notes that iGPT-S (24 blocks, 76M parameters) lacks the second peak, while iGPT-M (36 blocks, 455M) and iGPT-L (48 blocks, 1.4B) exhibit it. However, depth and width are confounded in this comparison—the larger models are both deeper and wider. A controlled experiment that isolates depth (e.g., training iGPT models of varying depth at fixed width, or analyzing intermediate representations from a single deep model with different numbers of blocks activated) would clarify whether the second peak is a depth phenomenon or a capacity phenomenon.

The relationship between peak ID and downstream performance is not quantified. For iGPT models, larger models achieve both higher peak IDs (~32 for iGPT-L vs. ~25 for iGPT-S) and higher peak semantic content (χ^{l,gt} ~0.35 vs. ~0.15). But the causal chain is unclear: does the higher expansion enable the higher semantic content, or are both independent consequences of having more parameters? A partial correlation analysis (controlling for parameter count) could distinguish these possibilities but is not performed.

The training dynamics analysis (Figure 3) is conducted on only one model size per domain. ESM-2 650M and iGPT-L are both the largest or near-largest models in their families. Whether the two-stage training dynamic (expansion first, compression later) holds for smaller models is unknown. The paper does not discuss whether the smaller models might converge to their ID profiles through a different temporal trajectory.

Domain-Specific Confounds Are Acknowledged but Not Controlled

The paper is transparent about potential confounds (Section 4): "variations of the pixel brightness and details of data preprocessing, including the substantial reduction of the input resolution" could "artificially alter the ID significantly" in iGPT. The quantization of RGB colors into 512 clusters via k-means—a preprocessing step specific to iGPT—introduces a discrete structure that could affect nearest-neighbor distances in ways that differ from the continuous amino acid embeddings in ESM-2. The paper does not analyze whether the ID profiles are sensitive to this preprocessing (e.g., by comparing iGPT representations on images with and without color quantization, or by analyzing a continuous-valued image transformer).

Similarly, the protein and image datasets differ in their label structures. SCOPe superfamilies are hierarchical and evolutionarily grounded; ImageNet classes are categorical and visually grounded. The near-identical χ^{l,gt} across ESM-2 model sizes (~0.8 for all three) while iGPT shows strong scale-dependence (0.15 → 0.27 → 0.35) could reflect differences in the label semantics (remote homology may be easier to capture, or may saturate earlier with model capacity) rather than differences in how the geometric patterns relate to semantics. The paper does not discuss this possibility.

The Practical Gain Is Demonstrated for Only One Task-Domain Combination

The 6% improvement in remote homology detection (Appendix Figure S2, ProtT5-XL-U50) is the paper's most concrete practical result. However, it is reported only in the Appendix, for a single model (not from the main ESM-2 family), using a nearest-neighbor classifier with $k=1$ . The paper does not evaluate whether the plateau layer also improves performance when using more sophisticated downstream methods (e.g., training a classifier on top of the frozen representations, fine-tuning), nor does it demonstrate that the ID-minimum strategy outperforms alternative unsupervised layer-selection heuristics (e.g., using the middle layer, using the layer that maximizes some variance-based criterion). The claim that the ID profile "can be exploited as an unsupervised method to identify the most semantically rich representations" (Section 1) is validated in the sense that the ID minimum aligns with maximal semantic content, but the paper does not compare this method against any baseline layer-selection strategy.

Statistical Rigor and Reproducibility

The paper reports point estimates of ID values (peak IDs of ~20, 25, 32; plateau IDs of ~5–7) without confidence intervals, making it difficult to assess whether differences between model sizes or domains are statistically significant. The TwoNN estimator has known variance properties, and the multiscale decimation procedure provides some robustness, but the paper does not report standard errors or bootstrap confidence intervals for any of the ID estimates. The overlap profiles are shown as curves without error bars or shaded regions.

The test-set sizes are modest for the semantic validation: 10,256 sequences for SCOPe (288 superfamilies, with some superfamilies having fewer than 50 members as noted regarding $k=50$ in Appendix Figure S5), and 90,000 images from 300 ImageNet classes (300 images per class). The SCOPe filtering ("each superfamily consists of at least two families" and "at least ten members") ensures statistical meaningfulness, but with 288 superfamilies and $k=10$ , the χ^{l,gt} measurement's variance depends on the within-superfamily neighbor density, which the paper does not characterize.

The paper provides code and states experiments were conducted on "a machine with 2 Intel(R) Xeon(R) Gold 6226 processors, 256GB of RAM, and 2 Nvidia V100 GPUs with 32GB memory" and that "code to reproduce our experiments and our analysis" is available online. The compute requirements for the main analysis are dominated by pairwise distance calculations in high-dimensional spaces, which scale quadratically in the number of data points—a practical constraint that the paper acknowledges motivated the ImageNet subsampling.

Missing Experiments That Would Strengthen the Claims

Several experiments are conspicuous by their absence:

Causal intervention: If the expansion enables compression and semantic abstraction, then ablating or disrupting the expansion layers (e.g., by injecting noise, removing attention heads, or skipping blocks) should impair the semantic content of the compressed representations. No such intervention is attempted.
Cross-model transfer: If the compressed representation is genuinely universal, representations at the ID minimum of one model should show higher transfer performance to a different model's downstream task than representations at other layers. This is not tested—all semantic evaluations use the same model for representation extraction and nearest-neighbor search.
Comparison with supervised models: The paper notes that the iGPT ID minimum (~22) is "compatible with" the ID at the output of supervised CNNs (13–24, from Ansuini et al., 2013). A direct comparison—extracting representations from a supervised ImageNet classifier at various layers, computing the ID profile, and comparing the semantic content to iGPT's intermediate representations—would directly test the claim that the self-supervised compression spontaneously recovers representations similar to supervised learning.
Alternative geometric metrics: The paper relies exclusively on the TwoNN ID estimator. Alternative ID estimators (local PCA, MLE, Fisher information-based methods) might produce different profiles, and the paper does not validate that the observed pattern is robust to the choice of estimator.
Task diversity: The paper evaluates one semantic property per domain. For proteins, other annotations are available (enzyme commission numbers, Gene Ontology terms, binding site annotations); for images, other labeling schemes exist (object detection bounding boxes, scene categories, texture attributes). Testing whether the ID minimum consistently maximizes semantic content across multiple, diverse label sets would substantially strengthen the universality claim.
Larger-scale NLP evaluation: The preliminary Llama-2-70B result on SST is a single sentence-level sentiment classification task. The NLP probing literature (Tenney et al., 2019; Liu et al., 2019) has established that different linguistic properties (part-of-speech, constituents, dependencies, semantic roles, coreference) emerge at different layers. The paper's framework would be significantly strengthened by showing that each type of linguistic information peaks at the nearest ID minimum in its respective emergence layer range, or by characterizing the relationship between the multiple ID peaks/minima in Llama-2 and the hierarchy of linguistic abstraction.

Summary of Claim Support

The paper's claims hold conditionally:

The ID profile has a characteristic multi-phase shape (expansion → compression → variable decoding): Strongly supported across all models tested (ESM-2 × 3, iGPT × 3, plus additional pLMs in Appendix Figure S1). The shape is robust to model size, training objective, and data modality under the specific preprocessing conditions used.
The ID minimum identifies layers with maximal semantic content: Supported for the specific semantic properties tested (remote homology in SCOPe, ImageNet class labels, SST sentiment). The claim is correlational—geometry predicts semantics, but no causal mechanism is established, and the practical advantage over naive heuristics (e.g., "use the middle layer") is not benchmarked.
The pattern is universal across self-supervised transformers: Supported across two domains with caveats. The pattern appears in both pLMs and iGPTs despite different objectives and modalities, but the decoding phase differs (final ascent vs. second peak), the absolute ID values differ (~6 vs. ~22 at the minimum), and the single NLP data point (Llama-2-70B on SST) shows a more complex multi-peak structure that the paper acknowledges is not yet fully understood.
The ID profile provides an unsupervised layer selection strategy: Supported in principle but not benchmarked against alternatives. The correlation between ID minimum and semantic peak is demonstrated, but the paper does not compare ID-guided layer selection against other unsupervised selection methods (variance-based, reconstruction-based, random) on downstream task performance, with the single exception of the 6% improvement in Appendix Figure S2.

The autoencoder analogy is an interpretive framework, not an empirically tested hypothesis. It organizes the observations coherently and generates testable predictions, but the predictions (causal role of expansion, functional symmetry of encoder/decoder, information-theoretic optimality of compression) remain untested in this paper.

6. Limitations and Trade-offs

The Cost of Difficulty Estimation Is Not Amortized

The assumption or constraint. The entire compute-optimal framework rests on knowing which difficulty bin a prompt belongs to before allocating test-time compute. Estimating this difficulty requires generating 2048 samples per question and scoring them, which the paper explicitly acknowledges as prohibitively expensive and does not include in any budget calculation. Section 3.2 states: "our experiments do not account for this cost largely for simplicity" and frames it as "an exploration-exploitation tradeoff—compute spent assessing difficulty versus compute spent solving the problem—flagging it as a key avenue for future work."

The consequence. In any realistic deployment, the total cost is difficulty assessment plus strategy execution, and the former can dominate the latter. The 2048-sample estimation step is comparable to or larger than the largest test-time budgets studied (256–512 generations). The reported 4× efficiency gains over best-of-N are computed after difficulty is known, without amortizing the cost of learning it. For a single prompt, the total cost would be 2048 + N generations, making the headline efficiency improvement unachievable in practice. Only in a setting where difficulty can be estimated once and reused for many similar prompts—a batch inference scenario with homogeneous difficulty—would the amortized cost approach the reported gains. The paper provides no strategy for such amortization.

What evidence exists in the paper. The limitation is discussed in Section 3.2, where the authors call out the exploration-exploitation tradeoff explicitly. Figures 4 and 8 show that predicted difficulty bins (using PRM scores rather than ground-truth labels) track the oracle bins closely, demonstrating that the approach works without access to answers—but the 2048-sample cost is still incurred in the predicted case. No experiment varies the number of estimation samples or attempts to estimate difficulty from fewer than 2048 samples. The paper does not report what happens with 64, 256, or 512 estimation samples, leaving open whether cheaper estimation is viable.

Mitigation status. The paper does not mitigate this limitation. It suggests future work on "pretraining or finetuning models to directly predict difficulty of a question" from its text (Section 8) and on adaptive approaches that interleave difficulty assessment with problem solving. These are identified as open problems, not attempted in the current work. Until solved, the compute-optimal strategy's practical advantage over simpler heuristics (e.g., always running a fixed beam search configuration, or using the PRM's average score over a small pilot batch of samples to guide allocation) remains unproven in deployment settings.

The 14× Larger Model Baseline Is Not Compute-Optimally Trained and Uses No Test-Time Compute

The assumption or constraint. The FLOPs-matched comparison in Section 7 scales model parameters while holding training data fixed, following the LLaMA paradigm rather than compute-optimal pretraining where both data and parameters are scaled equally (the Chinchilla recipe from Hoffmann et al., 2022). The paper acknowledges this in Section 7: "We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." Additionally, the larger model uses only greedy decoding—no majority voting, no best-of-N, no search, no revisions. The test-time compute budget is allocated entirely to the smaller model, while the larger model gets no inference-time augmentation at all.

The consequence. The comparison systematically favors test-time compute over pretraining. A Chinchilla-optimal model trained with 14× more total FLOPs (scaling both parameters and data) would likely outperform a parameter-only-scaled model, making the pretraining baseline artificially weak. The reported advantages—e.g., +27.8% on easy questions at R ≪ 1 for revisions (Figure 1, top-right bar chart)—may shrink or reverse against a properly compute-optimal larger model. More critically, giving the larger model even a modest test-time compute budget (say, best-of-8 or a short sequential revision chain) would create a much stronger baseline: a 14× larger model with best-of-8 has a total FLOPs cost that the smaller model would need to match using its own test-time budget, making the comparison genuinely about where to spend compute (pretraining vs. inference) rather than about whether to spend inference compute at all. The current setup confounds these two questions.

What evidence exists in the paper. The FLOPs-matched results appear in Figure 9 and Figure 1 (right bar charts). The parameter-only scaling choice is disclosed in Section 7 with the explicit caveat about Chinchilla-optimal training being left to future work. The greedy-only decoding assumption for the larger model is stated implicitly—the x-axis stars in Figure 9 represent the larger model's performance with no test-time compute sweep, meaning it uses a single greedy generation. No ablation tests the larger model with any flavor of test-time augmentation.

Mitigation status. Not mitigated. The paper frames the comparison as "representative of a canonical approach" but does not test sensitivity to the baseline strength. A reader should treat the FLOPs-matched advantages as upper bounds on the benefit of test-time compute relative to pretraining, not as precise quantitative estimates. The paper's own finding that hard problems (bin 5) show essentially zero benefit from test-time compute regardless of budget (Figures 3, 7, 9) is the more robust result—it cannot be explained away by a weak baseline—and suggests that the qualitative pattern (test-time compute helps on easy-to-medium problems, fails on hard problems) would persist even with a stronger pretraining comparison.

Single Model Family, Single Benchmark Per Domain

The assumption or constraint. All main results on search and revisions use PaLM 2-S* as the base model and the MATH benchmark (Hendrycks et al., 2021) as the evaluation dataset. The paper states in Section 4 that the authors "believe this model is representative of the capabilities of many contemporary LLMs," but this claim is not empirically verified within the paper. The choice is deliberate—MATH requires multi-step logical reasoning where test-time compute is expected to help, and PaLM 2-S* sits in a useful regime (roughly 10–19% pass@1) with room for improvement—but it means every quantitative finding is conditional on this specific model-dataset pair.

The consequence. Several aspects of the results could be model-specific or benchmark-specific. The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution; a model with different calibration or error patterns might exhibit different difficulty-dependent scaling curves. The revision model's ability to learn from in-context incorrect examples depends on the base model's in-context learning capabilities, which vary substantially across model families (e.g., GPT-4, Claude, Llama, Gemini). The MATH benchmark consists of competition-level problems requiring symbolic reasoning with ground-truth answers that can be checked with exact string matching. It is unclear whether the difficulty-dependent patterns—beam search hurting easy problems at high budgets, revisions dominating on easy problems but needing parallelism on hard ones—generalize to other reasoning domains (code generation, logical deduction, scientific QA) or to tasks where correctness is ambiguous or unverifiable (open-ended generation, dialogue, summarization). The finding that verifier over-optimization is the primary bottleneck may be specific to the quality of PRMs trainable on MATH-style step-by-step solutions; on tasks where step-level supervision is harder to obtain, the over-optimization threshold may be reached at different budgets.

What evidence exists in the paper. All main experiments use PaLM 2-S* on MATH (500 test questions, split into five difficulty quintiles). The paper does not include experiments on other model families, other benchmarks, or other task types. The Appendix tests additional protein language models (ProtBert, ProtT5, ESM-1v, ESM-1b; Figure S1) and iGPT image transformers (Figure 1, right), but these are separate domains entirely—they validate the geometric analysis framework, not the test-time compute scaling results that form the paper's primary contribution. The NLP preliminary result on Llama-2-70B (Appendix Figure S6) uses the SST sentiment dataset, which is also a single benchmark on a single model.

Mitigation status. Not mitigated. The paper does not claim generalization beyond MATH and PaLM 2-S*, but it also does not caution readers about the limits of single-benchmark evaluation. The test set of 500 questions is further split into five quintiles of approximately 100 questions each, then split again by two-fold cross-validation for the compute-optimal policy selection, meaning strategy choices are based on roughly 50 questions per fold per bin. This is a small sample, and the paper does not report confidence intervals on the compute-optimal scaling curves (Figures 4, 8), making it difficult to assess whether the observed gains are statistically reliable. A practitioner considering adopting the compute-optimal strategy for a different model, benchmark, or task domain would need to replicate the full analysis pipeline (PRM training, revision model training, difficulty estimation, strategy selection) on their specific setting before trusting the quantitative efficiency claims.

Verifier Over-Optimization Is a Hard Ceiling, Not a Solved Problem

The assumption or constraint. The paper documents that search-based test-time compute strategies suffer from verifier over-optimization: aggressive search against the PRM finds solutions that score highly under the PRM but are actually incorrect. This is most visible in Figure 3 (right), where beam search degrades easy-problem performance at high budgets—the clearest signature of reward hacking. Lookahead search, the most powerful optimizer, paradoxically performs worst overall (Figure 3, left), and qualitative examples in Appendix M show degenerate outputs (repetitive low-information steps, overly short solutions) that score highly under the PRM. The compute-optimal policy mitigates this by routing easy problems away from aggressive search (using best-of-N instead) and reserving beam search for medium-difficulty problems where the PRM signal has more room to provide genuine guidance, but it does not solve the underlying over-optimization problem.

The consequence. The compute-optimal approach is fundamentally bounded by verifier quality, and this bound is hard—increasing the test-time compute budget beyond a certain point yields no improvement, or even causes harm, regardless of how cleverly the budget is allocated. On medium-difficulty problems where beam search is deployed, over-optimization still limits the scaling ceiling: the beam search curves in Figure 3 flatten and sometimes decline well before the maximum budget is exhausted. This means the strategy cannot be expected to scale to arbitrarily large test-time budgets—there is a verifier-determined saturation point beyond which additional compute is wasted or counterproductive. Improving the PRM (through better training data, adversarial robustness, ensemble methods, or calibration techniques) would shift the difficulty thresholds and change the optimal policy, but the paper provides no guidance on how verifier quality interacts with the compute-optimal allocation—does a 2× better PRM produce a 2× higher saturation budget? A 10×? The relationship is unknown.

What evidence exists in the paper. Figure 3 (left) shows beam search with M=4 plateauing around 34% accuracy at 256 generations, while best-of-N weighted continues to improve to approximately 38% at 512 generations. Figure 3 (right) shows beam search on easy questions (bin 1) decreasing from roughly 78% to 77% as budget increases from 4 to 256—a direct measurement of over-optimization causing degradation. Appendix M provides qualitative examples. The paper discusses over-optimization in Section 5.3 and identifies it as a key phenomenon in Section 8, calling it "the primary bottleneck preventing unbounded improvements from additional compute."

Mitigation status. The compute-optimal policy mitigates but does not solve this. By routing easy problems to best-of-N (which is less aggressive and thus less susceptible to over-optimization) and reserving beam search for medium problems, the policy avoids the worst degradation. However, on medium problems where beam search is the optimal choice, over-optimization still caps the achievable accuracy. The paper suggests in Section 8 that "improving verifier robustness is the key bottleneck" for further scaling and proposes future directions including adversarial training of PRMs and constrained search with KL penalties to stay close to the base model's output distribution. None of these are implemented in the current work. A practitioner deploying this method should expect to hit a verifier-determined performance ceiling that cannot be exceeded by allocating more test-time compute, and the paper provides no method for estimating where that ceiling lies for their specific model and task.

Sequential Revision Strategies Incur Latency Costs Not Captured by Generation Counts

The assumption or constraint. The paper measures test-time compute in "generations"—the number of complete solutions sampled from the model. This is a reasonable proxy for total FLOPs and enables fair comparison between methods in terms of computational work. However, it ignores wall-clock latency. Sequential revisions (Figure 6) are inherently serial: each revision depends on the output of the previous one, so a chain of length N requires N sequential forward passes through the model. Parallel best-of-N can execute all N generations simultaneously with sufficient hardware (or in batches that amortize the per-token cost). A strategy that allocates a budget of 128 generations as 64 sequential revisions × 2 parallel chains takes approximately 64× longer in wall-clock time than one that runs 128 parallel samples concurrently.

The consequence. For latency-sensitive applications—interactive assistants, real-time decision-making, user-facing chatbots—the sequential-heavy strategies favored by the compute-optimal policy on easy-to-medium problems may be impractical regardless of their FLOPs efficiency. The paper's 4× improvement in generation efficiency could correspond to substantially worse wall-clock latency, depending on the sequential-to-parallel ratio selected. A practitioner choosing between a fully parallel best-of-256 (high throughput, moderate latency) and a compute-optimal strategy involving long revision chains (low throughput, high latency) faces a tradeoff the paper does not characterize. This is particularly acute because the optimal strategies for the problems where test-time compute is most beneficial (easy-to-medium difficulty) tend to favor sequential revisions over parallel search (Figure 7, right: easy and medium problems benefit from higher sequential-to-parallel ratios), meaning the practical utility of these strategies in latency-constrained deployments is unclear.

What evidence exists in the paper. Figure 7 (left) sweeps the sequential-to-parallel ratio at fixed total generation budgets, showing that at low budgets (8–32 generations), fully sequential is optimal, while at higher budgets (128–256), a balanced ratio with some parallelism performs best. The paper never discusses latency in the main text. Section 8 acknowledges that "combining PRM tree-search with revisions" is left to future work, but the latency implications of sequential dependencies are not mentioned as a limitation or design consideration. The FLOPs accounting in Section 7 measures total computational work (pretraining + inference FLOPs) but does not model wall-clock time or throughput.

Mitigation status. Not addressed. The paper provides no latency analysis, no throughput measurements, and no guidance on how to trade off latency against generation efficiency. Future work on combining search and revisions could potentially pipeline some operations (e.g., running revision chains on multiple parallel beams simultaneously where hardware allows), but the fundamental serial dependency of autoregressive revision generation remains. A practitioner deploying these strategies in a latency-sensitive setting would need to measure the wall-clock implications for their specific hardware and workload, and may find that fully parallel strategies with higher total FLOPs but lower latency are preferable to compute-optimal strategies with lower FLOPs but higher latency. The paper's efficiency numbers should be interpreted as FLOPs-optimal, not latency-optimal.

The Revision Model Exhibits a 38% Correct-to-Incorrect Reversion Rate, Patched Rather Than Solved

The assumption or constraint. The revision model is trained exclusively on sequences where all in-context answers are incorrect, followed by a correct target (Section 6.1). This means the model never sees a training example where the current answer is already correct and should be preserved. At test time, when a revision chain produces a correct answer at step t, the model has no learned behavior for "recognize that the current answer is correct and leave it unchanged," and will instead apply its learned transformation—revising an "incorrect" answer—potentially turning the correct answer into an incorrect one. The paper reports that "approximately 38% of correct answers get converted back to incorrect ones" using a naive approach (Section 6.1).

The consequence. Without mitigation, revision chains are self-undermining: each step has a non-trivial probability of destroying previously achieved correct answers. The paper's mitigation is to use majority voting or verifier-based selection across the entire chain—picking the best answer from any point in the chain rather than assuming the final revision is best. This works (Figures 6, 7, 8 show sequential revisions outperforming parallel sampling), but it is an imperfect patch. It means that compute spent on revisions after a correct answer has been produced is largely wasted—those generations could have been allocated to exploring alternative solution paths (parallel sampling) or applied to a different prompt entirely. More fundamentally, the revision model lacks the metacognitive capability to assess whether its current output is correct, a capability that would be necessary for autonomous self-improvement loops where the model must decide when to stop revising.

What evidence exists in the paper. The 38% reversion rate is reported in Section 6.1, though the paper does not detail exactly how this number was measured (e.g., on which dataset, at which revision step, with what verification method). The ReST^{EM} experiment in Appendix K (Figure 16) provides additional evidence that revision training is fragile: attempting to optimize the revision model with RL-style on-policy data collection caused performance to degrade substantially with sequential revisions, with fully sequential performance dropping to roughly 33.5% compared to 38.5% at the optimal ratio. The paper hypothesizes that "on-policy data collection exacerbates spurious correlations in revision data," suggesting the positive results depend on specific offline data construction choices that may not transfer robustly.

Mitigation status. Partially mitigated through chain-level selection (majority voting or verifier-based selection), but not solved. The mitigation recovers most of the benefit of sequential revisions despite the reversion problem, as evidenced by Figures 6 and 7. However, it does not address the root cause—the training data distribution mismatch—and it wastes compute on post-correct revisions. The paper does not explore training the revision model with correct-to-correct examples (teaching it when to stop) or incorporating an explicit confidence or stopping criterion. Section 8 mentions that combining revisions with PRM search is left to future work; one could imagine using the PRM to detect when a revision is correct and terminate the chain early, but this is not implemented. The ReST^{EM} negative result (Appendix K) suggests the revision training recipe is more brittle than the main paper's positive results might suggest, and a practitioner attempting to replicate the revision model on a different base model or dataset should expect to invest non-trivial effort in tuning the training data construction to avoid the reversion problem.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper changes the conversation around transformer representations from a layer-by-layer probing exercise into a principled geometric diagnostic — the intrinsic dimension profile — that reveals where semantic abstraction emerges without requiring any labels. The shift is methodological (the ID profile as an unsupervised navigation tool) and conceptual (the encoding–expansion–compression–decoding model of self-supervised transformer computation), and it reframes a scattered set of observations about intermediate representations into a coherent, testable framework.

The magnitude of this shift is best characterized as a reframing with practical diagnostics, not a paradigm overthrow. The paper does not introduce a new architecture, training objective, or optimization method. What it does is provide an organizing principle. Before this work, knowing which layer of a pretrained transformer to use for a downstream task required either labeled probing data (Chen et al., 2020 for iGPT, training linear classifiers at every layer) or a default heuristic ("use the last layer," as in Rives et al., 2021 for protein homology). The paper demonstrates that both approaches are suboptimal in measurable ways: the last layer of ESM-2 achieves a neighborhood overlap with remote homology superfamilies of ~0.4, while plateau layers achieve ~0.8 (Figure 4, left) — a factor-of-two degradation from the default choice, and a 6% absolute accuracy improvement available simply by switching layers (Appendix Figure S2). The ID profile provides an immediate, unsupervised answer: measure the geometry, find the compression minimum after the first expansion, and extract representations there. This converts layer selection from an empirical sweep requiring labels into a property of the data manifold that can be measured without any downstream task definition.

The field impact is threefold:

A unified explanation for conflicting probing results. The NLP literature had accumulated contradictory findings about where semantic information resides in transformers — some studies finding it in middle layers, others in final layers, others with no clear peak (Belinkov et al., 2017; Blevins et al., 2018; Liu et al., 2019; Tenney et al., 2019; Rogers et al., 2020, all cited in Section 4). The geometric framework resolves this by identifying the ID minimum as the relevant independent variable: semantic content peaks at the compression bottleneck, but the bottleneck's location varies across model architectures, sizes, and training objectives. A study probing only final layers on a model where the bottleneck is at relative depth 0.4 would naturally find weak semantics; a study probing mid-layers on a model with a broad plateau would find strong semantics. The contradiction dissolves when the ID profile is measured — it predicts where each study should have looked. The paper's preliminary NLP result on Llama-2-70B (Appendix Figure S6) directly supports this: the ID profile is more complex (three peaks, two minima), but "the highest overlap with class partition, determined by sentence sentiment, occurs in correspondence with the first local minimum." The framework generates falsifiable predictions for any transformer and any semantic label set, converting a descriptive post-hoc literature into a hypothesis-driven research program.
The autoencoder analogy as a new computational model for self-supervised transformers. Prior work on transformer internals focused on what individual attention heads attend to or what linear probes can decode. The paper's phase decomposition — expansion (early layers, rapid neighbor rearrangement, rising ID), compression (intermediate layers, stable neighbors, low ID, maximal semantics), and decoding (final layers, renewed rearrangement, rising ID, degraded semantics) — provides a functional-level description of what the stack of blocks collectively accomplishes, not just what individual components do. This is analogous to the shift in neuroscience from single-neuron recordings to population-level dynamics: the ID profile characterizes the manifold geometry of the entire population of representations, revealing computational phases that are not visible at the level of individual attention patterns. The finding that this structure emerges spontaneously from self-supervised training in architecturally symmetric networks (no bottleneck layer, identical blocks throughout) is particularly significant — it suggests that the encoding–compression–decoding decomposition is a convergent computational strategy that gradient-based optimization discovers, not an architectural constraint that must be explicitly designed. This reframes how we think about transformer depth: layers are not simply "more computation," but are functionally specialized into encoder, bottleneck, and decoder regimes, even when the architecture provides no explicit pressure to do so.
A practical method that changes how practitioners should extract representations. The 6% improvement in remote homology detection from switching to a plateau layer (Appendix Figure S2) is obtained "without any further training" — it costs nothing beyond measuring the ID profile once per model (which requires only unlabeled data and can be amortized across all downstream uses). This means the paper's findings are immediately actionable for any lab using pretrained protein language models for homology, structure prediction, or function annotation tasks. More broadly, the paper establishes that the standard practice of extracting final-layer representations from self-supervised models is a category mistake — using the decoder output when the encoder bottleneck carries the abstract content. This insight applies beyond the specific models studied, and the paper provides the diagnostic tool (TwoNN ID estimation) to identify the correct extraction layer for any new model or domain.

Some research directions become more attractive as a consequence of this work. Geometric analysis of representations — measuring intrinsic dimension, neighborhood overlap, manifold curvature — graduates from a niche descriptive tool to a practical diagnostic with demonstrated utility. Studies that probe why the expansion–compression pattern is optimal (or whether it is optimal at all) become tractable because the ID profile provides a clear quantitative phenotype to track. Cross-domain comparisons (do vision transformers, protein language models, and text LLMs all exhibit the same phases? what about multimodal models?) become straightforward: measure the ID profile on a new model and check for the characteristic peak–minimum–ascent shape. The paper's preliminary NLP result (Llama-2-70B, Appendix Figure S6) suggests that language models may exhibit more complex multi-peak profiles, opening the question of whether each peak–compression cycle corresponds to a distinct level of linguistic abstraction — a hypothesis that is now directly testable by measuring overlap with syntactic, semantic, and discourse-level annotations at each ID extremum.

Some directions become less attractive. Simply scaling up probing experiments — training linear classifiers at every layer across every possible task — is revealed as an inefficient way to find semantically rich layers when the ID profile can provide an unsupervised shortcut. The default assumption that "deeper is better" for representation extraction is refuted with quantitative evidence: deeper layers in self-supervised models are performing decoding, not semantic encoding, and their representations are objectively worse for downstream tasks (factor of ~2× degradation in χ^{l,gt} from plateau to final layer in ESM-2). Architecture designs that focus solely on final-layer output quality without considering intermediate representation geometry may be missing the most useful representations the model produces.

Follow-Up Research This Work Enables

Causal interventions to test whether expansion enables compression. The paper observes that the ID expansion precedes compression both in layer order and during training (Figure 3: the peak emerges first, compression develops later), and speculates that the expansion resembles the kernel trick — mapping data to a high-dimensional space where patterns become more separable before compressing. This is an untested causal hypothesis. A direct experiment would ablate the expansion layers in a fully trained model (e.g., by replacing early self-attention blocks with identity mappings, reducing the number of attention heads, or injecting noise into early-layer representations) and measuring whether the compression minimum still forms and whether semantic content at the bottleneck degrades. If expansion is causally necessary for effective compression, ablating it should reduce χ^{l,gt} at the ID minimum. If expansion and compression are independent phenomena that happen to occur in sequence, the compression quality should be robust to expansion ablation. A stronger variant would train models from scratch with constrained early-layer capacity (limited attention heads, reduced embedding dimension in early blocks) and compare the resulting ID profiles and downstream task performance to unconstrained baselines. The paper's training-dynamics analysis (Figure 3) provides the temporal baseline: the expansion develops by 10^4 steps in ESM-2 650M, compression develops between 10^4 and 3×10^4 steps. Ablations at these intermediate checkpoints could isolate whether compression requires the prior expansion or simply co-occurs with it.

Geometric characterization of the decoder phase and its functional role. The paper identifies that the final layers of self-supervised transformers perform a decoding function — the ID rises again (pLMs) or forms a second peak (larger iGPT models), neighbors rearrange rapidly, and semantic content degrades — but provides no detailed characterization of what these layers are computing. What information is recovered in the decoder that was lost during compression? Are decoder layers recovering positional information, token-level detail, or sequence-level features? A direct experiment would measure the overlap between a compressed (ID minimum) representation and the output representation using the same χ^{l,m}_k metric: if the decoder performs a systematic inversion of the encoder, there may be a detectable symmetry in the cross-layer overlap matrix — e.g., representations at symmetric depths in the first and second peaks might have high overlap. This could be tested on iGPT-L, where the second peak is most pronounced. An alternative approach would use mutual information estimators or reconstruction probing: given the compressed representation at the ID minimum, can a separately trained decoder reconstruct the input as accurately as the model's own decoder? If the compressed representation is an information bottleneck, reconstruction from it should be possible but lossy. If it primarily discards task-irrelevant detail, reconstruction quality from the bottleneck should correlate with ID at the bottleneck. The paper's iGPT models are particularly suitable because the next-pixel prediction objective provides a natural reconstruction metric: compare the model's pixel-level accuracy when forced to generate from the compressed representation versus from early-layer or final-layer representations.

Systematic probing of linguistic abstraction layers in large language models. The preliminary result on Llama-2-70B (Appendix Figure S6) shows "tree peaks and two local minima" in the ID profile, with sentiment overlap peaking at the first minimum. This suggests that NLP transformers may go through multiple expansion–compression cycles, potentially corresponding to the classical NLP pipeline hierarchy (morphology → syntax → semantics → pragmatics/discourse). A systematic study would measure the ID profile on Llama-2 or a comparable open model across all layers, then compute χ^{l,gt}_k for a diverse set of linguistic annotations: part-of-speech tags (morphosyntax), dependency parse depth (syntax), semantic role labels (semantics), coreference chains (discourse), and sentiment or natural language inference labels (task-level semantics). The prediction from the paper's framework is that each type of linguistic information should peak at the ID minimum nearest to its respective depth in the pipeline — syntactic information at an earlier minimum, semantic information at a later minimum, discourse information potentially at a third extremum. The paper's citation of Hewitt and Manning (2019), who found syntactic tree depth reconstruction peaks in middle layers of BERT, provides a starting point: does that syntactic peak align with the first ID minimum in BERT's profile? This experiment would directly test the universality of the geometry–semantics relationship across linguistic abstraction levels, and would either validate the multiple-compression-cycle interpretation or reveal that the relationship is more complex (e.g., different linguistic properties might have different optimal geometric conditions, with syntax preferring higher ID for structural encoding and semantics preferring lower ID for abstraction).

Cross-model and cross-task transfer from ID-minimum representations. The paper demonstrates that ID-minimum representations carry maximal semantic content within a single model evaluated on its own representations (nearest-neighbor overlap with labels). An open question is whether these compressed representations are more universal — do they transfer better across models, tasks, or domains than representations from other layers? A direct experiment would extract representations from the ID-minimum layers of multiple independently trained models (e.g., ESM-2 3B, ProtT5-XL-U50, ESM-1b) and evaluate nearest-neighbor transfer: use the representations from model A to retrieve homologs for sequences represented by model B. If compression produces model-invariant abstractions, the cross-model overlap should be highest when both models' ID-minimum layers are used, compared to using final-layer or expansion-layer representations. For the image domain, one could test whether the ID-minimum representations of iGPT-L transfer better to a supervised ResNet classifier's representations than final-layer iGPT representations do — measuring whether the self-supervised compression spontaneously aligns with supervised representations. The paper's observation that the iGPT ID minimum (~22) is "compatible with" the ID at the output of supervised CNNs (13–24, from Ansuini et al., 2013) makes this a natural test. If confirmed, it would position the ID minimum as not just a model-specific sweet spot but a universal representational format that emerges across training paradigms.

Training models with explicit geometric objectives to strengthen or reshape the ID profile. The paper is purely observational — it characterizes the ID profile that emerges from standard self-supervised training but does not intervene to modify it. A natural extension is to ask whether the ID profile can be engineered. For example, adding an auxiliary loss that penalizes high ID in intermediate layers (encouraging earlier compression) or that enforces a target ID trajectory during training. Would such a model converge faster, produce better downstream representations, or exhibit different semantic localization? Alternatively, one could test whether the ID-minimum representations can be distilled into smaller models: train a student model to directly predict the ID-minimum representations of a large teacher, bypassing the expansion and compression phases entirely. The paper's training-dynamics analysis (Figure 3) shows that compression develops relatively late in training (between 10^4 and 3×10^4 steps for ESM-2 650M, out of 5×10^5 total steps). An auxiliary geometric loss might accelerate this, effectively "teaching" the model to compress earlier. The Denti et al. (2022) higher-order ID estimator mentioned in Appendix B could provide a more stable training signal than TwoNN for such an auxiliary loss, since it has lower variance.

Understanding the failure modes: why does iGPT-S lack the second peak, and what does that imply for model scaling? The paper identifies a boundary condition: iGPT-S (24 blocks, 76M parameters) does not develop the second ID peak seen in iGPT-M (36 blocks, 455M) and iGPT-L (48 blocks, 1.4B). But depth and width are confounded — iGPT-S is both shallower and narrower. A controlled experiment would train a series of transformer models varying only depth (e.g., 12, 24, 36, 48 blocks at fixed width) or only width (e.g., embedding dimensions of 256, 512, 1024, 1536 at fixed depth) on the same next-pixel prediction task, and measure whether the second peak emerges at a specific depth threshold, a specific width threshold, or a specific total capacity threshold. This would clarify whether the full encoding–compression–decoding structure requires sufficient serial depth (enough layers to functionally specialize into encoder, bottleneck, and decoder regimes) or sufficient representational capacity (enough parameters to support a distinct decoder computation). The training dynamics of such a scaling series would also reveal whether the second peak emerges abruptly (a phase transition at a critical depth/width) or gradually (a continuous strengthening as capacity increases). The paper's observation that ESM-2 models do not exhibit a second peak at all — even the 3B, 36-block model shows only a final ascent — raises the further question of whether the decoder structure is objective-dependent (MLM vs. autoregressive), modality-dependent (proteins vs. images), or architecture-dependent, which a controlled cross-objective comparison could disentangle.

Practical Applications and Downstream Use Cases

Improved remote homology detection in protein biology, available immediately with no retraining. The paper's most directly actionable result is in protein remote homology — identifying proteins that share evolutionary ancestry and 3D structure despite having highly diverged sequences. The standard approach in the field, following Rives et al. (2021) and Schütze et al. (2022), uses the last hidden layer of a pretrained protein language model for nearest-neighbor search. The paper demonstrates that switching to a plateau layer (identified by the ID minimum) improves nearest-neighbor accuracy by approximately 6% on the ProtT5-XL-U50 model evaluated on SCOPe remote homologs (Appendix Figure S2). This improvement requires no model fine-tuning, no additional training data, and no architectural changes — only extracting representations from a different layer of the already-trained model. For a structural biology lab or a protein function annotation pipeline, the cost of adoption is: (1) run the TwoNN ID estimator once on a representative corpus of unlabeled protein sequences (e.g., a subset of UniRef50) to identify the plateau layers, (2) switch the representation extraction point from the final layer to the plateau for all subsequent queries. The 6% gain translates directly to more accurate functional annotation, better detection of distant evolutionary relationships for drug target identification, and improved sensitivity in metagenomic analysis where novel proteins with no close sequence homologs are common. The finding that the plateau ID (~5–7) is consistent with independent sequence-alignment-based estimates from Facco et al. (2019) adds confidence that this is not a model-specific quirk but reflects a genuine signal about the underlying biology.

Layer selection for transfer learning without labeled validation data. A common workflow in applied machine learning is: take a large pretrained transformer, freeze it, and train a task-specific classifier on top of the frozen representations. The standard practice is to use the final hidden layer, or to sweep across layers using a held-out validation set with labels to choose the best one. The paper's ID profile provides an unsupervised alternative: measure the ID curve on the (unlabeled) target dataset, identify the compression minimum, and extract representations from that layer for the downstream classifier. This is particularly valuable in domains where labeled data is scarce — which are precisely the domains where self-supervised pretraining is most valuable. For example, a rare disease diagnosis project using protein language model embeddings (as mentioned in the paper's Acknowledgments) could use the ID profile on the unlabeled patient protein sequences to select the optimal extraction layer before training a classifier on the limited labeled cases. The paper's finding that the ID-minimum layer generalizes across model sizes within a domain (all ESM-2 variants have a plateau at ID ~5–7, all iGPT variants have a minimum at ID ~22) means the ID profile measured once on a large model can inform layer selection even when deploying a smaller model, reducing the computational cost of the diagnostic. The practical recipe is: run TwoNN on a representative unlabeled corpus, plot ID vs. layer, select the layer at the first local minimum after the first peak, and use that for all downstream tasks until domain shift is suspected.

Guiding architecture design and training recipes for self-supervised models. The paper's phase decomposition suggests that self-supervised transformers should not be treated as monolithic depth-scalable units — early layers are encoders, intermediate layers are compressors, and final layers are decoders, and these functional roles are not interchangeable. For practitioners designing or fine-tuning transformers, this has concrete implications. If the downstream task requires semantic abstraction (classification, retrieval, zero-shot transfer), the model should be deep enough to reach the compression phase — a model that is too shallow may never exit the expansion phase, and its representations will carry less semantic content (consistent with the paper's finding that smaller models show weaker semantic peaks: χ^{l,gt} ~0.15 for iGPT-S vs. ~0.35 for iGPT-L). When fine-tuning with limited data, it may be counterproductive to update the compression layers, which encode task-general abstractions, while fine-tuning the decoder layers (which are task-specific to the pretraining objective) could be beneficial. Conversely, if the downstream task requires fine-grained reconstruction or generation, the decoder layers should be preserved or even expanded. The paper's training-dynamics analysis (Figure 3) further suggests a curriculum: early training establishes the expansion, and later training develops the compression. A practitioner could monitor the ID profile during training as a convergence diagnostic — if the compression minimum has not yet formed, the model may benefit from continued training even if the pretraining loss has plateaued, because the semantic content of intermediate representations may still be improving.

Quality assessment and comparison of pretrained models without downstream benchmarks. Currently, comparing two pretrained models (e.g., ESM-2 vs. ProtT5, or two different iGPT checkpoints) requires evaluating them on a suite of downstream tasks with labeled data — a process that is expensive, task-specific, and sensitive to the choice of evaluation protocol. The ID profile offers a complementary, label-free quality metric: a model with a well-defined compression phase (sharp ID minimum, stable plateau, high peak-to-minimum ratio) may produce more useful intermediate representations than a model with a flatter profile. The paper provides some evidence for this: larger models develop more pronounced ID peaks and sharper semantic localization (Figure 4, right: iGPT-L has a narrow, high χ^{l,gt} peak at the ID minimum, while iGPT-S has a broad, low peak). However, the paper does not systematically correlate ID profile sharpness with downstream performance — this is a hypothesis generated by the work that practitioners could test. If validated, the ID profile could serve as a rapid screening tool during model development: train candidate architectures, measure the ID profile on unlabeled data, and select the one with the most pronounced compression minimum for further evaluation on labeled benchmarks. This would reduce the iteration cycle for architecture search and hyperparameter tuning in settings where labeled benchmarks are expensive to run.

When to Prefer This Method

The paper frames the ID profile primarily as a descriptive and diagnostic tool, not as a method that competes directly against named alternatives for a specific predictive task. However, the paper does articulate an implied tradeoff between (a) using the ID profile for unsupervised layer selection versus (b) using labeled probing data to select layers empirically, and between extracting final-layer representations versus extracting ID-minimum representations for downstream tasks. The decision rules, grounded in the paper's results, are:

Prefer ID-guided layer selection when labeled downstream data is scarce or absent. The ID profile requires only unlabeled data to measure (Euclidean distances + TwoNN estimation) and predicts which layers carry maximal semantic content without any task-specific tuning. The paper demonstrates this quantitatively: χ^{l,gt} peaks at or near the ID minimum for three different semantic properties (remote homology, ImageNet classes, SST sentiment) across three model families (ESM-2, iGPT, Llama-2). In domains like rare disease genetics, metagenomics, or low-resource language tasks — precisely where self-supervised pretraining is most valuable — labeled data to sweep layers empirically may not exist, making the ID profile the only viable layer-selection strategy.
Prefer empirical layer sweeping with labeled probes when the downstream task is sufficiently different from the semantic properties validated in the paper, and sufficient labeled data exists. The paper tests only one semantic property per domain (remote homology for proteins, class identity for images, sentiment for text). If the downstream task involves a fundamentally different type of abstraction — e.g., predicting protein binding affinity rather than homology, segmenting objects rather than classifying images, or performing coreference resolution rather than sentiment analysis — the ID minimum may not be the optimal layer. The paper's framework predicts it will be (the geometric compression should capture any abstraction that is structurally important for the data manifold), but this prediction is untested for task types beyond those evaluated. A practitioner with ample labeled data should validate the ID-minimum layer against empirical sweeping before committing to it.
Prefer ID-minimum representations over final-layer representations for any downstream task requiring semantic abstraction from a self-supervised model. The paper provides direct evidence that final-layer representations of self-supervised transformers are degraded by the reconstruction objective: χ^{l,gt} drops from ~0.8 to ~0.4 in ESM-2 (Figure 4, left), and the 6% absolute improvement in remote homology from switching layers (Appendix Figure S2) is obtained with zero additional cost. This is not a tradeoff — there is no scenario described in the paper where the final layer outperforms the ID-minimum layer on semantic tasks, and the mechanism (decoder decompression) explains why this should hold generally for self-supervised models. The exception would be tasks that directly require the output of the pretraining objective — e.g., masked token prediction accuracy, next-pixel prediction quality — where the final layers are specifically optimized. For any task involving abstraction, classification, retrieval, or transfer, the ID-minimum layer is strictly preferable based on the paper's evidence.
The ID profile is most reliable as a layer-selection guide for models above a minimum capacity threshold. The paper shows that the ID minimum is broad and shallow in iGPT-S (76M parameters), making precise layer selection less critical but also less informative, and that the second peak (decoder structure) is absent in the smallest model. For very small transformers (fewer than ~12 blocks, fewer than ~100M parameters), the expansion–compression–decoding phases may not be fully developed, and the ID profile may not provide a clear signal. The paper does not systematically characterize this threshold — it is an observation across the three sizes tested — but a practitioner working with small models should validate that a clear ID minimum exists before relying on it for layer selection.