O N
E MERGENCE OF A CTIVATION S PARSITY IN T RAINED T RANSFORMERS
Anonymous authors
Paper under double-blind review
A BSTRACT
This paper reveals a curious observation that modern large-scale machine learning models with Transformer architectures have sparse activation maps.
By activation map we refer to the intermediate output of the multi-layer perceptrons (MLPs) after a ReLU activation function, and by “sparse” we mean that on average very few entries (e.g., 3.0% for T5-Base and 6.3% for ViT-B16) are nonzero for each input to MLP.
Moreover, larger Transformers with more layers and wider MLP hidden dimensions are sparser as measured by the percentage of nonzero entries.
Through extensive experiments we demonstrate that the emergence of sparsity is a prevalent phenomenon that occurs for both natural language processing and vision tasks, on both training and evaluation data, for Transformers of various conﬁgurations, at layers of all depth levels.
We discuss how sparsity immediately implies a way to signiﬁcantly reduce the FLOP count and improve efﬁciency for Transformers.
Moreover, we demonstrate perhaps surprisingly that enforcing an even sparser activation via Top- k thresholding with a small value of k brings a collection of desired but missing properties for Transformers, namely less sensitivity to noisy training data, more robustness to input corruptions, and better calibration for their prediction conﬁdence.
I NTRODUCTION
The great success of modern machine learning for tasks in computer vision, natural language processing, game playing and beyond is driven primarily by the computational model known as deep neural networks (DNNs) (LeCun et al., 2015).
With inspirations drawn from biological intelligent systems, DNNs are massive systems of distributed computational nodes (a.k.a. neurons) with learned inter-connections, which possess the capacity of accomplishing complex real-world tasks.
Although originally motivated from biological brains, there are differences at very fundamental levels on how DNNs work compared to biological neural networks.
One of such differences is in the sparsity of neural activities.
Evidence from neuroscience suggests that neural activity in biological brains is sparse , namely, only a small percentage of all neurons ﬁre at each time (Ahmed et al., 2020;
Barth &
Poulet, 2012; Kerr et al., 2005; Poo & Isaacson, 2009).
Sparse ﬁring suggests that despite having billions of neurons, only a small fraction of the brain participates in computation at each time, which may explain why brains can sustain at a very low energy cost.
In contrast, learning and inference with
DNNs rely primarily on dense computations where all neurons are involved for any input.
In fact, modern computational hardware for deep neural networks, such as GPUs and TPUs, are designed to facilitate massive scale dense computations.
Even with such dedicated hardware, DNNs are still notoriously resource-demanding to train and deploy.
Aside from computation efﬁciency, DNNs also lag far behind biological brains in terms of robustness to input perturbation, error correction for erroneous training labels, conﬁdence calibration for the predictions, etc.
1. A N I NTRIGUING O BSERVATION : A CTIVATIONS ARE S PARSE IN T RAINED T RANSFORMERS
This paper reveals a surprising observation that despite performing dense computations, DNNs produce very sparse activation in their intermediate layers once trained 1 .
Speciﬁcally, we study
[Figure]
Transformer (Vaswani et al., 2017), a DNN model architecture that has become a workhorse for modern applications.
Transformers are constructed by interleaving a self-attention module and a multi-layer perceptrons (MLPs) of depth 2, and the focus of this paper is on the activation map of the ﬁrst MLP layer.
Figure 1 shows the sparsity of the activation maps, measured by the percentage of nonzeros, in all MLP layers of a T5-Base model (Raffel et al., 2020) computed on the training set of C4.
We see that the percentage of nonzero entries is around 50% at initialization, which is expected: randomly initialized weights produce roughly equal numbers of positive and negative entries in the pre-activation map, resulting in about 50 % non-zeros after the ReLU.
However, at the end of training the percentage of nonzero entries reduces drastically: the average value across all encoder-decoder layers is 2.7% with the largest one being 12.0% and the smallest one being only 1.1%.
The emergence of sparse activation in Transformers bears a similarity to the sparsity of neural activities in biological brains, revealing an interesting connection between artiﬁcial and biological networks.
Moreover, unlike classical sparse methods where such a connection is established via explicit sparse regularization (Olshausen & Field, 1996), the sparsity observed in Transformers is emergent without any explicit design.
1. P REVALENCE , B ENEFITS , AND C AUSES OF S PARSITY
This paper studies the aforementioned phenomenon of sparse activation in trained Transformers, with a focus on the following two questions.
First, is the phenomenon shown in Figure 1 a corner case or does it occur broadly?
Second, why should we care about the sparsity in DNNs, other than the appeal of its similarity to biological brains?
Our main results along these two lines are summarized below.
Sparsity is a prevalent phenomenon.
We show in Section 2 that the emergence of sparse activation reported in Figure 1 is not an isolated and cherry-picked case.
Rather, sparsity is prevalent, and occurs broadly in Transformer models: it emerges in all layers of a Transformer, for Transformers trained on both vision and natural language data, for Transformers of various conﬁgurations, and for activation maps computed on both train and test data, etc.
Moreover, through controlled experiments on the width and depth of Transformers, we reveal that larger models are sparser, as measured by percentage of nonzero entries.
We also show in the Appendix B that sparsity emerges with many other architectures and with different optimizers.
Sparsity improves efﬁciency.
Sparsity of activation map in trained Transformers implies that a large proportion of the computation during inference is spent on multiplying values by zero.
Hence, FLOPs can be drastically reduced by avoiding all such computations, which we discuss in Section 3.1.
Motivated by this observation, and to obtain reduced FLOPs not only after training but throughout training, we introduce Top- k Transformer in Section 3.2, a simple modiﬁcation of Transformers where a Top- k thresholding is applied to the activation maps 2 . We show that Top- k Transformers with a reasonable sized k has on par performance with vanilla Transformers. To
demonstrate the computation beneﬁts of Top- k Transformers we provide proof-of-concept results on wall time reduction for the task of unbatched decoding on TPUv4 with a large Top- k T
Meanwhile, we emphasise that this result is far from fully realizing the beneﬁt of sparse activation,
due
to
a lack of
hardware support
for sparse computation Sparsity improves robustness and calibration.
We further show in Section 3.3 that enforcing explicit sparsity via Top- k
Transformers improves model performance in terms of less sensitivity to noisy training data, less sensitivity to input corruptions, and better conﬁdence calibration.
In addition, we provide a study on the causes of sparsity in the Appendix D, showing that sparsity is likely not an artifact of the training data, and may be attributed to the training dynamics in the optimization process.
1. E XPERIMENTAL S ETUP
We study the sparsity in activation maps of Transformers with two commonly used Transformer models, namely Text-to-Text Transfer Transformer (i.e., T5) and Vision Transformer (i.e., ViT).
• T5 is an encoder-decoder model for natural language processing tasks (Raffel et al., 2020).
We train T5 on the Colossal Clean Crawled Corpus (C4) using the span corruption task.
• ViT is an encoder model for vision tasks (Dosovitskiy et al., 2021).
Unless speciﬁed otherwise, we train ViT on ImageNet-21k (Deng et al., 2009), an image classiﬁcation dataset with 14M images and 21k classes.
For certain cases we also use ImageNet-1k which is a subset of ImageNet-21k with 1.3M images and 1k classes.
We measure the sparsity level (computed on training set unless speciﬁed otherwise) at the intermediate output of the two-layer MLPs in a Transformer .
Recall that an MLP performs the following mapping f
( x ; K , V ) .
= d
[Equation]
(cid:88) i = (cid:16) σ ( (cid:104) k
i , x (cid:105) ) · v i (cid:17) , or equivalently, f ( x ; K , V ) . = V σ ( K (cid:62) x ) , (1) where x ∈ IR d model is the input, K =
[ k 1 , . .
. , k d ff ] ∈ IR d model × d ff and V =
[ v 1 , . .
. , v d ff ] ∈
IR d model ×
d ff are learnable layer parameters, and σ () is a nonlinear activation function.
We use ReLU as the activation function σ () for both T5 and ViT 3 .
A two-layer MLP may be regarded as having d ff neurons in the hidden layer, where the i -th neuron performs the computation σ ( (cid:104) k
i , x (cid:105) )
· v i , and the ﬁnal layer output is the sum of the output of all neurons.
Each neuron is called activated
if σ ( (cid:104) k i , x (cid:105) ) is strictly positive.
Hence, the sparsity of neuron activation can be measured by the number of nonzero entries in the feature map
[Equation]
Both T5 and ViT come with several conﬁgurations for d model , d ff , number of layers, etc.
Unless speciﬁed otherwise, we will use the Base models (i.e., T5-Base and ViT-B/16) which have d model = 768 , d ff = 3072 , and 12 layers (for ViT) and 12 encoder layers +12 decoder layers (for T5).
Training details of T5 and ViT are provided in Appendix A.
P REVALENCE OF S PARSITY IN L EARNED T RANSFORMERS
This section shows thorough experiments on commonly used Transformers that sparsity in activation maps is a prevalent phenomenon.
We also show through some controlled experiments that deeper and wider Transformers tend to be sparser measured by percentage of nonzero entries in activation maps.
2. S PARSITY IS A U BIQUITOUS P HENOMENON
We start by providing experimental evidence that the emergence of sparse activation in trained
Transformers is a ubiquitous phenomenon.
To this end, we plot the percentage of nonzero entries of activation maps in different Transformers, and present the results in Figure 2.
These results demonstrate the following.
• Sparsity emerges for both Vision and NLP tasks.
Figure 2a shows the percentage of nonzero entries of trained T5 and ViT models evaluated on their respective training datasets.
We see that both encoder and decoder of T5, as well as the ViT, all exhibit sparsity.
• Sparsity emerges on both training and evaluation data.
Figure 2b shows the percentage of nonzero entries in a trained T5 model measured on both the training data and the evaluation data.
We see
[Figure]
that the property of sparsity generalizes very well to evaluation data as the curves for training and evaluation data align very closely with each other.
• Sparsity emerges on datasets of varying scale.
Figure 2c shows the percentage of nonzero entries in ViT trained on both ImageNet-21k and ImageNet-1k, where the former is a superset of the later with approximately 10 × more images and 21 × more classes.
We see that the scale of data does not affect much of the sparsity level.
• Sparsity emerges on Transformers of varying conﬁgurations.
Figure 2d shows the percentage of nonzero entries for ViT of varying conﬁgurations in model size.
Figure 2e and 2f show the percentage of nonzero entries for encoder and decoder, respectively, of T5 with varying conﬁgurations in model size.
We see that sparsity persists for all cases
• Sparsity emerges across all layers of a Transformer.
Finally, all plots in Figure 2 show that sparsity emerges in all layers of a Transformer.
Moreover, in all cases the ﬁrst few and last few layers tend to be denser than intermediate layers.
The presence of sparsity in activation maps does not rule out the possibility that a small percentage of the neurons are always activated for all inputs, whereas the rest of the neurons are never activated.
To illustrate that this is not the case, we experiment with a pretrained T5 base model 4 to plot the percentage of layer inputs for which each of the d ff neurons is activated when evaluated on 800 examples taken from C4 dataset with span corruption task.
Note that there are 800 ×
= 409600 samples as MLP activation is computed per token.
The results are presented in Figurewith x-axis being indices of neurons in the ﬁrst encoder layer of T5 sorted in descending order according to percentage of layer inputs on which they are activated.
It can be seen that while a few neurons are activated for around 50% of the time, the vast majority of neurons (around 93.5%) are activated less than 10% of the time.
Moreover, there are no dead neurons that are never activated, and the least activated neuron is activated for around 0.001% of the time, and 99% of neurons are activated over 1% of the time.
Finally, while the
[Figure]
2. T HE L ARGER , THE S PARSER
We next examine the effect of model size on the sparsity level of activation maps.
Note that Figure 2e and Figure 2f provide evidence with T5 of varying conﬁguration that larger models tend to be sparser.
Here we perform controlled experiments to examine the effect of model depth, measured by the number of Transformer layers, as well as the effect of model width, measured by the dimension of activation map of MLPs (i.e., d ff ), separately.
Towards that, we take a standard T5 model and vary the depth and width, respectively while keeping the rest of the conﬁguration ﬁxed, and examine their sparsity level after training.
The results are presented in Figure 4 for the encoder, whereas we omit the results for the decoder as they are qualitatively the same as those for encoder.
It can be seen from Figure 4a that deeper Transformers are arguably sparser.
For example, many of the middle layers of the 32-layer model have less than 1% nonzero entries while all shallower models have more than 1% nonzero entries across all layers.
For comparing networks of different widths, we measure the sparsity with the percentage and the count of nonzero entries in Figure 4b and Figure 4c, respectively.
It can be seen that wider models have a lower percentage of nonzero entries, though a higher count of nonzero entries.
E FFICIENT , R OBUST , AND C ALIBRATED : S PARSITY IS A LL Y OU N EED ?
In this section we show that activation sparsity provides several practical beneﬁts.
In Section 3.1 we discuss how the free sparsity in trained Transformers brings us free computation efﬁciency in terms of FLOPs count during inference .
In order to obtain sparsity hence FLOPs reduction throughout training , in Section 3.2 we introduce Top- k Transformers, a simple modiﬁcation of Transformers where a top- k thresholding operation is applied to the activation maps in all MLPs. While existing hardware cannot well support sparse computation and fully realize the beneﬁt of FLOPs reduction, we provide a proof-of-concept experiment on preliminary beneﬁts of Top- k Transformer.
Finally, in Section 3.
we show that sparsity in activation is a good regularization for Transformers. Namely, enforcing sparser activation with smaller values of k in Top- k Transformer (without any other hacks, tweaks and hyperparameter tuning) bestows Transformers several desired properties, namely, robustness of training with erroneous annotations, less sensitivity to input noise/perturbation, and better conﬁdence calibration of the predictions.
3. E FFICIENCY FOR F REE
Given an embedding dimension d model and an MLP intermediate dimension d ff , the computational complexity of a Transformer for an input sequence of length N is O ( Nd 2 model + N 2 d model + Nd model
d
ff ) , where the ﬁrst term comes from computing the key, query, and value matrices, the second term comes from computing the self-attention matrix, and the third term comes from the MLP.
For a ﬁxed sequence length N , and considering the fact that d ff is often much larger than d model , it is arguable that MLP poses the computational bottleneck in large Transformers.
In the following, we explain how sparsity in activation map of MLP can be leveraged to signiﬁcantly reduce its computational cost, without affecting the model performance.
Efﬁciency for the Second MLP Layer.
The sparse activation immediately suggests that a lot of the computation for inference with Transformers is not needed at all.
That is, while doing dense matrix-matrix multiplications, much of it is about multiplying a vector by a value of zero, which can be avoided to save computation.
Speciﬁcally, we consider the second layer of the MLP in (1) which performs the computation [Equation]
[Equation]
where a ∈ IR d ff is the intermediate activation map of MLP (see (2)) and V ∈ IR
d model ×
d ff is the layer parameter.
Eq.
(3) involves a simple matrix-vector multiplication which has a FLOP count of 2 d model ×
d ff .
However, if a is sparse with, say s nonzero entries, then the FLOP count for (3) reduces to 2 d model × s .
Hence,
FLOP in the second MLP layer is reduced by a factor of 1 − sd ff .
Note that sd ff is exactly the percentage of nonzeros plotted in the y-axis of e.g. Figure 1, which is 2 .
7% averaged across all layers.
Hence, the computational cost of the second MLP layer can be reduced by a signiﬁcant amount.
More excitingly, the reduction factor 1 − sd ff is likely to be even bigger for larger Transformer models (see Figures 4a and 4b), pointing to a greater reduction in computation.
Efﬁciency for the First MLP Layer.
The sparsity in the intermediate activation map of MLP does not immediately suggest a reduction in computation for the ﬁrst MLP layer.
Nonetheless, it is possible to signiﬁcantly reduce the computation in the ﬁrst MLP layer by leveraging approximate nearest neighbor search, which we explain next.
Recall from (1) that the computation in the ﬁrst MLP layer is given by
[Equation]
with K =
[ k 1 , . .
. , k d ff ] ∈ IR d model × d ff being the layer parameter and x being the layer input.
If the output is sparse with k nonzero entries, then the calculation in (4) may be formulated as ﬁnding k points from the set { k i } d ff i =1 that are “closest” to the input x measured by values of inner product.
Such a problem is well-known as the nearest neighbor search (NNS) problem or the maximum inner product search problem.
While naive solution of the NNS problem has linear complexity in d ff , there exists approximate algorithms (Guo et al., 2020; Johnson et al., 2019; Shrivastava & Li, 2014) that are of sublinear complexity, and using them in Transformers means that
FLOP in the ﬁrst MLP layer may be reduced to have sublinear complexity in d ff .
There are of course the questions of whether such approximate NNS algorithms could hurt Transformer performance, which we leave for future study.
3. S PARSITY IN T RAINING VIA T OP - k T RANSFORMERS
The beneﬁt of efﬁciency from sparsity in Section 3.1 comes with caveats.
First, while the activation maps are sparse on average, there is the possibility that some of the activation maps for certain inputs are denser hence cannot beneﬁt from sparse computation.
Second, sparsity occurs only in trained
Transformers while the computation is dense during and particularly at the beginning of training.
Here we present Top- k Transformer, a simple modiﬁcation to Transformer architecture that allows us to control sparsity level for all model inputs, and throughout training.
Top-
k Transformer is built upon a regular Transformer with the only modiﬁcation being the MLP layers, where at the output of the activation function σ () (see (1))
we add a Top- k thresholding operator.
That is, the MLPs of
Top-
k Transformers perform the following computation
[Equation]
(5) where Top k ( · ) performs a thresholding that all entries other than those of the largest k values are set to zero with k being a hyper-parameter subject to design choices.
Note that Top-
k Transformer reduces to a regular Transformer if we set k = d ff .
By using a small value of k , the beneﬁt of efﬁciency in terms of reduction in FLOP as discussed in Section 3.1 applies to Transformer training as well.
[Figure]
The immediate question for Top- k Transformer is whether it offers training sparsity at the cost of a reduced performance.
Here we conduct experiments with Top- k T5 and Top- k ViT, and evaluate their performance measured by prediction accuracy for C4 span corruption and ImageNet-21k classiﬁcation tasks, respectively.
The results are provided in Figure 5.
We see that with the Top- k T5-{Base,
Large, 3B} (resp., Top- k ViT) Transformer, taking k to be 128 (resp., 256) is sufﬁcient for closely matching the test performance of the vanilla T5-{Base, Large, 3B} (resp., ViT).
Note that this is achieved without any other hyper-parameter tuning for the Top- k Transformers upon those used for a regular Transformer, and other hyper-parameter choices may further improve the performance of
Top-
Transformers.
We now provide experimental results with Top- k Transformers on wall-time beneﬁts from FLOPs reduction discussed in Section 3.1.
In particular, we evaluate the inference time latency reduction of Top- k Transformer.
In our experiment, we add a Top- k thresholding to T5X (Roberts et al., 2022) 5 . We gain efﬁciency in the second MLP layer by an implementation that avoids all multiplication by zero as described in
Section 3.1. The decoder per-token wall time for unbatched greedy decoding during inference on a single TPUv4 chip is presented in Figure 6.
We observe that larger models have more wall time reduction, due to the fact that they have larger d ff hence more FLOPs reduction. In particular, for
T5-11B we observe around 10% wall time reduction with k ≤ 128 , though this amount becomes smaller with a larger k = 256 . Finally, we emphasize that the sparsity in Top- k
Transformers is unstructured and data-dependent, which is not well supported on existing computation hardwares such as TPUs and GPUs.
Hence, the results in Figure 6 are for proof-of-concept purposes, and are far from fully realizing the beneﬁt of
FLOPs reduction via sparsity.
We leave a study of better implementation of sparse computation for obtaining higher wall time reduction to future work.
3. B ONUS ! I MPROVED R OBUSTNESS AND C ALIBRATION
Despite not being explicitly designed for such purposes, inducing sparse activation via Top- k Transformer has the beneﬁts of improving model robustness 6 and conﬁdence calibration.
We demonstrate this using the image classiﬁcation task with the ImageNet-1k dataset, and present the results in
Table 1.
All results for Top- k ViT are obtained without any model and training hyper-parameter tuning upon those for ViT. Contexts and details are presented below.
More results are presented in
Appendix C.
[Table]
Robustness to Label Noise.
An important challenge for DNNs is that they are highly susceptible to label noise, the problem where a certain percentage of training labels are corrupted or erroneously generated.
This may be attributed to the fact that DNNs are often over-parameterized, hence too “capable” that they tend to overﬁt, or “memorize” the noisy labels without generalizing to test data.
While many dedicated techniques exist (see e.g., Algan & Ulusoy (2021); Song et al.
for a review), here we show that a simple Top- k Transformer can effectively address the label noise issue.
We conduct experiments using the ImageNet-1k dataset for which we replace p % of the labels in the training set with a random label drawn uniformly from the set of all possible labels.
The evaluation performance under p ∈ { 40% , 80% } label noise is presented in Table 1.
It shows that Top- k offers a consistent performance gain with label noise.
Conﬁdence Calibration.
Aside from label noise, another symptom of over-parameterization of
DNNs is that they tend to be overly conﬁdent in their predictions.
In the context of classiﬁcation problems, they tend to assign a high (i.e., close to 1) probability to the class of its prediction, while it is more desirable that they produce a probability that is commensurate with its conﬁdence level (Guo et al., 2017).
A commonly used metric for conﬁdence calibration is the expected calibration error (ECE) (Naeini et al., 2015), which is the discrepancy between the probability to the class of a model’s prediction and the probability that its prediction is actually correct.
Here we measure the calibration of Top- k ViT via ECE and report the results in Table 1.
It shows that Top- k with k = 128 enables the Transformer to be more calibrated when compared to a vanilla Transformer.
Furthermore, results reported in Appendix C show that ECE monotonically decreases as k is decreased from 128 to 32.
Robustness to Input Perturbation.
Another important challenge with DNNs is that their outputs tend to be sensitive to naturally occurring image corruptions, which limits their application to mission critical tasks (Bhojanapalli et al., 2021).
Here we evaluate the robustness of Top- k ViT to three types of additive noises, namely Gaussian noise, impulse noise, and shot noise.
For that purpose, we train Top-
k ViT on standard ImageNet-1k training data and report their classiﬁcation accuracy on
ImageNet-C (Hendrycks & Dietterich, 2019), a benchmark that contains algorithmically generated
Gaussian, impulse, and shot noise (among many others types) applied to the ImageNet-1k test dataset.
For each noise type, there are ﬁve severity levels.
We report the averaged performance over all severity levels of each corruption type in Table 1 for k = 128 , and in Appendix C for a few other values of k . We see that robust accuracy is the highest with k = 64 , while taking k = 128 or k =also provides beneﬁts compared to the vanilla Transformer.
R ELATED W ORK
Prior efforts on introducing sparsity in deep neural networks abound, though often with diverse motivations and objectives.
Here we provide a brief overview of several popular lines of work.
Sparsity for Efﬁciency.
Sparsity in either model weights or activation maps is often used for improving training and inference efﬁciency (see e.g. Hoeﬂer et al.
(2021) for a review).
For activation sparsity in particular, sparsity for efﬁciency is explored perhaps ﬁrst in ConvNets (Georgiadis, 2019;
Kurtz et al., 2020; Rhu et al., 2018) before subsequently becoming a key design component in many of the largest Transformer based language and vision models (Du et al., 2022; Fedus et al., 2022a;b;
Rajbhandari et al., 2022).
The Top- k thresholding that we use in Top- k Transformer has also been previously used in Gupta et al.
(2021) to improve memory efﬁciency of Transformers.
However, it has been unclear a priori whether sparsity hurts model performance, hence the practice often relies on wishful design, trial-and-error, and post-hot justiﬁcation (Baykal et al., 2022).
Our discovery that
Transformers naturally produce sparse activation maps, and that larger models are even sparser, may provide principled perspectives towards efﬁciently training future large models.
Sparsity for Robustness.
Many works ﬁnd that smaller and sparser networks obtained by model compression are more robust to adversarial perturbation (Chen et al., 2022; Guo et al., 2018; Jordao & Pedrini, 2021) and label noise (Xue et al., 2022).
Another line of work that uses sparsity for robustness leverages the property that practical data corruption is often sparse (Ghosh et al., 2017; Liu et al., 2022; You et al., 2020).
None of the work mentioned above is based on sparsity in activation maps.
More closely related to ours is the work of Ahmad & Scheinkman (2019) where sparsity in activation map of convolutional DNNs is shown to improve robustness to input perturbation, and
Muthukumar & Sulam (2022) that leverages sparse activation to derive robust generalization error bounds.
Sparsity for Explainability.
Work on leveraging sparsity for interpreting deep learning models long exist but often in a post-hoc fashion for examining the semantic meanings encoded by a neuron of a trained model (Dalvi et al., 2019).
For Transformers, evidence suggests that the learned knowledge is encoded mainly in its MLPs with individual neurons expressing speciﬁc factual knowledge (Dai et al., 2022).
Moreover, enforcing neuron activation sparsity in MLPs helps to improve the percentage of neurons that are interpretable (Elhage et al., 2022).
Hence, our discovery may point to new directions towards developing more interpretable DNNs (Cuadros et al., 2022; Sajjad et al., 2021).
Sparsity for Data Modeling.
Following the seminal work of Olshausen & Field (1996), there are a lot of interests in sparsity as an effective modeling of natural signals (Mairal et al., 2014).
With the close resemblance of the computational structure of ReLU networks and sparse encoding algorithms (Gregor & LeCun, 2010), it became natural to study a DNN as a multi-layer sparse modeling of the data (Papyan et al., 2018).
Along with substantial theoretical understanding of such a modeling are obtained (Papyan et al., 2017; Sulam et al., 2018), there are also experimental results on their practical beneﬁts (Sun et al., 2018) though less often on modern large-scale data.
Sparsity for Theory of Over-parameterized Models.
Because of its simplicity and well-develped theory in classical machine learning (Candès & Wakin, 2008; Vidal et al., 2015; Wright & Ma, 2022), sparse modeling is often used to provide theoretical understanding of modern large and overparameterized models.
This include works on implicit regularization (Chou et al., 2021;
Nacson et al., 2022; Vaskevicius et al., 2019; Woodworth et al., 2020; Zhao et al., 2019), nonconvex optimization (Buhai et al., 2020; Sulam et al., 2022), noise interpolators (Chinot et al., 2022; Donhauser et al., 2022; Koehler et al., 2021), etc.
However, the aforementioned work uses sparsity as a testbed or toy model to gain insights, without implication of existence of sparsity in DNNs.
D ISCUSSION
This work demonstrates the natural emergence of sparse activation in commonly used Transformer models (Section 2).
The notion of sparsity pertains to the law of parsimony, a.k.a.
Occam’s razor , where among all possible explanations of observed data, the simplest ones are preferred.
It is a fundamental scientiﬁc principle broadly used in various scientiﬁc and engineering subjects (Domingos, 1999; Epstein, 1984), including classical machine learning (Tibshirani, 1996).
Hence, our discovery may be suggesting that the law of parsimony is playing a role in Transformers even though they are not explicitly designed so, resonating with recent view on the role of sparsity for intelligence systems (LeCun, 2022; Ma et al., 2022; Roberts, 2021; Vasudevan et al., 2021).
More importantly, we back such a perspective by providing evidence of improved robustness and calibration via enforcing sparsity using Top- k thresholding (Section 3), which indicates that sparsity is indeed a pertinent prior for good generalization.
We hope that our work may motivate future effort on introducing sparsity in deep learning models in a more principled way for obtaining more efﬁcient, robust, and calibrated models.
Finally, while our motivation of studying sparse activation in Transformers comes (partly) from study of biological brains, establishing such a connection may reciprocally beneﬁts efforts on applying artiﬁcial intelligence to the study of biology and neuroscience (Richards et al., 2022).
Appendices
The appendices are organized as follows.
In Section A we provide the implementation details for experiments conducted in this paper.
In Section B we demonstrate the emergence of sparse activation in other architectures and with other optimizers than those used in Section 2.
In Section C we provide additional experiments upon those in Section 3 to demonstrate the beneﬁts of sparsity.
In Section D we explore the potential causes of sparsity, with a focus on the effect of training data.
In Section E we present a derivation to show that during early training, the ﬁnal MLP layer’s intermediate activation tends to get sparse.
Finally in Section F we present insights on the emergence of activation sparsity from experiments on 2-layer MLP models.
A I MPLEMENTATION D ETAILS A. T
For most of the experiments, except the Top- k transformer, we used vanilla T5 architecture (Raffel et al., 2020).
We trained model with Adafactor optimizer, an inverse square root learning rate schedule, and no dropout.
For the ﬁrst 10,000 steps we also use a ﬁxed learning rate of 0.01 as warm-up.
The training task is span corruption without any mixture, and unless speciﬁed otherwise, we train the model for 100,000 steps with batch size of 256 to save compute and time, as the sparsity or accuracy trend is already clear by then.
We used 512 tokens on the encoder side and 114 tokens on the decoder side.
A. V I T
Following Dosovitskiy et al.
(2021), we train ViT using ADAM (Kingma & Ba, 2015) as the optimizer with β 1 = 0 .
9 , β 2 = 0 .
Other training details such as weight decay, dropout rate, and learning rate all follow the description in (Dosovitskiy et al., 2021, Section B.1) except that we train forepochs (as opposed to 300) on ImageNet-1k. A. T5 / V I T C ONFIGURATIONS
For the reader’s convenience, we summarize the conﬁguration of varying T5 / ViT models used in our paper in Table A.1.
[Table]
B A DDITIONAL R ESULTS O N P REVALENCE OF S PARSITY B. S PARSITY AND N ETWORK A RCHITECTURE
We evaluate the sparsity level of activation map in several commonly used network architectures beyond T5 and ViT.
This includes BERT which is also a Transformer based architecture, as well as non-Transformer based architectures such as MLP-Mixer and ConvNets.
We also examine whether residual connection accounts for the emergence of sparsity [Table] [Figure] [Figure]
BERT.
We evaluate the sparsity level of BERT models (Devlin et al., 2019).
We speciﬁcally consider BERT Base (12 layers) and BERT Large (24 layers) Transformer models, with ReLU activation in the MLP layers.
We follow the same training receipe as Devlin et al.
and pre-train these models on
Wikipedia and Books dataset using Masked Language Modelling (MLM) objective.
We train for 450000 steps with a batch size of 1024 using AdamW optimizer with 1 e − 4 learning rate.
In Figure B.2 we plot the sparsity levels of both BERT models for all the intermediate MLP layers (plots a and b).
We observe that both these models exhibit high levels of sparsity (< 10%) as other Transformer models.
We further visualize the preactivation values of the MLP layers as histograms in plots c and d.
We observe that while they have mean 0 at initialization, the mean quickly becomes negative as training progresses, resulting in high sparsity levels.
Finally, in Figure B.3 we provide a visualization of pre-activation of values with several popular activation functions, including GeLU, Sigmoid, and Tanh. We observe a similar distribution of preactivation values as ReLU with GeLU and Sigmoid activations. However Tanh activation has a different distribution of preactivation values. With Tanh activation, the network does not show sparsity and the accuracy is signiﬁcantly worse in comparison to ReLU/GeLU.
MLP-Mixer.
We evaluate the sparsity level of the MLP-Mixer (Tolstikhin et al., 2021), an all-MLP architecture constructed from cascading token-mixing and channel-mixing MLPs.
Speciﬁcally, we use Mixer-B16 as the architecture, ADAM with β 1 = 0 .
9 , β 2 = 0 .
999 as the optimizer, and train on ImageNet-21k for 300 epochs.
While Tolstikhin et al.
(2021) sweeps over a product set of hyper-parameters, here for simplicity we use a ﬁxed set of hyper-parameters with weight decay of 0.03, gradient norm clipping at 1.0, base learning rate of 0.003, RandAugment magnitude of 10, no mixup, no stochastic depth, and no dropout.
Figure B.1 shows the sparsity level at the intermediate layer of both token mixing and channel mixing
MLPs of Mixer-B16.
We also plot the sparsity level of ViT (i.e., the plot in Figure 2a) to Figure B. [Figure] [Figure] for a comparison.
It can be seen that the ﬁrst four layers of channel mixing MLP and ViT have almost identical sparsity levels, while the rest of the layers (other than the last one) of channel mixing MLP are denser than the corresponding layer of ViT. On the other hand, the token mixing MLPs produce dense activation maps with more than 50% nonzero entries, probably because the dimension of the activation maps (384) is too small.
Convolutional Neural Network (ConvNet).
Sparsity in activation maps has been studied for ConvNets such as the AlexNet (Krizhevsky et al., 2017) at least as early as in the work of Rhu et al.
There are also follow-up work (Georgiadis, 2019; Kurtz et al., 2020) on how enforcing sparse activation maps can help to gain computation efﬁciency.
For completeness, we evaluate and present results for the sparsity level of residual networks (ResNets)
(He et al., 2016), which is one of the most commonly used ConvNets, trained on ImageNet-1k.
In particular, we focus on ResNetand ResNet-50 which are constructed from stacking 8 standard residual blocks and 16 “bottleneck” residual blocks, respectively, where each block has two and three convolutional and ReLU layers, respectively.
We examine the sparsity of activation maps after each of the ReLU layers in each residual block.
The results for ResNet-18 and ResNet-50 are reported in Figure B.4a and Figure B.4b, respectively.
Here, the x-axis is the index of the residual block, and the sparsity of different layers in the residual blocks are plotted with separated curves in each ﬁgure.
It can be observed that
• Layers near the network output tend to produce sparser activation maps than layers near the network input.
This is aligned with the observation with ViT trained on ImageNet-1k (see Figure 2b).
• For each residual block, the intermediate layers (i
e., the 1st layer for ResNet-18 and the 1st & 2nd layers for ResNet-50) produce sparser activation maps than the output layer (i.e., 2nd layer for ResNet-18 and 3rd layer for ResNet-50).
In addition, all residual blocks are divided into four stages that have different output feature map sizes.
For ResNet-50, the four stages are composed of blocks 0 - 2, 3 - 6, 7 - 12, and 13 - 15.
Figure B.4b
[Figure]
• For the 1st layers, percentage of nonzeros decreases within each stage, and jumps up from the last layer of each stage to the ﬁrst layer of next stage.
• For the 2nd layers, percentage of nonzeros decreases quickly at the beginning of each stage then becomes stable.
• For the 3rd layers, percentage of nonzeros tend to increase slightly within each stage, and jumps down from the last layer of each stage to the ﬁrst layer of next stage.
Such observations may help to understand the role of each stage in ResNets.
Comparing the percentage of nonzero entries in ResNets (shown in Figure B.4a and Figure B.4b) and for Transformers (shown in Figure 2b), both of which are trained on ImageNet-1k, we see that
ResNets produce much denser activation maps with more than 10% nonzero entries in all layers.
One possible explanation is that ResNet uses batch normalization (BN) before each activation function, while Transformer’s MLP does not have BN before the activation function.
To understand the effect of BN on sparsity, we conduct an experiment with BN in ResNet removed.
Because ResNet cannot be effectively trained without BN, we decrease the learning rate from standard ResNet training by a factor of 10.
Moreover, we add a learnable scalar multiplier that is initialized as 0 to all the residual branches, following the study in Bachlechner et al.
; Qi et al.
The results for comparing with standard ResNet are reported in Figure B.5, where to separate the effect of using a smaller learning rate, we also compare with the method of training a regular ResNet but with a small learning rate compared to standard training. The two subﬁgures of Figure B.5 show the effect of width on sparsity of the ﬁrst and second layers in each residual block, respectively. It can be observed that, removing BN does not signiﬁcantly change the sparsity level, except for small set of layers. Meanwhile, the trend that larger models are sparser for Transformers (see Section 2.2) holds for
ResNets as well, as seen in Figure B.6.
Here, we vary the width of ResNet-18 by multiplying the number of output channels of each convolutional layer by a factor of 1 (for width = 64), 2 (for width = 128), and 4 (for width = 256).
The two subﬁgures show the effect of width on sparsity of the ﬁrst and second layers in each residual block, respectively.
In both cases, wider models have smaller percentage of nonzero entries across all layers, except for the very last layer (i.e., the 2nd layer in block #7 shown in Figure B.6b).
Sparsity and Residual Learning.
We provide a study on the effect of residual connections on activation sparsity.
Each Transformer block contains two types of residual connections: the one that is in parallel with the attention blocks, and the one that is in parallel with the MLP blocks.
We focus on the residual connection parallel to the MLP blocks.
We perform two different studies. • Effect of shortcut connection. Towards that, we train two T5-Large models, one using the vanilla Transformer block and the other with residual connection removed for the Transformer block on encoder layer 6 (i e., the 7 th encoder layer, as we count from 0).
There is a 1.6% evaluation accuracy drop with the latter model compared to the former model.
The percentage of nonzero entries of these two Transformers are presented in Figure B.7 for the encoder layers and in Figure B.8 for the decoder layers.
It can be seen that in encoder layer 6 for which the residual connection is removed, the sparsity has a very different trend during training compared to the corresponding layer of the vanilla Transformer.
Moreover, the sparsity level at all other layers also changes, though to a much smaller extend. • Effect of
initialization
scale of the residual branch
Many works have found that having the residual branch initialized at a smaller scale helps with stabilizing and accelerating the training of residual (convolutional) networks (Goyal et al., 2017; Zhang et al., 2019) and Transformers (Touvron et al., 2021).
Here for simplicity we consider the idea from Bachlechner et al. ; Qi et al. (2020) where a trainable scalar multiplier that is initialized at zero is applied to the residual branch (a.k.a., ReZero). We consider ViT trained on ImageNet-21k with ReZero added to the MLP modules.
We ﬁnd that this increases the training accuracy from 46.15% to 46.85% but reduces the validation accuracy from 47.58% to 46.75%.
We plot the sparsity level of ViT with ReZero and compare it with the vanilla ViT in Figure B.9.
It can be seen that ReZero reduces the percentage of nonzeros in layers near the network output.
[Figure]
[Figure]
[Figure]
[Figure]
B. S PARSITY AND O PTIMIZER
Transformers are usually trained using ADAM or its variants as the optimizer (Kingma & Ba, 2015).
It may be curious to ask whether the emergence of sparsity is speciﬁc to such optimizers and whether other optimizers, such as stochastic gradient descent (SGD), also leads to sparse activation maps.
However, we ﬁnd that SGD cannot effectively train Transformer architectures such as T5 and ViT.
Hence, we study the effect of optimizer on activation sparsity by looking at ResNet trained on
ImageNet-1k following the setup in Section B.1, since both SGD and ADAM can effectively train the network.
To train ResNet with ADAM, we use the same hyper-parameters as those used in SGD, with the only difference being that the optimizer is ADAM with β 1 = 0 .
9 , β 2 = 0 .
To make the comparison with SGD fair, we tune the base learning rate for ADAM and select 3 e − 3 , which is the one that gives the highest training accuracy among the set of { 1 e − 4 , 3 e − 4 , 1 e − 3 , 3 e − 3 , 1 e − 2 } .
The training accuracy obtained by ADAM with base learning rate 3 e − 3 is similar to that obtained by SGD, namely, 67.8% by ADAM vs 69.3% by SGD with ResNet-18, and 75.0% by ADAM vs 78.5% by SGD with ResNet-50.
The results for ResNet-18 and ResNet-50 are presented in Figure B.10 and Figure B.11, respectively.
For ResNet-18, we see that ADAM leads to a smaller percentage of nonzero entries particularly towards the output of the network for the ﬁrst layers of each residual block.
In contrast, ADAM and
SGD have very similar sparsity level at the second layers of each residual block.
Similar observation holds for ResNet-50, where the percentage of nonzero entries is smaller with ADAM for the ﬁrst and second layers of each residual block, while for the third layer the sparsity level does not change much. B. S PARSITY IN F INETUNING
In this section we show that activation sparsity not only occurs after model pretraining but persists after further ﬁnetuning on downstream tasks.
Here we take a T5 that has been pretrained on C4 as described in Section 1.3, and ﬁnetune the model on a open domain Natural Question (Kwiatkowski et al., 2019) QA task.
We follow the set up in Li et al.
(2022), where the retrieved passages are independently encoded by the encoder, and then passed to the decoder via cross attention.
The decoder takes the question as the preﬁx and produces the answer. The decoder is a standard auto- [Figure] regressive setup, despite passing the question as a preﬁx.
We used the same Wikipedia passages retrieved by FiD (Izacard & Grave, 2020) which uses DPR (Karpukhin et al., 2020) as retriever.
For this experiment, we used 20 passages for each question.
When calculating the sparsity level, we ignored the activation produced from paddings, for both encoder and decoder.
The sparsity level of different encoder and decoder layers after ﬁnetuning is reported in Figure B.12 and is compared to those before the ﬁnetuning.
It can be seen that ﬁnetuning does not drastically change the sparsity level.
C A DDITIONAL R ESULTS ON B ENEFITS OF S PARSITY C. T OP - k D OES NOT S IGNIFICANTLY A FFECT T RAINING C ONVERGENCE
One may argue that the convergence of Top- k Transformer can be slower than that of a vanilla Transformer because, with fewer neurons being activated, the amount of parameters that have nonzero gradient associated with each training sample is smaller.
Here we plot the training curves for ViT as well as
Top-
k ViT with k ∈ { 64 , 128 , 256 } on ImageNet21K (this is the same setup as the experiments in
Figure 5b), and report the results in Figure C.1.
It can be observed that taking Top- k does not signiﬁcantly reduce convergence speed, particularly when k is relatively large (e.g., k = 256 ).
C. B ENEFITS OF S PARSITY P ERSISTS WITH (cid:96) 1 -N ORM I NDUCED S PARSITY
While Top- k thresholding is used in Section 3.3 to demonstrate the beneﬁt of sparsity, we show that other means of obtaining sparsity, such as an explicit (cid:96) 1 norm regularization, also provides such beneﬁts.
We experiment with ViT for ImageNet-1k classiﬁcation under the same setup as in Section 3.3.
Here instead of the Top- k ViT, we train a regular ViT but with an additional loss term, which is the sum of the (cid:96) 1 norm of all activation maps of ViT across all layers.
We refer to the method as L1-ViT.
We vary the weight λ on the (cid:96) 1 loss in the set λ ∈ { 0 . 1 , 0 . 5 , 1 . 0 } to control the strength of the regularization, and denote the corresponding methods as L1-ViT- { 0 . 1 , 0 . 5 , 1 . 0 } .
The sparsity level, natural accuracy, robust accuracy under input perturbation, and ECE of L1-ViT are reported in Table C.1.
We see that with λ = 0 .
1 or 0 . 5 , the averaged percentage of nonzero entries do not change much, but already demonstrates performance gain in terms of accuracy under input perturbation and calibration without hurting the natural accuracy. Using a λ = 1 . 0 drastically reduces [Figure]
[Table] [Table] the percentage of nonzero entries but it offers a better performance in terms of accuracy under input perturbation and calibration compared to ViT [Table]
C. A DDITIONAL R ESULTS FOR C ALIBRATION
Under the same setup as in Section 3.3, we report the calibration of Top- k ViT during training and with varying values of k , and report the results in Figure C.2.
At the beginning of training the model has a low ECE because the output probabilities are mostly uniformly distributed across all classes hence the model is not conﬁdent, and that its prediction is purely random hence wrong with high probability.
The model tends to become overly conﬁdent as the training progresses, hence the ECE increases particularly towards the end of training.
What we can observe is that Top-
k enables the
Transformer to be more calibrated when compared to a vanilla Transformer, particularly for small values of k.
The results with k = 128 and its comparison with the vanilla ViT is also presented in Table 1.
C. A DDITIONAL R ESULTS FOR R OBUSTNESS TO I NPUT P ERTURBATION
Under the same setup as in Section 3.3, we report the performance of Top- k ViT with varying value of k in Figure C.3.
We can see that Top- k ViT offers a performance gain over the vanilla ViT, particular with k = 64 or k = 128 and the severity level is high.
The averaged performance over all severity levels of each corruption type is reported in Table 1.
[Figure]
D S PARSITY FROM T RAINING D YNAMIC ?
In this section we study the causes of sparsity in activation maps of trained Transformers.
Towards that, in Section D.1, D.2, and D.3 , we present a set of hypotheses and design corresponding experiments to validate or disprove them.
We discuss the observation from the experiments and draw a tentative conclusion in Section D.4.
D. S PARSITY FROM L ABELS ?
Transformers are usually trained via supervised learning (e.g., using the ImageNet dataset for ViT) or self-supervised learning (e.g., using the span corruption task for T5).
In both cases, the training labels provide a pertinent and meaningful description of the corresponding training data (e.g., image for
ViT and text for T5).
Sparsity may arise because the label set provides a structured organization of the massive training data, hence the training dataset admits a compact representation.
This motivates us to formulate the following hypothesis.
Hypothesis D.1 (Sparsity from labels) .
Sparsity in trained Transformers arises from the labels for
Transformer training, e.g., human annotations in supervised training or generated labels from data itself in self-supervised learning.
We use a random label experiment with ViT for image classiﬁcation to test Hypothesis D.1.
Speciﬁcally, we generate a new training dataset by replacing p % of the labels in the ImageNet-21k dataset with random labels drawn uniformly at random from the set of all possible labels, where p is varied to examine the effects.
With such a dataset, the labels for a certain percentage of images do not provide a meaningful description for the content of the image.
Hence, if Hypothesis D.1 is valid, then the activation map will become dense.
The sparsity level of ViT trained on the random label datasets is shown in Figure D.1a.
It can be seen that the percentage of activated neurons decreases with an increasing percentage of label noise up to 70% .
An even higher label noise level at 100% changes the sparsity level across layers as the shallow layers (i.e., layers 0 - 4) becomes sparser, while the deep layers (i.e., layers 5 - 11) becomes denser.
Nonetheless, even with 100% label noise, all layers have < 10% activated neurons.
D. S PARSITY FROM D ATA ?
While modern image and text data are often of high-dimensional, their intrinsic degree of freedom is much smaller, i.e., they are low-dimensional and admit compact representations (Vidal et al., 2015; Wright & Ma, 2022).
Hence, even if the labels do not provide meaningful descriptions of the data, it may still be possible that Transformers extract low-dimensional structures from data and produce compact representations in the form of sparse activation maps.
This motivates the following hypothesis.
Hypothesis D.2 (Sparsity from natural data) .
Sparsity in trained Transformers arises from natural training data (e.g., images for ViT and texts for T5).
We use a random image experiment to test Hypothesis D.2.
With the ImageNet-21k dataset, we replace each image with a random image generated by drawing pixel values from an i.i.d.
Uniform distribution in the range of [0, 255], and use these images (instead of the original images in ImageNet21k) for model training.
Such random images do not contain any low-dimensional structures nor compact representations.
The percentage of nonzero entries of a ViT trained on random image dataset is shown in Figure D.1b.
It can be seen that the ﬁrst four layers become sparser while the last few layers become relatively denser compared to training with natural images in ImageNet-21k.
Nonetheless, all layers have < 10% activated neurons.
D. S PARSITY FROM D ATA - FITTING ?
Modern deep neural networks are often over-parameterized , with sufﬁcient capacity to ﬁt practical training datasets and obtain close-to-zero training error.
There is evidence suggesting that this result holds true even if the data and label are generated in random (Zhang et al., 2021).
Hence, there is the possibility that sparsity arises because the training data, even if generated in random, is scarce relative to the scale of modern over-paremeterized models.
Hypothesis D.3 (Sparsity from data-ﬁtting) .
Sparsity in trained Transformers arises from the fact that models have more than sufﬁcient capacity to ﬁt training data of practical scale.
To test Hypothesis D.3, we design an inﬁnite data experiment where the amount of training data is inﬁnitely large so that any practical Transformer becomes under-parameterized relative to the data and cannot ﬁt the data.
The way we generate inﬁnite training data is to sample images with random pixels as in the random image experiment, and for each image we sample a random label as in the random label experiment.
Moreover, we generate sufﬁcient amount of such training data to make sure that the model never sees the same data point twice during the training.
The number of training iterations in the inﬁnite data experiment is kept the same as that of the random image and random label experiments.
The result of this experiment is presented in Figure D.1c.
It can be seen that the ﬁrst four layers produce sparser activation maps, while middle layers with index 4 - 7 are considerably denser compared to the baseline with near 10% to 20% nonzero entries.
D. D ISCUSSION
The results of random label, random image, and inﬁnite data experiments in Figure D.1 show that labels, data, and data-ﬁtting as conjectured in Hypothesis D.1, D.2, and D.3, respectively, all affect the sparsity level of the activation map.
Nonetheless, none of them fully explains the emergence of sparsity since for all results in Figure D.1, the percentage of nonzero entries is considerably smaller than at the initialization (i.e., 50%).
Our results point to the possibility that sparsity comes from the training dynamic.
Namely, at early training stage with any data, the descending direction of the gradient on the Transformer parameters tends to point to a regime where their MLPs produce sparse activation maps.
This is because at the start of training, any non-zero neuron activity in the MLP is more likely to cause the cross entropy loss to go up rather than go down.
We demonstrate this negative gradient effect at initialization on the last MLP layer in Appendix E, as well as an analysis on two-layer MLP models in Appendix F.
Empirically, we ﬁnd in Figure B.2 that the suppression of activations indeed happens at the early stages of training in Bert.
We leave it for the future work to fully investigate the effect of training dynamics on sparsity for Transformer models.
E N EGATIVE GRADIENT ON ACTIVATIONS AT INITIALIZATION
In this section, we provide theoretical insights into the emergence of sparsity in MLP hidden layers.
In particular, we consider a classiﬁcation model with a ﬁnal MLP layer that maps a feature layer of neurons to the classiﬁcation logits.
Under the following assumptions:
• The loss function is the standard cross entropy loss.
• The ﬁnal MLP layer has no bias terms.
• The feature layer uses the ReLU activation function.
• Each element of the weight matrix of the ﬁnal MLP layer follows an i.i.d.
random uniform initialization: W ki ∼ U ( − 1 , 1) .
(This can be easily relaxed to any mean 0, symmetric distribution.)
We show that the expectation value of the gradient on the feature layer activations, when averaged over different initializations of the last layer MLP weights, are negative at initialization:
[Equation]
Where A ni is the i -th feature layer activation resulting from the n -th datapoint in the training set.
Consider a neural classiﬁer with its ﬁnal layer being an output of a MLP:
[Equation]
where y nk are classiﬁcation logits of the n-th training example on the k-th class, k ∈ { 1 . .
N c } .
W ki are the ﬁnal layer weights, R denotes ReLU activations and A i ( x ) with i ∈ 1 . . .
N e are the pre-activations.
N e is the embedding width.
The pre-activations can be computed by an arbitrarily complex network.
The cross entropy loss is given by:
[Equation]
where n ∈ 1 . .
N d denotes different training examples and t nk is the ground truth label distribution, p nk is the predicted distribution computed by the model:
[Equation]
in particular,
[Equation]
We compute the gradient on the pre-activations of the i-th neuron:
[Equation]
The key observation of this derivation is that at initialization, there is no statistical correlation between the true label distribution and the model’s randomly initialized weights, therefore the ﬁrst term is zero on average.
But because p nk is computed using W ki , there are some correlation between them, making the second term positive.
As a result, the preactivations A ni will receive a negative gradient due to mismatch between the predictions p nk and the true distribution t nk .
To see this effect more clearly, we compute the average gradient with respect to all random initializations of the last layer weights.
Consider the cases where A ni is positive so that the gradient is non-trivial:
[Equation]
Assuming W ki follows a uniform random distribution from say ( − 1 , 1) , then [Equation]
[Equation]
where,
I
[Equation]
[Equation]
where,
[Equation]
Note that A ni > 0 from E.6.
In addition, both C ki and B ki are positive constants independent of
W ki :
[Equation]
With these positive constants, it is easy to see that
[Equation]
because for each positive value of W ki ∈ (0 , 1) we have f ( W ki ) >
f ( − W ki ) .
Therefore:
[Equation]
We have shown that the gradient at initialization on the pre-activations of the feature layer neurons are negative, assuming cross entropy loss and that the weights follow a uniform distribution at initialization.
Even though we only show this rigorously for the feature layer neurons, the intuition we are establishing here could be more general: if a neuron’s activation causes an increase in the loss, then it will receive a negative gradient.
We have shown that, at the beginning of the training, on average non-zero activations of randomly initialized neural networks are more likely to cause the loss function to go up rather than down.
Therefore all activations receive a downward pressure from the gradients on average.
This provides some hints to why sparsity may arise across different data types, including realistic data, random input, random labels as well as inﬁnite data.
More generally, we conjecture that for most losses, the entropy of larger losses is always bigger compared to smaller losses.
This is perhaps made clear for cross entropy loss in the derivation above.
To provide another example, at any given value of a L2 loss, the volume of the space with larger loss is always inﬁnite while the volume of the space with smaller loss is ﬁnite (the volume of a d-dimensional ball).
Thus at initialization a non-zero neuron activation is more likely push the result towards the higher entropy regime, resulting in higher loss.
This in turn causes gradient descent to push the neuron activations to negative values at early stages of training.
For each training example, we speculate that a small subset of neurons’ activations happen to reduce the loss, then this neuron will be strengthened by this gradient and is more likely to survive the general downward pressure demonstrated above, thus causing the emergence of sparse activations in neural networks.
F I NSIGHTS FROM S PARSITY IN MLP S
We study the sparsity of one of the simplest neural networks: 2-layer MLPs.
As with transformers, we train the network with different types of data and study their effect on sparsity.
We uncover interesting interplay between over/under-ﬁtting and sparsity.
In particular, we show that 1) in the extreme under-ﬁtting limit, the model becomes completely sparse.
This is consistent with our analysis in Section E; 2) as the model capacity increase and the model can ﬁt more data, the sparsity level decreases; 3) the sparsity level peaks roughly when the model is just large enough to memorize the data; 4) in the overﬁtting regime, larger model leads to more sparsity.
We setup the experiment using 3 different datasets, the standard MNIST dataset, the random dataset where each pixel of MNIST is replaced by an i.i.d.
uniform random number and each label is replaced by a random class amongst 10.
However the input-label correlation is ﬁxed for each training data entry so a large enough model trained long enough can memorize this correlation.
We also have the inﬁnite data regime where the random inputs and labels are generated on-the-ﬂy, represeting a random dataset of inﬁnite size.
We train MLPs with different hidden sizes: 32, 128, 512, 2048, 8192, 32768 and 131072 using three different optimizers: SGD, SGD with momentum and Adam for 200 epoches (as measured by the number of training data iterations in the MNIST and random data case), which is enough time for most of our metrics to converge.
We report training accuracy and non-zero rate in the hidden dimension at the end of the trainings in F.1.
We ﬁrst observe that all models becomes completely sparse in the inﬁnite data case.
During training we observe a smooth and monotonic decay of non-zero rates as training goes on.
This phenomenon is consistent with the intuition developed by our theoretical analysis in Section E, where discrepancy between the model’s prediction and the labels will put negative gradients on activations and drive sparsity.
At the other end of the spectrum, after the model enters the overﬁtting regime, for both MNIST and random data, increasing widths further will lead to more sparsity.
This is consistent with our ﬁndings in transformer models.
These two observations indicate an interesting transition at the regime where the model is just able to memorize the training data. From this point, increasing the model size leads to more sparsity, descreasing the model size leads to more underﬁtting, which also leads to more sparsity, therefore the non-zero rate has to peak around this regime. This is exactly what we found empirically with random data in Figure F.1.
[Figure]
hidden layer width
[Figure]
[Figure]
hidden layer width
[Figure]