Title: SS1: Accelerating Inference with Fast and Expressive Sketch Structured Transform
Venue: NeurIPS 2024
Accepted: yes
Keywords: fast linear transform

parameter sharing

latency improvement

deployment
Abstract: Tensor multiplication with learned weight matrices is the fundamental building block in deep learning models. These matrices can often be sparsified, decomposed, quantized, or subjected to random parameter sharing without losing accuracy, suggesting the possibility of more efficient transforms. Although many variants of weight matrices exist, unstructured ones are incompatible with modern hardware, slowing inference and training. On the other hand, structured variants often limit expressivity or fail to deliver the promised latency benefits. We present Sketch Structured Transform (SS1), an expressive and GPU-friendly operator that accelerates inference. SS1 leverages parameter sharing in a random yet structured manner to reduce computation while retraining the rich expressive nature of parameter sharing. We confirm empirically that SS1 offers better quality-efficiency tradeoffs than competing variants. Interestingly SS1 can be combined with Quantization to achieve gains unattainable by either method alone, a finding we justify via theoretical analysis. The analysis may be of independent interest.
Moreover, existing pre-trained models can be projected onto SS1 and finetuned for efficient deployment. Surprisingly, these projected models can perform reasonably well even without finetuning. Our experiments highlight various applications of the SS1:
(a) Training GPT2 and DLRM models from scratch for faster inference. (b) Finetuning projected BERT models for 1.31× faster inference while maintaining GLUE scores. (c) Proof of concept with Llama-3-8b, showing 1.11× faster wall clock inference using projected SS1 layers without finetuning. We open source our code :https://github.com/apd10/Sketch-Structured-Linear/

SS1: Accelerating Inference with Fast and Expressive Sketch Structured Transform

Abstract

Tensor multiplication with learned weight matrices is the fundamental building block in deep learning models. These matrices can often be sparsified, decomposed, quantized, or subjected to random parameter sharing without losing accuracy, suggesting the possibility of more efficient transforms. Although many variants of weight matrices exist, unstructured ones are incompatible with modern hardware, slowing inference and training. On the other hand, structured variants often limit expressivity or fail to deliver the promised latency benefits. We present Sketch Structured Transform(SS1), an expressive and GPU-friendly operator that accelerates inference. SS1 leverages parameter sharing in a random yet structured manner to reduce computation while retraining the rich expressive nature of parameter sharing. We confirm empirically that SS1 offers better quality-efficiency tradeoffs than competing variants. Interestingly SS1 can be combined with Quantization to achieve gains unattainable by either method alone, a finding we justify via theoretical analysis. The analysis may be of independent interest. Moreover, existing pre-trained models can be projected onto SS1 and finetuned for efficient deployment. Surprisingly, these projected models can perform reasonably well even without finetuning. Our experiments highlight various applications of the SS1: (a) Training GPT2 and DLRM models from scratch for faster inference. (b) Finetuning projected BERT models for 1.31× faster inference while maintaining GLUE scores. (c) Proof of concept with Llama-3-8b, showing 1.11× faster wall clock inference using projected SS1 layers without finetuning. Our code is open-source. 3

1 Introduction

Tensor-matrix multiplication is one of the fundamental operations in deep learning models across various domains. Linear transformation are especially crucial in transformer architectures, which form the backbone of foundational models responsible for the advanced capabilities of Large Language Models (LLMs). A significant portion of the computational load in LLMs comes from linear layers. For example, with a batch size of 4 at full sequence length, the MLP workload of even the smallest Llama-3-8B model [1] involves (32768 × 4096 × 4096) operations in attention and (32768 × 4096 × 14336) in the MLP. Larger models have even larger workloads. Given these demands, finding efficient and expressive alternatives to standard linear layers is a crucial research direction.

Deep learning models often have significant redundancies that can be removed using various techniques. Popular approaches include Sparsification [2, 3, 4, 5, 6], Quantization [2, 7], Randomized

Parameter Sharing (RPS) [8, 9, 10], and Low-rank decomposition [11, 3]. Unstructured sparsity, while expressive, is inefficient on modern hardware. Structured sparsity [12, 13, 14] and low-rank are efficient but often compromise quality. Recently, it was shown in [15] that sparsity is theoretically weaker than RPS methods regarding the quality of learned models under compression. However, traditional RPS methods only focus on reducing the parameter memory footprint and do not affect FLOPs or latency. This paper proposes Sketch Structured Transform (SS1), an RPS scheme that reduces FLOPs while maintaining quality and improving wall-clock inference for machine learning models.

Randomized parameter sharing (RPS) reduces the model's memory footprint by randomly tying parameters across the model. While these methods offer a superior memory-quality tradeoff, they do not alter the computational graph, resulting in similar FLOPs and latency. Addressing this issue - reducing computation and latency while retaining the memory-quality tradeoff - has been a key challenge with RPS methods. In this paper, we propose a random yet structured parameter-sharing method that maintains superior quality, reduces the number of FLOPs, and improves the wall-clock latency of linear layers which are compute bound in most workloads.

The key idea in SS1 is to tie parameters inside a single neuron weight. This weight-tying can be equivalently implemented as the input being reduced in dimension first and then multiplied with the compressed weight vector which reduces FLOPs and memory movement in a neuron computation. We show how to devise this tying in a GPU-friendly manner using what we call K -and N -coalescing. Additionally, SS1 can be integrated with existing RPS methods to obtain independent control of parameter memory and computation during training and inference. While SS1 layers can be used to build a model itself, we can also obtain SS1 models from pre-trained models. This is especially important since many useful models are pre-trained once, and weights are open-sourced for public usage. For this purpose, we provide a projection function to project full matrices onto SS1 matrices, which can transfer knowledge from pre-trained models. Moreover, we demonstrate that SS1 can be effectively combined with post-training quantization to harness the advantages of both approaches. Additionally, we present theoretical insights explaining why the integration of SS1 and quantization leads to performance improvements that neither method can achieve independently.

We evaluate SS1 layers on a broad set of settings and use cases. A summary of our findings is below.

SS1 has better quality-efficiency tradeoff than competing methods like Monarch[16] (state-of-theart structured sparsity) and LowRank[2] across various domains. With SS1, we can build better models at lower parameter counts while delivering superior inference latency. For example, we achieve up to a 1 . 30 × improvement in GPT2 [17] model inference throughput.

Pretrained models can be projected onto SS1 and further finetuned to deploy fast models. We show that we can maintain the GLUE[18] score of BERT[19] while speeding up inference by 1 . 31 × .

SS1 projected models can be used even without finetuning with reasonable accuracies. We show proof-of-concept benefits in the Llama-3-8B model[1] with 1 . 11 × faster inference.

Quantization[2] is a highly effective technique for improving efficiency that can be combined with SS1. Not only do we see this empirically, but we can also see it in theory.

SS1 also impacts CPU ML workloads significantly. For instance, we reduce the MLP workload of DLRM[20] MLPerf Model, which contributes over 70% inference latency, by approximately 2 × using SS1 layers without compromising model quality.

2 Related Work and Background

Unstructured Sparsity: The redundancy in the deep learning model can be removed by sparsifying the model using iterative procedure [2, 3, 4]. Another related line of work is that of Lottery ticket hypothesis [5, 6, 21], which tries to find a sparse model at the start of training. Apart from expensive procedures to find these subnetworks, the unstructured sparsity still needs to deliver on the promise of latency improvements for inference.

Structured Sparsity: Linear transformations that are efficient and expressive has been an active line of research for over a decade. The general direction here is to create a combination of sparse, diagonal, permutation, and sub-linear transformations such as FFT, DCT, and Hadamard [14, 22, 13, 23, 24]. Some of them also used fixed random sparse/diagonal matrices. Structured matrices proposed in this line of research for transformation of size K × K have O ( K ) parameters and O ( Klog ( K )) FLOPs. The actual speed-up obtained on modern hardware using these methods is limited. More importantly, the expressivity of these matrices is severely restricted due to very few learnable parameters. A

recent line of work exploits butterfly matrices [25, 16] and their variations, such as Monarch to obtain expressive transformations. These matrices also fit the general recipe of structured matrices specified above. However, a critical distinction in Monarch matrices is the presence of many more parameters. Specifically, the pure monarch matrices are supposed to contain O ( K √ K ) parameters and FLOPs when using √ K blocks or factors (more blocks imply fewer parameters). However, practically, the matrices that have enough representative power in deep learning context generally use two or four blocks [16], i.e. O ( K 2 ) parameters and FLOPs. Nevertheless, Monarch shows great latency benefits for training with scheduled different-sized monarch decomposition. However, to our knowledge, the inference benefits for these matrices are limited. Our paper compares SS1 against Monarch since it is the SOTA representative for this class of matrices.

