Title: The geometry of hidden representations of large transformer models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Representations

transformers

geometry

interpretability
Abstract: 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 taskand 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.

The geometry of hidden representations of large transformer models

Abstract

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.

1 Introduction

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,

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 Methods

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 i 1 , r i 2 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 µ i = r i 2 / r i 1 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 χ l,m k 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 l i of an example at layer l . Let A l be the adjacency matrix with entries equal A l ij = 1 if x l j belongs to the first k -nearest neighbors of x l i in Euclidean distance and 0 otherwise. Then the neighborhood overlap between representations l and m can be written as χ l,m k = 1 N ∑ i 1 k ∑ j A l ij A m ij , where A l ij A m ij 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 gt ij = 1 if y i = y j and 0 otherwise, and we call χ l,gt k overlap with the ground truth classes. Both χ l,m k and χ l,gt k 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 ≃ 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 R l × 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 R l × 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 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 χ 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, 256 GBof RAM, and 2 Nvidia V100 GPUs with 32 GBmemory. We provide code to reproduce our experiments and our analysis online at github.com/diegodoimo/geometry_representations.

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0

5

10

15

20

25

30

35

ID

ESM-2(3B)

ESM-2(650M)

ESM-2(35M)

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0

5

10

15

20

25

30

35

ID

iGPT-L(1.4B)

iGPT-M(455M)

iGPT-S(76M)

3 Results

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 35 M, 650 M, and 3 B 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 76 M, 455 M, and 1 . 4 B 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 χ l,l +1 k , 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, χ 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 ( χ l,l +1 ∼ 0 . 7 ) and a second stage where the rearrangements are less pronounced ( χ l,l +1 ∼ 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, χ l,l +1 is lower in shallower models that need faster rearrangements to achieve consistent results in fewer blocks.

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.0

0.2

0.4

0.6

0.8

1.0

l

,

l

+1

ESM-2(3B)

ESM-2(650M)

ESM-2(35M)

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.0

0.2

0.4

0.6

0.8

1.0

l

,

l

+1

iGPT-L(1.4B)

iGPT-M(455M)

iGPT-S(76M)

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) 3 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] · 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 · 10 4 and 3 · 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 · 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] · 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 · 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

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0

5

10

15

20

25

ID

0

1

10 4

3

10 4

1

10 5

5

10 5

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0

5

10

15

20

25

30

35

ID

0

1.3

10 5

2.6

10 5

5.2

10 5

1

10 6

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 χ l,gt k 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 χ l,gt k ∼ 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, χ 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 χ l,gt around a relative depth of ∼ 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 χ 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.

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.0

0.2

0.4

0.6

0.8

1.0

l

,

gt

ESM-2(3B)

ESM-2(650M)

ESM-2(35M)

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

l

,

gt

iGPT-L(1.4B)

iGPT-M(455M)

iGPT-S(76M)

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 × 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 [3337] 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.

Acknowledgments and Disclosure of Funding

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.

References

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 30, 2017.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina N. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , 1:4171-4186, 2019.

Christopher D. Manning, Kevin Clark, John Hewitt, Urvashi Khandelwal, and Omer Levy. Emergent linguistic structure in artificial neural networks trained by self-supervision. Proceedings of the National Academy of Science , 117(48):30046-30054, December 2020.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res. , 21(1), jun 2022.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 1877-1901. Curran Associates, Inc., 2020.

Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C. Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences , 118(15):e2016239118, 2021.

Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Wang Yu, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, Debsindhu Bhowmik, and Burkhard Rost. ProtTrans: Towards Cracking the Language of Lifes Code Through SelfSupervised Deep Learning and High Performance Computing. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2021.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with AlphaFold. Nature , 596(7873):583-589, 2021.

Yanrong Ji, Zhihan Zhou, Han Liu, and Ramana V Davuluri. DNABERT: pre-trained Bidirectional Encoder Representations from Transformers model for DNA-language in genome. Bioinformatics , 37(15):2112-2120, 02 2021.

