Exploring the Impact of a Transformer’s Latent Space Geometry on Downstream Task Performance

Transformer-based large language models, such as BERT Devlin et al. (2018) and the GPT models Radford et al. (2018), have outperformed other language models on linguistic benchmarking tasks time and again. These models largely rely on a pre-training and fine-tuning training protocol, where the base model is trained with a very large corpus on a generic language task, such as masked language modeling, and then the weights are updated in a second round of training with a more specific task, such as text similarity or entailment. The assumption motivating this process is that the pre-training task provides the model with general linguistic information, which allows it to better learn a specific linguistic task during fine-tuning.

However, this assumption has been challenged in several recent works in which the BERT model training process is altered such that we would expect the model to perform poorly on benchmarking tasks. Examples of these pathological changes include using a non-linguistic training task or shuffling the model’s input embeddings Alajrami and Aletras (2022); Sinha et al. (2021); Zhang (2021). If the outstanding benchmarking performance of the standard models relied on the presence of general linguistic knowledge in the underlying model, then these alterations should reduce that performance to a level similar to an untrained model. However, these works consistently find benchmarking performance that is closer to that of a standard model than to that of an untrained model.

We propose that at least some of the benchmarking performance seen with this pre-training/fine-tuning protocol can be attributed to geometric characteristics of the latent space that result from the transformer architecture, divorced from any generic linguistic knowledge. Many common downstream tasks for language models, such as classification or paraphrase, use measures of similarity or separability in their loss functions. These measures rely on the way that latent representations are organized in this space, so it stands to reason that there could be a measurable characteristic of this organization that would be advantageous for learning these fine-tuning tasks. If the latent space is already organized with data placed in various high-dimensional nooks and crannies, it may be easier for a powerful transformer model to learn how to access the data appropriately during fine-tuning, even if the organization scheme is non-linguistic.

Here we explore a variety of measures that could be used to describe the geometry of a high-dimensional contextual language model latent space, and we examine the relationships between these measures and the GLUE benchmarking task performance on a range of BERT-type contextual language models. We introduce a quantized cell density measure that appears to have a strong linear relationship with GLUE performance when applied to latent space representations. We also find that this measure can be used to predict the surprising GLUE performance of many of the non-standard models presented in the aforementioned works, though the performance of a few of these models remains unexplained.

When Vaswani et al. (2017) introduced the transformer model, it was as an end-to-end autoencoder that relied on the attention mechanism to provide contextual information for sequential data modeling. That architecture was adapted to be used in large language models by isolating either the encoder component (as in the BERT models from Devlin et al. 2018) or the decoder component (as in the GPT models from Radford et al. 2018) and training with a very large dataset on a generalized linguistic task, such as masked language modeling. These models are intended to be fine-tuned on specific downstream tasks, generally with a much smaller, labeled dataset, by adding an additional layer and jointly training the pre-trained and new layer weights on the new task.

It has been shown that this pre-training/fine-tuning procedure produces improved performance on downstream tasks, both on transformer models Devlin et al. (2018); Radford et al. (2018) and on more traditional neural models Dai and Le (2015); Howard and Ruder (2018). The advantage of this dynamic is generally understood to fall under the theory of transfer learning, whereby the model learns general information during the pre-training task, and this can then be specialized during fine-tuning. This means that many fewer parameters need to be trained from scratch on the fine-tuning tasks, allowing for large models to be effectively applied to these tasks without drastically overfitting to the (often small) supervised training datasets.

Benchmarking datasets provide the research community with the ability to compare performance on the same exact tasks, and generally aim to encapsulate some knowledge or skill that would be useful for the type of model being evaluated. The pre-training/fine-tuning process lends itself nicely to benchmarking measures, as the limited-size benchmarking dataset can be used to fine-tune a very large model. The GLUE benchmarking suite is a collection of tasks for language models, and each of the nine datasets is focused on evaluating a linguistic skill, from sentiment to paraphrase to inference Wang et al. (2018). Although other language model benchmarks have been proposed, GLUE remains a key set of measures for evaluating language model performance, and, at the time of writing this, nine of the ten top performing models are transformer-based models, and all ten follow the pre-training/fine-tuning paradigm NYU-MLL (2024b).

However, several recent works bring into question the source of this impressive benchmarking performance. Sinha et al. (2021) showed that RoBERTa models perform comparably to the standard model when trained with randomly shuffled word order. Zhang (2021) found that using English-trained input embeddings with BERT models trained in another language has surprisingly little effect on the GLUE performance, as does shuffling the indices of trained input embeddings. Alajrami and Aletras (2022) found that RoBERTa models with non-linguistic pre-training tasks, described in Section 3.4, still perform surprisingly well on GLUE tasks.

The unexpected performance on the GLUE tasks found in these papers suggests that the benefit these models get from pre-training may not be due to generalized linguistic knowledge. Sinha et al. (2021) suggest that the models may simply be learning complicated distributional co-occurrence statistics, while Zhang (2021) suggest that the linguistic information may actually be contained largely within the trained input embeddings. We instead propose that the benefit of the pre-training process may be related to the geometric characteristics of a latent space resulting from training a model with the transformer architecture, divorced from any specific content knowledge.