Quantization: Post training quantization [2] is currently one of the most successful tools for improving large models' efficiency. The basic idea is to reduce the precision of the weights and activations to reduce the memory footprint of the model and also exploit integer arithmetic [7] for better compute throughput. The literature on quantization is vast, and more details can be found in the survey [26]. In this paper, we show that we can combine SS1 with quantization techniques to further improve latency. We also explore the theory of combining quantization with SS1 approximation in section 4, which can be of independent interest.

Randomized Parameter Sharing: Randomized parameter sharing was first introduced in [8] as a general model compression tool. RPS is primarily used to reduce the parameter memory footprint of models. The parameter memory is separated from the model's actual computational graph. Each model's weight is mapped to the parameter memory using a random hash function. The value of the weight is then the value from the parameter accessed via hash functions. Essentially, if two weights are mapped to the same value in parameter memory, they are tied together and share a single learnable parameter.

The randomness of weight tying that leads to theoretical guarantees concerning projection quality was also the reason behind the extremely slow systemic performance of the proposal. The systemic performance was fixed using block-based hash mappings showing practical applications in various domains [10, 9, 27]. Surprisingly, the projection quality improves with block-based projections[10] leading to a strictly superior RPS system. RPS quality was further enhanced using global parameter sharing (parameter sharing across modules) which is superior to module-specific parameter sharing [9]. This, however, leads to additional challenges regarding the stability of training the model. These challenges were resolved in [15]. This work also proved a missing link in the theoretical analysis of RPS methods - it showed that the quality of dimensionality reduction (alt. projection quality) directly correlates with the quality of models learned under projection for linear models. Further, it was shown that random dense projections (which underlie RPS) are superior to random sparse projections (which underlie Pruning) which justifies why RPS methods convincingly outperform pruning methods, especially at high compressions in [15]. Interestingly, most of the RPS literature is strictly focused on reducing the memory footprint and does not affect the computational workload. In this work, we further improve the utility of RPS techniques by deploying them to reduce computation.

2.1 Background on Randomized Parameter Sharing required for SS1

Tying of parameters using hash functions ( h, g ): Under RPS, there is a single parameter memory M , |M| = m , and each weight inside the model is mapped to one of the parameters in the memory using hash functions. Let the flattened weight vector of the entire model be θ , | θ | = n . A weight, say θ [ i ] , is uniquely identified using a set of integers (module number, location inside the module, etc.), say id ( i ) ∈ N k for some k , and then the value of the weight is

Sketch representation of recovery: We can write the entire recovery of the vector θ as a linear projection from M using a sketch matrix S ∈ R n × m which is defined as

̸

Then the weight θ and the parameter memory M is related by, θ = S M

Dimensionality reduction problem: Traditionally, random projections are used as a dimensionality reduction technique for data points. The projection is considered better if it can maintain the structure of the dataset (inter-point distances or equivalently inner products between points). The

W:K×N

2D super group

Rotationofweightsinagroup

using hashfunction

K-coalescing:

01234

３40１２

shareparameter

chunks

1D super

BK

dnob

BN

N-coalescing:share the same hashfunction

(denotedbyrotationoffsetposition)

standard theoretical setup is as follows. Consider two arbitrary vectors x, y ∈ R n , Consider a projection function PROJ (or compression technique in general), then ˆ x, ˆ y are projected vectors, ˆ x = PROJ( x ) , ˆ y = PROJ( y ) . Then, the quality of the projection is measured by

where the expectation is over the randomization scheme of the projection. This problem is of interest to evaluate RPS-based learning since the quality of the learned model under projection f ( S M ) correlates with the quality of data projection S ⊤ for the standard dimensionality reduction problem with ˆ x = S ⊤ x and ˆ y = S ⊤ y [15].

3 Sketch Structured Transform(SS1)

The SS1 layer falls into the category of RPS methods. We will use small case letters for scalars, boldface small cases for vectors, and boldface capitals for matrices. We use numpy notation of indexing. For instance z [(range( n ) + 5)% n ][ i : j ] first rotates the vector z by 5 places and then selects a subarray from i to j .

3.1 Parameter sharing in SS1

In standard RPS, where the weights are mapped into a single memory M , with a high probability, weights corresponding to a single neuron are not tied together. While this seems advantageous for expressivity, it is also the reason why RPS methods cannot reduce computation since we need to perform O ( K ) multiplications for a single neuron. In SS1, we perform restricted parameter sharing where parameters are tied inside a single neuron only. Consider a linear transform y ⊤ = x ⊤ W where x ∈ R K , y ∈ R N and W : R K × N matrix.

Single Neuron RPS: Let us consider a single neuron y = x ⊤ w . Under SS1, the weights w ∈ R K come from compressed parameter vector z ∈ R K//c where c is the compression factor. We use // to denote integer division. For simplicity, we assume throughout the section that c is an integer and c | K . Each neuron has its own z . Weights w are recovered from z using standard RPS. i.e.,

where S : K × K//c is sparse matrix according to Equation 1. When parameters are shared in this manner, the computation can be reduced.

Thus computing y needs only K//c multiplications since x ⊤ S , a sketch of input x , can be implemented using only additions and subtractions. As is, this randomized mapping will not be cache-efficient and thus is not GPU-friendly. We show how to make it GPU-friendly using K and N -coalescing in the subsequent parts of this section.

K -coalescing: The single-neuron computation shown above is not GPU-friendly for arbitrary choice of RPS mapping hash function h . We now explain the hash function h used in SS1.

Equivalently, w can be represented as

The parameter tying is illustrated in Figure 1. We first divide the weight vector w into chunks or groups of size B K , a hyper-parameter. The chunk-id and offset inside the chunk of weight can be written as,

Then c ( recall that c is the integral compression factor ) chunks are grouped together into 1D supergroups. The 1D super group ID of the weight can be written as,

We want to restrict the hash function in SS1 to ensure the following,

Each chunk of size B K is contiguously located in z and do not share weights.

chunks that belong to the same 1D supergroup share weights among each other.