Mark Chen, Alec Radford, Rewon Child, Jeff Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pretraining from pixels. In Proceedings of the 37th International Conference on Machine Learning , ICML'20. JMLR.org, 2020.

Alexander Kolesnikov, Alexey Dosovitskiy, Dirk Weissenborn, Georg Heigold, Jakob Uszkoreit, Lucas Beyer, Matthias Minderer, Mostafa Dehghani, Neil Houlsby, Sylvain Gelly, Thomas Unterthiner, and Xiaohua Zhai. An image is worth 16x16 words: Transformers for image recognition at scale. Proceedings of the International Conference on Learning Representations , 18-24 Jul 2021.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 8748-8763. PMLR, 18-24 Jul 2021.

Alessio Ansuini, Alessandro Laio, Jakob H Macke, and Davide Zoccolan. Intrinsic dimension of data representations in deep neural networks. Advances in Neural Information Processing Systems , 32, 2019.

Diego Doimo, Aldo Glielmo, Alessio Ansuini, and Alessandro Laio. Hierarchical nucleation in deep neural networks. Advances in Neural Information Processing Systems , 33, 2020.

Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Salvatore Candido, and Alexander Rives. Evolutionary-scale prediction of atomiclevel protein structure with a language model. Science , 379(6637):1123-1130, 2023.

B. Scholkopf and A.J. Smola. Learning with Kernels: Support Vector Machines, Regularization, Optimization, and Beyond . Adaptive Computation and Machine Learning series. MIT Press, 2018.

Hariharan Narayanan and Sanjoy Mitter. Sample complexity of testing the manifold hypothesis. Advances in neural information processing systems , 23, 2010.

J Teenbaum, DD Silva, and JA Langford. global geometric framework for nonlinear dimensionality reduction [j]. Science , 290(5500):2319-2323, 2000.

Daniele Granata and Vincenzo Carnevale. Accurate estimation of the intrinsic dimension using graph distances: Unraveling the geometric complexity of datasets. Scientific reports , 6(1):31377, 2016.

R. Bennett. The intrinsic dimensionality of signal collections. IEEE Transactions on Information Theory , 15(5):517-525, 1969.

Elena Facco, Maria d'Errico, Alex Rodriguez, and Alessandro Laio. Estimating the intrinsic dimension of datasets by a minimal neighborhood information. Scientific Reports , 7, 2017.

Emily Cheng, Corentin Kervadec, and Marco Baroni. Bridging information-theoretic and geometric compression in language models. arXiv preprint arXiv:2310.13620 , 2023.

Lorenzo Basile, Nikos Karantzas, Alberto D'Onofrio, Luca Bortolussi, Alex Rodriguez, and Fabio Anselmi. Relating implicit bias and adversarial attacks through intrinsic dimension. arXiv preprint arXiv:2305.15203 , 2023.

Henry Kvinge, Davis Brown, and Charles Godfrey. Exploring the representation manifolds of stable diffusion through the lens of intrinsic dimension. arXiv preprint arXiv:2302.09301 , 2023.

Aldo Glielmo, Iuri Macocco, Diego Doimo, Matteo Carli, Claudio Zeni, Romina Wild, Maria d'Errico, Alex Rodriguez, and Alessandro Laio. Dadapy: Distance-based analysis of datamanifolds in python. Patterns , page 100589, 2022.

Mohammed AlQuraishi. ProteinNet: a standardized data set for machine learning of protein structure. BMC Bioinformatics , 20, 2019.

John-Marc Chandonia, Lindsey Guan, Shiangyi Lin, Changhua Yu, Naomi K Fox, and Steven E Brenner. SCOPe: improvements to the structural classification of proteins - extended database to facilitate variant interpretation and machine learning. Nucleic Acids Research , 50(D1):D553D559, 12 2021.

