Uncovering Layer-Dependent Activation Sparsity Patterns in ReLU Transformers

While attention mechanisms are the feature that differentiate Transformers from other sequence models, the Multi-Layer Perceptrons (MLPs) within a Transformer block seem to be key to how they learn and store information. MLPs contain the majority of parameters within most modern language models, and it is believed that they are central to how models store factual information[4]. Given this, we believe that rigorous work to better understand the internal machinery of Transformer MLPs is of potentially high value to our understanding of this important class of models.

Previous work demonstrated that MLPs that use ReLUs tend to be sparse in their activations - that is, for any given token vector that makes a forward pass through a MLP, only a small fraction of the available hidden units will have values greater than 0 after a ReLU activation is applied [2]. This work additionally focuses on sequence and batch measures of sparsity. We believe these metrics are important both for understanding the degree to which sparsity can or can’t be computationally exploited, and also for understanding what sparsity might be able to tell us about features being learned.

To investigate this, we define and calculate three core metrics of sparsity - or, inversely, of “hidden unit use”. (We decided that focusing on the positive measure was more intuitive in this setting, but this is fundamentally just sparsity from a different perspective). These metrics are calculated on the hidden dimension activations between the two dense layers that make up a Transformer MLP, and all of them are calculated independently for each layer in the network.

1. 1.

   Per-Token Hidden Unit Use - The number of hidden units that are nonzero post-ReLU, calculated per token, and then averaged over sequence and batch.

2. 2.

   Per-Sequence Hidden Unit Use - The number of hidden units that are nonzero post-ReLU anywhere in a sequence. This is fundamentally an “or” over all token in a sequence - if a hidden unit is used in even one token, it is included in this count. Experiments in this paper are done with a standard sequence length of 512.

3. 3.

   Per-Batch Hidden Unit Use - The number of hidden units that are nonzero post-ReLU anywhere in a batch. This is similar to Sequence Hidden Unit Use in that a hidden unit dimension is counted here if any token in any sequence in the batch has a value greater than zero. Metrics are calculated per-device, with a device batch size of 16, meaning a batch for this purpose contains 16*512 = 8192 tokens.

We additionally calculate percentile metrics of how frequently over the course of a sequence a hidden unit neuron is used, when it is used at all. These are calculated by first taking the hidden units that have a nonzero value anywhere in the current sequence, and, for each of them, counting what fraction of tokens in that sequence each hidden unit is used on. We then take different percentiles of that distribution. This can intuitively be thought of as taking a histogram of how frequently in a sequence a hidden unit is used, removing values of zero, and taking different percentiles of that histogram, for the sake of making a metric that can aggregate more easily over training.

For all the plots shown in this paper, we chose to focus on the early stages of training, since behavior tends to stabilize by about 70,000 steps, and showing a full plot makes the early-in-training dynamics harder to see. Full plots are available in the appendix whenever a shortened version is shown.

The primary subject of investigation in this paper is a 6-layer decoder-only Transformer model, with a hidden dimension size of 32768 and ReLU MLP nonlinearity, trained on C4 next token prediction [6]. LeCun Normal initialization is used for the dense kernels within the MLPs. An AdamW optimizer is used, with a cosine decay learning rate schedule after a warm-up of 5000 steps up to a peak LR of 3e-3. In section 5 we discuss the extent to which our observations transfer to different model depths, hidden dimensions, and learning rates, but outside of that section all results are shown on the model described above. Each model was trained with GCP TPU-v2 slice of type v2-32, and in this environment took about 2.5 days of wallclock time to reach convergence.

To understand how sparsity develops during training, we need to first understand the sparsity characteristics of an untrained model. In Table 2, we can see that, at initialization, higher layers of the network use monotonically fewer of their dimensions at in the first training batch. Given that weight initialization is shared across layers, this suggests that the transformations being applied even by a random network shift the inputs to each successive MLP in a way that induces sparsity.