We use the following mapping h ′ : [ K ] → [ K//c ] to satisfy both conditions. Let h : N × N → N be a universal hash function. Then,

Note that many hash functions which follow the two conditions exist. However, we will stick with this hash function since this is the function we implement. With this hash function, the j th chunk of i th 1D super group in w can be written as,

where R (0 : n ) = (0 , 1 , ..., n -1) is the range function. In matrix form using Equation 1, let S i,j : B K × B K be the matrix representation of above hash function restricted to j th chunk of i th 1D supergroup. Then,

Then, we can write the single neuron computation as,

We can implement the above computation as follows: Bring in one chunk of z ; then bring in c chunks belonging to the 1D super group of x one at a time and aggregate them into a single chunk of size B K using sketch matrices S ; perform the dot product on the chunk of z and aggregated chunk of x and move on to the next chunk of z and 1D supergroup of x . Note that we can perform this operation in a block manner and using exactly one read of z and x due to the nature of the hash function.

N -coalescing: Now let us move our discussion from single neuron computation in SS1 to computing all the neurons, i.e. the entire matrix multiplication. In this case we have the complete weight matrix W : K × N which is derived from a compressed matrix Z : K//c × N . Apart from K - coalescing specific to each neuron, to fully utilize the GPU capabilities, we further restrict parameter sharing along the N dimension. The weight mapping is illustrated in Figure 1. We divide the weights along the N dimension in column blocks of size B N . Each neuron that belongs to the same block will have the same set of hash functions. This allows us to do a block-based computation as,

Note that each neuron has its own parameter space in Z . The complete algorithm is presented in Algorithm 1. It considers batched multiplication and uses B M hyper-parameter, which is standard in matrix multiplication. Also, it uses a slightly different notation for ease of expression. However, it maintains the principles introduced in this section. The kernel implementation and optimization details are deferred to Appendix E for lack of space.

3.2 SS1 projection from the pre-trained model

A pre-trained model can be projected into SS1 and fine-tuned for downstream tasks. The complete model projection boils down to projecting each full matrix in a linear layer, say W : K × N into parameter tied SS1 matrix Z : ( K//c ) × N . Recall that SS1 has parameter sharing independent in each neuron. Thus, we just have to find the projection for each neuron's weight. Consider the neuron

Algorithm 1 SS1(Z, X)

y = x ⊤ w and corresponding compressed with z and recovery sketch matrix S . Given w , we want to find z that minimizes

The solution is that of linear regression, z ∗ = ( S ⊤ S ) -1 S ⊤ w . Note that ( S ⊤ S ) is a diagonal matrix with non-zero diagonal elements. Once we solve for the value of z for each neuron in each weight matrix, we have our overall projection onto SS1. The algorithm for projection can also be implemented in a blocked manner and is given in Algo 2 in Appendix G

3.3 SS1 combined with standard RPS.

SS1, which reduces the computation of the linear layers, can be combined with RPS methods, which reduce the parameter memory to create a holistic efficiency system - one that lets you control both memory and computation independently. The only change to Algo 1 is in where we read the weight tile. In the current algorithm, we maintain Z : K//c × N , the compressed weight matrix, in a row-major format. Under RPS, we can locate the tile Z [ k B K : ( k +1)B K , j B N : ( j +1)B N ] in the RPS memory, say M , using another hash function, say h ′ , at M [ h ′ ( k, j )% |M| B K B N ] .

4 Theoretical aspects of SS1

Quantization + SS1 can beat the individual methods: In this section, we analyze a combination of quantization and SS1 (i.e., projection) in the standard dimensionality reduction setup. We consider stochastic integer quantization for our analysis. Consider two vectors x , y ∈ R n . We aim to reduce the memory by a factor of c . In case of quantization, assuming initial precision as F , c = b/F if we are using b-bit quantization. Similarly, for projection, c = m/n if projecting the data to R m . In case of projection the compressed vector can be represented:

Let the component values be restricted to ( -D/ 2 , D/ 2) for some D . The compressed vector in case of stochastic quanti-

n=128

Cg=0.3

Variance of unit normvectors

0:5

ProjectionVariance

Quantization-variance

Quantization+Projection

0!1

0!2

Compression

zation can be written as,

The range D is divided into 2 b buckets and ⌊ . ⌋ q rounds the value down in the buckets. The variance of the inner product between ⟨ ˆ x , ˆ y ⟩ for projection, quantization, and combination that first performs b-bit quantization followed by projection is given below,

Theorem 1 Consider two arbitrary vectors x , y ∈ R n such that || x || ≤ 1 , || y || ≤ 1 in F precision. For some fixed value k in [1 , 2] ,

The variance of quantization ( V ( ⟨ ˆ x , ˆ y ⟩ q ) ), projection ( V ( ⟨ ˆ x , ˆ y ⟩ p ) ) with compression c , can be tightly bounded by

The variance of combination with compression c = c p c q with quantization with compression c q followed by projection with c p is bounded by

The best way to understand this theorem is to plot the variance as a function of compression as shown in the Figure 2. In the plot, we use c q = 0 . 3 , n = 128 and k = 1 . 5 for illustration. As we can see with the combination, the variance is better than individual methods in certain regions. The reason is that both methods do quite well in the low compression regime, and explode steeply at very high compression. Thus, when we combine the two methods, we can remain in relatively low error zones for both methods while obtaining greater overall compression.

Analyzing parameters of SS1: This section discusses the new aspects of SS1 and how it affects the quality of SS1. Specifically, we discuss the effect of parameters B K and B N . From a latency viewpoint, it is natural that these parameters should be autotuned to optimize for latency, which is standard for matrix multiplication. However, since these parameters affect the weight tying, it is important to investigate their effect on the quality of models. We consider the standard dimensionality reduction setup to evaluate the changes in mapping performed by SS1. We find that B K is a pure latency parameter with no impact on learning quality. We present this as the following result.

Theorem 2 Consider two vectors x , y ∈ R n in higher dimensional space. Also, let h and g be the SS1 hash functions used to project these vectors to ˆ x , ˆ y ∈ R m ( m<n ). Let B K be the K-coalescing factor of h . Then, under compression c = n/m and assuming c B K | m | n , the inner product estimation is unbiased and has variance,

̸

where I ( i, j ) is a Kronecker delta function indicating if i and j belong to the different groups in the super group of size c B K . Under permutation p of elements of x and y before projection,

̸

In the above theorem, we find that the quality of projection is unaffected by the parameter B K . This can be understood by considering a specific element in x . The number of other elements with which it can interact under projection depends only on the compression factor c = n/m and not on B K . We know from [15], that the quality of projection is directly related to the quality of learned linear models under compression. Thus, B K does not have any impact on learning quality as well. We analyse the effect of B N as well. But the analysis is deferred to the appendix (see section F.2) for lack of space.

(a) GPT models on wikitext-103 (standard deviation : around 0.2PPL measured for baseline model)

(b) MLPMixer models on CIFAR datasets (Stdev: 0.001 (2x), 0.002 (4x), 0.02 (8x))

5 Experiments

5.1 Accuracy vs. Latency evaluation of SS1

We first evaluate the expressiveness and latency of our SS1 against popular baselines of Monarch, SOTA structured sparsity-based transformation, and LowRank transformation. We benchmark various shapes for the "nn.Linear", Monarch, LowRank, and SS1 kernels to evaluate the standalone kernel latency. We use the official Monarch implementation by the authors 4 with number of blocks (nb) as the compression controlling factor. The results are deferred to the appendix due to space constraints. We also measure end-to-end model latency improvement for various sizes of GPT2 and MLPMixer (MM) [28] models. We show the latency of three sizes of models (small: S, medium:M,

SS1 projection followed by finetuning

SS1 projection without finetuning

and large:L)for both architectures. Exact details are provided in appendix I) Additionally, we perform quality experiments in end-to-end training using these matrices for two domains: (a) NLP: GPT2-S model on wikitext-103 dataset [29] using test perplexity (PPL) and loss as the metric (b) Vision: MLPMixer on CIFAR (C10, C100)[30] and Tiny-Imagenet datasets [31] using accuracy as the metric. The results are presented in Table 1. The details of these experiments can be found in the Appendix I. We make the following observations.

In our kernel latency evaluation, SS1 consistently outperforms Monarch and LowRank under similar parameter budgets (Table 4 in appendix). Monarch is generally worse than full multiplication except in larger shapes. LowRank is competitive with SS1 for higher compression.

The kernel-level latency benefits also translate to end-to-end model latency (Table 1 and Table 1b). This can be seen across both Transformer and MLPMixer architectures, two of the SOTA architectures for deep learning. Although, at times, LowRankLinear gives better performance at higher compression, the quality of the model obtained is inferior to SS1. With larger models, we see an increase of around 1.3 × in end-to-end throughput with SS1. It also enables us to run larger batches.

In terms of quality, SS1 consistently outperforms Monarch and LowRank across NLP and Vision domains on GPT and MLPMixer architectures under similar parameter budgets.

Overall, we see that SS1 allows better quality per unit compute than its structured competitors.

Building expressive and faster models: Should we build and train models with SS1 layers instead of standard linear layers? Although answering this question requires a broader evaluation, we cautiously believe that the answer might be yes. We can indeed obtain better quality models by using SS1 layers instead of "nn.Linear" layer. This is demonstrated in Table 2. The quality of the SS1 model with 74 M parameters is better than that of standard models with similar or more parameters.

5.2 SS1 + Quantization

Quantization is an important efficiency technique to improve inference time memory and latency. Although SS1 already reduces the memory footprint and improves inference latency, it can be combined with quantization to obtain further inference benefits. One might wonder if combining two approximations (SS1 and quantization) can lead to worse-quality models. We show that SS1 can be combined with quantization without significantly impacting quality (also supported by theory) In fact, the effect of quantization on SS1 is similar to the full standard model. The quality details before

PPL(Latency)

and after quantization are in Table 2. To measure the latency benefits, we would have to implement quantized SS1 kernels, which is out of the scope of this paper.

5.3 Dense-to-structured finetuning of pre-trained models for efficiency

We show that pre-trained models can be projected onto SS1-structured space and finetuned for downstream tasks. We demonstrate this by projecting selected linear transformations in BERT to SS1 and training the resulting model on the GLUE Benchmark. We obtain a 1 . 31 × faster model which gives a similar GLUE score in downstream tasks. More details are mentioned in I. The results are summarized in Table 3

We find that we can improve inference of Llama-3-8B model for downstream tasks if we project carefully selected layers onto SS1. The selection is done on a calibration dataset that corresponds to the task at hand. We show that, even without finetuning, we can reduce the computation involved in the model while maintaining reasonable accuracy. Results can potentially improve if we finetune post-projection as we see in BERT. The results are summarized in Table 3.

5.4 Improving CPU workloads e.g. DLRM MLPerf Benchmark

While this paper focuses on providing GPU kernel for SS1, the algorithm can also be utilized to improve the inference performance of CPU deployments. Specifically, recommendation workloads of models such as the Deep learning recommendation model (DLRM) are often run on CPUs due to humungous embedding tables. We consider the DLRM MLPerf Benchmark, which has 100GB embedding tables and 10MB MLP layers. Nevertheless, on CPUs, the latency bottleneck is that of matrix multiplications. Specifically, 71.3% of the time is spent in the top MLP component (see Table 4). We show that we can train from scratch a SS1 version of DLRM, which has half the parameters in MLP and maintains the quality of the full model. While CPU-kernel implementation of SS1 is out of the scope of this paper, we expect that benefits shown on GPUs will also translate to CPUs, and thus, with SS1, we can improve the throughput of DLRM MLPerf Benchmark on CPUs without compromising the quality of the model.

5.5 Comparison with standard RPS method

We also compare SS1 with standard RPS method. Since SS1 focuses on compressing linear layers, we compare it with ROAST-MM compression with each matrix having its individual separate memory. The results are presented in Table 5. As expected, RPS methods do not improve the latency.

Limitations: Our work is mainly limited on two aspects. On the theoretical side, our analysis of parameters and the combination of quantization with SS1 is conducted mainly over linear models, which is a reasonable starting point. However, these results may not necessarily extend to deep learning models, and we plan to explore this in future research. Additionally, regarding our SS1 proposal, the efficiency gains are marginal beyond an 8 × compression of parameters. In future work, we aim to develop efficient structures where the efficiency gains are more closely aligned with the degree of parameter reduction.

6 Conclusion

In this paper, we introduce an efficient structured linear transformation termed SS1. We demonstrate that SS1 outperforms state-of-the-art structured baselines in terms of quality per unit of computation. Additionally, we show that SS1 provides immediate benefits across diverse applications, including language understanding, generative AI, and recommendations. Furthermore, we illustrate how SS1 can be combined with the popular quantization approach to achieve further improvements. We also theoretically explain the underlying reasons that may be of independent interest.

7 Acknowledgments

This work was supported by National Science Foundation SHF-2211815 and grants from Adobe, Intel, Total, and VMware.

References

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. ICLR , 2016.

Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems , 28, 2015.

Song Han, Jeff Pool, Sharan Narang, Huizi Mao, Enhao Gong, Shijian Tang, Erich Elsen, Peter Vajda, Manohar Paluri, John Tran, et al. Dsd: Dense-sparse-dense training for deep neural networks. arXiv preprint arXiv:1607.04381 , 2016.

Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv:1803.03635 , 2018.

Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M Roy, and Michael Carbin. Stabilizing the lottery ticket hypothesis. arXiv preprint arXiv:1903.01611 , 2019.

Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer-arithmetic-only inference. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2704-2713, 2018.

Wenlin Chen, James Wilson, Stephen Tyree, Kilian Weinberger, and Yixin Chen. Compressing neural networks with the hashing trick. In International conference on machine learning , pages 2285-2294. PMLR, 2015.

Aditya Desai, Keren Zhou, and Anshumali Shrivastava. Hardware-Aware Compression with Random Operation Access Specific Tile (ROAST) Hashing. In International Conference on Machine Learning , pages 7732-7749. PMLR, 2023.

Aditya Desai, Li Chou, and Anshumali Shrivastava. Random Offset Block Embedding (ROBE) for compressed embedding tables in deep learning recommendation systems. In D. Marculescu, Y. Chi, and C. Wu, editors, Proceedings of Machine Learning and Systems , volume 4, pages 762-778, 2022.

Tara N Sainath, Brian Kingsbury, Vikas Sindhwani, Ebru Arisoy, and Bhuvana Ramabhadran. Low-rank matrix factorization for deep neural network training with high-dimensional output targets. In 2013 IEEE international conference on acoustics, speech and signal processing , pages 6655-6659. IEEE, 2013.

Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint arXiv:1902.09574 , 2019.

Marcin Moczulski, Misha Denil, Jeremy Appleyard, and Nando de Freitas. Acdc: A structured efficient linear layer. arXiv preprint arXiv:1511.05946 , 2015.

Quoc Viet Le, Tamás Sarlós, and Alexander Johannes Smola. Fastfood: Approximate kernel expansions in loglinear time. arXiv preprint arXiv:1408.3060 , 2014.

Aditya Desai and Anshumali Shrivastava. In defense of parameter sharing for model compression. In The Twelfth International Conference on Learning Representations , 2024.

Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference on Machine Learning , pages 4690-4721. PMLR, 2022.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461 , 2018.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sundaraman, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G Azzolini, et al. Deep learning recommendation model for personalization and recommendation systems. arXiv preprint arXiv:1906.00091 , 2019.

Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In International Conference on Machine Learning , pages 3259-3269. PMLR, 2020.

Yu Cheng, Felix X Yu, Rogerio S Feris, Sanjiv Kumar, Alok Choudhary, and Shi-Fu Chang. An exploration of parameter redundancy in deep networks with circulant projections. In Proceedings of the IEEE international conference on computer vision , pages 2857-2865, 2015.

Vikas Sindhwani, Tara Sainath, and Sanjiv Kumar. Structured transforms for small-footprint deep learning. Advances in Neural Information Processing Systems , 28, 2015.

Zichao Yang, Marcin Moczulski, Misha Denil, Nando De Freitas, Alex Smola, Le Song, and Ziyu Wang. Deep fried convnets. In Proceedings of the IEEE international conference on computer vision , pages 1476-1483, 2015.

Tri Dao, Beidi Chen, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. arXiv preprint arXiv:2112.00029 , 2021.

Amir Gholami, Sehoon Kim, Zhen Dong, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. A survey of quantization methods for efficient neural network inference. In Low-Power Computer Vision , pages 291-326. Chapman and Hall/CRC, 2022.

Aditya Desai and Anshumali Shrivastava. The trade-offs of model size in large recommendation models: 100GB to 10MB Criteo-tb DLRM model. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 33961-33972. Curran Associates, Inc., 2022.

Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mlp-mixer: An all-mlp architecture for vision, 2021.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Ya Le and Xuan S. Yang. Tiny imagenet visual recognition challenge. 2015.

Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning , pages 38087-38099. PMLR, 2023.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771 , 2019.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, highperformance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'AlchéBuc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

Philippe Tillet, Hsiang-Tsung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages , pages 10-19, 2019.

Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXiv:1710.03740 , 2017.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 16344-16359. Curran Associates, Inc., 2022.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding, 2021.

Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019.

Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023.

Felipe Maia Polo, Lucas Weber, Leshem Choshen, Yuekai Sun, Gongjun Xu, and Mikhail Yurochkin. tinybenchmarks: evaluating llms with fewer examples, 2024.

A Additional Details for checklist

1. Resource Details:

MLPMixer experiments: Each run takes 2 hours on a single RTX Quadro 8000 machine

GPT2-Small experiments: Each run takes around 13 hours on four 32-GB V100 GPUs

BERT experiments: QQP, the largest dataset for finetuning with BERT-Large, takes around 7-8 hours on RTX-8000 Quadro.

Llama experiments: on MMLU takes 1 hour on 1 40GB A100

2. License details:

MLPMixer Experiments uses: CIFAR-10 , CIFAR-100 datasets [30] (license: could not find) with MLPMixer GitHub https://github.com/omihub777/MLP-Mixer-CIFAR (MIT License)

GPT2 Experiments uses: Wikitext-103 dataset [29] (Creative Commons AttributionShareAlike License (CC BY-SA 4.0) License) the code is obtained from fly GitHub repo: https://github.com/HazyResearch/fly presented in [16] (Apache-2.0 license), the quantization experiments on GPT2 uses the quantization method employed in [32] with GitHub repo: https://github.com/mit-han-lab/smoothquant (MIT License) Codes are generally obtained from the HuggingFace Transformers library[33] including GPT2 tokenzier the repo is at https://github.com/huggingface/transformers with Apache-2.0 License. The backend of experiments is Pytorch[34] at https://github.com/pytorch/pytorch

Llama3 [1] pre-trained weights and code are obtained from https://github.com/metallama/llama3 (META LLAMA 3 COMMUNITY LICENSE AGREEMENT)

The kernels are implemented with Triton [35] with GitHub repo: https://github.com/triton-lang/triton (MIT License)

Massive Multitask Language Understanding Dataset - MIT License. Obtained via HuggingFace Dataset library.

Winogrande Dataset - Apache 2.0 License. Obtained via HuggingFace Dataset library.

The Microsoft Research Paraphrase Corpus - License unknown. Obtained via HuggingFace Dataset library.

The Multi-Genre Natural Language Inference Corpus - MIT License. Obtained via HuggingFace Dataset library.

The Quora Question Pairs2 dataset -Licensed by Quora. https://www.quora.com/about/tos. Obtained via HuggingFace Dataset library.

The Corpus of Linguistic Acceptability - Obtained via HuggingFace Dataset library.

The Recognizing Textual Entailment (RTE) datasets - Unknown license. Obtained via HuggingFace Dataset library.

The Stanford Sentiment Treebank - CC0: Public Domain license. Obtained via HuggingFace Dataset library.

The Semantic Textual Similarity Benchmark - Unknown license. Obtained via HuggingFace Dataset library.

The Winograd Schema Challenge - Unknown license. Obtained via HuggingFace Dataset library.

Question-answering NLI dataset - CC BY-SA 4.0 license. Obtained via HuggingFace Dataset library.

B Predicted Common Questions (PCQ)

1. Q: How do you choose datasets?

We begin with what Monarch evaluated w.r.t deep learning - train GPT and finetune BERT. Additionally, we wanted to do a quality tradeoff analysis on wider data. So, we chose the vision datasets that LTH uses. Apart from that, we choose the DLRM MLPerf benchmark as a representative of recommendation models.

2. Q: Do you see any training latency benefits?

The computation in the backward pass is not reduced in SS1. So, in general, the backward

pass costs the same as the full model. In this sense, SS1 is geared towards inference benefits. We might see some improvements due to the lower overall memory footprint and metadata that the optimizer stores. But beyond that, we do not expect to see training latency benefits. Under the assumption that inference is more costly than training, this seems to be an okay thing.

Q: Why have you not measured latency with Quantization?

To measure latency with Quantization, we would have to write a quantized SSL kernel. For this paper, we want to restrict our code development to a single kernel. Thus, we only show the quality preservation with SS1 + Quantization. Having said that, we expect the gains from Quantization to be present in SS1 + Quantization.

Q: Does this mean standard linear layers should always be replaced by the SS1 ? We would like to believe there is an "additional" representative value in SS1 layers that can push the boundary of quality per unit compute. However, we want to be cautious while making this claim, and we will be rigorously evaluating SS1 in the future.

5. Q: Why not block-sparse as a baseline? Most of our experiments were done in FP16. Also, our kernels are optimized for FP16. Pytorch block sparse does not have a stable implementation when it comes to FP16. We see segmentation faults. Thus, we skip Block sparse as a baseline. Additionally, with recent work on sparsity vs. RPS, we believe that RPS should beat Block sparsity in quality per unit parameters.

Q: Monarch performance seems lower than what was reported in the original paper. We were surprised as well. We expected Monarch to do better. However, the paper also spoke only about the training times. In our paper, we measure inference times. Additionally, we confirmed with the authors that we were benchmarking the correct code.

Q: Why was expressivity of SS1 not analyzed the way it has been done for other methods such as Monarch? While the analysis of Monarch is interesting and a valuable piece of knowledge, it is debatable if the expressivity analysis is representative of deep learning requirements. Indeed, we show that SS1 beats Monarch in quality vs. parameters ( and compute), emphasizing that deep learning may not align with the notion of common transformations.

Q: Will the results on GPT translate to bigger models such as Llama? It is hard to say. But we believe that whenever the model is at capacity with the information it can learn, having these layers in a slightly bigger model ( so it has a similar computation as a smaller model) can help boost the quality.

Q: Where is the code? Visit our landing page: https://github.com/apd10/Sketch-StructuredLinear for all the details.

Q: What care needs to be taken for dense-to-structured finetuning? Different layers have different sensitivity to compression, especially when we want to preserve the accuracy of the pre-trained model. Thus, we have to identify layers that cannot withstand the approximation. This can be done using some calibration data.

Q: Isn't this just feature hashing? No. Feature hashing (followed by linear transformation) is a low-rank operation, with the first matrix being the fixed sketch matrix. Also, feature hashing is not GPU-aware and has bad latency performance. SS1 applies feature hashing but for each neuron separately. The independence makes SS1 full rank with high probability. Also, due to the coalescing factors, we make SS1 GPU aware and efficient.

C How does learning rate tuning affect the quality?

In our experiments, we always use the hyperparameters that are optimal for the original model. However, these hyperparameters are not optimal for different methods. We perform learning rate tuning for Monarch, a competitive baseline, and find that, while the perplexity improves for all methods with finetuning, the findings remain unchanged. The results are shown in 6

★

SS1

190

+

Monarch

LowRank

Latency (ms)(lower is better)

base model

180

170

160

150

140

19.5

20.0

20.5

21.0

21.5

22.0

22.5

Perplexity (lower is better)

D Different representation of latency quality results

The plot of latency vs. quality for GPT-2 is presented in figure 3.

E Latency vs. Parameters and Kernel Details

We have implemented three kernels for SS1, including forward , backward-weight-gradient , and backward-input-gradient . Our implementations can be used on GPUs with or without tensor cores and support multiple precisions, such as float16 and float32 (TF32).

We implemented permutation by adjusting the addresses to load instead of permuting the data in shared memory. If permuting is done in shared memory, we have to permute the data again to adjust the value distribution to fit the tensor core operands, thus wasting instructions on shared memory accesses and time in synchronizing threads within a CTA (Cooperative Thread Array).

In the forward kernel, we launch ⌈ M/ B M ⌉ × ⌈ N/ B N ⌉ CTAs. Each CTA, while accumulating along the K dimension, permutes the input block (i.e., X ) with each data load. In the backward-inputgradient kernel, each CTA permutes the weight block while accumulating along the N dimension. In the backward-weight-gradient kernel, each CTA only permutes the result before storing it in DRAM using atomicAdd .

We have introduced the following optimizations to improve the performance of SS1 kernels:

Vectorized Memory Load . Rather than permuting individual elements, we permute larger chunks of V EC elements along the K dimension to maximize the GPU's memory bandwidth efficiency. The size of V EC can be selected from 1, 2, 4, or 8.

Pipelining . Utilizing asynchronous memory load instructions supported by the GPU architecture, we preload N stages blocks prior to the main accumulation loop to overlap compute and memory transactions.

Autotuning . To optimize performance across different tensor shapes, we implemented an autotuner to select the ideal B M , B N , B K , and N stages . We also prune configurations that exceed the GPU's shared memory capacity before compilation and benchmarking.

Fusion . We have fused bias computation directly into the forward kernel to reduce kernel launch and memory access overhead.

SS1's forward kernel reduces the number of load instructions and computational overhead compared to full matrix multiplications. However, the backward kernels of SS1 maintain the same level of load instructions and computational demand as those of full matrix multiplications, leading to similar performance outcomes.

Shape =(12288,768,2304)

Shape =(12288,768,50256)

nnLinear

LowRankLinear

Monarch

SSL

nnLinear

LowRankLinear

Monarch

SSL

0.6000

10.0000

8.0000

0.4000

6.0000

4.0000

0.2000

2.0000

0.0000

0.0000

16

16

compression

compression

Shape =(36864,768,3072)

Shape =(36864,3072,768)

nnLinear  LowRankLinear

Monarch

SSL

LowRankLinear

Monarch

SSL

2.0000

1.2500

1.5000

1.0000

0.7500

1.0000

0.5000

0.5000

0.2500

0.0000

0.0000

16

16

compression

compression

F Theoretical Results

F.1 Claim: B K is a pure latency parameter

We will evaluate this under the standard dimensionality reduction setup (which we now know is related to learned model quality). Consider two vectors x and y in R n space. Also, consider the smaller dimension m to which these vectors are being projected using hash functions defined in section 3 with K-coalescing factor B K . Let the final mapping be h . The inner product can be written as,

Assume that B K | m | n for simplicity. Then, the inner product can be written as,

̸

The estimator is unbiased since E ( g ( i )) = 0 . The variance is,

̸

̸

̸

Nonzero terms appear if either of the two cases holds:

̸

Note that

Let I ( i, j ) be the delta function that i, j belong to different groups inside same 1D supergroup

̸

Under a permutation p of the initial input x and y , the following is the probability since every element can interact with all the elements from other groups inside the same supergroup.

Thus, under permutation,

̸

̸

Thus, B K is a pure latency parameter and should be chosen for optimizing shared memory usage.

F.2 How does B N affect the quality?

To investigate the effect of B N , let us consider how neurons in the same 2D super group and different 2D super groups affect each other. We present the result in the following theorem,

Theorem 3 Consider two neurons y 1 , y 2 ∈ R output from a linear transformation under SS1. Let h 1 , h 2 , g 1 , g 2 be the hash functions used for the neurons. Let the weights for these two neurons be p, q ∈ R m . Let the input be x ∈ R n . Then, under the assumption that each signal in x ∼ D is independent, the Covariance between the values of neurons can be written as,

̸

̸

̸

̸

̸

The above theorem points out that if two neurons belong to the same group, i.e., they share the hash functions, then the neuron outputs have unintended correlations and, thus, mutual information. So the larger value of B N implies more neurons having interdependency and, consequently, lesser expressivity. Fortunately, in our experiments, we find that larger values of B N also work well, and we optimize it for latency.

details: Intuitively, sharing the mapping functions across neurons will restrict the representative power of neurons since they will sketch the same parts of the input vector x with the same hash functions. We verify this intuition here. Consider the following two neurons under SS1 with weight vectors p, q

Consider the mutual information between y 1 and y 2 . Since we are dealing with linear computations, we measure the Covariance of the two neurons,

Assuming independent signals in x

Case 1: y 1 and y 2 belong to the same 2D super group

Only non-zero terms when i 1 = i 2 since h ( j ) can only have one value

Case 2: y 1 and y 2 belong to the different 2D super group

Thus, if two neurons are in the same 2D group, the structure imposes the two neurons to be correlated and thus reduces the expressive power of the model.

F.3 Quantization + SS1: How combining them results in a better method for compression.

Analysing quantization under standard dimensionality reduction Consider two vectors x, y with values from the range [ -D/ 2 , D/ 2] . Under stochastic Quantization for b bits, we will have

The estimator is unbiased.

Convince yourself that covariance terms are 0 .

the inner product

Note that the inequality is tight and occurs when all the values of x are exactly at the midpoint of the bin.

where w = D/ 2 b

It follows that this inequality is also tight.

The equation for projection is

̸

we use ≤ t to show that it is tight. The tight upper bound can be represented using some 1 ≤ k ≤ 2 Under same budget m = nb F where F = 32 or F = 16 . for whatever precision width we are using for the original vectors.

if || x || ≤ 1 , || y || ≤ 1

F.3.1 Combined

Let us consider x and y with norm 1. Let us consider the procedure where we first apply projections to get ˆ x 1 and ˆ y 1 , and then Quantization is applied to obtain ˆ x 2 and ˆ y 2 . Let us assume that the bit width for Quantization is b and the projection is into m .

Let's apply Quantization and then projection. The resulting estimator can be written as

̸

The estimator is unbiased. Now let us look at the variance w.r.t all stochastic elements.

̸

̸

Note that each of the term on the right side has covariance 0 with each other since E ( AB ) = E ( A ) E ( B ) = 0 , so Cov ( A,B ) = 0 Hence,

Thus,

effective bits per element is B = ( mb ) /n = F ( m/n )( b/F )) = c p c q F .

