SS1: Accelerating Inference with Fast and Expressive Sketch Structured Transform

We thank you for the suggestions to improve the paper. Incorporating the suggestions, we have made several writing changes to the paper that are listed in the common rebuttal. We hope to have addressed your concerns and incorporated your suggestions to your satisfaction. Specific concerns are addressed below,

 Comparison with RPS:

Existing RPS methods focus on reducing the parameter memory but not the FLOPs. Thus the latency of RPS methods generally is slightly worse than full model latency no matter the amount of compression. Even so, we perform RPS experiments using ROAST for matrix compression. The results are as follows and as expected. We will add these results to the final paper.

the performance of SS1 compared to other methods on large models like LLM is unclear.

While we do not have the computing power to train 1B+ LLM models, our rigorous tests on different domains, model architectures, and settings encourage us to believe that the results we see are indeed general and will translate even to bigger models.

Writing quality improvements We apologize for the issues with our presentation. We have fixed all the issues listed here and more. The detailed summary of major changes is presented in the common rebuttal.

How does Algorithm 2 guarantee the minimization of the Frobenius norm error in Equation 10? We have improved our explanation of projection from full matrix to SS1 compressed matrix. It should help avoid a lot of confusion. We will provide a brief explanation here. ( we have also added this explanation to appendix and have moved algorithm 2 to appendix)

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 \times N$ weight matrix is independently projecting each neuron. Consider a single neuron under RPS $y^\top = x^\top w$ and $w = S z$. Given $w$ , the solution to get best $z$ which minimizes $|| w - Sz||_2$ is nothing but solution of linear regression, i.e. $z^* = (S^\top S)^{-1}) S^\top w$. Also, since S is a sparse matrix with exactly one non-zero in each row, $S^\top 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^\top w$) and then normalized by the total number of elements from $w$ that map to this element in $z$ ( computed via $S^\top 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.

B_K and B_N parameters : implementation, do they vary with GPU, does it affect throughput The GEMM operation has three block parameters $B_M, B_K, B_N$. In the implementation of SS1, our coalescing parameters B_K and B_N behave like standard GEMM block parameters. i.e. they determine what block sizes to bring into shared memory for performing matrix multiplication. What is interesting is that algorithmically, they also determine how the hash function and SS1 recipe in general should behave, which is the novelty in SS1 that ensures that the algorithm is hardware-friendly. These parameters are specific to each GPU and need GPU-specific optimization. Also, the choice of these parameters has a huge impact on throughput, as is true for GEMM operations.

  Llama-3-8b experiments We were able to run quality experiments for lowrank. We could not find a ready implementation for monarch projection in original repository for rectangular matrices. We will add complete results in the final paper. The results are very similar across different methods.

We want to note that this is just a proof of concept in support of the fact that SS1 can be used in the era of pretrained models as well. The main experiment still is when we train SS1 from scratch to evaluate quality vs. latency.

We hope to have addressed your clarifications and incorporated your suggestions to your satisfaction. If you are satisfied with our changes, please consider giving us an updated score.

We thank the reviewer for the support of our paper. We have made writing changes to improve the manuscript and additional experiments as suggested by other reviewers. They are listed in common rebuttal. Specific clarifications are answered below,

K- and N-coalescing, add complexity to the implementation

All CUDA kernels optimize block parameters for performance. The K and N coalescing is similar to block size choosing generally required for matrix multiplication. Libraries such as CUBLAS do this optimization and cache the parameters for use, hiding the complexity of the kernels from the end user. Similarly, we do the optimizations in the code and the end user does not need to worry about these parameters.

We hope to have addressed your concerns to your satisfaction. We are happy to provide further clarifications if needed.

We thank the reviewer for their support of our paper. We apologize for writing issues and assure the reviewer that we have taken due care to fix them both in the main paper and appendix. The list of changes made are presented in the common rebuttal. We clarify some specific concerns below

Notation overload of c The $c$ in the text always refers to the compression factor ( a factor of c=4 means compression of 4x). In the text we have said that the supergroup contains “c” groups. Having said that we have refactored the section 3 with standardized notations and explanations to improve the readability of the section and to avoid confusions.

Algorithm 2 We do not need to perform brute force to find the solution to equation ||WX-SS1(Z,I(k))X||. The solution of the equation is exactly shown in algorithm 2. We found a easier way to explain the projection mechanism and have updated the manuscript accordingly. We hope that alleviates the concerns regarding algorithm 2. We also explain it here,

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 \times N$ weight matrix is independently projecting each neuron. Consider a single neuron under RPS $y^\top = x^\top w$ and $w = S z$. Given $w$ , the solution to get best $z$ which minimizes $|| w - Sz||_2$ is nothing but solution of linear regression, i.e. $z^* = (S^\top S)^{-1}) S^\top w$. Also, since S is a sparse matrix with exactly one non-zero in each row, $S^\top 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^\top w$) and then normalized by the total number of elements from $w$ that map to this element in $z$ ( computed via $S^\top 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.

Other presentation issues We have refactored section 3, added detailed comments in algorithm 1, simplified the projection explanation, and made the notation clear and consistent. We hope the reviewer finds the new version friendlier.

Section 4 explanation

Recent literature on RPS shows that the quality of the learned linear model under compression is directly related to the quality of inner product preservation under compression. Thus, inner product preservation is an important problem to analyze in this regard. [1]

The goal in section 4 Theorem 1 is to understand whether projection ( which is the basis for RPS) and quantization ( basis for quantization based compression) can be combined together to obtain better compression rates while preserving quality. While there is general consensus from empirical observations that compression methods can be combined, theorem 1, to our knowledge, is the first theoretical exposition on this topic which also provide guideline as how much compression to perform from individual methods.

The goal of section 4 theorem 2 is to understand the impact of the coalescing-factors have on the “inner product preservation” ( and hence quality of the learned model).

The unbiasedness and variance is under the randomness of permutations ( implemented via hash functions). It is true that we fix these permutations at the start, the randomness is still important to guarantee that we do not choose a “bad” permutation that will compromise the quality of the learned model with a high probability.

[1] Desai, A. and Shrivastava, A.. In defense of parameter sharing for model-compression. International conference on learning representations 2024

We hope to have addressed your clarifications and incorporated your suggestions to your satisfaction. We are happy to provide further clarifications if needed.

Dear reviewer, we would be happy to perform comparative experiments on banded matrix while we prepare for rebuttal as time permits.

Can you please share links to the fast implementation of Banded Weight matrices? In our preliminary search, we have not found a good implementation for banded weight matrices. The closest we found was this genbmm package (https://blog.rush-nlp.com/visualizing-banded-sparse-matrices.html) but it seems to only multiple two banded matrices whereas the setting requires us to multiply one banded matrix (weight) and full matrix (input).

We appreciate the reviewer rating our paper well on counts of soundness and contribution. We have made writing changes to improve the manuscript and additional experiments as suggested by other reviewers. They are listed in the common rebuttal.

Plots of latency vs. perplexity: Thanks for the suggestion. We plot latency vs. perplexity for the GPT experiment in figure 3 ( page 15) and we can see that indeed SS1 curve is the best lying below other curves. We do not use such plots in the main paper since we want to show more latency results for larger models for which we do not have quality numbers in the table in the main paper.

About changing dimensions along with compression rates

The direction of increasing hidden dimension is definitely meaningful not just for baselines such as lowrank, monarch etc. but also our proposed SS1. It makes the experimental space of models explode. Our choice of experimental setup is primarily to set up a practical and fair experimental setting to compare different alternatives. In fact, sticking to base model dimensions and evaluating various compression methods is quite standard in papers that try to fundamentally compare different methods. Some examples being [1,2,3].

[1] Tanaka, H., Kunin, D., Yamins, D.L. and Ganguli, S., 2020. Pruning neural networks without any data by iteratively conserving synaptic flow. Advances in neural information processing systems, 33, pp.6377-6389.

[2]Frankle, J., Dziugaite, G.K., Roy, D.M. and Carbin, M., 2020. Pruning neural networks at initialization: Why are we missing the mark?. arXiv preprint arXiv:2009.08576.

[3] Desai, A. and Shrivastava, A.. In defense of parameter sharing for model-compression. International conference on learning representations 2024

Learning rate tuning You are correct. Our learning rate is tuned for base models and used across structured models. This again is a commonplace way to compare against different methods [1,2,3]. While it is not optimal for different methods, it avoids entering into hyperparameter tuning for each method x setting. Having said that, we understand the concern. We quickly tested if learning rate tuning would change the findings. We find that learning rate tuning improves results for Monarch and SS1 and does not change the relative ordering.

The findings do not change with learning rate tuning.

suggestion on structure aware tuning The work cited Qui et. al. is indeed interesting work and will be useful for our future research. Thanks for pointing it out.

Banded Matrix Comparison We would love to compare against a banded matrix. However, as mentioned in our official comment, we did not find a good implementation of a banded matrix that multiplies with a full matrix. Please let us know if you have any references. We would be happy to include comparative study in our camera ready. Conceptually, SS1 and the banded matrix are very different. Banded matrix is motivated via “sparsity” which drops the input signals for neuron computation whereas SS1 is motivated via parameter sharing which sketches the input to reduce dimensionality. In previous literature we have seen that sketching generally is superior in quality than sparsity. [1]

[1] Desai, A. and Shrivastava, A.. In defense of parameter sharing for model-compression. International conference on learning representations 2024

We hope to have addressed your clarifications and incorporated your suggestions to your satisfaction. we kindly ask you to consider giving us an updated score.

We thank the reviewer for the support of our paper. We have made writing changes to improve the manuscript and additional experiments as suggested by other reviewers. They are listed in common rebuttal.

Please find the clarifications requested below:

Line 285: In table 2 (right), we can see that the SS1 model with 74M parameters has better (lower) perplexity than GPT-S-6x with 76M parameters. Thus SS1 models with less parameters can outperform standard model.

Would the benefits of SS1 over other approaches hold at larger model sizes (1B+)? We believe so. While we do not have the computing power to train 1B+ LLM models, our rigorous tests on different domains, model architectures, and settings encourage us to believe that the results we see are indeed general and will translate even to bigger models.

Line 123: The equivalence is due to the relation between inner products and norms. Inner product => norms. Since ||x||^2 = <x, x> norms are just inner product with itself. Norms => inner products. Since 2<x, y> = ||x - y ||^2 - ||x||^2 - ||y||^2 Thus, if norms are preserved, then inner products will also be preserved and vice versa. More discussion can be found in the book[1] (Page 12, first paragraph).

[1] Woodruff, D.P., 2014. Sketching as a tool for numerical linear algebra. Foundations and Trends® in Theoretical Computer Science, 10(1–2), pp.1-157.

Line 155: Let us clarify the line here. Claim: If we do not perform K-coalescing, then we might end up requiring to load input multiple times to compute the value of the single neuron. Argument: Consider a neuron that is computed as <z, Sx>, where the dimensions of z and x are large enough so that they cannot be cached or stored on shared memory. In such a case, under certain hash functions that generate S, x[i] and x[i+1] ( for any i) can get multiplied with z[j] and z[k] where j and k are far apart ( In fact under the randomness of the hash function, this event is highly likely). In such a case, if we take a single pass on z, we will have to read the cache line containing x[i] and x[i+1] at least twice.

Taking cue from the need for clarification along with notation and presentation issues pointed out, we have refactored our explanation in section 3 in manuscript while standardizing the notation. We believe the current version is clearer and will help the readers.

Line 153, Equation 9. Thanks for pointing out the inconsistency in handling W. We have fixed the presentation such that we always use W as a K x N matrix and all expressions are written as vector and matrix multiplications (We have removed the inner product notation) for clarity.

Line 202 : B_K B_N is subtracted in the indexing in accordance with the hash function treatment of the ROAST[2] -- recent RPS method. It is a simple solution used to avoid overflows when reading a block of size B_K x B_N from the indexed location.

[2] Desai, A., Zhou, K. and Shrivastava, A., 2023, July. Hardware-aware compression with random operation access specific tile (ROAST) hashing. In the International Conference on Machine Learning (pp. 7732-7749). PMLR.

Line 263: fixed.

We hope to have addressed your clarifications and incorporated your suggestions to your satisfaction. We kindly ask you to consider giving us an updated score as per your indication in the review. We are happy to provide additional clarifications if needed.

Thanks again for the support of our paper.