Johannes Söding and Michael Remmert. Protein sequence comparison and fold recognition: progress and good-practice benchmarking. Current opinion in structural biology , 21 3:404-411, 2011.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Elena Facco, Andrea Pagnani, Elena Tea Russo, and Alessandro Laio. The intrinsic dimension of protein sequence evolution. PLoS computational biology , 15(4):e1006767, 2019.

N.S. Detlefsen, S. Hauberg, and W. Boomsma. Learning meaningful representations of protein sequences. Nature Communications , 13, 2022.

John Hewitt and Christopher D. Manning. A structural probe for finding syntax in word representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4129-4138, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

Yonatan Belinkov, Lluís Màrquez, Hassan Sajjad, Nadir Durrani, Fahim Dalvi, and James Glass. Evaluating layers of representation in neural machine translation on part-of-speech and semantic tagging tasks. In Proceedings of the Eighth International Joint Conference on Natural Language Processing (Volume 1: Long Papers) , pages 1-10, Taipei, Taiwan, November 2017. Asian Federation of Natural Language Processing.

T. Blevins, O. Levy, and L. Zettlemoyer. Deep rnns encode soft hierarchical syntax. Proceedings of the 56th Association for Computational Linguistics (ACL) , 2018.

Nelson F. Liu, Matt Gardner, Yonatan Belinkov, Matthew E. Peters, and Noah A. Smith. Linguistic knowledge and transferability of contextual representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 10731094, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

Ian Tenney, Dipanjan Das, and Ellie Pavlick. BERT rediscovers the classical NLP pipeline. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pages 4593-4601, Florence, Italy, July 2019. Association for Computational Linguistics.

Anna Rogers, Olga Kovaleva, and Anna Rumshisky. A primer in BERTology: What we know about how BERT works. Transactions of the Association for Computational Linguistics , 8:842-866, 2020.

Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Sal Candido, et al. Language models of protein sequences at the scale of evolution enable accurate structure prediction. bioRxiv , 2022.

Bo Pang and Lillian Lee. Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales. arXiv preprint cs/0506075 , 2005.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023.

D. Moltchanov. Distance distributions in random networks. Ad Hoc Networks , 10(6):1146-1166, 2012.

Francesco Denti, Diego Doimo, Alessandro Laio, and Antonietta Mira. The generalized ratios intrinsic dimension estimator. Scientific Reports , 12(1):20005, 2022.

Alex Rodriguez, Maria d'Errico, Elena Facco, and Alessandro Laio. Computing the free energy without collective variables. Journal of Chemical Theory and Computation , 14(3):1206-1215, 2018. PMID: 29401379.

Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. MSA Transformer. Proceedings of the 38th International Conference on Machine Learning , 139:8844-8856, 18-24 Jul 2021.

Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alex Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. Advances in Neural Information Processing Systems , 34, 2021.

The UniProt Consortium. UniProt: the universal protein knowledgebase in 2021. Nucleic Acids Research , 49(D1):D480-D489, 11 2020.

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.

Appendix

A Architecture details