̸

̸

̸

̸

̸

G Projection Details

Firstly, note that parameter sharing in SS1 is only inside a single neuron. So, there are independent RPS instances for each neuron. Thus, the projection of the K × N weight matrix is independently projecting each neuron. Consider a single neuron under RPS y ⊤ = x ⊤ w and w = Sz . Given w , the solution to get best z which minimizes || w -Sz || 2 is nothing but solution of linear regression, i.e. z = ( S ⊤ S ) -1 ) S ⊤ w . Also, since S is a sparse matrix with exactly one non-zero in each row, S ⊤ S is a diagonal matrix. The hash function defined in the paper ensures that all elements of the diagonal are non-zero. Thus, it is invertible. In fact, if you view S as a mapping matrix that maps each element of w (i.e., range K) to some element of z ( range K/c ), then the value of z is just an aggregation of all those elements from w that map to this element in z ( computed via S ⊤ w ) and then normalized by the total number of elements from w that map to this element in z ( computed via S ⊤ S ) -1 ). This is straightforward to implement and can be done in a blocked manner for the entire matrix Z . This is presented in the Algo 2.

Algorithm 2 SS1-project(W)

H Backward kernel algorithms

We implement both the forward and backward kernels with Triton [35]. The backward pass computes the compressed weight gradients and also the input gradients with respect to the hash functions used in the forward pass. It includes two kernels to output the SS1 weight gradients and input gradients.