Several works have shown that increasing the isotropy of a language model’s latent space, or causing the resulting representations to be more uniformly spread in all directions, may increase performance on benchmarking tasks Mu et al. (2017); Liang et al. (2021); Kaneko and Bollegala (2020); Gao et al. (2019). The motivation behind this work rests on the finding that the latent spaces of these language models are highly anisotropic, with all of the data concentrated in a narrow “cone”, rather than spread throughout the available space and dimensions Mimno and Thompson (2017); Mu et al. (2017); Ethayarajh (2019). By adjusting the models to have more isotropic latent spaces through training or retrofitting methods, these works have attempted to cause the models to use more of the available space, producing some improvement on selected benchmarking tasks. However, there is a growing body of evidence that an isotropic latent space may not be as beneficial as initially thought. Ding et al. (2021) suggest that the benefits found in earlier work may be isolated to static word embedding models, like Word2Vec and GloVe. Rudman and Eickhoff (2023) find that increasing isotropy is detrimental to contextual model performance when measured with a more reliable measure than used in previous work Marbut et al. (2023); Rudman et al. (2021). Cai et al. (2020) and Rajaee and Pilehvar (2021) find that increasing overall isotropy is not beneficial in contextual latent space models, but that increasing isotropy within local clusters may improve model performance.

Similarly, it has been found that applying quantization methods during model training or fine-tuning may improve benchmarking performance as well Liao et al. (2020); Sablayrolles et al. (2018); Hu et al. (2022). Quantization is a compression method in which high-dimensional data are typically clustered within some number of subspaces that can be recombined to recover the full-dimensional data with minimal loss Gholami et al. (2022). Minimizing this loss during some part of the model training could result in a latent space that is more easily separable, which may prepare data to be more easily partitioned during fine-tuning on downstream tasks.

To explore the question of how the geometric characteristics of a contextual language model’s latent space may be related to benchmarking performance, we produced a series of 170 BERT-type models with varying levels of noise added to the pre-trained weights of all the encoder layers, leaving the pre-trained input embeddings un-noised. We then examined the relationship between the GLUE benchmark performance of these models and the values of measures intended to approximate characteristics of data spread and separability in the models’ latent spaces. Finally, we explored the possibility of using these measures to predict the surprising GLUE performance of non-standard models from the literature.

In this work, we address the question of whether there is some measurable geometric characteristic to the latent space of transformer models that could be used to predict performance on downstream tasks, specifically the GLUE benchmarking tasks. We produced models with gradually decreasing GLUE scores by transitioning the weights of pre-trained language models to random noise, and then we applied a set of measures chosen to approximately capture the amount of data spread and separability present in the resulting models’ latent spaces.

In examining the relationship between these measures and GLUE benchmarking performance, we identified a quantized cell density measure, based on the patchiness index Wade et al. (2018), that has a strong linear relationship to the GLUE score. We also considered whether any of these measures could be used to predict the unexpected GLUE performance found in a body of work exploring non-standard training tasks and model manipulations. Although the evidence was less clear, we once again found that the patchiness measure showed promise as a predictive measure for the GLUE score of most of the non-standard models that we examined.

This research is potentially relevant to the growing body of work attempting to reduce the resource requirements for pre-training large language models. If a measure of separability within a model’s utilized latent space could be used to initialize a transformer model before any linguistic training begins, it is possible that the pre-training requirements could be greatly reduced without sacrificing performance on downstream tasks.

Our experiments also uncovered interesting non-linear relationships between many of the measures that we examined and average GLUE score. Further exploration of these non-monotonic patterns could shed light on the inner workings of transformer models, contributing to work on increasing the general interpretability of transformer models.

Our experiments were limited by two major difficulties: (1) The practically infinite nature of a contextual model’s latent space, and (2) the lack of established measures for describing high dimensional data distributions. To overcome these obstacles, we had to accept approximate solutions to both.

Our measure calculations are based on a sample of the latent space that we showed to be representative of a sizable corpus, but is still necessarily an imperfect approximation of the theoretical complete latent space. Additionally, and perhaps more relevant to our findings, the contextual nature of the latent spaces required manipulation of the transformer model itself to approximate a gradual change in the geometry of the latent space. As briefly discussed in Appendix B, some of our results clearly indicate that adding noise to the model weights is not the same as adding noise to the latent space representations, which would perhaps provide a clearer signal for our intended experiment.

While the data spread measures that we examined are fairly well understood, our use of quantization algorithms to approximate some measure of the clustered nature of the data introduces a bit of a “black box” into the measure calculations. The algorithms are designed to effectively segment the data for optimal reconstruction, and it is unclear how this goal interacts with some of the measures that we apply. For example, the quantized cell density measure (for which we found the most compelling results) is based on the distribution of the number of points assigned to each centroid in the quantized space. However, it may be the case that the quantization algorithms themselves are in direct conflict with this measure, specifically attempting to adjust the location of the centroids and subspaces to even out the distribution of points assigned to each cluster. Also, the number of quantized components may not adequately describe the complexity of space utilization – imposition of too many or too few components may play a role in some of the observed non-monotonic relationship between density measure and GLUE score. We do not believe these to be issues for our experimental results, but it does hinder our understanding of how our measures are related to the latent space geometry.

Additionally, we remain unsatisfied that any of our measures quite captured the concept of separability that we intended, somehow quantifying the extent to which data is located in “nooks and crannies” in the latent space. We point this out not to discredit our work and results, but to again highlight the fact that all of these measures are approximations of complex geometric relationships. It is certainly possible that a different measure exists that would better capture the relationship between GLUE performance and latent space organization, and it is perhaps equally possible that the relationship is too complex to be fully approximated by a single measure.

Finally, with the future goal of using this paper’s results to reduce pre-training requirements for transformer models, we come up against the major issue that existing quantization methods are non-differentiable. This means that none of the quantization-based measures explored here can be used as an additional or initial loss function. However, we do intend to explore the possibility of using differentiable k-means algorithms in place of the quantization methods used here to address this problem in future work.

We thank Katy McKinney-Bock, Conner Copeland, and Daniel Olson for enlightening conversations during the development of this project. Computational analysis was performed on Jetstream2 resources at Indiana University through allocation BIO220047 from the NSF ACCESS program, which is supported by OAC Grants Nos. 2138259, 2138286, 2138307, 2137603, and 2138296.