Transformer models for protein sequences. The input data points, corresponding to proteins, are variable-length sequences of l letters s = a 1 , a 2 , . . . a l , chosen from an alphabet of n a ( ≃ 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 := f 0 ( x ) of size l × d . A model with B blocks transforms a data point x ∈ R l × d into a sequence of representations: f 0 ( x ) → f 1 ( x ) → . . . → f B ( x ) → f out ( x ) , where f i , i = 1 , . . . , 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 × n a . The size of each hidden layer does not change across the model and is equal to l × d ; therefore, the action of the model is a sequence of mappings R l × d → R l × d . For each layer i , we choose to perform global average pooling across the row dimension f i ( x ) → 1 l ∑ l j =1 ( 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 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 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 ∈ R l × 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 ∈ R l × d . The final head f out projects the output of the last attention block to R l × d in , which encodes at each sequence position i the conditional distribution

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 µ i = r i 2 / r i 1 between the distances to the second and first nearest neighbor follow a Pareto distribution p ( µ i | d ) = dµ -d -1 i . Assuming that the dataset likelihood factorizes over the N single-point likelihoods, the ID can be easily inferred given the empirical values of the µ i as the p ( µ i | d ) depend only on d . Facco et al. propose to estimate the ID from the cumulative distribution function F ( µ ) = 1 -µ -d with a linear regression fit. Indeed, the ID is the slope of the line through the origin obtained regressing log( µ σ ( i ) ) against -log(1 -F emp ( µ σ ( i ) )) where µ σ ( i ) are the µ i sorted in ascending order and F emp ( µ σ ( i ) ) := i N is the value of empirical cumulative distribution function at µ σ ( 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 , ... ] . The relevant data ID is chosen as the value of ˆ d where the graph ˆ 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.

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0

5

10

15

20

25

ID

ProtBert

ProtT5-XL-U50

ESM-1v

ESM-1b

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.0

0.2

0.4

0.6

0.8

1.0

l

,

gt

ProtT5-XL-U50

D Impact of the implementation details and hyperparameter choices.

Cifar10 - WR28-8

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0

5

10

15

20

25

ID

avg_pooling

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0

10

20

30

40

50

60

70

ID

full_id

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0

5

10

15

20

25

30

35

ID

post_ln

post_attn

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

l

,

gt

post_ln

post_attn

iGPT-L

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.0

0.2

0.4

0.6

0.8

1.0

l

,

l

+1

k = 1

k = 2

k = 5

k = 10

k = 20

k = 50

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.0

0.2

0.4

0.6

0.8

1.0

l

,

gt

k = 1

k = 2

k = 5

k = 10

k = 20

k = 50

ESM-2(650M)

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.0

0.2

0.4

0.6

0.8

1.0

l

,

l

+1

k = 1

k = 2

k = 5

k = 10

k = 20

k = 50

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.0

0.1

0.2

0.3

0.4

0.5

l

,

gt

k = 1

k = 2

k = 5

k = 10

k = 20

k = 50

iGPT-L

E Preliminary analysis of the intrinsic dimension and overlap in LLama2

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

30

35

40

45

50

55

ID

sst

0.0

0.2

0.4

0.6

0.8

1.0

relative depth

0.525

0.550

0.575

0.600

0.625

0.650

0.675

0.700

l

,

gt

sst

llama2-70B

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work investigates the geometry of hidden representations of transformer models trained via a self-supervised task on either amino acid prediction in proteins or pixel prediction in images. The work uses two tools to understand this geometry: intrinsic dimension (ID) (estimated via the TwoNN algorithm) and neighborhood overlap. The paper shows that as data passes through the layers of a transformer, both the ID and neighborhood overlap change in characteristic ways that the work connects with the extent to which the representation is organized by the semantic content in the data. The paper validates this claim by looking at the extent to which each data point shares neighbors that belong to the same semantic class, showing that this peaks at layers when ID is low and changes in neighborhood overlap is also low. The work then speculates that these observations could provide a way of identifying the best representations to use for downstream tasks.\n\nSTRENGTHS:\n- While there is a rich literature exploring hidden representations of deep learning models, most works continue to focus on CNNs or MLPs with supervised training on medium sized datasets. Given the growing importance of transformers to NLP and vision and the increasing use of self-supervision for large-scale training, there is a need for works that explore how these approaches impact a model’s hidden representations. Thus, this is a welcome work that will doubtless be of value to researchers training large transformers via self-supervision. \n- The paper contains careful analysis of the experiments (as opposed to other works which all too frequently just list summary statistics). The conclusions which are reached are all fairly-well supported within the scope of the experiments that were performed. Crucially, the two metrics that are used, intrinsic dimension and neighborhood overlap, reinforce each other’s conclusions. This increases the believability of the results significantly.\n- The paper is able to suggest some practical value (identifying layers of the transformer that capture the most semantic content) in their scientific observations (patterns in representation geometry in large-transformers), helping to connect practice with theory.\n\nWEAKNESSES:\n- The experiments in this work all use transformer models trained via self-supervised tasks based around data reconstruction (as opposed, for example, to a contrastive learning task). While most of the phenomena observed in this paper is attributed to the “large transformer” architecture, this reviewer wonders if some of the conclusions are also contingent on the self-supervision task. For example, does one see similar behavior in a large vision transformer that has been trained in a supervised manner. Disentangling which of the conclusions are due to architecture and which are due to training method would make the work significantly stronger.\n- Many of the claims in the work rely on comparison between the shapes of curves which plot intrinsic dimension of data between layers. While the claims seem mostly reasonable given the figures, it would make the work stronger if more quantitative measures were used to, for example, compare ID curves. This sort of automation might also make it easier to compare against a broader range of models and datasets. While this reviewer thought that the use of both protein and image datasets was a strength of this work, comparing against other types of transformer architectures and datasets would reinforce the conclusions.\n- The main content of the paper consists of discussion of several figures. This reviewer felt that the way this discussion was written/organized, it was easy for the reader to get disoriented and forget the main points already established. Possibly making the discussion more concise or better highlighting the main takeaways from each section would help the reader to mentally organize the primary findings of the work. Being more succinct might also allow further experiments to be included.\n\nNitpicks\n\n- The abstract describes a transformer as being composed of a “sequence of functionally identical transformations”. To this reviewer, functionally identical transformations would be transformations that behave the same way on the level of functions (for the same inputs, they give the same outputs), but may be parametrized in different ways. The layers of a transformer are instead transformations that all belong to the same functional family, though each layer is generally a different function.\n- Some sentences in this work have an over-abundance of commas. For example, “In this work, we systematically investigate, in some self-supervised models, fundamental geometric properties of representations, such as their intrinsic dimension (ID) and neighbor composition.” To improve the flow of the paper, it would be good to find such sentences and remove some of the commas.\n- Line 66: “…where the annotation is very scarce” -> “…where the annotations are very scarce.”\n- Reading this work would have been easier if the figures were closer to their corresponding analysis. As it is, the reader must flip between pages frequently to validate claims.\n- Line 110: “More in detail,” -> “In more detail,”.\n\nQUESTIONS:\n- Given the success of transformers trained with self-supervised learning in the NLP domain, I was curious why the authors did not study this setting? \n- Line 76: Is the assumption of locally constant density really “mild”? This reviewer wonders if image datasets, which have moderately high intrinsic dimension ([1] estimates ImageNet as having intrinsic dimension between 26 and 43) might be fairly sparsely sampled even when the dataset is large.\n- What happens to the results when $k$ is varied (rather than being fixed at 30)?\n- The paper says that hidden representations are extracted after the first normalization layer of the attention blocks. Do you have a sense of whether results change if you use representations from other parts of attention blocks?\n- The results show the intrinsic dimension of the data increasing at times between layers. We know that mathematically, maps $f: M \\rightarrow \\mathbb{R}^n$ that map a $k$-dimensional manifold $M$ into $\\mathbb{R}^n$ such that $\\dim(f(M)) > k$ exist (e.g., space filling curves) but it is this reviewer's understanding that the space of such $f$ has measure zero with respect to many reasonable probability measures on the space of such maps. Given this, is it really possible that the intrinsic dimension of the representations is increasing? Or is their some other change that is changing the ID estimation?\n\n[1] Pope, Phillip, et al. \"The intrinsic dimension of images and its impact on learning.\" arXiv preprint arXiv:2104.08894 (2021).\n\nLIMITATIONS:\nThis reviewer believes that some limitations could have been discussed, including:\n- The use of intrinsic dimension estimators, which can provide misleading feedback in certain situations.\n- A limited number of different model types and datasets.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents an analysis of the internal representations of transformers trained with self-supervised learning, such as masked language modeling, from two perspectives: intrinsic dimension and adjacency structure. Experiments on two datasets - protein sequences and image data - revealed that the intrinsic dimension within the transformer has two peaks, one in the early layers and another in the latter layers. This result is qualitatively different from previous convolutional networks trained with supervised learning. Notably, the intermediate layer, where the intrinsic dimension is minimized, is the most suitable for categorical discrimination in downstream tasks.\n\nSTRENGTHS:\nThis paper presents a notable analysis of the intrinsic dimension of transformer models trained with self-supervised learning. While many approaches have been used to analyze the internal dimension of deep neural networks, most have focused on convolution-based networks. The discovery of two peaks in the internal dimension is intriguing, offering insightful contributions to the community. The results are clearly presented, and the paper is well-structured, making it easy for readers to follow the logical development. Of particular interest is the finding that such representations can spontaneously appear in models trained with masked modeling, even without an explicit bottleneck structure in the model.\n\nHowever, this interesting result also raises many considerations and discussions. There are points in the current paper where sufficient evidence to support the authors' claims is lacking and further room for discussion remains. These points will be discussed in detail below.\n\nWEAKNESSES:\nThe paper's experimental comprehensiveness raises two issues. The first issue revolves around the focus on internal representations of transformer models trained with self-supervised learning, with the assertion that the evolution of the internal dimension displays two peaks. However, it appears that sufficient experiments have not been performed to identify the causes of this occurrence definitively. An ablation study would be incredibly beneficial in distinguishing the elements that stem from the model architecture (transformer) and those arising from the training protocol (masked/autoregressive modeling).\n\nSecondly, the paper state in the introduction that the adjacency structure is “rearranged” at the peak of the representation’s internal dimension, regardless of the dataset domain. Yet, upon the qualitative comparison of Figures 1 and 2, it is challenging to confirm a consistent change in the internal dimension and the overlap of the adjacency structure between layers in the image domain experiments. Therefore, the current portrayal of the contributions gives the impression of over-claiming. As asserting a clear correlation in qualitative behavior is difficult, a comparison with some control conditions should be made. Additionally, it would be desirable to have a discussion concerning the influence of differences in the data domain.\n\nQUESTIONS:\nI have several questions for the authors:\n\n1. In interpreting Figure 1, the main text mentions that the changes in the internal dimension qualitatively match the experiments for the two domains. However, there seems to be a qualitative difference between domains concerning whether there is concentration or expansion near the final layers in the latter half. What could be the reasons for such a difference?\n2. Regarding the analysis of the overlap of the adjacency structure, the paper mentioned that the authors used $k=10$ and $k=1$ in the protein domain experiments, and $k=30$ in the image domain experiments. Generally, the experimental results would depend on this hyperparameter. How robust are the authors' claims with respect to these hyperparameters, and what is the rationale or justification for the chosen hyperparameter values?\n3. This paper mentioned that the authors used the GAP type averaging procedure along the token direction when analyzing the internal representation of the transformer model. While this is not a problem when comparing results within the same model architecture, it seems caution is needed when comparing the results with other models. For instance, when comparing the results with convolution-type neural networks, how can we guarantee the qualitative robustness of the statements depending on whether this averaging operation is performed or not?\n\nLIMITATIONS:\nGiven that this paper is not proposing a new method but focusing on analyzing the internal representation of existing methods, it may not strictly apply to a discussion on limitations. The estimation of the internal representation is based on a method proposed in previous research, and the applicability of this method has been sufficiently demonstrated. However, as mentioned in the Weakness and Question sections, there are concerns about the comprehensiveness of the experiments in this paper. Improvements in these points would lend greater significance to the authors' claims.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work focuses on characterizing the geometrical and statistical properties of data representations across the hidden layers of large transformer models. Specifically, it demonstrates the similarity in the evolution of geometric properties, such as the intrinsic dimension, between image-based and protein-based transformers. Additionally, the work proposed an intuitive, unsupervised strategy to identify the most semantically rich representations of transformers. The effectiveness of this strategy is demonstrated by the SOTA performance achieved by leveraging previous work and replacing the last layer with the layer identified in this study.\n\nSTRENGTHS:\n- This work presents a novel strategy to select the layers that produce the most semantic representation to be used in downstream tasks, opening many possible applications. Moreover, the method seems to be general and could be exploited in other architectures where the layer choice is currently arbitrary, such as classifiers.\n  \n- This work demonstrates state-of-the-art results in identifying protein relations by leveraging existing methods. It simply involves swapping the previously employed last layer with the layers that maximizes the semantic content identified through the strategy proposed in this work (Figure 5, right).\n\n- The writing is extremely clear and easy to follow, presenting intuitive ideas, a solid experimental setup and convincing results.\n\n- This work does not involve training models but instead analyzes publicly available models, completely avoiding any possible bias that could have been introduced in the training process.\n\n- Among the others, the insight that large transformers behave essentially like sophisticated autoencoders (lines 337-338) is insightful and may serve as a source of ideas for future research.\n\nWEAKNESSES:\nThis work presents a novel idea with exceptional clarity in its writing, accompanied by a robust experimental setup that yields convincing results. While no major weaknesses were identified, a more comprehensive discussion on the current limitations of the study would enhance its overall quality even further.\n\nMinor:\n\n- In line 36 (and others), the paper mentions the term \"semantically rich representation.\" However, it is important to clarify how this term is defined. It can be argued that the richness of semantic content in a representation is dependent on the specific downstream task being addressed. In other words, a representation can be considered more or less semantically rich based on its effectiveness in a particular downstream task.\n\n- To enhance readability, it would be beneficial to include the ID (or at least the minimum ID) directly in the plot in Figure 4. This additional information (even though repeated from Figure 3), would facilitate the reader to compare the peak categorical information with the ID minimum.\n\nQUESTIONS:\n- In lines 131-132, the paper mentions a method to compress the resolution and color channels of images to address memory constraints. However, it raises the question of why a methodology similar to the one described in [1] was not adopted. In [1], a \"first-stage\" autoencoder was trained specifically to compress the images into lower-dimensional latent representations while maintaining perceptual equivalence (and kept frozen in subsequent stages). Is there a specific reason why this approach was not considered in the current work?\n\n- In line 286, it is reported that the plateau extends beyond half of the layers. This raises the question of whether this observation could be an indication that the model is over-parameterized or, at the very least, has more layers than necessary.  Further investigation or discussion could shed light on this question and provide valuable insights in the minimum network complexity to produce representations with the same semantic information.\n\n\n[1] Rombach, Robin, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. “High-Resolution Image Synthesis with Latent Diffusion Models.”, Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022\n\nLIMITATIONS:\nAlthough not extensively, Section 4 explores the limitations and potential areas for future research.", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
The paper studies the hidden representation of pretrained transformers via the ID (intrinsic dimension) on protein language tasks and image reconstruction tasks . The papers show that on protein LM, from input to output layer, the ID first increase to a peak, then decrease to an elbow, and finally increase to near its ID value at input layer; this is robust from small to largest models. For image reconstruct, the picture is quite different; the input & output layers have the smallest ID and there two peaks near the input and output layers, respectively. The results on iGPT is less robust compare to protein LM; (e.g., small iGPT does not seem to have two peaks). 


Overall, I think the paper is well written, presentations and experimental setup are clean, the results are interesting. However, I am not very familiar with the field and could not judge the novelty of the paper.

STRENGTHS:
A well written paper; structure is clean; presentation is good; authors spent decent amount of efforts to introduce / motivate / explain key concepts using in the paper, e.g. ID etc. 

The results seem interesting even to folks who are not working this area.

WEAKNESSES:
The results on iGPT is less robust (compare to protein LM). In particular, for the small model (I don’t see two peaks). Naively, I also expect the shape of ID of iGPT to be similar to pLM, encode (decreasing ID, smaller than input) and decode (increasing ID), just like the elbow in the pLM. Is there an explanation? 

Why not also including a third task: NLP  (casual language model) in the main text?

QUESTIONS:
See above.

LIMITATIONS:
Have several discussion about further extension of the current approach.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08