The algorithms used in the backward kernels are presented in Algo 3 and Algo 4.

Algorithm 3 SS1-backward: weight grad

I Experiment settings

I.1 GPT experiment settings

We follow the standard GPT2-Small implementation of monarch paper [16] from the HazyResearch group at fly repo (https://github.com/HazyResearch/fly.git).

To optimize training the GPT model on our GPUs, we replaced Huggingface's transformers[33] 'Conv1D' layers with PyTorch[34] 'nn.Linear' in the feedforward blocks, as the transformation

remains the same. The architecture of the GPT2 models used in these experiments is detailed in 7. All models are trained using mixed-precision-training [36] with FP16 precision and implemented using PyTorch [34] AMP. We use V10032 GB GPUs to train the models from scratch.

Hyperparameters: For training we use the AdamW optimizer with α =6 e -4 , β 1 =0 . 9 , β 2 =0 . 999 , ϵ = 1 e -08 and weight _ decay = 0 . 1 we employ a linear schedule and wramup 1% of steps. The effective batch size is 512 , which is not achievable on our hardware memory; thus, we perform gradient accumulation every 32 step to reach that. All models are trained for 100 epochs. The hyperparameters are adopted from the Monarch paper [16].

Dataset: We use wikitext-103 [29] to train, evaluate, and test the model. The reported perplexity numbers are based on an evaluation of the test dataset. The length of the generated sequence is 1024.

I.1.1 Baselines

SS1: To build SS1 models, we substitute the transformer's feedforward linear layers with our SS1 layer. We use a fixed block size of Block_Size_K = 32 , Block_Size_N = 32 , and Block_Size_M = 64 for both the forward and backward kernels across all layers during training and testing. Additionally, each SS1 layer is assigned a different seed to ensure reproducibility of the hash functions used. In addition to the coalescing parameters, the vectorization parameter E V EC is set to 4.

Monarch: : For Monarch baseline, we convert the same layers as SS1 from the original Monarch implementation [16] and change the number of blocks to achieve different compression rates.

Lowrank: To implement the Lowrank layer, we take the linear layer matrix and parameterize it into two consecutive smaller matrices (linear layers) with reduced intermediate dimensions. The intermediate dimension in the replaced layer will be:

Where K s and N s are the input and output channels of the original matrix respectively, c ∈ N is the compression factor, and K lr is the intermediate dimension of Lowrank matrices. Afterward, like the other baselines, we swapped the linear layers of feedforward and replaced them with the Lowrank layer.

Smallmodel: The Smallmodel baseline aims to assess if the mlp layers of feedforward blocks of Large Language Models (LLM) are overparameterized and the impact on model quality by choosing smaller matrices for those layers. To achieve this, we take the same modules of the standard model and adjust the inner_dimension of the linear layers by a compression factor c where c ∈ N .

We were surprised to find that the smaller version of the model had moderate accuracy degradation. We also experimented with combining the small model and SS1 by training the small model with SS1 matrices for the FFN layers. We observed improved performance with the same number of parameters for 'Smallmodel 4 x+SS1 2 x' compared to 'Smallmodel 8 x'. The conversion process is the same as what we did in SS1-conversion.

BlockSparse: A pruning baseline that we sought to conduct the same experiment for was pytorch block sparse 5 . However, pytorch block sparse is constrained to fp32 precision training, and for the large models such as GPT2, it does not fit our available GPU memory space.

FlashAttention [37] suggests a significant latency of the inference phase lies in applying the softmax function on attention scores, which is a memory-bound operation. Using FlashAttention kernels shifts the workload from memory-bound to compute-bound; we know from the [37] that linear transformations are compute-bound because of the extensive matrix multiply operations. Considering all that, the performance contribution of SS1 is in the compute-bound setting because of reducing

the inner dimension of the linear matrices. Thus, in GPT2 latency measurement experiments, we use the Pytorch kernel abstraction of FlashAttention, the scaled-dot-product-attention (SDPA). As the FlashAttention kernels are yet to be released for older GPUs such as V 100 , we were not able to apply them in the training phase on V 100 GPUs; thus, we only replaced them in the latency measurement experiments for all the baselines above. For the latency measurement, we employ a single A 100 -40 GB GPU.

I.1.2 Quantization

To unravel the potential of SS1 in the quantization domain, we utilize uniform post-training Quantization. The method is as follows:

Where X is a floating point tensor, and ˆ X is its quantized version to an integer with n number of bits. S is the scale and ˆ S is the quantization level. ⌊·⌉ denotes rounding to the nearest integer. We have chosen the scale to be max( | X | ) as suggested by [32] due to its effectiveness in mitigating the impact of large activations in LLMs.

In our application, we implement the proposed method to reduce the precision of activation and weight tensors from 16 to 8 bits during the inference mode. For weight quantization, we apply the per-channel activation, meaning the scale is computed based on values of the tensor for each output channel; for activation quantization, we do per-token Quantization based on [32]. The quantization process is dynamic in our case. The results are obtained from testing the pre-trained GPT2-standard and SS1 models that we trained.

I.2 BERT finetuning settings

We conducted experiments to show that pre-trained linear layers can be projected onto SS1 layers while retaining previous knowledge. To validate this, we project the encoder layers of BERT-base and BERT-large models [19] onto SS1. BERT (Bidirectional Encoder Representations from Transformers) consists of a stack of Transformer encoder layers, each containing a self-attention mechanism and a feedforward neural network. In this work, we focus on two BERT model sizes: BERT-base (12 encoder layers, 768 hidden size, 12 attention heads, 110M parameters) and BERT-large (24 encoder layers, 1024 hidden size, 16 attention heads, 340M parameters). We then finetune these models on the GLUE (General Language Understanding Evaluation) benchmark [18]. Hyperparameters used during the finetuning stage are provided in Tables 10 and 11. GLUE consists of a collection of diverse natural language understanding tasks, such as textual entailment, question answering, and sentiment analysis. For our experiments, we utilize the GLUE dataset provided by the Hugging Face Datasets library.

The results are presented in Tables 8 and 9. We are able to achieve 50% compression of the encoder layers on both BERT-base and BERT-large models, with only a minimal drop in model quality across the GLUE tasks. This reduction in parameters is not achieved by uniformly compressing all layers in the encoder. Instead, we only project layers 1, 6, 8, 9, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24 (14 out of 24) for BERT-large and layers 1, 7, 8, 9, 10, 11, 12 (7 out of 12) for BERT-base at 8x compression. This leads to a net compression ratio of 1 -14 / 8+10 12 = 0 . 51 for BERT-large and 1 -7 / 8+5 12 = 0 . 51 for BERT-base, approximately 50% compression for both models.

These specific layers are selected using the RTE (Recognizing Textual Entailment) task from GLUE as calibration data. We treat RTE as a development set to determine which layers are less sensitive to compression. Layer sensitivity is measured by the degree of performance degradation on RTE when applying SS1 compression to individual layers in the model. The layers chosen by this method are able to generalize well to all other GLUE tasks, achieving evaluation metrics that closely match those of the full-sized BERT models with linear layers.

In addition to parameter reduction, BERT models with SS1 layers are able to achieve up to 1.31x increase in inference throughput. Detailed latency measurements are provided in Table 12. On the GLUE task, our BERT-large model with SS1 compression attains an accuracy of 79.76 (2.6 drop compared to the full model), while our compressed BERT-base model reaches 79.9 accuracy (2.2 drop) - demonstrating that our SS1 projection method can maintain strong performance with significantly reduced parameters and runtime.

batch

Latency ms (median)

I.3 Llama experiments

Meta's Llama family has emerged as one of the most powerful open-source Large Language Models. They are designed to be powerful and efficient models for various natural language processing tasks. In this work, we focus on the Llama3-8B model, which has 32 decoder layers, a hidden size of 4096, and 32 attention heads, resulting in approximately 8 billion parameters.

We conduct experiments to demonstrate that the pre-trained linear layers in Llama can be directly projected onto SS1 layers without any finetuning while maintaining model performance. This is a significant finding, as it highlights the effectiveness of our SS1 compression method in preserving the knowledge captured by the pre-trained model. By avoiding the need for finetuning, our approach offers substantial computational savings and facilitates more efficient deployment of compressed models. To validate our approach, we project the transformer layers of the Llama3-8B model onto SS1 layers. Specifically, we use the 'meta-llama/Meta-Llama-3-8b' model provided by

the Hugging Face Transformers library. We then evaluate the compressed model on the MMLU (Massive Multitask Language Understanding) [38] and Winogrande [39] tasks without any further training or finetuning using Language Model Evaluation Harness [40].

For calibration purposes, we utilize two small datasets: tinyMMLU and tinyWinogrande [41]. These datasets consist of 100 curated examples from the original MMLU and Winogrande datasets, respectively. By using these tiny datasets, we aim to identify the layers that are less sensitive to compression while minimizing the computational overhead of the calibration process. The results of our experiments are presented in Table 9. We achieve a 25% compression ratio of the decoder layers in Llama3-8B without any finetuning. This compression is obtained by projecting layers 18, 19,20, 21, 22, 23, 24, 25, 26, 27 in the case of MMLU and layers 18, 19, 20, 21, 22, 23, 24, 25, 26, 30 in the case of Winogrande onto SS1 at 8x compression, resulting in a net compression ratio of approximately 0.25. The selection of layers for compression is based on their sensitivity to compression, as determined by the performance on the tinyMMLU and tinyWinogrande datasets. The layers chosen using the tiny calibration datasets generalize well to the full MMLU and Winogrande datasets. This observation suggests that the compressed model retains the essential knowledge captured by the pre-trained Llama model, enabling it to perform competitively on the complete datasets without any finetuning.

Table 9 presents the evaluation metrics of the compressed Llama model on the full MMLU and Winogrande datasets. On MMLU, our compressed model achieves an accuracy of 61.26%, which is only a slight decrease from the 65.05% accuracy of the original Llama3-8B model. Similarly, on Winogrande, the compressed model obtains an accuracy of 69.93%, compared to the 76.1% accuracy of the uncompressed model. These results demonstrate the effectiveness of our SS1 compression method in maintaining the performance of the Llama model while significantly reducing its parameter count. Furthermore, the compressed Llama model exhibits improved inference speed, achieving up to 1.1x increase in throughput compared to the original model.

The ability to compress large language models like Llama without finetuning opens up new possibilities for efficient deployment and utilization of these models in various applications. Our SS1 compression method offers a practical solution for reducing the computational and memory requirements of large models while preserving their performance on downstream tasks.

I.4 Vision

We use the following MLPMixer models We use the repository: https://github.com/omihub777/MLP-

Mixer-CIFAR with default settings for all runs. We do not use hyperparameter tuning for compression methods.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We give our main kernel implementation.(link in PCQs in appendix) The experiment section is created using set of patched github repositories forked from original implementations. We will release our experimentation repository after acceptance. We detail our experiment methodology in appendix

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have added resource details to the appendix section A

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: To the best of our knowledge there is no specific societal impact of our work apart from that we are contributing to faster AI.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have provided all the information at sec A

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a structured randomized parameter sharing scheme (SS1) for computational complexity reduction. A couple of coalescing techniques are suggested, where the chunk size does not affect the approximation error. The proposed method is easy to combine with the quantization method to achieve further gain. Experimental results demonstrate that SS1 improves the throughput of GPT-2.\n\nSTRENGTHS:\nOverall, the paper suggests good additional features to Random Parameter Sharing (RPS).\n- The idea of reducing the computational complexity with RPS is interesting and adequately tackles a limitation of RPS.\n- The proposed method achieves actual latency improvement in some cases.\n- The paper focuses also on the GPU kernel implementation. The details about Forward and Backward Kernels are discussed well.\n\nWEAKNESSES:\nHowever, a few points should be addressed to improve the presentation. \n- Experimental results lack an important baseline--the comparison between SS1 and the existing RPS methods should be considered since SS1 is built upon RPS. \n- The latency and accuracy comparison of SS1 with other structured matrices is limited to small-sized models. Hence, the performance of SS1 compared to other methods on large models like LLM is unclear. \n- Writing quality needs to be improved, especially for the mathematical expressions and algorithms. I am willing to increase the score if the presentation of the algorithms, theorems, and proofs meet the quality of the top-tier ML conference.\n  - Some mathematical expressions are confusing, e.g., a vector and its element are sometimes denoted by the same alphabet $x$. \n  - Algorithms 1 and 2 are hard to follow. Adding comments to each important line or block might help readers understand what they are for. \n  - The abbreviations used in Table 1 are not properly introduced in the text. \n  - Proofs in Appendix contain typos in many places and should be written with a more professional nuisance.\n  - It is hard to follow how Algorithm 2 finds the solution of Equation 10. More details regarding Algorithm 2 should be followed.\n\nQUESTIONS:\nMajor questions\n- How does Algorithm 2 guarantee the minimization of the Frobenius norm error in Equation 10? More details should be provided. \n- How are $K$- and $L$-coalescing used in the actual SS1 implementation?\n- What are the latency and accuracy of Llama-3-8b with low-rank or Monarch matrices?\n- How much is SS1 better/worse than the existing RPS methods in terms of accuracy and latency?\n\nMinor questions\n- In Equation 10, I suggest denoting the Forbenius norm by $\\|\\|\\cdot\\|\\|_F$, not $\\|\\|\\cdot\\|\\|_2$, where the latter is usually reserved for the operator 2-norm induced by the vector 2-norm.\n- Do the hyperparameters $B_K$ and $B_N$ vary from GPU to GPU? How does the inference throughput change over those parameters?\n- Typo in Line 220 “[1,2])”\n- I suppose “SSL” in Table 1 is a typo and should be SS1.\n\nLIMITATIONS:\nSS1 requires a specialized kernel to be effectively utilized on GPUs. Different types of layers (e.g., convolution layer) may require different kernel implementations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-08"}

