Learning Neural Networks with Sparse Activations

In recent years, transformer based deep neural networks (Vaswani et al., 2017) and the subsequent development of large language models have marked a paradigm shift in the fields of natural language processing and computer vision (Brown et al., 2020; Chowdhery et al., 2022; Chen et al., 2022b; Dosovitskiy et al., 2020). These models have significantly improved performance across various tasks, setting new benchmarks and enabling previously unattainable breakthroughs. However, the computational cost of training and deploying these models, especially the largest variants, presents a significant challenge. A notable portion of these models’ computational and parameter overhead is attributed to the Multi-Layer Perceptron (MLP) layers. These layers are integral to the transformer architecture, playing a crucial role in its ability to solve many different tasks.

Despite their efficacy, the resource-intensive nature of these models has spurred a wave of research focused on enhancing their efficiency (Banner et al., 2019; Frankle and Carbin, 2018; Gholami et al., 2022; Hinton et al., 2015; Anil et al., 2018; Harutyunyan et al., 2023). Among the various strategies explored for improving the inference efficiency of large transformers, attempting to sparsify the transformer is a promising approach.

A motivation for exploiting sparsity is rooted in an intriguing empirical observation made in recent works (Li et al., 2023) regarding the behavior of MLP layers within large transformer models. Post-training, these layers tend to exhibit a high degree of sparsity in their activations; often each input activates as low as 3% of the neurons in the MLP layers, suggesting a natural emergence of sparsity in activations. This leads to these MLP layers behaving like key-value lookups (Geva et al., 2020). The extremely low sparsity (3%) suggests that there might be significant room to sparsify the MLP layers leading to both training and inference efficiency. In addition, such sparsity also helps with interpretability of transformers by disentangling neurons corresponding to distinct concepts (Elhage et al., 2022). Moreover, through extensive ablation studies Li et al. (2023) observe that this phenomenon is highly prevalent. It occurs in convolutional networks (CNNs), as well as in vanilla fully connected feedforward networks.

Despite the potential benefits, effectively harnessing dynamic sparsity has proven challenging. Although, there have been many recent efforts (Li et al., 2023; Grimaldi et al., 2023; Liu et al., 2023; Dong et al., 2023; Csordás et al., 2023; Mirzadeh et al., 2023), they have led to limited success. None of the approaches achieve speedups (either in training or in inference) anywhere close to the the potential factor of 33x that is suggested by 3% sparsity. Moreover, by explicitly enforcing sparsity via methods such as choosing only the top-$k$ activations, the quality of the model degrades in some cases.

A key reason for the hardness in exploiting activation sparsity is that this form of sparsity is dynamic in nature and is input-dependent (i.e., not a fixed pattern). While each input example activates a small number of neurons, the overall sparsity pattern cannot be localized to a small subset of the model weights. For instance, the dynamic nature precludes the use of typical weight quantization or pruning based methods to exploit sparsity empirically. On the other hand, having a non-localized sparsity pattern is crucial in ensuring the model has rich expressiveness.

The above observations suggest that post-training, large transformer networks belong to an intriguing function class that is highly expressive yet exhibits high sparsity. Given the challenges in exploiting this behavior in practical settings, in this work, we initiate a theoretical study of the statistical and computational properties of such functions in the probably approximately correct (PAC) learning framework (Valiant, 1984).

We introduce the class of sparsely activated MLPs. We focus on the case of depth-$1$ MLPs with $n$ input units and $s$ hidden units with the standard ReLU activations. We define the class $\mathcal{H}_{n,s,k}$ as the class of depth-$1$ ReLU networks in $n$-dimensions with the promise that on each input in the support of the data distribution, at most $k$ of the $s$ hidden units are active:

Note that this sparsity differs from dead sparsity, where some neurons are never active on any of the inputs, and consequently, can be deleted from the network without impacting its functionality. The form of dynamic sparsity we study can be crucial for the networks to be more expressive. We provide a couple of examples of useful functions represented using sparsely activated networks here:

