This is the first entry in a sequence of posts which compare a mathematical theory of attention against trained transformers.
Links:[The code for these experiments.][GitHub repository]:[Geshkovski et al., the theory this investigation is built on][Original paper]:[a video going over the original paper][YouTube walkthrough]: These posts describe experiments building on the paper *"A Mathematical Perspective on Transformers," *by Geshkovski et al. This background is going to give the definition of a cluster and describes the structure of the activation space that the self attention mechanism builds by virtue of its architecture. Please check out the paper or the YouTube video.
The paper focuses on self attention. Layernorm between self attention layers means that our transformer's residual stream can be considered on the surface of the hypersphere of dimensionality . Attention (from equations 2.4 and 2.5 in the paper) is:
With matrices set to identity:
The inverse temperature sets how sharply attention concentrates. So, instead of seeing the QK and VO circuits as a set of pipes that move information around, we now see attention as a set of particles from some initial conditions being placed on a hypersphere and then interacting as the attention progresses through the layers of the network. Our transformer is a particle system.
The paper's perspective is that the interacting particles are tokens from the prompt passed through the embedding layer onto the hypersphere, then interact with each other on the surface of that hypersphere. This is the structure of our activation space. Initially, the particles are going to be distributed based on whatever the embedding layer has learned and from adding on the positional encoding. The paper proves that if you had infinite layers, all of these particles would collapse down to a singular point (if all matrices are identity) somewhere on this surface. In the paper's Proposition 3.4, we show that we have an energy that we want to maximize defined as:
Maximizing this term -> maximizing the inner product in the exponential -> the particles have a more similar inner product -> the particles are approaching each other as they evolve on the surface of the hypersphere. The particles cluster.
The interesting thing happens between placing the particles onto the surface of the hypersphere and passing them through an infinite amount of layers. The particles do not isotropically collapse down to that singular point after infinite layers, they clump and form discrete clusters in the activation space. This intermediate phase of 'metastable states' we call clusters is the setting we investigate here.
The architecture of the Transformer drives tokens to cluster and then collapse as they move through the layers of the network.
The weights of the neural network learn to resist clustering and prevent collapse during training.The Transformer functions by resisting its own architecture.
Experimental Setup
[machine generated]
**GPT-2:**
Chosen for its breadth across scale, its familiarity, and its expressive power. GPT-2 small is extremely popular in the mechanistic interpretability field and is regarded as a model small enough to be easily experimented with but large enough to be seen as a 'real' model (not a toy model). The entire family was used to have a variety of depths and dimensionalities.
GPT-2 Large was used twice: once with trained weights, and once with random weights. Large was used instead of small/medium because large has different dynamics than medium.
GPT-2: | Layers | Dim | Small | 12 | 768 | Medium | 24 | 1024 | Large | 36 | 1280 | XL | 48 | 1600 |
**Albert: **
Chosen because it is the model used in the original paper, and because its weight-sharing architecture makes the dynamical analysis structurally cleaner. Albert is a single transformer block applied iteratively: there is one shared weights matrix, so the properties governing clustering are the same at every depth. It is also the model used in the paper.
Albert v2 Base (L=12, 24, 36, 48) was used twice: once with trained weights, and once with random weights. No version of Albert other than v2 is used in the post.
Albert: | Layers | Dim | | v2 Base | 12*, 24, 36, 48 | 768 | | v2 XL | 12, 24*, 36, 48 | 2048 |
The * represents the length the model was trained on.
Bert:
Chosen because Bert's bidirectional masked-LM pretraining produces a different routing structure. This makes Bert a good test of whether the theoretical framework is architecture-agnostic or specific to causal models.
Bert (uncased): | Layers | Dim | Base | 12 | 768 | Large | 24 | 1024 |
see this link for full config.
One important question to ask is whether token clusters exist outside of the toy models given in the paper. As we will soon show, the answer is a definite
yes.The clusters are a real phenomenonin the activation space of the transformer, although the space is nottotallydefined by clustering. Tokens are clustered ~50% of the time (see figure 8), and that clustering is persistent across layers and prompts.
We analyzed models with trained weights and with randomly initialized weights with norms equal to the trained values. We ran experiments across the GPT-2 family, Albert v2 XL and v2 base, as well as Bert. We will primarily focus on GPT-2 Large for the comparison, as it was the most interesting model in preliminary analysis, and focusing on one single model helps with brevity.
We lead in with some visualization of the activation space given by PCA, TSNE, and UMAP. We define the clusters via HDBSCAN. Under certain situations models will cluster particles together and eventually collapse. However, this assumes that the weight matrices are set to identity. Are we able to measure clustering in models with their matrices set to random values, or with trained weights? Do the dynamics of the model reflect the mathematical guarantees outlined in the paper?
We start with the strongest visualization to persuade the reader that we do, in fact, see clusters in real models.
Figure 1: The projections struggle to meaningfully spread the activation space in the random model. The colors represent different clusters labeled by HDBSCAN (a discrete cluster labelling tool). The gray 'x' marks denote unclustered particles labeled by HDBSCAN. Colors are not consistently the same cluster through depth.
Figure 2: A similar figure, but for the trained GPT-2 Large. There are visibly more clusters containing fewer particles.
HDBSCAN Clustering Explanation (Optional)
[Machine Generated] Clustering asks the token cloud a specific question: are these particles organized into distinct groups, and if so, which particles belong to which group? We don't want to fix the number of groups in advance (that should be discovered, not assumed), and we need "this token doesn't belong to any group yet" to be a valid answer rather than an error. HDBSCAN gives us both.
The starting point is DBSCAN. A cluster is a dense region separated from other dense regions by sparse space. A point is a "core point" if it has at least minPts
neighbors within distance ; clusters grow outward by linking neighboring core points; points reachable from a core point but not core themselves are border points; everything left over is noise (label −1). This handles irregular cluster shapes for free, since density is a local property, and it leaves room for tokens that are mid-transition instead of forcing every particle into a group.
Plain DBSCAN uses one fixed density threshold everywhere on the sphere, and that's the problem for us: a threshold that resolves a sparse cluster will merge two nearby dense ones, and a threshold tight enough to keep dense clusters apart will erase the sparse one. Since we don't know in advance how tight or loose a metastable cluster of tokens will be at a given layer, one fixed threshold is guaranteed to be wrong somewhere.
HDBSCAN (McInnes and Healy, 2017) removes the fixed threshold by building the full density hierarchy instead of committing to one cut:
This leaves dense-region distances alone, since points there are already closer than their core distances, and inflates sparse-region distances outward, pushing isolated points further from everything.
The result is cluster assignments for tokens in stable dense regions, a noise label (−1) for tokens in low-density regions or mid-transition, and a cluster count that was discovered, not specified.
We run this directly on the sphere, using cosine distance () on each layer's L2-normed activations, with min_cluster_size=2
, permissive enough that even a pair of tokens can register as its own cluster. That's deliberate: we're not hunting for a handful of large semantic categories, but instead trying to detect the metastable states from the paper's theory wherever they show up. The noise label matters here too. During metastability some tokens have committed to a cluster and others haven't, and a layer where 30% of tokens land in noise is a layer where clustering is mid-transition. Tracking that noise fraction layer to layer is one of the more useful diagnostics in this analysis.
Here is a statquest on DBSCAN, and another video of full HDBSCAN. We notice that the structure of the trained model under these projections is far more isotropic than the random model. This will be a common theme we see across the experiments, as **the random model will exhibit strong tendencies towards collapse, while the trained models resist this collapse. **
The clusters of the random model are far closer together, while the trained model separates the clusters to be more even across the activation space. The unclustered particles only appear in a specific subspace of the random model, while they are more evenly distributed across the trained model.
Similar projections with Albert Base (Optional)
The clustering is visually more clear in Albert base, and in the random model we see near total collapse. The projections make the space appear more spread than other figures will later: we could not use the Albert base random for many analyses because its utilized activation space collapses down to under 2 dimensions. GPT-2 Large will be the primary model we analyze going forward.
Figure 3: Mass near 1 across several models. Mass near 1 is the fraction of inner products between particles greater than 0.9 in each layer. Layer depth is normalized because each model may have a different number of layers. A value pinned to 1 means that the model's activation space has collapsed to a single point, and the unembedding matrix will only resolve whatever token that degenerate point is associated with. Mass near 1 does not tell the entire story, but we will expand on inner products below.
Mass Near 1 Explained (Optional)
[Machine Generated]
Layer normalization via root mean square (RMS) is used by every model. It divides each token's residual-stream vector by its L2 norm before passing it to the next layer. After this operation, *every token vector ** satisfies *. The token cloud lives on the unit sphere:
On the sphere, the natural measure of similarity between two points is their inner product:
where is the angle between the two vectors. The inner product is the cosine similarity when both vectors are unit length.
The range is :
For a layer with token particles, there are distinct pairs. Compute every pairwise inner product and put them in a histogram. The x-axis runs from to . The y-axis is density and integrates to 1. Reading the histogram across layers:
Early layers build Gaussian about 0.In high dimensions, random vectors on the sphere concentrate near the equator. By concentration of measure, most pairwise inner products between randomly-oriented high-dimensional unit vectors land close to 0.
Clustering in progress: values growing at 1.As layers push similar tokens together, a subset of pairs achieves high inner product. The histogram develops a second peak migrating toward 1 while the main mass stays near 0. Each such peak describes clustering/collapse: the pairs at high inner product are tokens that have been pulled together.
Full collapse: all values at 1.When all tokens have merged into a single point on the sphere, every pairwise inner product equals 1. The histogram is a single spike at 1. This is what the theory's long-run prediction (consensus / a single Dirac mass) looks like empirically.
Tracking the full histogram across all layers, models, and prompts produces many plots. We simplify the histogram into a scalar: Mass near 1 is the fraction of pairs with inner product above 0.9.
This is a direct, threshold-based read on "how much of the token cloud has clustered so far." A layer where mass near 1 goes from 0 to 0.4 in one step is a layer where 40% of pairs snapped into high agreement at once, signaling a merge event.
The inner product denotes particle proximity. A consistent value means that some number of tokens are consistently grouped together, although we cannot immediately tell if they are maintaining consistent groups over the layers from this value alone. Clustering that is stable across layers would lead to this consistent value over time, as we see (more or less) in the GPT models. Clustering that merges with other structures and leads to total collapse, as is outlined in the paper, is more reminiscent of the Albert models.
For Albert base, we see at a depth of ~0.8 (layer 20 or so) that the model experiences total collapse. All inner products rise to 1 and do not recover, meaning that they all rest in the same location. When the model is unembedded, the only output will be a repeated token or two, over and over again. All Albert models will eventually collapse if run for long enough, as they do not have a predetermined length. They are simply models with the same attention block applied again and again, so all particles are attracted to the same point forever. The model can only resist this for some length of time. They were trained to a normalized depth of 0.5, so going beyond this is out of distribution for these models.
We see that the random cases collapse nearly immediately, while the trained models resist that collapse. This is the primary result of the experiments: **the architecture is driven to collapse the activation space, and the weights learn to resist that collapse. **The clustering we see is a result of the architecture as well as the collapse. The weights are explicitly trained to minimize loss, so the expansion/resistance to collapse must be driven by the model's need to accurately predict the next token. Any space not dedicated to the computation of the model is automatically collapsed, and has the capacity to be reclaimed. GPT random does not completely collapse because there is some that each particle is attracted to per layer, and each layer is likely to have a different . This moving target prevents total collapse, where Albert maintains only a single .
Figure 4: Histogram of the random model's inner products. We can see that the first layer for this initialization has a rather tight gaussian, and by layer 5 has pushed the majority of inner products past the 0.9 threshold. Under random initializations, the model is driven to push all particles together.
Figure 5: Histogram of the trained model's inner products. We can see that the initialization has far greater variance to begin, and does not send all inner products (demarcating particle distance) to very high values like the random weights did.
These are the inner products of the trained vs random GPT-2 Large designed to isolate the dynamics of the random model. We can see that the random model, similar to the model with weights set to identity in the paper, drives the particles to cluster and fall into a degenerate state of very low (if not singular) dimensionality. **The weights of the trained model resist the collapse, while the randomly initialized models drive towards collapse. **
Figure 6: The inner product of particles divided by interaction type: whether the interaction is between particles of one cluster, two clusters, cluster/unclustered, and unclustered/unclustered. IP means inner product. Noise means an unclustered particle. Notice the energy given beta is set to 1. We defined energy in the summary above, and will return to it later as well.
Figure 7: The same figure given for the trained GPT-2 Large instead of the random model. We can see that the graph above (mass near 1) was dealt primarily with particles in the same cluster in the trained case.
In this view of the inner products, we measured the relationship between particles in as well as out of clusters as defined by HDBSCAN. The first graph given is the mean of our sets of inner products, the second is very similar to the mass near 1 above but focused on whether particles are or are not clustered, and the third shows the population size.
One point of interest is the second highest set of inner products in the second subgraph: the random model has noise/noise interactions as the second highest value, while the trained model has cluster/cluster particles interacting second most. Although the values are not directly comparable, it appears that the cluster/cluster particles maintain a similar value while the noise/noise particles become farther apart in this trained setting.
We again see that the random case has tighter clusters and remains closer to all other particles than the trained case. The model has learned to reduce inner products within its cluster as well push all other particles away compared to the random models.
Figure 8: We analyze whether the trained or random models have more particles in clusters. The HDBSCAN output is compared with in cluster versus out of cluster particles and analyzed in GPT across layers.
trained vs random
We can see that the trained model tends towards the lower end of variance of particles in clusters, compared to the random model. We also see the largest variance in the random model out of all of the other visualizations that we have seen so far. This highlights the discrete nature of HDBSCAN and the fact that clusters technically do not have a hard edge as we are using in our visualization to discern clusters.
The trained model appears to have fewer particles within clusters according to the clustering method that we use, although this is within variance of the random model. This aligns with figures 6 and 7 above showing that the clusters have a higher inner products in the random case compared to the trained case. This also aligns with what we see from effective rank with the expansion of rank as the model continues forward.
This follows with the idea that we have seen so far that the trained network wants to resist the architecture's drive towards collapsing particles into clusters. The clusters are expanded outwards and are generally farther apart from other clusters in the trained case. We will also see later (in cluster count across all models) that we simply have a larger number of clusters with smaller numbers of particles within those clusters in the trained case.
Figure 9: the cluster counts across all models. The number of HDBSCAN clusters are plotted for all models across depth.
Interestingly, all models have a similar number of clusters. One might initially think that some models might have more clusters than others due to different tokenization methods, model size, prompt length, etc., but that is not the case. It is consistently near 50 for all trained models. Are the clusters doing something similar across models? Is there some universality to the clusters in a similar way that we see a quasi-universality to the dictionary entries of a sparse autoencoder?
We can see that the random models have fewer clusters, and we saw from the inner product values above that those clusters are more tightly packed. We also know that a similar number of particles are in clusters, so we can deduce that there are a greater number of particles per cluster in the random case. The model learns to reduce the particles in the clusters and creates numerous small clusters with greater internal distance in the trained case.
Albert random appears to be a special case, but upon inspection one sees that the effective dimensionality plummets to zero. The rise and subsequent fall of the total cluster count is associated with a collapse in the activation space. HDBSCAN is having to discern what is and is not a cluster as the space falls to essentially a degenerate point mass. Before collapse, we can see that the cluster count sits around 35, just above the norm of GPT random.
Figure 10: The Effective rank across several models. The effective rank is a measure of how many dimensions the model uses within the residual stream. We can see that the models like to use ~250 dimensions according to this measurement despite the models having various activation space dimensionalities from 700-1600. Interestingly, Bert base generally maintains the highest rank of all the models. The GPT models tend to increase their values as the layers progress, and the Albert models decay downwards. The models with random weights collapse quickly down, signifying that the raw architecture wants to compress the space as it attempts to place all particles into the global attractor .
Effective Rank (Optional)
[Machine Generated]
Inner products tell you about pairs. Effective rank tells you about the whole cloud at once: how many independent directions is the token cloud actually using, out of the dimensions available? You can have plenty of pairs at high inner product without the whole cloud collapsing: two well-separated clusters give zero cross-cluster pairs near 1, but the cloud is still only using about two directions.
Effective rank globally measures whether the whole cloud is collapsing onto a low-dimensional subspace. We use the entropy form (Roy–Vetterli 2007): we normalize the singular values to a distribution and exponentiate the Shannon entropy. One direction -> rank 1 (a point). A flat spectrum -> rank d. It is smooth and threshold-free.
Take the activation matrix (n tokens, d dimensions per token). Compute its singular values . Convert them to a probability distribution by normalizing, then take the exponential of the Shannon entropy:
This is the effective rank of Roy and Vetterli (2007). Entropy measures how spread-out a distribution is. A distribution concentrated on one value has entropy 0; a uniform distribution over values has entropy . Exponentiating maps entropy back to a "number of effective components."
Applied to the singular value spectrum:
All weight on one direction(): , entropy , effective rank . The cloud is essentially 1D. This is a line through the origin, meaning all tokens are nearly identical.
Spectrum flat(): entropy , effective rank . The cloud uses all directions equally.
In between: smooth interpolation. A cloud where a few directions dominate gets an effective rank near that count, not some artifact of whatever threshold you happened to pick.
The key advantage is that it's smooth and weights directions by their share of the spectrum. A direction carrying a negligible fraction of the total barely contributes, even though it would count as a "nonzero singular value" in a threshold-based count.
Effective rank tends to start high at early layers (many directions contribute roughly equally) and fall as clustering progresses (token representations align, the spectrum concentrates in fewer directions). The endpoint depends on architecture:
Effective rank tracks the full distribution of the spectrum, not just the high-agreement pairs like the inner products from before. A model with two perfectly separated clusters of equal size will show mass near 1 = 0 (no cross-cluster pairs near 1) but effective rank ≈ 2 (two principal directions). They're measuring different things and can complement each other from their local and global views.
When effective rank drops low enough, the token cloud is nearly a point-mass. At this point:
Rather than report these unhelpful quantities in the analysis, we gate them on effective rank and suppress them below the threshold.
The fact that all models have similar effective ranks despite having very different dimensionalities is very strange. This may be an artifact of the embedding matrix's structure, although no investigation has yet taken place.
We can still see that the GPT family tends to take an immediate drop as soon as the tokens enter the model. This is likely the movement of positional encoding information and the routing of particles into the regular dynamics of the model learned during training, and caps out by layer 5. Curiously, all members of the GPT family take a slight rise in rank, then drop down again around a third of the way through all layers. The rank then slowly increases until it reaches the end of the layers, then abruptly drops.
Albert takes the opposite route: The effective rank rises sharply, then nearly plateaus through the first half of the model. The effective rank then begins to plummet. The trained length for Albert base was 12 layers, which here correlates with the moment the rank begins to fall. This shows that the weights maintain the rank for the entire length it was trained upon, and only after does collapse begin. Albert XL is a similar story. The collapse we see in both Albert models is an artifact of them running longer than intended.
The random case follows the theory almost exactly. The particles are placed onto the activation space (our hypersphere from the theory), then begin to collapse towards a single point. Albert seems to rise before experiencing total collapse, which leads one to wonder if the embedding matrix, the tokenization, or the positional encoding play some part in the effective rank of the model. The trained models maintain effective rank, preventing the collapse the architecture is driving. The random models have not learned to resist this collapse.
Figure 11: The Fiedler values of the attention matrix mapped out for GPT. The value seems to dive low, then rise steadily as the network progresses through the layers. Notice that the random network maintains a relatively high, stationary value.
Figure 12: The Fiedler values mapped out for Albert. This graph also shows a great deal more separability in the trained case. The raw values are not comparable across architectures. We can see that Albert trained maintains a much lower Fiedler value than its random counterpart. The value only rises once the model is run longer than it was trained for.
The values given here represent a global view of who attention routes to, and are calculated from the routing graph . We are visualizing the second smallest eigenvalue, which signifies how densely connected the graph is.
* A low Fiedler value* for our circumstance means attention has organized into small groups of tokens that mostly attend within their own group.
* A high Fiedler value* means that every token is allowed to attend to every other token, similar to uniform mixing. We can see that the trained models have lower Fiedler values than the random models.
Fiedler Value Definition (Optional)
[Machine Generated]
A Fiedler value measures how hard it is to split a graph into two disconnected pieces. The graph here is the attention matrix at a given layer, not the residual stream: nodes are tokens, edge weight is attention paid from one token to another, after Sinkhorn-normalizing that layer's attention into a doubly stochastic matrix: row-normalize, then column-normalize, repeat until both hold:
This is computed separately for each attention head, then averaged as "Fiedler value (mean across heads)" in the plots below.
Ask this of that graph: what is the cheapest way to cut it into two groups, where "cheapest" means breaking the fewest and weakest connections? If attention has already organized into two groups that mostly attend within themselves, the cut is nearly free. If attention crosses every boundary with roughly equal strength, any cut is expensive.
The Fiedler value is the number that captures this cost. First symmetrize the attention graph:
then build the normalized Laplacian from its degree matrix:
The Laplacian's eigenvalues, ranked smallest to largest, describe the graph's connectivity:
is always zero: , the constant vector is always an eigenvector — and is a fixed property of any graph, not a measurement. The Fiedler value is , the first number in the list that actually reflects the graph's shape.
Read on the model: a low Fiedler value means attention has arranged itself into at least two token groups that barely attend to each other. A high Fiedler value means attention is close to uniform with no fault line anywhere. A high value has two possible causes that look identical on the plot: a model that never learned a fault line (the random baseline), or a model whose fault line existed and then collapsed into uniformity. Worth keeping separate when reading the trained curves below.
This gives a fourth, independent way to ask what the rest of this section has been asking: does the network resist the collapse the architecture pushes it toward? HDBSCAN and the inner-product histograms geometrically describe where tokens sit on the sphere. Effective rank describes how many directions that geometry occupies. The Fiedler value describes something else: not where tokens sit, but how they're allowed to talk to each other.
These two graphs tell a more interesting story than what one may see at first glance: the architecture with random weights allows a more general attention (anything can attend to anything), while the trained weights cut away this global view and limit who can attend to who. The trained weights learn to deny which tokens attend to one another.** **They do not learn which tokens should attend to which other ones. This would mean that training a network is closer to subtractively carving a sculpture from a large stone rather than assembling attention brick by brick. The algorithm wants to have uniform global connections, but the network learns to resist the raw architecture seen in the random case by learning where to cut connections and make the subgraphs more easily divisible.
Figure 13: CKA seen across Albert and GPT. Albert base stops appearing when effective rank falls under 3. CKA measures kernel similarity between layers and shows the persistence of clusters across the network as opposed to their existence in a single layer. Notice that the absolute difference in all models is small.
Centered Kernel Alignment (CKA), Optional
[machine generated] CKA asks whether the layer-to-layer map is approximately the identity over some set of layers, meaning the pairwise-similarity structure of all tokens is preserved from one layer to the next. It is the primary plateau signal because it is basis-free and scale-invariant: it does not care how the residual stream rotates or rescales, only whether the relational geometry is the same. A flat run near 1 is a plateau; a sharp single-step drop is its end.
We should expect the particles to be placed on the hypersphere and then quickly cluster. These clusters will persist through the layers, and because there was the period of rapid change into clusters before settling in, we should expect to see a plateau form. This signifies the persistence of clusters.
At layer you have an activation matrix (sphere-projected). At layer you have , same shape. You want** one number saying how similar these two representations are** geometrically.
The natural object is the representational similarity matrix (RSM): the Gram matrix , whose entry is . Two representations are considered the same geometry if their RSMs agree, regardless of how the underlying coordinates are oriented. This is exactly the inner-product matrix from Group A. CKA is built on top of the same object.
The Hilbert-Schmidt Independence Criterion measures whether two sets of features are statistically related. For linear kernels and , the unnormalized linear HSIC reduces to
after centering. F is the Frobenius norm. This gives us the alignment between the two RSMs. it is large when pairs of tokens that are close in are also close in .
HSIC on its own scales with the magnitudes of and , so it is not comparable across layers. CKA normalizes it by the self-similarities:
The result is in : 1 means the two RSMs are identical up to an orthogonal transform and isotropic scaling. 0 means they are orthogonal.
This is 1 when the two RSMs are identical up to rotation and uniform scaling, and 0 when they're orthogonal. This is the property that makes CKA the right instrument:
So a CKA plateau near 1 means the model is rotating and rescaling the token cloud from one layer to the next but not reorganizing which tokens are close to which. That's precisely the metastable picture the whole post is looking for: clusters fixed in place, dynamics idling.
We see a large drop in Albert around 0.5, where it normally would end during training. Close to the end, we see both trained models falling ~5%. The changes are small in magnitude and the CKA is nearly 1 for all models. Each layer is very consistent with the one before, implying that the clusters are consistent as well.
Figure 14: the energy landscape of Albert. We can see that the random case monotonically rises, and then peaks at some value. This consistently happens near layer five. Looking at the scale of the Y axis, we can see that the inverse temperature beta plays a strong role in the total amount of energy that the network would have during training. Troughs can be seen where the network would stop during its training.
Figure 15: the energy landscape of GPT. We can see once again that the inverse temperature beta controls the total amount of energy that the network contains. Similar to the Albert case, we see that the energy has a similar shape across all temperatures tested, and the random case monotonically increases in energy.
Energy monotonically increasing is of paramount importance in the paper. In the paper, if V is set to identity, then all interactions will be attractive. The energy is maximized whenever all points fall into a degenerate state and minimized in the uniform case, so the progression of the model is necessarily associated with a monotonic rise in energy. This is explicitly true for when , and is empirically true from what we can see in the random case.
However, **the trained case does not have monotonically rising energy. **This is true for every model, prompt, and temperature tested. This means that there are more dynamics in the model than purely attractive forces that will collapse the model to a single degenerate point. These dynamics are the starting point for the model to learn how to minimize the loss and simulate natural language. The architecture is designed in such a way that it will experience collapse, but during training the weights learn to avoid monotonically increasing the energy in order to reduce loss and prevent collapse.
For the bulk of both models under both conditions, the energy is only slowly ascending or descending, meaning that energy is more or less preserved as the layers progress. This leads to interesting questions about symmetries within the activation space. If we force energy to be preserved, what symmetries will arise? Figure 16: we analyze the magnitudes of attention to see whether particles inside of clusters receive a different amount of attention than those outside. The first graph shows the proportion of attention that each group receives per layer (mean attention per particle summed for group). The second graph compares the amount of punctuation or Other tokens, which are extremely frequent per prompt. The third one shows the unclustered population.
Figure 17: we see a very similar graph to figure 16, but for the trained case instead of the random case above. It appears that there is a meaningfully higher noise fraction than the random case, but in figure 8 it is shown that the trained case is below the random mean but within random variance for unclustered particles. Interestingly, the probability of particles being associated with punctuation is similar in both.
The trained case is interesting: The values are about equal up until layer seven, but then the unclustered particles ascend in attention weight to ~1.8 (representing almost 90% of attention despite composing ~50% of tokens) until the end of the network. Attention is allocated differently in trained models.
In the random case, we can see that clusters tend to receive a slight majority of the attention in the layer. This could be a reflection of the fact that the randomly initialized model will have certain particles which have a higher than average V, and they are more likely to attract other particles towards them. This would result in the clusters being preferred by the attention mechanism. The trained case may be different if a fixed position token (e.g., [sep] or [bos]) receives a disproportionate amount of attention and is trained to be unclustered.
This opens the question of what role the clusters play semantically/syntactically, and will be the subject of future posts.
We have shown that the theory of the paper holds for both random and trained matrices of the Transformer, but the trained case has interesting caveats. The energy does not monotonically increase in the trained case and does not experience collapse unless driven out of distribution, as shown with Albert under long runs it was not trained for.
The architecture drives the particles to collapse into subspaces which we have called clusters, then collapse everything to a single point. We have shown that the trained weights resist total collapse, and the clusters formed are smaller and more numerous.
Many questions remain. Do the important computations of the language model correspond to clusters? Is there a universality to clusters across trained models? We know that induction heads and bigram statistics exist: will these structures be found in the unclustered particles of the activation space? Can we find symmetries in the latent space from the near preservation of energy?
Next post coming soon!