A Theoretical View on Sparsely Activated Networks
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Deep and wide neural networks successfully ﬁt very complex functions today, butdense models are starting to be prohibitively expensive for inference.
To mitigatethis, one promising direction is networks that activate a sparse subgraph of thenetwork.
The subgraph is chosen by a data-dependent routing function, enforcinga ﬁxed mapping of inputs to subnetworks (e.g., the Mixture of Experts (MoE)paradigm in Switch Transformers).
However, prior work is largely empirical,and while existing routing functions work well in practice, they do not lead totheoretical guarantees on approximation ability.
We aim to provide a theoreticalexplanation for the power of sparse networks.
As our ﬁrst contribution, we presenta formal model of data-dependent sparse networks that captures salient aspectsof popular architectures.
We then introduce a routing function based on localitysensitive hashing (LSH) that enables us to reason about how well sparse networksapproximate target functions.
After representing LSH-based sparse networks withour model, we prove that sparse networks can match the approximation power ofdense networks on Lipschitz functions.
Applying LSH on the input vectors meansthat the experts interpolate the target function in different subregions of the inputspace.
To support our theory, we deﬁne various datasets based on Lipschitz targetfunctions, and we show that sparse networks give a favorable trade-off betweennumber of active units and approximation quality.
Introduction
Overparameterized neural networks continue to yield performance gains as their sizes increase.
Thistrend has been most prominent with large Transformer based language models [6, 12, 32].
However,using large, dense networks makes training and inference very expensive, and computing a forwardpass may require trillions of ﬂoating point operations (FLOPs).
It is an active area of research toimprove the scalability and efﬁciency without decreasing the expressiveness or quality of the models.
One way to achieve this goal is to only activate part of the network at a time.
For example, the
Mixture of Experts (MoE) paradigm [22, 37] uses a two-step approach.
First, each input is mappedto a certain subnetwork, known as an expert.
Then, upon receiving this input, only this particularsubnetwork performs inference, leading to a smaller number of operations compared to the totalnumber of parameters across all experts.
Some variations of this idea map parts of an input (calledtokens) individually to different experts.
Another well-studied generalization is to map an input to agroup of subnetworks and aggregate their output in a weighted manner.
Switch Transformers [14]successfully use a reﬁned version of the MoE idea, where the input may be the embedding of a tokenor part of a hidden layer’s output.
Researchers have investigated many ways to perform the mapping,such as Scaling Transformers [21] or using pseudo-random hash functions [34].
In all cases, sparsityoccurs in the network because after the input is mapped to the subnetwork, only this subset of neuronsneeds to be activated.
Moreover, the computation of the mapping function, sometimes referred to as a‘routing’ function, takes signiﬁcantly less time than the computation across all experts.
The success of these approaches is surprising.
Restricting to a subnetwork could instead reduce theexpressive power and the quality of the model.
However, the guiding wisdom seems to be that itsufﬁces to maximize the number of parameters while ensuring computational efﬁciency.
Not allparameters of a network are required for the model to make its prediction for any given example.
Ourgoal is to provide a theoretical explanation for the power of sparse models on concrete examples.
Our Results.
Our ﬁrst contribution is a formal model of networks that have one or more sparselyactivated layers.
These sparse layers are data dependent: the subset of network activations andtheir weight coefﬁcients depend on the input.
We formally show that our model captures populararchitectures (e.g., Switch and Scaling Transformers) by simulating the sparse layers in these models.
We next prove that the above class of sparsely activated models can learn a large class of functionsmore efﬁciently than dense models.
As our main proof technique, we introduce a new routingfunction, based on locality sensitive hashing (LSH), that allows us to theoretically analyze sparselyactivated layers in a network.
LSH maps points in a metric space (e.g., R d ) to ‘buckets’ such thatnearby points map to same bucket.
The total number of buckets used by a LSH function is referred toas the size of the hash table.
Prior work has already identiﬁed many routing functions that work wellin practice.
We do not aim to compete with these highly-optimized methods, but rather we aim to use
LSH-based routing as a way to reason about the approximation ability of sparse networks.
In Theorem 4.1, we show that LSH-based sparse models can approximate real-valued Lipschitzfunctions in R d .
Although real data often lives in high dimensions, the underlying manifold where theinputs are supported is often assumed to be low-dimensional (e.g. the manifold of natural images vs 58 R 224 × 224 for a 224 × 224 image).
We model this by assuming that our inputs lie in a k -dimensionalmanifold within R d .
To get (cid:15) approximation error, we need an LSH table of size approximately
O (( √ dk
/(cid:15) ) k ) but a forward pass only requires time O ( dk log(1 /(cid:15) ))
as only one of the O (( √ dk/(cid:15) )
k )non-empty buckets are accessed for any given example
In Theorem 4.3, we complement our upper bounds by proving a lower bound of Ω((2 √ d/(cid:15) )
k ) on thesize needed for both dense and sparse models (when points live in k dimensions).
This also transformsinto a lower bound on the number of samples required to learn these functions using a dense model.
This lower bound implies that a forward pass on a dense model takes time Ω( d (2 / √ k(cid:15) ) k ) , whichis exponentially worse than the time taken by the sparse model.
Altogether, we show that for thegeneral class of Lipschitz functions, sparsely activated layers are as expressive as dense layers, whileperforming signiﬁcantly fewer ﬂoating point operations (FLOPs) per example.
To support our theory, we perform experiments in Section 5 on approximating Lipschitz functionswith sparse networks.
We identify several synthetic datasets where models with data-dependentsparse layers outperform dense models of the same size.
Moreover, we achieve these results withrelatively small networks.
This provides a few minimum working examples of when sparsity isbeneﬁcial.
Our results complement the existing work that has already shown the success of verylarge MoE models on real-world NLP and vision tasks [14, 21, 33].
Our experiments also generalizethe intuition from Figure 1a, where the model outputs a different constant value for each bucket.
Apuzzling empirical result was observed in [34] where a pseudo-random uniform hash function is usedto route input tokens to experts.
Intuitively, this might map very nearby or similar points to distinctexperts, leading to a potentially non-smooth behavior.
We show that on our synthetic, Lipschitzdatasets, the LSH model outperforms uniform random hashing, which further corroborates our theory.
Related Work.
Sparsely activated networks have had enormous empirical success [1, 5, 13, 25,29, 33, 37, 40].
The Switch Transformer [14] is one of the ﬁrst major, contemporary applicationsof the sparsely activated layers.
Follow-up works such as Scaling Transformers [21] and otherhash functions
[34] aim to improve the sparse layers.
These papers build upon seminal MoEworks [19, 20, 22], and other uses of the MoE paradigm [27, 36].
There has been work on usingstatistical mechanics to analyze the generalization of MoE models [23].
However, to the best of ourknowledge, there is no systematic theoretical study of modern sparsely activated networks.
The above work on dynamic sparsity builds on previous static sparsity efforts, e.g., weight quantization [28], dropout [38], and pruning (see the survey [18] and references).
Static sparsity means thatthe subnetwork activation does not change in a data-dependent way.
The focus is on generalizationand compression, instead of achieving fast inference time with a huge number of parameters.
           \MRTYXJIEXYVIW            ]   J Y R G X M S R  Z E P Y I W   (a) 1D example       2YQFIVSJEGXMZEXIHYRMXW