• •

  Junta functions: The class of functions on $n$ variables which depend on only a $p$-sized subset ($p<n$) of the variables is known as $p$-junta functions. Sparse parities are a canonical example of junta functions. We show in Theorem 4.2 that we can represent $\log(s)$-juntas using $\mathcal{H}_{n,s,1}$.

• •

  Indexing function: Consider the function $\mathsf{Index}_{b}:\{-1,1\}^{b+2^{b}}\to\{0,1\}$, where $\mathsf{Index}_{b}(z)$ is the $x$-th bit of $y$ ($-1$ mapped to $0$), where $x$ is the integer represented by the first $b$ bits of $z$ in binary representation, and $y$ is the remaining $2^{b}$ bits vector. This can be represented as a $1$-sparse activation network of size $2^{b}$ (i.e., in $\mathcal{H}_{b+2^{b},2^{b},1}$): $\mathsf{Index}_{b}((x,y))=\sum_{\alpha\in\{-1,1\}^{b}}\sigma(\langle{w_{\alpha},z}\rangle-b+\frac{1}{2})$ where the first $b$ coordinates of $w_{\alpha}$ are $\alpha$ and the $\alpha$-th coordinate among the last $2^{b}$ coordinates is $\frac{1}{2}$. On input $z=(x,y)$, only the neuron corresponding to $\alpha=x$ is activated, and the output is precisely $\frac{1}{2}y_{x}+\frac{1}{2}$.

In both the examples presented above, removing any of the $s$ neurons will change the functionality of the network. However, each weight vector $w_{i}$ is quite sparse. In Appendix A, we present an example of a sparsely activated network where even the weight vectors $w_{i}$ are not sparse. Hence, in general, it is not clear if sparsely activated networks can be represented with fewer neurons or sparse weight vectors.

In order to provide learning guarantees, we have to assume an upper bound on the scale of $u$, $w_{j}$’s and $b_{j}$’s. We will use the following natural scaling for the paper:

We then consider the problem of learning sparsely activated networks efficiently. We consider the domain to be the Boolean hypercube $\mathcal{X}=\{1,-1\}^{n}$ as a natural first-step and as a domain where sparsely activated networks can compute non-trivial functions. The Boolean hypercube provides a setting where the function can be sparse everywhere in the domain while maintaining expressiveness; this appears harder in the continuous setting. For instance, if the inputs are Gaussian over $\mathbb{R}^{n}$, one likely needs the biases in the ReLU units to be very large to enforce $1$-sparsity. This suggests that, in the continuous domain, more non-standard distributions are likely necessary to obtain a rich class of functions which are sparse everywhere in the domain. Hence for theoretical simplicity we focus on functions on the Boolean hypercube.

Even with the sparsity assumption, the class $\mathcal{H}_{n,s,1}$ is likely hard to learn in polynomial time (or even quasi-polynomial time) under an arbitrary distribution on the hypercube. In particular, we show that parities on the hypercube on $k$ variables can be computed by $\mathcal{H}_{k^{2},2k,1}$, with coefficient vectors of norm at most $O(k)$. Thus, $\mathcal{H}_{n,O(\sqrt{n}),1}$ need $2^{\Omega(\sqrt{n})}$ queries in the powerful Statistical Queries (SQ) model (see \Crefsec:lb-uniform for details). We also show cryptographic hardness results for learning $\mathcal{H}_{n,s,1}$ under generic distributions on the hypercube.

We consider the problem of learning real-valued functions over the input space $\mathcal{X}=\{-1,1\}^{n}$, to small expected $\ell_{2}$-squared error, namely for the underlying distribution $\mathcal{D}$ over $(x,y)\in\mathcal{X}\times\mathbb{R}$, our goal is the minimize the population loss of a predictor $f:\mathcal{X}\to\mathbb{R}$ given as $\mathcal{L}_{\mathcal{D}}(f):=\operatorname*{\mathop{\mathbb{E}}}_{(x,y)\sim\mathcal{D}}\ell(f(x),y)$ where $\ell(\hat{y},y):=\frac{1}{2}(\hat{y}-y)^{2}$. For any dataset $S\in(\mathcal{X}\times\mathbb{R})^{*}$, we denote the empirical loss as $\mathcal{L}_{S}(f):=\frac{1}{|S|}\sum_{(x,y)\in S}\ell(f(x),y)$.

