SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lucrezia Valeriani $^{1,2*}$
Diego Doimo $^{1,3*}$
Francesca Cuturello
Alessandro Laio<sup>3,4</sup>
Alessio Ansuini $^{1\dagger}$
Alberto Cazzaniga $^{1\dagger}$
<sup>1</sup> AREA Science Park, Trieste, Italy
$^{2}$  University of Trieste, Trieste, Italy
$^{3}$  SISSA, Trieste, Italy
4 ICTP, Trieste, Italy
Large transformers are powerful architectures used for self-supervised data analysis across various data types, including protein sequences, images, and text. In these models, the semantic structure of the dataset emerges from a sequence of transformations between one representation and the next. We characterize the geometric and statistical properties of these representations and how they change as we move through the layers. By analyzing the intrinsic dimension (ID) and neighbor composition, we find that the representations evolve similarly in transformers trained on protein language tasks and image reconstruction tasks. In the first layers, the data manifold expands, becoming high-dimensional, and then contracts significantly in the intermediate layers. In the last part of the model, the ID remains approximately constant or forms a second shallow peak. We show that the semantic information of the dataset is better expressed at the end of the first peak, and this phenomenon can be observed across many models trained on diverse datasets. Based on our findings, 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.
Deep learning has significantly changed the landscape of scientific research and technology advancements across multiple disciplines in recent years. A particular class of deep learning models, transformers trained with self-supervision, combine high predictive performance and architectural simplicity. These transformer models consist of a stack of identical self-attention blocks trained in a self-supervised manner using masked language modeling (MLM) or auto-regressive objectives [1, 2]. It has been shown that the features learned by these models can be used to solve a wide range of downstream tasks in natural language processing [3-5], biology [6-9], and computer vision [10-12].
Previous studies analyzing convolutional architectures have shown that data representations in deep learning models undergo profound changes across the layers [13, 14]. In transformers, each module maps the data into a representation, and it has already been observed that the organization of representations in the last hidden layer can reflect abstract, domain-specific properties [10]. However,
†Correspondence: alessio.ansuini@areasciencepark.it, alberto.cazzaniga@areasciencepark.it
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.
In this work, we study the intrinsic dimension and neighbor composition of the data representation in ESM-2 [15] and Image GPT (iGPT) [10], two families of transformer architectures trained with self-supervision on protein datasets and images (see Sec. 2.3). We find consistent within-domain behaviors and highlight similarities and differences across the two domains. Additionally, we develop an unsupervised strategy to single out layers carrying the most semantically meaningful representations.
Our main results are:
- representations in large transformers evolve across the layers in distinct phases, revealed by simultaneous changes in global ID and local neighborhood composition (Sec. 3.1);
- a common trait of the ID profiles is a prominent peak of the ID, in which the neighbor structure is profoundly rearranged, followed by relatively low ID regions in which the representations encode abstract semantic information, such as the class label for images, and remote homology for proteins (Sec. 3.2). This abstract information emerges gradually while the ID decreases (Sec. 3.1);
- in the case of protein language models (pLMs) and iGPT, our experiments show that the ID profile can be used to identify in an unsupervised manner the layers encoding the optimal representation for downstream learning tasks, such as protein homology searches or image classification (Sec. 3.2).
These findings suggest a general computational strategy learned by these models, in which the reconstruction task is performed in three phases: 1) a data expansion in a high intrinsic dimension space. This expansion resembles the strategy followed in kernel methods [16], where one implicitly expands the feature space by introducing non-linear functions of the input features; 2) a compression phase that projects the data in low-dimensional, semantically meaningful representations. Strikingly, this representation emerges spontaneously in all the self-supervision tasks we considered; 3) a decoding phase where the model addresses the minute decision-making needed to reconstruct the data from a compressed representation. The layers performing this task behave similarly to the decoder layers in a "vanilla" autoencoder. In particular, understanding this behavior can be exploited as an unsupervised method to identify the most semantically rich representations of transformers trained on large datasets where the annotation is scarce.
2.1 Intrinsic dimension.
Learning complex functions in high-dimensional spaces is challenging because as the dimension of the space increases, the required number of samples must grow exponentially [17]. Fortunately, real-world datasets, such as images and protein sequences exhibit regularities that result in strong constraints and correlations among their features [18, 19]. Consequently, although these datasets may have many features, they lie around low-dimensional, often highly non-linear manifolds. A dataset's intrinsic dimension (ID) is the dimensionality of the manifold approximating the data and can be described as the minimum number of coordinates that allow specifying a data point approximately without information loss [20, 21]. We measure the ID with the TwoNN estimator [21], which requires only the knowledge of the distances  $r_{i1}$ ,  $r_{i2}$  between each point  $x_i$  and its first two nearest neighbors. It can be shown that under the assumption of locally constant density, the ratio  $\mu_i = r_{i2} / r_{i1}$  follows a Pareto distribution with a shape parameter equal to the ID. The intrinsic dimension can be inferred by maximum likelihood or linear regression from the cumulative distribution function (see Appendix Sec. B). The algorithm is robust to changes in curvature and density of the data and has been used effectively to analyze representations in deep neural networks in [13, 14] and more recently in [22-24]. We adopt the TwoNN implementation of DADApy [25] and discuss the technical details of our analysis together with the validation of the constant density assumption in Appendix B.
2.2 Neighborhood overlap.
Changes in the data representation between successive layers of a model can be traced in the rearrangement of the neighbor structure of each data point. The neighborhood overlap  $\chi_{k}^{l,m}$  introduced in [14] measures the similarity between two representations  $l, m$  with the average fraction of common  $k$ -nearest neighbors of a data point. Consider the activation vector  $x_{i}^{l}$  of an example at layer  $l$ . Let  $A^{l}$  be the adjacency matrix with entries equal  $A_{ij}^{l} = 1$  if  $x_{j}^{l}$  belongs to the first  $k$ -nearest neighbors of  $x_{i}^{l}$  in Euclidean distance and 0 otherwise. Then the neighborhood overlap between representations  $l$  and  $m$  can be written as  $\chi_{k}^{l,m} = \frac{1}{N}\sum_{i}\frac{1}{k}\sum_{j}A_{ij}^{l}A_{ij}^{m}$ , where  $A_{ij}^{l}A_{ij}^{m}$  is the number of common  $k$ -neighbors of point  $x_{i}$  between these representations.
In the presence of a labeled dataset  $\{(x_i, y_i)\}$ , a generalization of the neighborhood overlap measuring the consistency of a label within the nearest neighbors of each data point can be used as a semantic probe. In this case, the "adjacency matrix" with respect to the dataset labels is  $A_{ij}^{gt} = 1$  if  $y_i = y_j$  and 0 otherwise, and we call  $\chi_k^{l,gt}$  overlap with the ground truth classes. Both  $\chi_k^{l,m}$  and  $\chi_k^{l,gt}$  lie in [0, 1] and depend on the choice of neighborhood size  $k$ . However, in the Appendix (see Fig. S5), we show that the qualitative trend of the overlap is robust to changes in  $k$ . We use  $k = 10$  when analyzing the overlap with the protein superfamily of the SCOPe dataset, and  $k = 30$  in the case of the transformers trained on the ImageNet dataset, consistently with the value used in Doimo et al. [14].
2.3 Models and Datasets
Transformer models for protein and image datasets. The ESM2 single-sequence protein language models (pLMs) and iGPT networks are characterized by similar architectures. The token vocabulary comprises an alphabet of  $\simeq 20$  amino acids for pLMs and 512 colors for iGPTs (see Sec. A and Chen et al. [10] for further details). An embedding layer encodes each sequence token into a vector of size  $d$ , which is then added to a learned positional embedding common to all sequences. A sequence of  $l$  embedded tokens of size  $\mathbb{R}^{l\times d}$  is thus processed by a stack of architecturally identical blocks composed of a self-attention map followed by a multi-layer perceptron (MLP) producing a series of data representations of size  $\mathbb{R}^{l\times d}$ . We extract the hidden representations of the sequences after the first normalization layer of each block and then average pool along the sequence dimension to reduce a sequence into a data point embedded in  $\mathbb{R}^d$ . This way, we can compute the Euclidean distances needed for the ID and overlap analyses between sequences of different lengths.
We study models pre-trained on large corpora of proteins and images with different self-supervised objectives: ESM2 were trained by Lin et al. [15] on the Uniref50 dataset with a masked language model objective and have been publicly released at github.com/facebookresearch/esm; the iGPT networks were trained by Chen et al. [10] with a next token prediction objective on the ImageNet dataset and are available at github.com/openai/image-gpt.
Datasets. We consider two benchmark datasets for our analysis of pLMs: ProteinNet and SCOPe. ProteinNet [26] is a dataset of 25, 299 protein sequences for evaluating the relationships between protein sequences and their structures. We use the ProteinNet training set to analyze the ID curves and compute the neighborhood overlap of consecutive layers. The Astral SCOPe v2.08 (SCOPe) dataset [27] contains subsets of genetic domain sequences hierarchically classified into fold, superfamily, and family. We first preprocess the sequences with the filtering procedure recommended in [28]. Secondly, to define the remote homology task, we select proteins that belong to superfamilies with at least ten members and ensure that each superfamily consists of at least two families. As a result, we obtain a dataset of 10,256 sequences and 288 superfamilies, used to evaluate the  $\chi^{l,gt}$  in Fig. 4.
For the analysis of the iGPT representations, we choose 90,000 images from the ImageNet training set [29]. We randomly select 300 classes and keep 300 images per class.
Reproducibility. All experiments are performed on a machine with 2 Intel(R) Xeon(R) Gold 6226 processors, 256GB of RAM, and 2 Nvidia V100 GPUs with 32GB memory. We provide code to reproduce our experiments and our analysis online at github.com/diegodoimo/geometry Representations.
Figure 1: Intrinsic dimension (ID) of the data representations in hidden layers of large transformers. (left) The ID profile of ESM-2 protein language models of small (35M, green), medium (650M, blue), and large (3B, purple) sizes show a peak in the first part of the network, a long plateau, and a final ascent. (right) The ID profile in iGPT models of small (76M, yellow), medium (455M, orange), and large (1.4B, red) sizes also present an early peak followed by a less pronounced second peak for the medium and large model. The ID is plotted against the relative depth: the block number divided by the total number of blocks.
3.1 The intrinsic dimension profile in large transformers is similar in different tasks and architectures
In large transformers, the input representation is processed in a sequence of identical blocks, creating a series of representations in vector spaces of the same dimension but radically different in their geometrical properties and semantic content. We start by estimating the ID of each representation for large transformers trained by self-supervision on protein language and image reconstruction tasks. The ID profiles provide a glimpse into the geometry of the manifolds in which the data are embedded. As we will see, the ID changes through the layers in a highly non-trivial manner, with three phases that can be recognized in all the transformers and a fourth phase that we observed only in the iGPT models. In the following sections, we will see how these phases signal the development of structured semantic content in intermediate representations, which are hidden instead in the input and, by construction, in the final embeddings.
ID curve of protein language models. We extract representations of the protein sequences in the ProteinNet dataset as explained in Sec. 2.3 and plot their ID against the block number divided by the total number of blocks, referred to as relative depth in Fig. 1. This section describes our findings for pLMs in the ESM-2 model family with 35M, 650M, and 3B parameters. Analogous results hold for other models trained on different datasets, as shown in the Appendix (see Fig. S1).
Figure 1 (left) shows the typical shape of the ID profile in pLMs with three distinct phases: a peak phase, a plateau phase, and a final ascent phase. The peak develops in early layers, spanning roughly the first third of the network. During this phase, the ID expands rapidly, reaching a maximum of approximately 20, 25, and 32 for ESM-2 35M, 650M, and 3B, respectively. The ID then contracts quickly, stabilizing at notably low values, between 5 and 7, that define the plateau phase. While the ID at the peak grows with the model size, 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. Notably, the ID observed in the plateau phase is also consistent with the values between 6 and 12 measured by
Facco et al. [30]. Their study involved different metrics for computing distances directly applied to pairwise alignments of protein sequences. This remarkable correspondence between different analytical approaches suggests that hidden representations in the plateau region can effectively gauge the underlying degree of variability arising from evolutionary changes in protein sequences. In the final ascent, the ID grows again, returning progressively to values close to the ID computed on the input representation after the positional embedding. This is a consequence of the masked language modeling objective, which focuses on using the context information to recover the missing tokens.
ID curve of image transformers. Figure 1 (right) shows the ID profiles of iGPT models of increasing size. Specifically, we examine the small (S), medium (M), and large (L) versions of iGPT trained on ImageNet, which have 76M, 455M, and 1.4B parameters, respectively.
In all cases, the IDs of the output are similar to those of the input, which aligns with the findings from protein language models. In the iGPT-S network, the ID profile has a hunchback shape quite similar to that in convolutional models trained on ImageNet [13] but with a peak value significantly smaller (around 25). This discrepancy could be partly attributed 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, typically falling within the range of  $O(10^{4})$  to  $O(10^{6})$ . As we increase the model size, similar to what is observed (pLMs), the maximum intrinsic dimension (ID) grows, reaching 28 for iGPT-M and 32 for iGPT-L. However, after an initial decrease between 0.3 and 0.4 of the relative depth, the ID shows a local minimum of 22 common between the architectures. In the last part of the network, in contrast with what is observed in pLMs, the ID experiences another, more gradual increase, forming a second, shallower peak near the end of the network. Interestingly, the ID of the most compressed representations at around 0.4 of relative depth 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. In their study, the ID varied from 13 to 24 and was found to be correlated with the classifier's accuracy (refer to Figure 4 in Section 3.3 of [13]).
The results shown in Fig. 1 hint at a scenario that will be validated in the following: the ID is an indicator of semantic complexity, and the representations that resemble more closely those extracted by supervised learning on the same datasets are expected to correspond to the relative minimum of the ID curve, situated after the initial peak. This minimum is well defined in iGPTs, while in pLMs, the ID profile forms a plateau.
The neighborhood rearrangement across the layers mirrors the ID profile. We will now examine the evolution of the data distribution across the layers, specifically by analyzing the neighborhood overlap  $\chi_k^{l,l+1}$ , defined in Sec. 2.2. This parameter quantifies the rate at which the composition of the neighborhood of each data point changes between nearby blocks. In deep convolutional models, this rate is significant only in the very last layers [14], where the representation is forced to comply with the ground-truth classification of the images.
In pLMs, we evaluate the neighborhood overlap on the representations of the ProteinNet dataset. The results of our analysis, presented in Fig. 2 (left), reveal a relationship between changes in neighborhood composition and the ID. This observation points to an interesting connection between ID and neighborhood overlap, which differs from what is typically found in convolutional architectures. In the first  $40\%$  of the layers,  $\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. This corresponds with the peak phase of the ID. However, during 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. 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.
In image transformers, Fig. 2 (right), the profiles are qualitatively similar to those observed in protein language models. However, the transition between the initial stage, coinciding with the first ID peak and relative depth  $< 0.4$ , with rapid changes in neighborhood composition  $(\chi^{l,l+1} \sim 0.7)$  and a second stage where the rearrangements are less pronounced  $(\chi^{l,l+1} \sim 0.9)$  is more gradual than pLMs. A significant neighborhood rearrangement is always observed in the last layers, where the reconstruction task is carried on. Similarly to pLMs,  $\chi^{l,l+1}$  is lower in shallower models that need faster rearrangements to achieve consistent results in fewer blocks.
Figure 2: The neighborhood rearrangement across the layers in large transformers. (left) Neighborhood overlap of consecutive layers  $\chi^{l,l+1}$ , computed for ESM-2 (35M, green), ESM-2 (650M, blue), and ESM-2 (3B, purple), shows major local rearrangements in the layers corresponding to the peak and in the final layers and minor changes in the plateau region, mirroring the behavior of the ID. (right) The trend of  $\chi^{l,l+1}$  in iGPT-S (yellow), iGPT-M (orange), and iGPT-L (red) is similar to pLMs: the neighborhood composition changes more in early layers where the ID reaches a peak and in the final layers.
The evolution of the representations during training. Until now, our focus has been on observing changes in some geometric quantities across layers in fully trained models. We will now focus on understanding how these quantities evolve during training. Specifically, we investigate how representations of the data manifold transform during training for two models: the protein language model ESM-2 (650M) and the image transformer iGPT-L.
For the ESM-2 (650M) model, as shown in Figure 3 (left), we analyze checkpoints corresponding to  $[0, 0.1, 0.3, 1, 5] \cdot 10^5$  training iterations, with the last checkpoint representing the fully converged model. 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. Between  $1 \cdot 10^4$  and  $3 \cdot 10^4$  training steps, the ID of the plateau layers substantially decreases. In contrast, the ID measured in the last layers progressively increases towards the input ID. From  $3 \cdot 10^4$  training steps to convergence, the ID curve takes its final shape, with a slight increase of the ID measured at the peak and a minor compression at the plateau. 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 compression is tightly related to the emergence of semantic information that we discuss in Sec. 3.2.
For iGPT-L, we consider the representations at checkpoints corresponding to  $[0, 0.13, 0.26, 0.52, 1] \cdot 10^6$  training iterations (see Fig. 3 (right)). The ID dynamics is similar to that of pLMs: before training, the ID curve is flat, and in early training iterations ( $< 1.3 \cdot 10^5$ ), the first ID peak emerges in the initial layers of the network. In the later stages of training, the ID slightly increases at the first peak and decreases at 0.4 of relative depth, forming the local minimum. Additionally, a second peak emerges in the last third of the model's hidden layers.
3.2 The intermediate representations of large transformers are the most semantically rich
A common characteristic among the ID profiles shown in Fig. 1 is the presence of a distinct peak in the first part of the network, observed consistently in all the models we examined. In protein language models, this initial peak is followed by an extended plateau that spans more than half of the layers. In contrast, in iGPT models, a second shallow peak in the ID is observed, with the exception of the iGPT-S model, which does not have one. In section 3.1, we found that the ID at the minimum point between the two peaks is quantitatively similar to the one measured at the output of convolutional networks trained with supervision on the same dataset [13]. This points to a scenario where the representations at this local minimum encode semantic information about the datasets. Chen et al. [10] have already demonstrated with linear probes that features extracted from the intermediate layers
3The weights of the checkpoints of ESM-2 (650M) analyzed in Fig. 3 were kindly provided by R. Rao.