Importantly, we can see that a hidden unit being zero in the first batch of training does not mean that it never has the possibility to shift outside of that regime - in the “Turned On“ column of 2 we see that all layers turn on some hidden units that began training unused. The tendency for layers to "turn on" hidden units is monotonic in the opposite direction, with higher layers turning on more hidden units. This could be indicative of the model learning features at higher layers that benefit from a larger number of hidden units, but could also mean that, while all layers would benefit from turning on hidden units, later layers in the network are more able to do so because of the distribution shift happening underneath them in the earlier layers.

Beyond converging to different levels of sparsity, there seem to be meaningful patterns of when in training each layer becomes sparse. Figure 2 shows two metrics early in training: the fraction of hidden units active per token (again, simply the inverse of a traditional sparsity metric) and the median fraction of tokens within a sequence that use a hidden unit (excluding hidden units that never activate in the sequence). These metrics are explained in more detail in 2.1. Both of these values drop early in training, and in particular both drop first in the final layer of the network, then the penultimate layer, and then the lower layers in nearly-monotonic succession from highest to lowest.

Another important note from Figure 2(b) is that unlike the other layers, Layer 0 reverses its trend and gradually uses each of its hidden units more frequently over the course of a sequence, where other layers are moving in the opposite direction. This is especially striking since, in Figure 1, we can see that during this period of training the total number of hidden units being used per batch drops precipitously for Layer 0, meaning that it is using fewer hidden units overall but using the ones it does use more frequently. We do see some tendency for other layers to undergo a drop in batch dimension starting at a similar point (in particular layers 1 and 2), but the degree is much less, and tends to partially recover after an initial dip.

The accepted understanding of how neurons become “dead”, or drop out of active model use, is that the behavior emerges randomly or accidentally, due to outlying batches that contribute a highly negative gradient impulse and knock the dimension into a regime where it stops activating. We believe like these results add complexity to and fundamentally challenge that mental model.

First, as mentioned in the previous section, a neuron or hidden unit being empirically dead within a batch does not mean that it cannot be turned on by the training process. It seems, then, strange to refer to these hidden units as dead, if that death can be reversed. Intuitively, this aligns with the fact that even the earliest MLP in a Transformer sits on top of an embedding layer and an attention mechanism, both of which can change the input distribution a dense layer experiences, and thus shift the distribution of its outputs even if the weights of the vector within the dense layer itself do not change.

Additionally, if “neuron death” were a random or accidental process, one would expect it to happen with similar frequency over the course of training, or at least to happen to a degree that scales to learning rate. Instead, in Figure 1 we see very few dimensions turned off for the first 8000 steps of training - which for this model encompasses the period of highest learning rate before cosine decay begins - and then a dramatic increase after that. We believe these two facts are suggestive that sparsity is being driven by the model’s construction of certain feature spaces to operate in, rather than simply happening accidentally and pathologically.

We hypothesize that some of the sparsity behavior observed here can be explained by something we here call Feature Specificity. To explain the concept, imagine two ways features could be constructed. These are both extremes, and we don’t expect either represents the internal behavior of real models, but we find them helpful models to illustrate a spectrum of possibilities.

1. 1.

   Continuous Shared Subspace In this paradigm, all tokens have informative (nonzero) values along all dimensions, and information is communicated through the continuous differences in the values along each dimension.

2. 2.

   Binary Representation In this paradigm, there are many features, features take on binary values, and the information content of a token is represented by the set of hidden units which are on. In this world, you would expect the number of dimensions that are on (or 1) to differ dramatically from token to token, and would expect most dimensions to only be used in a small fraction of tokens, because if all dimensions were switched on all the time, there would be very little information to differentiate tokens. Intuitively, the lower information content of a 0/1 signal compared to a fully continuous one means that the full set of dimensions available to switch on or off would need to be larger in this paradigm to capture a similar amount of information.