For any hypothesis class $\mathcal{H}\subseteq\mathbb{R}^{\mathcal{X}}$, we say that $\mathcal{D}$ is $\mathcal{H}$-realizable, if there exists $h^{\star}\in\mathcal{H}$ such that $h^{\star}(x)=y$ holds with probability $1$ for $(x,y)\sim\mathcal{D}$. Following the standard definition of probably approximately correct (PAC) learning (Valiant, 1984), we say that a learning algorithm $\mathcal{A}$ $(\varepsilon,\delta)$-PAC learns $\mathcal{H}$ with sample complexity $m(\varepsilon,\delta)$ if for all $\mathcal{H}$-realizable distributions $\mathcal{D}$ over $\mathcal{X}\times\mathbb{R}$, and for $S\sim\mathcal{D}^{m(\varepsilon,\delta)}$, it holds with probability at least $1-\delta$ that $\mathcal{L}_{\mathcal{D}}(\mathcal{A}(S))\leq\varepsilon$. We say that a learning algorithm $\mathcal{A}$ $(\varepsilon,\delta)$-PAC learns $\mathcal{H}$ under distribution $\mathcal{P}$ (over $\mathcal{X}$) if the learning guarantee holds for all $\mathcal{H}$-realizable $\mathcal{D}$ with the marginal over $\mathcal{X}$ being $\mathcal{P}$. In particular, we use $\mathcal{U}$ to denote the uniform distribution over $\mathcal{X}$.

In this section we provide a learning algorithm for $k$-sparsely activated networks under the uniform distribution.

At a high level, we show that all hypotheses in $\mathcal{H}_{n,s,k}^{W,B}$ exhibit low-degree Fourier concentration and hence can be learned over the uniform distribution using the low-degree algorithm (\Creflem:low-degree-alg). To show Fourier concentration, we bound the noise sensitivity of sparse-activated networks by first showing a bound on the average sensitivity and then converting this to a bound on noise sensitivity.

Next, we can use the bound on average sensitivity to bound the noise sensitivity of functions in $\mathcal{H}_{n,s,k}^{W,B}$. To do so we use an argument attributed to Peres for converting bounds on average sensitivity to bounds on noise sensitivity, allowing us to get better low-degree approximations.

The proof of \Creflem:as-to-ns-generalk is provided in \appendixrefapx:as-to-ns.

[\theoremrefthm:generalk-uniform-ub] We combine \Creffact:fourier, \Creflem:low-degree-alg and \Creflem:as-to-ns-generalk. Fix an error parameter $\varepsilon$. Then, by \Creflem:as-to-ns-generalk, there is a constant $c>0$, such that for

any $h\in\mathcal{H}_{n,s,k}^{W,B}$, satisfies

Thus, we can choose a suitable $d=\Theta((k^{8}W^{4}\log(ns)^{4}+k^{6}B^{4}\log s)/\varepsilon^{2})$, such that by \Creffact:fourier,

Finally, note that $C_{\mathcal{H}_{n,s,k}^{W,B}}=k(W\sqrt{n}+B)$; since at most $k$ neurons are active on any input, and each neuron can at most contribute $W\sqrt{n}+B$. Thus, the theorem now follows from combining the above with \Creflem:low-degree-alg. The run-time and sample complexity will be $O(n^{d}\log(1/\delta)/\varepsilon)$ where $d$ is as above.

Note that the previous section implies a quasi-polynomial time learning algorithm for the class $\mathcal{H}_{n,s,1}$ of $1$-sparsely activated networks. We next show that a quasi-polynomial run-time is likely necessary for learning $\mathcal{H}_{n,s,1}$ under the uniform distribution and stronger lower bounds under arbitrary distributions.