SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 3: Evolution of ID curves during training. (left) ID curves for the pLM ESM-2 (650M) at checkpoints corresponding to  $[0,10^{4},3\cdot 10^{4},10^{5},5\cdot 10^{5}]$  training steps. Firstly, the peak develops; then, the plateau reaches low-ID values, and the ID in the final layers reaches values similar to the input. (right) ID curves for iGPT-L (1.4B) at checkpoints  $[0,0.13,0.26,0.52,1]\cdot 10^{6}$  training steps. Firstly, the ID peak followed by a plateau with a local low-ID minimum develops; later in training, a second peak emerges in the final layers.
of iGPT models encode this semantic information. In this section, we extend this analysis to protein language models, showing that, in this case, semantic information is maximally present in the whole plateau region of the network. Therefore, we show that the ID profile can be used to predict the specific layers where this information is most pronounced.
The representations in the plateau region code remote homology relationships. Proteins are considered remote homologs when they share a similar structure resulting from common ancestry despite having highly dissimilar amino acid sequences. Rives et al. [6] observed that the Euclidean distances among representations in the last hidden layer of pLMs encode information related to remote homology. We study how this biological feature emerges in pLMs, analyzing the hidden representations of the SCOPe dataset. For every layer  $l$ , we compute the neighborhood overlap with the ground truth labels  $\chi_{k}^{l,gt}$  with  $k = 10$ . Here, "gt" represents the classification by superfamily, with the exclusion of neighbors in the same family, allowing us to focus on remote homology specifically. Figure 4 (left) shows that structural homology information is absent in the positional embedding layer; it grows smoothly in the peak phase, reaching a stationary maximum  $\chi_{k}^{l,gt} \sim 0.8$  in the plateau phase, and it suddenly decreases in the final ascent phase. This analysis shows that homology information is highest in the plateau layers, and it appears early in the networks right after the peak of the ID. Notably, 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. Therefore, 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. The Appendix shows that the entity of such improvement is approximately 6% (see Fig. S2).
Image semantic features emerge in the layer where the ID has a relative minimum. It has been shown that predicting pixels in a self-supervised learning context can generate representations that encode semantic concepts [10]. Expanding upon these findings, we establish a connection between this remarkable property and the geometric aspects of representations. Similar to the case of pLMs, we quantify the semantic content of a representation by computing the overlap with the ground truth labels,  $\chi^{l,gt}$ . In Figure 4 (right), we plot for the iGPT models the overlap with the ImageNet labels for 300 classes of the training set (see Sec 2.3). In all the cases, we consistently observe a peak of  $\chi^{l,gt}$  around a relative depth of  $\sim 0.4$  where the ID in the larger models iGPT-M and iGPT-L is lower (see Fig. 1). This peak becomes sharper, and the overlap values increase as we scale up the model size. Specifically, the peak value of  $\chi^{l,gt}$  is 0.15 in iGPT-S, 0.27 in iGPT-M, and 0.35 in iGPT-L. Similar to our observations in pLMs, the representations where the semantic abstractions are better encoded are also those where the ID is low.
Figure 4: Local geometry and semantic relations in hidden representations. (left) Overlap with remote homologs  $\chi^{l,gt}$  for ESM-2 (35M, green), ESM-2 (650M, blue), and ESM-2 (3B, purple). Information about remote homologs is highest in the plateau layers, corresponding to low ID values. (right) Overlap with ImageNet labels  $\chi^{l,gt}$  for iGPT-S (yellow), iGPT-M (orange), and iGPT-L (red). The information related to ImageNet labels is highest in correspondence with the ID minimum.
4 Discussion and conclusions
In this work, we investigate the evolution of fundamental geometric properties of representations, such as the intrinsic dimension and the neighbor structure, across the layers of large transformer models. The goal of this analysis is twofold: 1) understanding the computational strategies developed by these models to solve the self-supervised reconstruction task, and 2) monitoring how high-level, abstract information, such as remote homology in proteins and class labels in images, is related to the geometry of the representations. In particular, we propose an explicit, unsupervised strategy to identify the most semantically rich representations of transformers.
The qualitative picture that emerges from our results is that 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. In the largest iGPT model, this similarity is particularly evident: 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)). 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.
Crucially, the compression is preceded by an intrinsic dimensionality expansion and a relatively fast-paced rearrangement of the neighbors during the encoding process, which is very similar across all the models we studied. The relation between a low intrinsic dimension at the end of the encoding part and the richness in abstract content of representations is robust and quantitative in all the models we considered. A vast rearrangement of the neighbors in the layers close to the output is also observed for all the architectures we analyze. This is devoted to the process of dealing with the decoding task, and it causes a degradation of the abstract content.
The analysis we performed could be further reinforced, for example, by analyzing the similarity of representations in the ID minima with those generated in models trained in a supervised setting. This can provide a deeper understanding of the role of the second peak in the large iGPT model and the reasons for its appearance in later training phases, as well as offer a clearer interpretation of low ID representations in the plateau phase of pLMs.
Impact of some implementation choices. In this study, we average pool the representations along the sequence dimension as a convenient method to compute the Euclidean distances between proteins of varying lengths. This averaged representation can serve multiple purposes, including measurement of representation quality using linear probes [10], and can provide sufficient biological information to address homology, structural, and evolutionary tasks directly or after fine-tuning. It's worth noting, however, that other approaches for comparing sequences of variable lengths have been explored, such as those in [31] and [22]. When the inputs have all the same sequence length as in
the iGPT case, reducing the sequence representation to a  $d$ -dimensional vector is also convenient from a computational perspective. Indeed, 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. In section D of the Appendix, we show with an experiment on CIFAR10 that the ID of a representation decreases after average pooling. Still, the fundamental shape of the ID does not change qualitatively (see Fig. S3). In the same section, we also show that the ID profiles do not quantitatively change if we extract the representations after the attention maps or after the MLP blocks (see Fig. S4).
Discussion of different ID shapes. In Sec. 3.1, we showed that the ID profile in pLMs has three phases and terminates with a final expansion, while in iGPTs, there is a fourth phase where the ID is compressed in the last layers. These differences can be attributed to various factors. For instance, in the case of ImageNet, variations of the pixel brightness [13] and details of data preprocessing, including the substantial reduction of the input resolution (see Sec. A and [10]) can artificially alter the ID significantly. In the middle layers, instead, the ID is connected to the semantic complexity of the dataset, and its value tends to be consistent not only within transformers of the same class but also with the last hidden layer of convolutional networks in the case of the ImageNet dataset [13] and with other methods relying on different distance metrics in the case of proteins [30] (see Sec. 3.1). This implies that 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. These considerations may explain why the final layers of the network can either reduce the ID (as seen in iGPT) or expand it (as observed in ESM models). However, our key finding remains consistent in both cases: the layers where the ID reaches a local minimum correspond to representations where a semantic property of the data is most pronounced. Moreover, other important factors, such as the distinct training objectives (MLM in protein language models and next pixel prediction in iGPTs) and the different data modalities used for training, also influence the ID's characteristic shape.
Preliminary exploration of NLP representations. The NLP literature has previously discussed the idea that hidden representations in the middle layers of large transformer models encode abstract information that can be leveraged to solve various tasks successfully. Syntactic information is more prominent in the middle layers of deep transformers, as shown by syntactic tree depth reconstruction in [32]. However, the localization of semantic information has yielded contrasting results [33-37] possibly due to the complexity of language data and the ongoing evolution of model size and architectures. These studies show that language data requires diverse tasks and probes for comprehensive understanding. In addition, there is a substantial distinction between pLMs and Large Langue Models (LLMs). While for pLMs, we almost reached an overparameterization regime on the UniRef dataset [38], this saturation level has not been observed in the context of language. Recently, we started analyzing the geometry of hidden representation in the NLP domain with our methods. Our early analyses on GPT-2-XL on the SST dataset [39] did not reveal a second peak similar to the one in iGPT. With the release of the Llama models [40], we repeated the experiment on the Llama-v2 with 70 billion parameters (see Appendix, Sec. E). In this case, the evolution of the ID profile is more complex, showing three peaks and two local minima across the hidden layers. Remarkably, our key finding remains consistent: the highest overlap with class partition, determined by sentence sentiment, occurs in correspondence with the first local minimum. After the second peak, the ID profile shows a more complex behavior, which we are currently investigating. We will explore these aspects further in future work dedicated to NLP tasks.
The authors acknowledge the AREA Science Park supercomputing platform ORFEO made available for conducting the research reported in this paper and the technical support of the Laboratory of Data Engineering staff. We thank Roshan Rao for kindly providing the weights of the ESB-2(650M) models at training checkpoints used to produce Fig. 3 (left).
A.A., A.C., and D. D. were supported by the project "Supporto alla diagnosi di malattie rare tramite l'intelligenza artificiale" CUP: F53C22001770002. A.A., A. C., and F.C. were supported by the European Union - NextGenerationEU within the project PNRR "PRP@CERIC" IR0000028 - Mission 4 Component 2 Investment 3.1 Action 3.1.1. L.V. was supported by the project PON "BIO Open Lab (BOL) - Rafforzamento del capitale umano" - CUP J72F20000940007.
[47] Konstantin Schütze, Michael Heinzinger, Martin Steinegger, and Burkhard Rost. Nearest neighbor search on embeddings rapidly identifies distant protein relations. Frontiers in Bioinformatics, 2, 2022.
Table 1: Characteristics of the models employed for extracting representations.
Model	#Blocks	Emb. dim.	#Heads	#Params	Dataset	Reference
ESM-2 (35M)	12	480	20	35M	UR50/D	[15]
ESM-2 (650M)	33	1280	20	650M	UR50/D	[15]
ESM-2 (3B)	36	2560	40	3B	UR50/D	[15]
iGPT-S(76M)	24	512	8	76M	ImageNet	[10]
iGPT-M(455M)	36	1024	8	455M	ImageNet	[10]
iGPT-L(1.4B)	48	1536	16	1.4B	ImageNet	[10]
Transformer models for protein sequences. The input data points, corresponding to proteins, are variable-length sequences of  $l$  letters  $s = a_{1},a_{2},\ldots a_{l}$ , chosen from an alphabet of  $n_a(\simeq 20)$  tokens corresponding to amino acids. Each token is encoded by an embedding layer into a vector of size  $d$  so that the generic protein  $s$  is represented as a matrix  $x\coloneqq f_0(x)$  of size  $l\times d$ . A model with  $B$  blocks transforms a data point  $x\in \mathbb{R}^{l\times d}$  into a sequence of representations:  $f_{0}(x)\rightarrow f_{1}(x)\rightarrow \dots \rightarrow f_{B}(x)\rightarrow f_{out}(x)$ , where  $f_{i},i = 1,\dots ,B$  stands for the self-attention module at the  $i^{th}$  block, and the final LM-head  $f_{out}$  is a learned projection onto dimension  $l\times n_{a}$ . The size of each hidden layer does not change across the model and is equal to  $l\times d$ ; therefore, the action of the model is a sequence of mappings  $\mathbb{R}^{l\times d}\to \mathbb{R}^{l\times d}$ . For each layer  $i$ , we choose to perform global average pooling across the row dimension  $f_{i}(x)\rightarrow \frac{1}{l}\sum_{j = 1}^{l}(f_{i}(x))_{j}$ , since this reduction retrieves sufficient biological information to solve, directly or possibly after finetuning, homology, structural and evolutionary tasks. For a given pLM, the action of the network on a dataset of  $N$  proteins can thus be described by  $B + 1$  collections of  $N$  vectors in  $\mathbb{R}^d$ : these are the data representations that we investigate.
Transformer models for image processing. The structure of the Image GPT (iGPT) transformers is very similar to that of pLMs. Due to the high memory footprint required by the attention layers, Chen et al. [10] reduce the image from the standard ImageNet size  $(r = 224^2)$  to  $r = 32^2$  (S),  $r = 48^2$  (M),  $r = 64^2$  (L) and the three color channels are encoded in an "embedding axis" with size  $d_{in} = 512$ . In practice, 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. An input image is then represented by a data point  $x \in \mathbb{R}^{l \times d_{in}}$  where the  $l$  pixels are placed along the sequence axis in raster order and  $d_{in}$  encodes the color of each pixel. Like in pLMs, an input data point, after the embedding layers, is processed by a stack of attention blocks that produce a sequence of data representations of identical size  $x \in \mathbb{R}^{l \times d}$ . The final head  $f_{out}$  projects the output of the last attention block to  $\mathbb{R}^{l \times d_{in}}$ , which encodes at each sequence position  $i$  the conditional distribution
$$
p \left(x _ {i}\right) = p \left(x _ {i} \mid x _ {0}, \dots x _ {i - 1}\right) \tag {1}
$$
Once the network is trained, an image can be generated by sampling from  $p(x_{i})$  one pixel at a time.
B Details about the intrinsic dimension estimation.
Facco et al. [21] showed that the ID can be estimated using only the distances  $r_{i_1}, r_{i_2}$  between each point  $x_i$  and its first two nearest neighbors. Under the assumption of the Poisson point process in space [41], the ratios  $\mu_i = r_{i_2} / r_{i_1}$  between the distances to the second and first nearest neighbor follow a Pareto distribution  $p(\mu_i|d) = d\mu_i^{-d - 1}$ . Assuming that the dataset likelihood factorizes over the  $N$  single-point likelihoods, the ID can be easily inferred given the empirical values of the  $\mu_i$  as the  $p(\mu_i|d)$  depend only on  $d$ . Facco et al. propose to estimate the ID from the cumulative distribution function  $F(\mu) = 1 - \mu^{-d}$  with a linear regression fit. Indeed, the ID is the slope of the line through the origin obtained regressing  $\log (\mu_{\sigma (i)})$  against  $-\log (1 - F^{emp}(\mu_{\sigma (i)}))$  where  $\mu_{\sigma (i)}$  are the  $\mu_i$  sorted in ascending order and  $F^{emp}(\mu_{\sigma (i)}) := \frac{i}{N}$  is the value of empirical cumulative distribution function at  $\mu_{\sigma (i)}$  [21].
In practice, due to its very local nature, the TwoNN estimator can be affected by noise, which typically leads to an overestimation of the ID. To identify the most plausible value of the data intrinsic dimension, Facco et al. apply the TwoNN estimator on the full data set and then on random subsets of decreasing size:  $n = [N,N / 2,N / 4,\dots]$ . The relevant data ID is chosen as the value of  $\hat{d}$  where the graph  $\hat{d} (n)$  exhibits a plateau, and the ID is less dependent on the scale. Recently, Denti et. al [42] proposed a better strategy to perform a multiscale analysis of the ID considering nearest neighbors of higher order instead of decimating the data set. This approach reduces the variance of the likelihood, giving more stable estimates.
These methods gave similar results in the experiments of this work: 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.
Validity of the local density assumption. To quantitatively validate the constant density assumption, we employed the Point Adaptive kNN (PAk) method introduced by Rodriguez et al. [43]. PAk 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. Applying PAk (as implemented in DADApy [25]) to our dataset revealed that, on average, the density can be considered constant within the first 6 neighboring data points. Our study measures the ID using the distances between a data point and its first two nearest neighbors. This analysis allows the conclusion that, at this scale, the assumption of local density holds.
C Intrinsic dimension and overlap profiles on other protein language models.
Figure S1: The ID shape for different pLMs architectures. The latest developments in the application of pre-trained pLMs for the solution of diverse biological tasks have been mainly fuelled by two families of models: Prot-Trans [7] and Evolutionary Scale Modelling (ESM) [6, 38, 44, 45]. During the last years pLMs with different architectures, number of parameters, and embedding sizes have been trained on several datasets obtained starting from the UniProt [46] database. The profiles in the figure complement the analysis done in Sec. 3.1, showing the ID curve of ProteinNet on pLMs trained on UniRef50 (ESM1b and ProtT5-XL-U50, red and orange) UniRef90 (ESM1v, green), Uniref100 (ProtBert, purple). Despite the significant differences of the pLMs considered in the analysis, the consistency of the three-phased behavior of the ID remains: an initial peak is followed by a plateau where the ID assumes low values, and the ID grows again to values close to the one measured after the positional embedding.
Figure S2: Nearest neighbor search in plateau layers improves identification of protein relations. It was recently shown in [47] that the first nearest neighbor searches for remote homologous protein domains based on the last hidden layer representations of the ProtT5-XL-U50 pLM outperform state-of-the-art methods based on sequence similarity. Adapting the approach of Sec. 3.2, we mimic the experiment performed in Sec. 2 of [47] by 1) considering protein domains in SCOPe belonging to a superfamily with at least 2 sequences, 2) setting the number of neighbors to  $k = 1$ . Considering representations in the plateau layer improves the accuracy of the 1-kNN homology search. In particular, in the figure, we observe an improvement of  $\sim 6\%$  performing the search on a plateau layer instead of the last layer before the output. Importantly, the performance gain of  $\sim 6\%$  is obtained without any further training.
D Impact of the implementation details and hyperparameter choices.
Cifar10 - WR28-8
Figure S3: Impact of the average pooling on the intrinsic dimension profiles. To study how the average pooling affects the ID profile, we tested the representations of CIFAR10 in the Wide-ResNet28-8 model. The left panel shows the intrinsic dimensionality (ID) computed in the full feature space, while the right panel displays the ID after applying the global average pooling along the spatial dimensions (width and height). Notably, 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 (see [13]). From a quantitative perspective, the ID profile after average pooling exhibits a downward shift as a consequence of the averaging procedure, particularly pronounced at the initial stages of the architecture. This effect is likely due to the low number of channels after the first block (16) compared to the later blocks (128, 256, and 512, respectively). Nevertheless, we can be confident about the qualitative robustness of the profiles as long as the number of "channels" (where channels are interpreted as embedding dimension) significantly surpasses the ID as in large transformer models. Indeed, in this case, the number of "channels" is substantially higher than in early CNN layers (ranging from 512 to several thousand in modern large language models) and remains constant throughout the architecture.
iGPT-L
Figure S4: Intrinsic dimension and neighborhood overlap after the self-attention blocks. The figure compares the ID (left) and the overlap with the ground truth labels  $(\chi^{l,gt}$ , right) of the ImageNet representations after the first normalization layer with those of the representations after the attention maps of each self-attention block for the iGPT Large model. The ID and profiles are consistent, indicating the robustness of our analysis with respect to the layer choice.
ESM-2(650M)
Figure S5: Neighborhood overlap profiles are robust w.r.t. the choice of the neighborhood size  $k$ . Overlap profiles in ESM-2 (650M, top) and iGPT Large (bottom) varying the neighborhood size  $k$  in the range  $\{1, 50\}$ . The curves of the overlap between consecutive layers  $\chi^{l,l+1}$  (left) and with the ground truth classes (right) are qualitatively unchanged. When considering the ESM-2 model (top), the lower values of  $\chi_k^{l,l+1}$  and  $\chi_k^{l,gt}$  when  $k = 50$  is the consequence of certain superfamilies having fewer than 50 elements.
E Preliminary analysis of the intrinsic dimension and overlap in LLama2
Figure S6: Intrinsic dimension and neighborhood overlap of sst in LLama2-70b. [Top] The intrinsic dimension (ID) of representations in hidden layers of Llama-v2 for the Stanford Sentiment Treebank (SST) dataset, plotted against block number normalized by the total number of blocks (relative depth). [Bottom] Overlap with ground truth labels, i.e. sentiment of the sentences, at each self-attention block, plotted against relative depth. In correspondence with the first local minimum of the ID profile, at 0.3 relative depth, the overlap with the class partition in the Bottom panel is the highest.