comments: SUMMARY:
The paper introduces the Sketch Structured Transform (SS1), a novel approach designed to enhance the efficiency of tensor multiplication in deep learning models. SS1 leverages structured yet random parameter sharing to reduce computational load while maintaining the expressive power of the models. The authors provide empirical evidence demonstrating SS1's superior quality-efficiency tradeoffs compared to existing methods like Monarch and LowRank. Key findings include significant improvements in inference speed for models such as GPT2, BERT, and Llama-3-8B, even without finetuning. The combination of SS1 with quantization further enhances efficiency, and SS1 also proves effective in reducing inference latency in CPU workloads, such as the DLRM MLPerf Model.

STRENGTHS:
- SS1 introduces a unique method of structured yet random parameter sharing that effectively reduces computational requirements while preserving model expressivity.
- Extensive experiments demonstrate SS1's superior performance across a variety of models and applications, including significant improvements in inference throughput for GPT2, BERT, and Llama-3-8B.
- SS1 can be applied to pre-trained models, allowing them to be projected onto SS1 and finetuned for faster deployment, which is highly practical for real-world applications.

WEAKNESSES:
- The details of GPU optimization, such as K- and N-coalescing, add complexity to the implementation, which might limit widespread adoption.

QUESTIONS:
None

LIMITATIONS:
None
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nthe authors propose the Sketch Structured Transform (SS1), a randomized parameter-sharing method that reduces FLOPs and latency while maintaining model quality. SS1 ties parameters within a single neuron weight and can be implemented to reduce input dimensions before multiplying with a compressed weight vector. This method is GPU-friendly and can be integrated with existing RPS methods, providing control over parameter memory and computation. SS1 can also be applied to pre-trained models using a projection function, preserving the benefits of existing models and enabling fast deployment.\n\nSTRENGTHS:\n1. The authors created a kernel that can work on available HW and reduce workload latency. Moreover, they open source the kernel!\n\n2. The basic idea is clear as described in section 3.1, the rest is kind of messy.\n\nWEAKNESSES:\n1. Very hard to follow. They abused notations many time (c is compression rate, size of super group etc.)\n\n2. The algorithms are unclear. For instance, algorithm 2. If we wish to find Z that minimize the frobenius norm between the compressed weight and the original weight we somehow need to search the different options to do it. If you are simply using brute force then just say it and push the algorithm to the appendix. Note that using only w and not checking the activation value when defining compression scheme will be sub-optimal as you might reduce the MSE of ||W-SS1(Z,I(k)||| but ||WX-SS1(Z,I(k))X|| (where x is the input to the linear layer) can be further reduced. \n\n3. Generally the paper is just poorly written and it is a shame as it seems is contains have very interesting ideas\n\nQUESTIONS:\n1. Can you explain section 4 in high level. What are you trying to prove and why? Specifically to the best of my understanding, you compressed\\quantize only the weights  and activation undergo reduction. So why should we care about the variance of dot product of compressed unit norm vectors. \n\n2. Same goes to theorem 2 – can you explain, in what sense you mean unbiased? Is it under permutation of the hash function? If so, I am not sure I understand why it is important if we pick one permutation and use that for the entire ( I mean the weight are tied once)\n\nLIMITATIONS:\nThe paper show promising results on small scale models. For llama3-8b the results are less promising as only 1.1x speedup was achieved by applying ss1 only on some of the layers (selected based on calibration set). The limitations are partially discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Sketch Structured Transform (SS1), a hardware-efficient structured matrix for improving the latency of linear layers in neural networks. The paper theoretically analyzes the effectiveness of SS1 in random projections, optionally combined with quantization. Experiments show favorable performance compared to alternatives such as Monarch and low-rank structures, as the compression ratio is varied.\n\nSTRENGTHS:\n1. Due to its use of the sketch matrix, SS1 is conceptually novel and distinct compared to other popular and performant structures such as Monarch and Low Rank, which only use batched matrix multiplies.\n2. SS1 show promising performance relative to Monarch and low rank.\n3. SS1 is designed to be hardware efficient, so that runtime overhead is low and latency improvements can be realized even at fairly small scale like GPT-2 Small.\n\nWEAKNESSES:\nThe experiment section does not convincingly demonstrate the superiority of SS1 over Monarch and Low Rank for the following reasons:\n1. The presentation of the results makes it hard to parse how SS1 compares to alternatives **as a function of latency**, which is the metric the paper aims to optimize for. For example, in Table 1, it would be much better to plot perplexity or loss as a function of latency for each structure. It would only be fair to claim SS1 outperforms alternatives if its loss vs latency curve stays below those of others.\n2. It does not make sense to control for the model's hidden dimension and only vary the compression ratio, as is done in the experiments. For example, one can increase the compression ratio (number of blocks) in Monarch or decrease the rank in low rank **while scaling up the hidden dimension and achieve the same latency**. It's possible this is a favorable trade off for the other structures, but the experiment section does not investigate this possibility. Indeed, recent work [1] has shown Monarch with more blocks can perform better for the number same parameters or FLOPs by simultaneously scaling up the hidden dimension. **The truly meaningful comparison would be a scaling law plot of loss vs latency with one curve for each pair (structure, compression ratio)**. \n3. I strongly suspect the learning rate is not well-tuned in the GPT experiments, making it difficult to trust the results. The learning rate used was $6e-4$ for all models. This value is well-tuned for the original dense GPT-2 model, but it is suboptimal for structured models as shown in [1]: Monarch with more blocks and low rank with lower ranks would generally require higher learning rates, otherwise their performance can be significantly underestimated.\n\nQUESTIONS:\n1. At a conceptual level, SS1 seems similar to a banded matrix with $N/c$ bands, which also have hardware efficient implementations. How does its performance compare to a banded matrix?\n2. Have you tried either tuning the learning rates for each model in the GPT-2 experiments? Alternatively, you can use the structure-aware learning rate rescaling introduced in [1] to ensure a good value is used for each model without having to tune it.\n\n[1] Qiu et al. 2024. Compute Better Spent: Replacing Dense Layers with Structured Matrices\n\nLIMITATIONS:\nThe authors discussed some limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the sketch structured transform (SS1) method for fast inference. SS1 is a form of randomized parameter sharing (RPS) with connections to low-rank/dimensionality-reduction methods. Theory is used to elucidate SS1’s compatibility with quantization -- it is also compatible with other RPS methods. In empirical experiments with popular NLP and CV setups, relative to competing methods, SS1 can preserve performance better as it reduces parameter count and inference latency. SS1 is also shown to have flexibility in the stage it is applied at -- it is applicable to models without any finetuning, before finetuning, and during training.\n\nSTRENGTHS:\nOriginality: The RPS method SS1 is introduced, and it is shown to reduce parameters, FLOPs, and latency. The method is made complementary with tiled matrix multiplication via schemes to coalesce values along multiple axes. Theorems are introduced to helpfully demonstrate SS1’s compatibility with quantization and its hyperparameter flexibility.\n\nQuality: The overall approach is straightforward and intuitive, and it is tested in a range of scenarios that vary the model, dataset domain, and application timing (before or after training). It is compared against strong benchmark methods.\n\nClarity: The SS1 method is clearly explained and positioned relative to prior work. The figures, equations, and algorithms motivate design choices and illustrate implementation details. The paper is mostly well written.\n\nSignificance: This paper addresses efficient inference for matrix-multiplication-focused models (e.g., LLMs), a crucially important topic. The proposed method is applicable to pretrained models, giving it broad relevance.\n\nWEAKNESSES:\nThe paper doesn't have any major weaknesses. I address some minor weaknesses below (see Questions). These mostly deal with improving the presentation. I would be happy to raise my score if some of them are addressed.\n\nQUESTIONS:\nLine 285: this seems incorrect. The quality with SS1 seems worse in Table 2 (right).\n\nWould the benefits of SS1 over other approaches hold at larger model sizes (1B+)?\n\nLine 123: I think this equivalence only holds for points with unit length.\n\nLine 155: Can you please rewrite this sentence to better clarify why the input would have to be read multiple times?  \n\nLine 153: should this be “$N \\times K$” instead of “$K \\times N$”? The latter notation is used on this line, which is inconsistent with the output dimension appearing first in the weight matrix notation on line 139.\n\nIs equation 9 missing the inner product notation?\n\nLine 202: why is $B_KB_N$ subtracted in the indexing?\n\nLine 263: this is a figure, not a table.\n\nLIMITATIONS:\nThe limitations are well articulated.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}