Examining tables 1 and 3, as you rise higher in the network, features appear to shift from being closer to the continuous model to being closer to the binary model. Layer 0 is the most extreme case – it converges to active use of a substantially smaller set of hidden units (1), but without a commensurate drop in number of hidden units used per token. We can also see in Table 3 that hidden units in layer 0 are used much more frequently over a sequence: the 65th percentile hidden unit activates for at least 31% of tokens in a sequence. The next highest layer on this metric, layer 2, activates for only of only 3.2% of tokens, nearly 10x smaller.

On the other extreme, we see layers 4 and 5 behaving in ways that more closely correspond with the binary feature model. They have the highest number of hidden units in active use per-batch (28.2K and 31.3K respectively), but use comparatively few for each token token. Looking at 3, we see that the hidden units in these final two layers activate only very infrequently in the course of a sequence; the 90th percentile hidden units only activate 3.7% and 6.9% of tokens in a sequence respectively. To make this more concrete, for layer 4, out of dimensions that are ever used in a batch, 90% of those dimensions only have nonzero values for 20 tokens or fewer out of a 512-length sequence.

On a very simplistic level, all of this batch-level sparsity seems like a waste: capacity that the model could be making productive use of that it appears to be ignoring. In one framing, the number of dimensions the model converges to for each layer could be a sort of revealed preference of capacity - perhaps for this width or depth or dataset, the capacity needs of some layers are bottlenecked by others.

To investigate this as a possible framing, observed the number of hidden units used at convergence in a standard experimental setting. We then started a new model run wherein each layer was given that number of hidden units as its maximum MLP hidden dim. If the previously observed sparsity was indicative of the capacity the model could effectively utilize, one would expect to see very little sparsity emerge in this second experiment, since the model would need to use every dimension available to it to reach the same capacity. Instead, we see the model trained in Round 2 uses a similar per-layer fraction of its available capacity to what we saw in Round 1, when every layer started with 32768 dimensions of available capacity. This is despite the fact that the model at round 2 has substantially worse performance, indicating this capacity reduction is costly.

From this we hypothesize that each layer has a certain fractional subset of initialization space that it is able to effectively use for the features that layer preferentially learns, and when you give the network a lower total hidden dim, you proportionally reduce the number of dimensions that fit this criteria. We weren’t able to more fully explore this hypothesis in this work, but believe it could be an important area for future study.

Our optimizer is structured to use a warm-up period of 5000 steps, up to a peak learning rate, and then cosine decay thereafter. Table 5 examines how the converged Per-Batch Hidden Unit Use of a 6-layer model varies based on peak learning rate. Note that the standard experimental setting of 3e-3 has the highest performance, since it was selected via hyperparameter search. Some things are clearly consistent over these different learning rates: in all cases layer 0 substantially fewer dimensions over the course of a batch compared to any other layer, and the two highest layers of the network still use the most dimensions, with the final layer continuing to use the most dimensions in all but one case where layer 4 overtakes it by a minute margin. Overall, aggregate sparsity levels, and particularly the extreme batch-level sparsity in layer 0, goes down as peak learning rate decreases.

In [2] the authors showed that wider models had generally higher levels of sparsity, and we see see the same directional result here in 3, both for hidden units used per token (the measurement closest to the one used in that paper) and for hidden units used per batch. This would be consistent with the idea that there is a minumum necessary absolute number of hidden units for good performance (and thus a higher fraction of a narrower model), but that after that the marginal hidden unit declines in value.

To explore the sensitivity of our observations to model depth, we trained a 6, 10, and 14 layer model with a hidden dimension of 16384 (so that all models could fit on available hardware) In 4 we can see that, when it comes to depth rather than width, we do not get the same directional result as above: there isn’t a clear trend in hidden units used per token, but hidden units per batch tend to increase in a deeper model. One explanation that could be at play here is that as we deepen the network we create new layers that are more similar in their behavior to the later layers of a six layer model. That does seem broadly supported by the results here, though we seem to also see some cases (layer 6, in the 14 layer model) which more closely follow the pattern of layer 0 in using a very small fraction of its hidden units per batch, suggesting that as models get deeper there may be value in additional instances of multiple of the layer archetypes so far discussed, not just those of the later layers.