PSKWGEPI       1 I E R
  U Y E V I H  ) VV S V HIRWI[MHXL!\E\MWZEPYIW (71XSXEP[MHXL! (71XSXEP[MHXL!
(71
XSXEP[MHXL! (71XSXEP[MHXL!

[Figure]
Our work builds on locality sensitive hashing (LSH), a well-studied technique for approximate nearestneighbor search (see the survey [4] or the book [16] and references therein for LSH background).
A seminal and popular LSH family uses random hyperplanes to partition the space [3, 7].
Anotheroption is to use balls centered at lattice points, which works well for (cid:96) 2 distance [2].
We use theseideas to design sparse networks.
For uses of LSH in deep learning, sketch-based memory improvesnetwork capacity [15, 30].
Other empirical work uses LSH to improve dense network training timeor memory [8, 9, 31].
Our work differs from the prior studies because we implement the LSH-basedapproach with sparsely activated networks, with the goal of reducing inference time and achievinglow regression error.
LSH-based sparsity may be most relevant for vision transformers [24, 33, 42],using distance metrics and routing functions tailored to images, or for time series analysis [41].
Preliminaries
Let f : R d → R be a multivariate real-valued function that we want to learn with a neural network.
We consider regression, and we aim to minimize the mean-squared error or the (cid:96) ∞ error.
For afunction f deﬁned on a set Γ and an estimator ˆ f n , we deﬁne (cid:107)
f − ˆ f n (cid:107) ∞
=
sup x ∈ Γ | f ( x ) − ˆ f n ( x ) | .
For some results, we approximate f on a subset V ⊆ R d .
For example, V may be the intersection of[ − 1 , 1] d and a k -dimensional subspace.
We also consider the case when f is L -Lipschitz, meaningthat | f ( x ) − f ( x (cid:48) )
| ≤ L · (cid:107) x − x (cid:48) (cid:107) 2 for all x, x (cid:48) ∈ R d .
We deﬁne [ n ] = { 1 , 2 , . . . , n } .
2.1 Data-Dependent Sparse Model
A dense neural network g with a fully connected ﬁnal layer can be expressed as g ( x ) =
A · φ ( x )where A ∈ R 1 × t is a matrix and φ : R d → R t is a function (e.g., φ captures the representationlearned up until the ﬁnal layer).
Here, t is the width of the ﬁnal layer and d is the input dimensionality.
Our focus is on the types of networks captured by what we call the Data-Dependent Sparse Model ( DSM ) .
This is a network with a sparsely activated ﬁnal layer.
Formally, let t be the width, and let s ≤ t be a sparsity parameter.
Then, we consider functions g of the form g ( x ) =
A x · φ ( x ) where
A x ∈ R 1 × t and φ : R d → R t .
The crux of the model is the ﬁnal layer.
The sparsity comes from letting A x =
A ◦ mask ( x ) , where mask ( x ) ∈ { 0 , 1 } 1 × t is an s -sparse indicator vector, and “ ◦ ” is the entry-wise product.
The mask zeroes out certain positions, and A contains the learned weights but no longer depends on x .
Intuitively, the mask is the “routing function” for the sparse activations.
To do so, we deﬁne mask ( x ) ◦ φ ( x ) to mean that we zero out element of the vector φ ( x ) .
Then, we have g ( x ) =
( A ◦ mask ( x )) φ ( x ) =
A ( mask ( x ) ◦ φ ( x )) .
Under the above deﬁnitions, let DSM ( d, s, t ) be the set of functions
g =
( A ◦ mask ( x )) · φ ( x ) .
Inwhat follows, we use A x as shorthand for A ◦ mask ( x ) .
In the DSM model, we wish to understand the effect of sparsity on how well the network canapproximate certain functions.
It is important to allow the ﬁnal layer A x ∈ R 1 × t to be data-dependent.
The values in A x depend on the target function f , and we learn this from training data.
Real-world transformers may have multiple sparse layers.
To capture this, we can compose DSMfunctions.
For example, two sparse layers comes from g ( x ) =
A ◦ mask 2 ( x ) ◦ φ 2 ( mask 1 ( x ) ◦ φ 1 ( x )) .
For simplicity, we focus on a single sparse layer in what follows, which sufﬁces for our main results.
2.2 Hash-based Routing
Prior sparse models compute hash functions of the input vector x to determine mask and the networkactivations [34].
Our main theorem uses this strategy, considering LSH families.
LSH has the propertythat nearby points are more likely to end up in the same hash bucket than far away points.
We can use
LSH to deﬁne a general class of efﬁcient regression functions.
In Section 3, we prove that the DSMmodel captures popular sparse architectures and the following LSH model.
LSH Model.
We review an LSH framework based on sign patterns.
Let h 1 , . .
, h m : R d → {− 1 , 1 }be m distinct hash functions.
Partition the space into 2 m buckets based on the m sign patternsz x =
( h 1 ( x ) , . . .
, h m ( x )) over all x ∈ R d .
Then, for each z ∈ {− 1 , 1 } m , we can specify a functionˆ g z ( x ) , where the goal is for g z to approximate the target function f in the part of space
associatedwith z (i.e., points x that have pattern z under h 1 , . . .
, h m ).
More generally, we can allow s setsof such hash functions ( h i 1 , . .
, h im ) , and s sets of these approximation functions (ˆ g 1 z 1 , . .
, ˆ g sz s ) fori = 1 , . . .
, s .
On input x , we compute the sign patterns z 1 , . . .
, z s and output g ( x )
= (cid:80) si =
α
i ˆ g
iz i ( x ) .
Further, we can restrict each ˆ g iz to be a degree ∆ ≥ 0 polynomial.
For a ﬁxed LSH family of possiblehash functions, we let LSH ( d, s, m, ∆) denote this class of functions.
In many cases, we only needˆ g iz to be a constant function, i.e., ∆ = 0 , and we shorten this as LSH ( d, s, m, 0) :
= LSH ( d, s, m ) .
Our goal is to understand the trade-offs of the “memory” parameter m and the “sparsity” amounts that we need to approximate large classes of d -dimensional functions (i.e., L -Lipschitz functionssupported on a k -dimensional manifold).
Euclidean LSH Model [11].
This is a popular LSH family for points in R d .
In this case, each hash function outputs an integer (instead of ± 1 above).
Each bucket in this model is deﬁned by a set of hyperplane inequalities.
There are two parameters ( D, (cid:15) ) .
We sample D random directions a 1 , . .
. , a D where each coordinate of each a i is an independent normal variable.
In addition we sample b i ∼ Unif[0 , (cid:15) ] independently for i ∈
[ D ] .
For a point x ∈ R d , we compute an index into a bucket via a function h
i : R d → Z deﬁned as
[Equation]
Here, the index i ranges over i ∈
[ D ] , leading to a vector of D integers.
Simulating Models with DSM
We formally justify the DSM ( d, s, t ) model by simulating other models using it.
We start with simpleexamples (interpolation and k nearest neighbor ( k -NN) regression), then move on to transformersand the LSH model.
For the simple examples, we only need one hidden layer, where φ ( x ) =
σ ( Bx )for a matrix B ∈ R t × d and non-linearity σ .
Interpolation.
We show how to compute f at t points x 1 , . . .
, x t .
We only need sparsity s =
1and σ is simply the identity function.
We set A i = f
( x i ) / (cid:104) b
i , x i (cid:105) , where b i is the i th row of B .
Further, we let mask ( x i ) have a one in the i th position and zeroes elsewhere.
Then, we deﬁneg ( x i ) =
( A ◦ mask ( x i ))
Bx i =
f
( x i ) , which computes f at the t input points.
Sparse networks perform k -NN regression .
We sketch how the DSM ( d, k, n ) model can simulate k -NN with g ( x ) =
( A ◦ mask ( x ))
σ ( Bx ) ,when the dataset contains vectors with unit (cid:96) 2 norm.
Let the rows of B be a set of n unit vectorsb 1 , . . .
, b n ∈ R d .
For the target function f , let A = 1 k ( f ( b 1 ) , . . .
, f ( b n )) .
Deﬁne σ ( Bx ) to have ones in the top k largest values in Bx and zeroes elsewhere.
For a unit vector x , these k positionscorrespond to the k largest inner products (cid:104) x, b i (cid:105) .
Since (cid:107) x − b i (cid:107) 2 = 2 − 2 (cid:104) x, b i (cid:105) , the non-zeropositions in σ ( Bx ) encode the k nearest neighbors of x in { b 1 , . . .
, b n } .
Thus, g ( x ) computes theaverage of f at these k points, which is exactly k -NN regression.
Moreover, only k entries of A areused for any input, since σ ( Bx ) is k -sparse; however, computing σ ( Bx ) takes O ( nd ) time.
Whilethere is no computational advantage from the sparsity in this case, the fact that DSM can simulatek -NN indicates the power of the model.
3.1 Simulating transformer models with DSM
Real-world networks have many hidden layers and multiple sparse layers.
For concreteness, wedescribe how to simulate a sparsely activated ﬁnal layer.
As mentioned above, we can composefunctions in the DSM model to simulate multiple sparse layers.
Switch Transformers [14].
The sparse activations in transformers depend on a routing function
R : R d → { 1 , . . .
, L } , where R ( x ) speciﬁes the subnetwork that is activated (for work on the bestchoice of R , see e.g., [34]).
To put this under the DSM ( d, s, t ) model, consider a set of trainablematrices A 1 , . .
A L ∈ R 1 × s , where the total width is t =
s · L .
On input x , we think of A x asa 1 × t matrix with s non-zero entries equal to A R ( x ) .
In other words, A is the concatenation of
A 1 , . . .
, A L , and mask ( x ) is non-zero on the positions corresponding to A R ( x ) .
Scaling Transformers [21].
The key difference between Switch and Scaling transformers is that thelatter imposes a block structure on the sparsity pattern.
Let t be the width, and let s be the numberof blocks (each of size t (cid:48)
=
t/s ).
Scaling Transformers use only one activation in each of the sblocks.
In the DSM ( d, s, t ) model, we capture this with A x as follows.
Let e i ∈ { 0 , 1 } t (cid:48) denote thestandard basis vector (i.e., one-hot encoding of i ∈ [ t (cid:48) ]) .
The sparsity pattern is speciﬁed by indices( i 1 , . .
, i s ) .
Then, A x = ( α 1 e i 1 , . . .
, α b e i s ) for scalars α 1 , . . .
, α s ∈ R .
3.2 Simulating the LSH model using DSM
We explain how to represent the LSH model using the DSM model.
The key aspect is that A x dependson the LSH buckets that contain x , where we have s non-zero weights for the s buckets that containeach input.
This resembles the Scaling Transformer, where we use the LSH buckets to deﬁne routingfunction.
In the LSH ( d, s, m, ∆) model, there are s sets of m hash functions, leading to s · 2 m hashbuckets.
We use width t = s · 2 m for the DSM ( d, s, t ) network.
The entries of A x are in one-to-onemapping with the buckets, where only s entries will be non-zero depending on the s buckets that xhashes to, that is, the values ( h i 1 ( x ) , . .
, h im ( x )) ∈ {− 1 , 1 } m for i
= 1 , 2 , . .
, s .
We now determine the values of these s non-zero entries.
We store a degree ∆ polynomial ˆ g ( x ) : 185 R d → R associated with each bucket.
For our upper bounds, we only need degree ∆ = 0 , but wemention the general case for completeness.
If ∆ = 0 , then ˆ g is simply a constant α depending on thebucket.
An input x hashes to s buckets, associated with s scalars ( α 1 , . . .
, α s ) .
To form A x , set sentries to the α i values, with positions corresponding to the buckets.
For degree ∆ ≥ 1 , we storecoefﬁcients of the polynomials ˆ g , leading to more model parameters.
Section 4 contains details onusing LSH to approximate Lipschitz functions with sparse networks.
Computing and storing the LSH buckets.
Determining the non-zero positions in A x only requires
O ( sm ) hash computations, each taking O ( d ) time with standard LSH families (e.g., hyperplane
LSH).
We often take m to be a large constant.
Thus, the total number of operations to compute aforward pass in the network O ( smd )
≈ O ( sd ) .
The variable m above determines the total numberof distinct buckets we will have ( 2 m ).
For an n point dataset, m =
O (log n ) is a realistic setting intheory.
Therefore, 2 m = poly( n ) is often a reasonable size for the hash table and moreover the densemodel requires width O (2 m ) as well.
In practice, the number of experts is usually 32–128 [14, 33].
The hash function typically adds very few parameters.
For example, to hash d dimensional datainto 2 m buckets, LSH requires O ( md ) bits and takes O ( md ) time as well.
In summary, the LSHcomputation does not asymptotically increase the number of FLOPs for a forward pass in the network.
We also later bound the number of non-empty LSH buckets.
The memory to store the hash table iscomparable to the number of samples required to learn the target function.
Data-Dependent Sparse Models are more Efﬁcient than Dense Models
For a general class of Lipschitz functions, LSH-based learners yield similar (cid:96)
∞ error as dense neuralnetworks while making inference signiﬁcantly more efﬁcient.
It is a common belief in the machinelearning community that although many of the datasets we encounter can appear to live in highdimensional spaces, there is a low-dimensional manifold that contains the inputs.
To model this, weassume in our theory that the inputs lie in a k -dimensional subspace (a linear manifold) of R d .
Herek (cid:28) d .
Theorem 4.1 shows that the LSH model we propose can learn high-dimensional Lipschitzfunctions with a low (cid:96) ∞ error efﬁciently when the input comes from a uniform distribution on anunknown low-dimensional subspace.
Theorem B.1 extends this result to when the input comes froman unknown manifold with a bounded curvature.
We present Theorem 4.1 here and defer Theorem B.1to Section B. All proofs are in the appendix.
Theorem 4.1.
For any f : [ − 1 , 1] d → R that is L -Lipschitz, and for an input distribution D that isuniform on a k -dimensional subspace in [ − 1 , 1] d , an LSH-based learner can learn f to (cid:15) -uniformerror with O ( kL √ d k log( L √ d/(cid:15) ) /(cid:15)
k ) samples using a hash table of size
O ( L √ d k /(cid:15) k )
with probability ≥ 0 . 8 .
The total time required for a forward pass on a new test sample is O ( dk log( L √ d/(cid:15) )) .
The key idea behind this theorem is to use LSH to produce a good routing function.
The locality ofthe points hashed to an LSH bucket gives us a way to control the approximation error.
By using alarge number of buckets, we can ensure their volume is small.
Then, outputting a representative valuesufﬁces to locally approximate the target Lipschitz function (since its value changes slowly).
Thenext lemma bounds the size of the sub-regions corresponding to each bucket, as well as bounding thenumber of non-empty buckets.
Lemma 4.2. Consider a Euclidean LSH model in d dimensions with Ck hyperplanes and widthparameter (cid:15) where C is a large enough constant.
Consider a region Γ deﬁned by the intersection of ak -dimensional subspace with [ − 1 , 1] d .
We have that the LSH model deﬁnes a partitioning of Γ intobuckets.
Let c be a constant.
Then, with probability ≥ 0 . 9 ,1.
Projecting any bucket of the LSH onto Γ corresponds to a sub-region with diameter ≤ (cid:15)/c .2.
At most (cid:16) 2 √ d (cid:15) (cid:17) O ( k ) buckets have a non-empty intersection with Γ .
The above construction assumes knowledge of the dimensionality of the input subspace k .
Fortunately,any upper bound on k would also sufﬁce.
The table size of the LSH model scales exponentially in kbut not d .
Thus, an LSH-based learner adapts to the dimensionality of the input subspace.
Theorem 4.1 shows that sparsely activated models (using LSH of a certain table size) are powerfulenough to approximate and learn Lipschitz functions.
Next, we show a complementary nearlymatching lower bound on the width required by dense model to approximate the same class offunctions.
We use an information theoretic argument.
Theorem 4.3.
Consider the problem of learning L -Lipschitz functions on [ − 1 , 1] d to (cid:96) ∞ error (cid:15) when the inputs to the function are sampled from a uniform distribution over an unknown k -dimensional subspace of R d ∩ [ − 1 , 1] d .
A dense model of width w with a random bottom layer requires w = Ω (cid:32) ( √ dL ) k ( C(cid:15) )
k (cid:33) , for a constant C > 0 to learn the above class of functions.
The number of samples required is Ω (cid:32)
[Equation]
Our approach for this theorem is to use a counting argument.
We ﬁrst bound the number of distinctfunctions, which is exponential in the number of parameters (measured in bits).
We then construct alarge family of target functions that are pairwise far apart from each other.
Hence, if we learn thewrong function, we incur a large error.
Our function class must be large enough to represent anypossible target function, and this gives a lower bound on the size of the approximating network.
Theorem 4.3 shows a large gap in the time complexity of inference using a dense model and an LSHmodel.
The inference times taken by a dense model vs. a sparse model differ exponentially in 1 /(cid:15) .
[Equation]
Overall, the above theorems show that LSH-based sparsely activated networks can approximate
Lipschitz functions on a k -dimensional subspace.
The size and sample complexity match betweensparse and dense models, but the sparse models are exponentially more efﬁcient for inference.
Experiments
To empirically verify our theoretical ﬁndings, we align our experiments with our proposed modelsand compare dense models, data-dependent sparse models ( DSM ), LSH models, and sparse modelswith random hash based sparse layers.
While DSM and LSH models are analyzed in the previoussection, random hash based layers introduce sparsity with a random hash function, as an alternativeof learnable routing modules and LSH in MoE models [34].
Our goal is to show that the DSM and
LSH models achieve small MSE while using much fewer activated units than dense models.
5.1 Experimental set-up
Dense, DSM and random hash
sparse models contain a random-initialized, non-trainable bottomlayer, (a Top-K layer for DSM and a random hash layer for random hash models to enforce sparsity),and a trainable top layer, with varying number of hidden units and sparsity levels.
LSH models havenon-trainable hyperplane coefﬁcients for hashing and a trainable scalar in each bucket (the scalardetermines the output of the network for points in that bucket).
We compare dense models and threesparse models ( DSM , LSH , and random).
We evaluate with synthetic data from two random, Lipschitz target functions that are commonly usedbasis functions for arbitrary continuous functions.
These random functions allow us to empiricallyevaluate the construction from Theorem 4.1, while comparing different routing functions.
Random polynomial.
p ( x ) of degree d for x ∈ R n with sum of coefﬁcient absolute values < 1 / ( d ) .
As a result, we ensure that the function is 1-Lipschitz over [ − 1 , 1] n .
Random hypercube function.
f : [ − 1 , 1] n → R which interpolates the indicator functions at eachcorner with random {− 1 , 1 } value at each corner.
Concretely, the function is deﬁned as follows: foreach corner point y ∈ {− 1 , 1 } n , its indicator function is I y ( x ) = (cid:81) ni
=1 1+ y i
x
i. Sample randomvalues v y ∈ {− 1 , 1 } with probability (0 , 5 , 0 . 5) independently for each y ∈ {− 1 , 1 } n , the randomhypercube polynomial function is f ( x ) = (cid:80) y ∈{− 1 , 1 } n v y I y
( x ) .
For the random polynomial (random hypercube) functions, we use n = 8 and d = 4 in the this sectionand provide more results for other parameter settings in the appendix (especially the case when thetarget polynomial functions concentrate on a low dimension subspace).
We also present results on areal dataset, CIFAR-10, which corroborates our ﬁndings from the synthetic data experiments.
5.2 Results
MSE for random functions.
Figures 2 and 3 show the scaling behavior of the DSM and LSH modelsfor a random polynomial function and hypercube function (lower is better).
Sparsity helps for the
DSM and LSH models, both achieving better quality than dense models using the same number ofactivated units.
In Figure 4, we compare the DSM /LSH models with the random hash sparse models,and we see random hash sparse models underperform dense models, suggesting data-dependentsparsity mechanisms (such as DSM /LSH) are effective ways to utilize sparsity.
FLOPs.
To further qualify the efﬁciency gain of sparse models, we compare the MSE at the same
FLOPs for sparse/dense models in Table
FLOPs eval MSE (dense) eval MSE (DSM 50% sparsity) eval MSE (DSM 25% sparsity) 0.
      2YQFIVSJEGXMZEXIHYRMXWPSKWGEPI      
[Figure]
R   U
[Figure]
E V I H  ) VV S V HIRWI[MHXL!\E\MWZEPYIW (71XSXEP[MHXL! (71XSXEP[MHXL!
(71
XSXEP[MHXL! (71XSXEP[MHXL!
       2YQFIVSJEGXMZEXIHYRMXWPSKWGEPI       1 I E R
  U Y E V I H  ) VV S V HIRWI[MHXL!\E\MWZEPYIW
07,FYGOIXW! 07,FYGOIXW!
07,FYGOIXW! 07,
FYGOIXW!
[Figure]
  
[Figure]
U Y E V I H  ) VV S V HIRWI[MHXL!\E\MWZEPYIW (71XSXEP[MHXL! (71XSXEP[MHXL!
(71
XSXEP[MHXL! (71XSXEP[MHXL!
       2YQFIVSJEGXMZEXIHYRMXW
PSKWGEPI   
  1 I E R   U Y E V I H  ) VV S V HIRWI[MHXL!\E\MWZEPYIW 07,
FYGOIXW! 07,FYGOIXW!  07,FYGOIXW! 07,
FYGOIXW!
CIFAR-10.
We also compare the scaling behavior of DSM and dense models on CIFAR-10 [26].
The baseline model is a CNN with 3 convolutional layers (followed by max-pooling), a dense layerwith varying number of units, and a ﬁnal dense layer that computes the logits (referred as CNN +dense).
For the data-dependent sparse models, we use the same architecture, except we change thepenultimate dense layer with a data-dependent sparse layer (referred as CNN + DSM).
Both modelsare trained with ADAM optimizer for 50 epochs and evaluated on the test dataset for model accuracywith
no data augmentation; see Figure 5 and Table 2 for the accuracy versus number of activated units.
As with the synthetic datasets, DSMs outperform dense models at the same number of activated units.
5.3 Discussion
Our experimental results (e.g., Figures 2, 3, and 4) show that sparsely activated models can efﬁcientlyapproximate both random polynomials and random hypercube functions.
Intuitively, the DSM and
LSH models employ the sparsity as a way to partition the space into nearby input points.
Then,because the target function is Lipschitz, it is easy to provide to local approximation tailored to thespeciﬁc sub-region of input space.
On the other hand, the uniform random hash function performspoorly for these tasks precisely because it does not capture the local smoothness of the target function.
      2YQFIVSJEGXMZEXIHYRMXWPSKWGEPI         1 I E R   U Y E V I H  ) VV S V HIRWI[MHXL!\E\MWZEPYIW VERHSQLEWLXSXEP[MHXL!
(71XSXEP[MHXL!
07,FYGOIXW!
(a) random polynomial       2YQFIVSJEGXMZEXIHYRMXWPSKWGEPI   
  1 I E R   U Y E V I H  ) VV S V HIRWI[MHXL!\E\MWZEPYIW VERHSQLEWLXSXEP[MHXL!  (71XSXEP[MHXL!
07,FYGOIXW!
[Figure]
function
[Figure]
'22HIRWI[MHXL!\E\MWZEPYIW'22(71XSXEP[MHXL!'22(71XSXEP[MHXL!'22(71XSXEP[MHXL!
Model \ # activated units
Dense 69.
DSM (50% sparse) 70.
[Table]
On CIFAR-10, we also see that the DSM model performs comparably or better than the dense network.
In particular, Figure 5 shows that the “CNN + DSM” approach with total width 1024 improves uponor matches the dense model.
In this case, the sparse activation allows the network to classify wellwhile using only a fraction of the number of FLOPs.
In Table 2, we see that DSM model outperformsthe dense model when we control for the number of activated units in the comparison.
Limitations.
Our experiments focus on small datasets and 2-layer networks as a way to align withour theoretical results.
Prior work on sparsely activated networks has shown success for large-scale
NLP and vision tasks. Our experiments complement previous results and justify the DSM and LSHmodels by showing their ability to approximate Lipschitz functions (consistent with our theorems).
It would be good to further evaluate the DSM and LSH models for large-scale tasks, for example,by using them as inspiration for routing functions in MoE vision transformers, such as V-MoE [33].
It would be interesting to evaluate on larger datasets, such as ImageNet as well.
We experimentedwith a handful of hyperparameter settings, but a full search may lead to different relative behaviorbetween the models (similarly, we only evaluated a few parameter settings for the random functionsand the synthetic data generation, which is far from exhaustive).
Conclusion
We provided the ﬁrst systematic theoretical treatment of modern sparsely activated networks.
To doso, we introduced the DSM model, which captures the sparsity in Mixture of Experts models, such as
Switch and Scaling Transformers.
We showed that DSM can simulate popular architectures as wellas LSH-based networks.
Then, we exhibited new constructions of sparse networks.
Our use of LSH to build these networks offers a theoretical grounding for sparse networks.
We complemented ourtheory with experiments, showing that sparse networks can approximate various functions.
For future work, it would be interesting to implement LSH-based networks in transformers forlanguage/vision tasks.
A related question is to determine the best way to interpolate in each LSHbucket (e.g., a higher degree polynomial may work better).
Another question is whether a densemodel is more powerful than a sparse model with the same number of total trainable parameters.
Theorem 4.1 only says that a sparse model with similar number of parameters as a dense modelcan more efﬁciently (fewer FLOPs) represent Lipschitz functions.
This does not say all functionsexpressible by a dense model are also expressible by a sparse model.
This is non-trivial questionas A x depends on the input (i.e., DSM ( d, t, t ) may be more expressive than the dense model withwidth t ).
We expect that dense networks can be trained to perform at least as well as sparse networks,assuming the width is large enough .
The dense networks should optimize the weights in the last layerto approximate the function, but they may not favor certain neurons depending on the input.
A Proofs of Main Upper and Lower BoundsA.1 Proof of the Lipschitz Upper Bound
Proof of Theorem 4.1.
We use the Euclidean- LSH construction of Lemma 4.2 with parameter (cid:15)/L .
In any sub-region of the k -dimensional subspace that has a small diameter, the Lipschitz natureof the function together with Lemma 4.2 will imply that we can approximate it by just a constantand incur only (cid:15) error in (cid:96) ∞ .
In particular, given a point x 1 belonging to an LSH bucket, we canset ˆ f ( x ) =
f
( x 1 ) everywhere in that bucket.
For any x 2 also mapping to the same bucket, from
Lemma 4.2, we have that (cid:107) x 1 − x 2 (cid:107) 2 ≤ (cid:15)/L .
Since f is L-Lipschitz,
[Equation]
Next we look at how many samples we need to obtain the guarantee (cid:107) ˆ f − f (cid:107) ∞ ≤ (cid:15) .
A rare scenariothat we have to deal with for the sake of completeness is when there exist buckets of such smallvolume that no training data point has mapped to them and consequently we don’t learn any valuesin those buckets.
At test time, if we encounter this rare scenario of mapping to a bucket with novalue learnt in it, we simply run an approximate nearest neighbor search among the train points.
For our prediction, we use the bucket value associated with the bucket that the approximate nearestneighbor maps to.
To control the error when doing such a procedure, we take enough samples toapproximately form an (cid:15)/ Γ L cover of Γ for a large enough constant Γ . The size of an (cid:15)/ Γ L cover ofΓ is O ((2Γ L √ d/(cid:15) ) k ) .
This implies that, via a coupon collector argument, when the input distributionis uniform over the region Γ , O ( k (2Γ L √ d/(cid:15) )
k log(2Γ L √ d/(cid:15) ))
samples will ensure that with veryhigh probability, for every test point x there exists a train example x i such that (cid:107) x − x i (cid:107) 2 ≤ 2 (cid:15)/ Γ L .
The test error is | f ( x ) − ˆ f
( x i ) | ≤ | f ( x ) − f ( x i ) | + | f
( x i ) − ˆ f
( x i ) |
= O ( (cid:15) ) .
Computing the exactnearest neighbor is a slow process.
Instead we can compute the approximate nearest neighbor using
LSH very quickly.
We lose another O ( (cid:15) ) error due to this approximation.
Choosing Γ appropriatelywe can make the ﬁnal error bound exactly (cid:15) .
This leads to our stated sample complexity bound. This implies that (cid:107) f
− ˆ f (cid:107) ∞ ≤ (cid:15).
Hence using an Euclidean LSH with O ( k ) hyperplanes we canlearn an (cid:15) -approximation to f .
The time to compute ˆ f ( x ) for a new example is the time required tocompute the bucket id where it maps to.
Since there are k hyperplanes and our input is d -dimensional,computing the projections of x on the k hyperplanes takes O ( dk ) time.
Then we need to perform adivision by the width parameter (cid:15)/L .
This takes time equal to the number of bits required to represent
L/(cid:15) .
Hence the total time taken would be O
( dk log( L/(cid:15) )) .
The above theorem uses a lemma about Euclidean LSH, which we prove next.
Proof of Lemma 4.2.
Let K
= Ck .
Let the random hyperplanes chosen by the Euclidean LSH bea 1 , . .
. , a K .
Let the width parameter used by the LSH be (cid:15) LSH .
The value of (cid:15) LSH we choose willbe determined later.
Since the distribution of entries is spherically symmetric, the projection of thevectors onto the k -dimensional subspace will also form a Euclidean- LSH model.
Henceforth in ouranalysis we can assume that all our inputs are projected onto the k -dimensional space Γ and that weare performing LSH in a k -dimensional space instead of a d -dimensional one.
Let A =
[ a 1 , . .
. , a K ] (cid:62)be the matrix whose columns are the vectors perpendicular to the hyperplanes chosen by the LSH.
Note that A ∈ R K × k .
Then we have, from tail properties of the smallest singular value distributionof Gaussian random matrices (e.g. see [10, 35]), for a large enough constant c ,
[Equation]
For two points x 1 , x 2 ∈ Γ to map to the same LSH bucket, (cid:107) A ( x 1 − x 2 ) (cid:107) ∞ ≤ (cid:15) LSH .
This impliesthat (cid:107) A ( x 1 − x 2 ) (cid:107) 2 ≤ (cid:15) LSH √ k .
Together with (3), this implies that (cid:107) x 1 − x 2 (cid:107) 2 ≤ (cid:15) LSH /c withprobability ≥ 9 / 10 .
At the same time, since we x 1 , x 2 ∈ [ − 1 , 1] d , the maximum distance alongany direction is at most the length of any diagonal, which is 2 √ d .
Moreover, along any hyperplanedirection sampled by the LSH, we grid using a width (cid:15) LSH .
Since the total number of hyperplanes is
Ck the maximum number of LSH buckets possible is (cid:16) 2 √ d (cid:15) LSH (cid:17) O ( k ) .
We set (cid:15) LSH = c(cid:15) .
Then, thediameter of any bucket ≤ (cid:15) LSH /
A.2 Proof of the Dense Lower Bound
Proof of Theorem 4.3.
Assuming B bits per parameter, in our dense layer model we have 2 Bwdistinct possible conﬁgurations.
We lower bound the width w by constructing a class of functions Fdeﬁned on a k -dimensional subspace within [ − 1 , 1] d such that three properties simultaneously hold:
each f ∈ F is L -Lipschitz,
the number of functions in F is at least Ω(2 (2 √ dL/C(cid:15) ) k )3.
for f 1 (cid:54) = f 2 ∈ F , we have (cid:107) f 1 − f 2 (cid:107) ∞ > (cid:15) .
These three properties together will imply that w ≥ 1 B ˙(2 √ d/C(cid:15) )
k as otherwise by there would have to be two functions f 1 (cid:54) =
f 2 ∈ F that are approximated simultaneously by the same dense network, which is impossible since (cid:107) f 1 − f 2 (cid:107) ∞ > (cid:15) .
We construct F as follows.
Given the d -dimensional cube [ − 1 , 1] d , we pick a subset of k diagonals of the cube such that they are linearly independent.
We consider the k -dimensional region deﬁned by the intersection of the subspace generated by these diagonals and the cube [ − 1 , 1] d .
Denote the region we obtain by G .
Let e 1 , . .
, e k form an orthonormal basis for the subspace G lies in.
We grid G into k -dimensional cubes of side length 2 (cid:15)/L aligned along its bases { e i } ki =1 .
For the center of every cube we pick a random assignment from { + (cid:15), − (cid:15) } .
Then we interpolate the function everywhere in G such that (i) it satisﬁes the assigned values at the centers of the cubes and (ii) its value decreases linearly to 0 with radial distance from the center.
That is, given the set of cube centers V
[Equation]
To understand the Lipschitz properties of such an interpolation, note that the slope at any given point in G is either 0 or L , which bounds the Lipschitz constant by L .
The total number of cubes that lie within G is at least ( √ dL/C(cid:15) )
k for some constant C and hence F contains a total of (2) ( √ dL/C(cid:15) )
k functions.
Moreover, given any f 1 , f 2 ∈ F such that f 1 (cid:54) = f 2 , there exists a cube center where their values differ by 2 (cid:15) giving us the third desired property as well.
Consequently, we get that to attain (cid:15) -uniform
error successfully on F we need
[Equation]
which implies that w = Ω(( √ dL ) k / ( C(cid:15) ) k ) .
B LSH Models Can Also Learn Lipschitz Functions on k -Manifolds
A k -dimensional manifold (referred to as a k -manifold) can loosely be thought of as a k -dimensionalsurface living in a higher dimensional space.
For example the surface of a sphere in 3-dimensions is a2-dimensional manifold.
We consider k -manifolds in R d that are homeomorphic to a k -dimensionalsubspace in R d .
We assume that our k -dimensional manifold M k is given by a transform f : R k →
R k applied on k -dimensional subspace of R d L k .
To control the amount of distortion that can occurwhen going from L k to M k , the Jacobian of f is assumed to have a constant condition number for allx ∈
L k .
We now state our main upper bound for manifolds, showing that LSH models can adapt andperform well even with non-linear manifolds of a bounded distortion from a linear subspace.
Theorem B.1.
For any f : [ − 1 , 1] d → R that is 1 -Lipschitz, and for an input distribution Dthat is uniform on a k -manifold in [ − 1 , 1] d , an LSH model can learn f to (cid:15) -uniform error with
O ( k √ dk k log( √ dk/(cid:15) ) /(cid:15) k ) samples using a hash table of size O ( √ dk k /(cid:15) k ) with probability ≥ 0 . 8 .
The total time required for a forward pass on a new test sample is O ( dk log(1 /(cid:15) )) .
Proof.
The main idea of the proof is to follow similar arguments from Theorem 4.1 on the subspace
L k and try to bound the amount of distortion the arguments face when mapped to the manifold
M k .
Since we are no longer dealing with a subspace (linear manifold), the argument that an LSH ind -dimensions can be viewed as an equivalent LSH in k -dimensions does not hold.
We use EuclideanLSH models with O ( d ) hyperplanes.
Furthermore, we will use multiple LSH models each deﬁned using O ( d ) hyperplanes.
The main challenge in the proof is to show that the total number of bucketsused in approximating f do not grow exponentially in d , which is a possibility now as we use O ( d )hyperplanes.
Lemma B.2.
For any x ∈ R d , a d -dimensional sphere of radius O ( (cid:15)/d ) centered at x is fullycontained in the bucket where the center of the sphere maps to with probability ≥ 0 .
Proof.
Along any hyperplane direction the gap between parallel hyperplanes is (cid:15) .
Since any pointis randomly shifted before being mapped to a bucket we get that with probability 1 − O (1 /d ) , x ismore than Ω(1 /d )
away from each of the two parallel hyperplanes on either side.
So with probability(1 − O (1 /d ))
O ( d ) = Ω(1) the entire sphere is contained inside the LSH bucket x maps to.
Lemma B.3.
Using O ( k log d ) Euclidean- LSH functions, we get that every x ∈ L k , there exists abucket in at least one of the O ( k log d ) buckets x gets mapped to such that the entire k -dimensionalsphere of radius O ( (cid:15)/d ) centered at x is contained within the bucket.
Proof.
We use a covering number argument.
The maximum volume of a k -dimensional subspacewithin [ − 1 , 1] d is (2 √ d ) k .
We cover this entire volume using spheres of radius (cid:15)/d .
The total numberof spheres required to do this are O
((2 d √ d ) k /(cid:15) k ) .
We now do a union bound over all the spherecenters in our cover above.
For a single sphere, the probability that it does not go intact into a bucketin any of the O ( k log d ) LSH functions is d − Ω( k ) .
By a union bound, we can bound the probabilitythat there exists a sphere center that does not go into a bucket to be d − Ω( k ) .
Hence the Lemmastatement holds with exceedingly large probability of 1 − d Ω( k ) .
Now, we only include buckets with volume at least (Ω( (cid:15)/ ( d √ k )))
k .
We can do this procedure usingapproximate support estimation algorithms [39].
This takes time and sample complexity S/ log Swhere S is the size of the support.
With constant probability all points in L k are mapped to somesuch high volume bucket in at least one of the LSH functions.
The total number of buckets withthis minimum volume is at most ( O (( d 2 √ k ) /(cid:15) ))
k , which is also an upper bound on the samplecomplexity and running time of the support estimation procedure.
Now, we lift all the above resultswhen we go to M k from L k .
Since the Jacobian of the manifold map f has a constant conditionnumber, its determinant is at most exp( k ) ; so the volume of any region in L k changes by at mostan exp( ± O ( k )) multiplicative factor when it goes to M k .
So all volume arguments in the previousproofs hold with multiplicative factors exp( ± O ( k )) .
This concludes our proof.
C Lower Bound for Analytic Functions
The functions described in the lower bound presented earlier are continuous but not differentiableeverywhere as they are piecewise linear functions.
In Theorem C.1 we show that we can makethe lower bound stronger by providing a construction of L -Lipschitz analytic functions (which aredifferentiable everywhere).
See Figure 6 for an example of analytic vs non-analytic functions.
Theorem C.1.
A dense model of width w with a random bottom layer requires
[Equation]
where C 1 is a large enough constant, to learn L -Lipschitz analytic functions on [ − 1 , 1] d to (cid:96) ∞ error (cid:15) when the inputs are sampled uniformly over a unknown k -dimensional subspace of R d ∩ [ − 1 , 1] d .
Moreover, the number of samples required to learn the above class of functions is Ω ( w log w ) , where w = Ω
[Equation]
( √ kπ(cid:15) )
k (cid:19) .
Proof of Theorem C.1.
We construct a family F of analytic functions that are L -Lipschitz and are described using the Fourier basis functions.
Each f ∈ F will be of the form
[Equation]
for x ∈ [ − 1 , 1] k .
We pick a small value of 0 < (cid:15) 1 < 1 .
We assume 1 /(cid:15) 1 is an integer for convenience.
If it is not, we can simply take (cid:100) 1 /(cid:15) 1 (cid:101) instead.
For a set of integers ( n 1 , n 2 , . . . , n k ) ∈ [1 /(cid:15) 1 ] k , letη n 1 n 2 ...n k ∈ {±} .
We use η n as a shorthand when it is not ambiguous.
The family F is deﬁned asthe set of functions f below
[Equation]
where each η n is chosen to be either ± L/ ( C √ kπ ) for a large enough constant C and α will bedetermined later.
There are (1 /(cid:15) 1 ) k Fourier bases in each f and the coefﬁcient of each is set to be± L(cid:15)
α 1 / ( C √ kπ ) .
Hence we have
[Equation]
Next we argue that a larger than 0 .
9 fraction of the functions in F are L -Lipschitz.
We have,
[Equation]
where the last expectation is over the uniform measure over functions in F .
To get a bound on(cid:107)∇ f
( x ) (cid:107)
we bound each ( ∇ f ( x ))
i with high probability.
Each ( ∇ f ( x ))
i is a sum of (1 /(cid:15) 1 ) k independent random variables, namely η n .
We saw above that E [( ∇ f ( x )) i ]
= 0 .
To bound | ( ∇ f ( x ))
i |with high probability we will use McDiarmid’s inequality.
An upper bound on how much the valueof ( ∇ f ( x ))
i can change when any one η n ﬂips in value is computed as 4 (cid:15) ( α − 1)
L/C √ k .
Then, an application of McDiarmid’s concentration inequality gives us that,
[Equation]
with probability ≥ 0 .
9 for a large enough constant C .
This implies that
[Equation]
with probability ≥ 0 .
9 for a randomly sampled f ∈ F .
Now, let η f denote the vector of η n values insequence for any f .
Using McDiarmid’s (or Hoeffding’s) concentration bound again, we also get that,with probability ≥ 0 .
9 , the Hamming distance between η f 1 and η f 2 for two f randomly sampled from
F is at least c (1 /(cid:15) 1 ) k for a small enough constant c < 1 .
This implies that for randomly sampledf 1 , f 2 ,
[Equation]
where η (cid:48) n is non-zero for at least c (1 /(cid:15) 1 ) k of the terms from the above argument about the Hammingdistance.
Parseval’s identity then implies that
[Equation]
Finally we note that by union bound, at least a 0 .
8 fraction of the functions in F satisfy both our
Lipschitzness property (9) and (11) simultaneously.
Setting α = k/ 2 + 1 and (cid:15) 1 = C √ kπ(cid:15)
L √ c 2 k/ 2 we get that to achieve a strictly smaller error than 2 (cid:15) in the (cid:107) .
(cid:107) ∞ sense
, one requires a dense model with a width of
[Equation]
D Experiment details
D.1 Details of learning random functions
In Section 5, we demonstrated experiments for randomly generated polynomial/hypercube functions.
Here we present the details for the experiment settings.
Random function generation.
For the random polynomial functions, we randomly generate coefﬁcients of the monomials by sampling from a uniform distribution U ([ − 1 , 1]) and scale the coefﬁcientsso that their absolute values sum up to 1 .
(this is to ensure the Lipschitz constant of the generatedfunction is bounded by a constant independent of dimension and degree of the polynomial).
For therandom hypercube function, we sample values of the function at each corner independently from auniform distribution on − 1 , 1 , and interpolate using the indicator functions.
Train/Test dataset generation.
For a given target function f (polynomial or hypercube), we sampleindependently from U ([ − 1 , 1] n ) (where n is the input dimension) to generate the input features x and compute target value y = f ( x ) .
The train dataset contains 2 16 ( x, y ) pairs and the test datasetcontains 2 14 ( x, y ) pairs.
Training setting.
All the models in Section 5 are trained for 50 epochs using the RMSProp [17]optimizer with a learning rate of 10 − 5 .
For the one dimension example in Section 1, the model istrained for 200 epochs using the RMSProp optimizer with a learning rate of 5 × 10 − 6 .
Random hash sparse model.
We discussed the design of DSM and LSH models in Section 2.
Herewe present the details of the random hash model, where the sparsity pattern is determined by a randomhash of the input data (i.e. the same input data would always have the same sparsity pattern).
Thefollowing code snippet shows the generation of a random mask that only depends on the input datausing TensorFlow 2.x.import tensorflow as tf648 # seed: a fixed random seed# inputs: the input tensor# mask_dim: size of the masked tensor# num_buckets: a large integer# k: the dimension after masking654 input_dim =
inputs.shape[-1]if input_dim !
= mask_dim:proj =
tf.random.stateless_normal(shape=(input_dim, mask_dim),seed=
seed)inputs =
tf.einsum(’...
i,io->...o’, inputs, proj)hs = tf.strings.to_hash_bucket_fast(tf.strings.as_string(inputs),num_buckets=
num_buckets)top_k_hash = tf.expand_dims(tf.nn.top_k(hs, k).values[..., -1],axis=-1)mask = hs >=
top_k_hash
D.2 Learning random polynomials under other parameter settings
We present experiment results for learning random polynomial target functions with low intrinsicdimensions.
To be precise, the target polynomial is p ( Ax ) , where p is a polynomial of degree dwith sum of coefﬁcient absolute value < 1 , x ∈ R n is a vector, A ∈ R k × n is a matrix with randomorthonormal rows, and n > k .
Note now the intrinsic dimension of the domain is k , while the inputsx has higher dimension n .
In Figure 7, we compare the mean squared loss for dense models and
DSM s for n = 64 , k = 8 , and d = 4 .
We observe similar behavior as Figure 2, where the inputdimension is the same as the intrinsic dimension.
This validates our analysis in Section 3.
[Figure]