Sparse Learning for State Space Models on Mobile

Thanks for the suggestion. Llama.cpp is currently the only publicly available inference engine supporting Mamba on mobile platforms, making it the most practical baseline for our comparisons. We are actively monitoring developments in mobile inference engines, including MNN-LLM, MLC-LLM, TVM, ollama, ExecuTorch, NCNN, PowerInfer, exLLaMa, fastLLM, and TinyChatEngine, yet none of them support Mamba on mobile platforms now. We agree that a broader comparison with additional baselines would strengthen the paper further. Meanwhile, we aim to show how our optimizations outperform llama.cpp while providing insights into fundamental differences in execution strategies. We hope this addresses your concerns and demonstrates the significance of our findings despite current limitations in baseline availability.

Response to Question 5:

Thanks for pointing this out. Our work focuses on accelerating State Space Models (SSMs), specifically the Mamba model, on resource-constrained mobile devices. We achieve this by introducing sparsity into the model's weight matrices, reducing computational complexity and memory usage. At the algorithm level, we propose a novel fine-grained structure pruning scheme called Cn4 Kernel Design, which offers high accuracy due to its flexibility and precision. We utilize a Sparse Learning Framework and Weight Compensation to further enhance the model's sparsity and accuracy. At the system/compiler level, we design Weight Reordering to benefit from regular memory access patterns and improve cache hit rates (increasing L1/L2/L3 cache hits by 5%, 4%, and 7% respectively; details will be added in our evaluation section). Our Efficient Sparse Weight Storage reduces storage costs compared to traditional CSR format by 1.5x.

In contrast, DNNFusion aims to accelerate deep neural network execution broadly without focusing on specific models or platforms. It applies to various DNNs including CNNs, RNNs, and Transformers. SmartMem introduces a framework that eliminates most layout transformations by enabling multiple operators to use the same tensor layout through strategic choices in layouts and operator implementations. However, neither approach focuses on sparse computation and cannot be directly applied to sparse Mamba models. We will clarify this in our revision.

Response to Question 6:

Our solution is specifically designed for mobile architectures, so we have not tested it on non-mobile devices. However, with appropriate adaptations, it could potentially be applied to non-mobile platforms, which we may explore in future work.

Response to Question 7:

Our approach can also be extended to NPUs, as our Cn4-specific storage format is compatible with various devices. However, since NPUs have different memory organizations, what needs to be done is to adopt specific layouts for different operators when eliminating layout transformations in NPUs.

Response to Question 8:

It is not a typo. By adopting calibration in our post-training pruning method, we effectively eliminate non-essential connections, streamlining the network to focus on the most relevant features. This refinement allows the pruned model to achieve slightly better results than the dense model.

Response to Question 1:

Compared with NVIDIA 2:4 pruning (and n:m pruning scheme in paper [1]), the Cn4 kernel is motivated by the observation that modern mobile processors often feature Single Instruction Multiple Data (SIMD) units, which are optimized for parallel processing of data in groups, typically four elements at a time. By aligning the pruning pattern with the SIMD architecture, our method aims to maximize hardware utilization without specialized hardware support. NVIDIA's 2:4 pruning, and more generally n:m pruning, typically fall under the category of structured pruning techniques. These techniques prune weights in a structured manner, removing entire groups of weights (e.g., rows, columns, filters) instead of individual weights. The paper's Cn4 kernel allows for adaptive sparsity, where the number of pruned weights (n) can vary across different kernels. This adaptability provides greater flexibility in tailoring the sparsity to the specific characteristics of the model and the data. Our methodology also emphasizes learning the optimal sparsity pattern through a dedicated sparse learning framework. This framework considers factors like accuracy loss, sparsity level, and latency constraints to guide the pruning process.

Response to Question 2:

We provide the perplexity results of LLaMA family on WikiText2 with 50% sparsity compared to the SparseGPT and Wanda in Table A4 at supplementary and also the following table. The results show that our method achieves better performance than the other two methods.

Response to Question 3:

We have focused our exploration on SSMs within the NLP research domain, as their architectures are quite similar. Regarding SSMs in the computer vision domain, such as Vision-Mamba, we recognize the significance of this work and plan to apply our method to Vision-Mamba in future research endeavors.

Response to Question 4:

Our optimizations do specifically consider the unique architecture of mobile devices. For example, in the high-level architecture of Adreno GPUs (see page 15 in [2]), mobile devices feature two different memory units: buffer and image. Buffers can only be read and written through the L2 cache, whereas images can be read through the L1 cache to achieve higher bandwidth but must still be written through the L2 cache. Therefore, in our layout transformation elimination design, we account for this specific architecture, though we omitted detailed explanations due to paper length limitations. Since we primarily focus on mobile devices and design optimizations tailored to their architecture, we did not compare our results with server devices.

[1] Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch

[2] Qualcomm® Snapdragon™ Mobile Platform OpenCL General Programming and Optimization

Dear Reviewer,

Thank you very much for your valuable questions and feedback. As the discussion period is coming to a close, we sincerely hope you’ve had the opportunity to review our detailed responses to your previous questions and comments. If you have any further inquiries or require additional clarification, please don’t hesitate to reach out. We will do our best to address them promptly before the discussion deadline.

Thank you again for your time and consideration.

Best regards, The Authors

Dear Reviewer,

Thank you again for your valuable time and thoughtful feedback on our submission. As the deadline for the Author/Reviewer discussion is approaching, we would greatly appreciate it if you could let us know whether our responses have addressed your concerns. Your confirmation or any further feedback would help us refine our work more effectively.

We are more than happy to provide additional clarifications or address any remaining questions you might have.

Thank you once again for your effort and consideration.

Best regards, The Authors

Response to Weakness 1:

We have included the latency results for 30% sparsity in Table 3 and also the following table. The accuracy performance under 30% sparsity is already shown in Table 1. With this level of sparsity, our framework achieves a 1.1x to 1.2x speedup on mobile devices compared to the dense model and a 3.5x to 4.7x speedup compared to llama.cpp, highlighting its efficiency in mobile deployment.

We have added additional hardware metrics, including memory usage and energy consumption, in Table A5, Table A6 in Appendix, and also the following tables, to provide a more comprehensive evaluation.

Response to Weakness 2:

Thanks for the suggestion, we revised this in our rebuttal revision.

Response to Question 1:

We add more hardware metrics, including memory and energy consumption in Table A5 and Table A6, respectively, and also the following tables. Our sparse models result in slightly decreasing in memory, varying between 2.4% and 5.5%, while achieving substantial energy savings of up to 45.8% and significant inference acceleration. This highlights the efficiency and practicality of our approach, particularly for energy-sensitive and performance-critical applications.

Peak memory consumption results.

Energy consumption results.

Dear Reviewer,

Thank you very much for your questions. Since the discussion will end soon, we sincerely hope that you have found time to check our detailed response to your previous questions/comments. If you have any further questions, please feel free to let us know. We will try our best to reply to you before the discussion deadline.

Thanks Authors.

Response to Weakness:

To make it more clear, we define our kernel design in Section 5.1, and discuss the rationality of this design. The weights are split into multiple groups and each group has four consecutive weights. Our kernel is designed as Cn4, which removes n elements (0<=n<=4) from each group, as we can achieve practical acceleration with our compiler optimization for each case of Cn4.

Response to Question 1:

In practice, we have to adopt the dampening technique in Remark 5.2 to ensure the computation of matrix inversion. Without the dampening technique, the error of unavailable matrix inversion happens quite often. We set the dampening ratio $\gamma$ to 0.001. We ablate the performance of different dampening ratios and the results are demonstrated in Figure A1 of Appendix. We test on Mamba-2.8B model with 50% sparsity and compute the perplexity on LAMBDA dataset. As observed, smaller dampening ratios typically result in better performance, and further reducing the dampening ratio only leads to marginal improvements. A dampening ratio $\gamma$ smaller than 0.001 still leads to difficulties of matrix inversion. So we set the dampening ratio $\gamma$ to 0.001.

Response to Question 2:

Although the computations in Equation (6) and (7) may seem to be complex, we can avoid the matrix multiplication with $e_{q_i}$ and $M_i$ by selecting rows or columns from a matrix with reduced complexity, due to their specific formation. For example, $e_{q_i}^T W$ just selects the corresponding rows of $W$. So most of matrix multiplications can be implemented very efficiently. The most computation intensive part is the matrix inversion, which does not cost too much with a small or moderate dimension in 3B models or smaller.

In the context of sparse learning, the original model weights are kept frozen, while only the sparse masks for each linear layer are trained. We explore optimal pruning strategies for each layer, leveraging the proposed specialized kernel while adhering to constraints such as sparsity and latency. As highlighted in Section 6.1, the sparse learning process for Mamba-2.8B requires just 50 minutes on an A6000 GPU, demonstrating its efficiency and relatively low computational overhead.

Response to Question 3:

Our work focuses on accelerating State Space Models (SSMs) with software-hardware co-design, specifically the Mamba model, on resource-constrained mobile devices. We achieve this by introducing sparsity into the model's weight matrices, reducing computational complexity and memory usage. At the algorithm level, we propose a novel fine-grained structure pruning scheme called Cn4 Kernel Design, which offers high accuracy due to its flexibility and precision. We utilize a Sparse Learning Framework and Weight Compensation to further enhance the model's sparsity and accuracy. At the system/compiler level, we design Weight Reordering to benefit from regular memory access patterns and improve cache hit rates (increasing L1/L2/L3 cache hits by 5%, 4%, and 7% respectively; details will be added in our evaluation section). Our Efficient Sparse Weight Storage reduces storage costs compared to traditional CSR format by 1.5x.

We developed our sparse acceleration optimizations using C++, Python, and OpenCL. Specifically, we use Python for sparse training and compiler weight storage optimizations, while C++ and OpenCL are used for code generation and GPU acceleration. This original work has not been submitted or accepted elsewhere. We appreciate the author's recognition of our contribution and will add a subsection in the evaluation to clarify this in the revision.

I appreciate the detailed response from the authors. I am convinced that the work makes a significant stride compared to the current SOTA works. I have increased the score to 6. However, I recommend that authors make some amendments to the paper to make the paper easier to read + make the details (various optimizations and the design) less abstract.

Thank you for your detailed rebuttal. I appreciate the additional experiments and insights. I have now raised my score.

Response to Question (Question 1):

The details about layout transformation elimination, as well as the process of searching for the optimal data layout for each operator are elaborated in Appendix B.2. Regarding different target devices, our proposed method is compatible with different devices. We focus on mobile devices because they have limited bandwidth. By eliminating the layout transformation operations and selecting optimal data layout, we can significantly reduce the cost of memory access, which is a major overhead in mobile devices.

Response to Question (Question 2):

We utilize sparse learning to train the sparse masks and determine the optimal pruning strategy, keeping the model's original weights frozen throughout the process. For Mamba-2.8B, sparse learning is highly efficient, requiring approximately 50 minutes on a single A6000 GPU. Once the optimal sparse masks for each layer are obtained, pruning is performed based on these masks. Following pruning, we apply a compensation method to enhance the performance of the pruned model. This compensation process is a post-training, forward-only procedure that does not involve any backward computation, ensuring simplicity and efficiency.

Response to Question (Omissions/Extensions 1):

Our method does not conflict with the quantization method, we can adopt GPTQ for direct weight quantization to reduce the model size. While for GPUs, the acceleration brought by activation quantization may not be that large compared to CPUs. We further benchmark the performance on other edge devices (Xiaomi 6 mobile phone, equipped with QualComm Adreno 540 GPU and 6GB of memory). The result of energy consumption is reported in Table A6 at supplementary and also the following table. Compared to the dense model, our sparse model demonstrates a significant energy saving ratio, varying from 2.6% to 45.8%. The reduction in energy consumption is primarily attributed to our hardware design, which optimizes the proposed framework to significantly reduce computational demands. An important observation is that energy consumption decreases as the sparsity increases, highlighting the efficiency of our method in leveraging sparsity to minimize resource usage.

Response to Question (Omissions/Extensions 2):

We also present the results for the LLaMA model family in Table A4 of the supplementary materials and the table above, demonstrating that our algorithm is well-suited for the LLaMA family. However, our framework currently does not support hardware deployment for the LLaMA models, as this would require additional optimizations tailored to the transformer architecture. We recognize the significance of this limitation and plan to address it in future research efforts.

Response to Question (Omissions/Extensions 3):

Quantization is fully compatible with our proposed method and can be leveraged for additional model compression. For example, post-training quantization techniques like GPTQ [1] can be applied directly to quantize the model weights. Subsequently, our method can be employed to produce a sparse model, enabling efficient acceleration on mobile devices with a 4-bit representation.

Response to Question (Omissions/Extensions 4):

There is no such limitations, our method is based on the linear layers, it can be extended to Vision Mamba.

Response to Question (Omissions/Extensions 5):

We show the peak memory consumption during inference in Table A5 at supplementary and also the following table. Notably, The sparse model requires only a slight increase in memory compared to the dense model, ranging from 2.4% to 5.5%, while delivering significantly faster inference. This demonstrates the efficiency and practicality of our approach.

Response to Question (Omissions/Extensions 6):

In our experiments, we adopt the uniform sparsity ratio for each SSM block.

[1] GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers

Response to Question (Evaluation 1):

There is a typo in the introduction; the value should be '2.5s' instead of ‘5.8s’. We will correct it in the revision.

Response to Question (Evaluation 2):

We have added the 30% sparsity results in Table 3 and also below to show the speedup for each 1% decrease in accuracy.

Response to Question (Evaluation 3):

Llama.cpp does not support GPU inference for Mamba on mobile devices during evaluation, so we only include CPU results. Currently, llama.cpp still lacks support because some SSM kernels are not implemented on mobile platforms yet.

Response to Question (Evaluation 4): Yes, we maintain the same accuracy presentation as other state-of-the-art works.

Response to Weakness 1:

The perplexity results of LLaMA family on WikiText2 with 50% sparsity compared to the SparseGPT and Wanda are added as Table A4 in supplementary and also in the following table,

We also conducted the evaluation on a low-end device, the Xiaomi 6, which is equipped with a Snapdragon 835 featuring an octa-core CPU and Adreno 540 GPU with 6GB of memory. The results are presented in Table A7 at supplementary, as well as the table below. Due to the memory limitations, we report the results for Mamba models with 130M, 370M, and 790M model scales. The results show similar trends: our dense version achieves a speedup ranging from 3.3x to 4.4x compared to llama.cpp on the mobile CPU. Meanwhile, compared to our dense version, our sparse method demonstrates considerable acceleration, achieving approximately 1.1x to 1.2x speedup at 30% sparsity, 1.1x to 1.3x speedup at 50% sparsity, and 1.2x to 1.8x speedup at 70% sparsity. These results demonstrate that our proposed method is both compatible and efficient across different edge devices.

Response to Weakness 2:

Thanks for pointing this out. We will discuss the on-device efforts as follows,

“Recently, several frameworks for on-device LLMs have emerged, including Llama.cpp, Power-Infer, Power-Infer2, ExecuTorch, MediaPipe, OpenLLM, and BentoLLM. Llama.cpp offers an efficient plug-and-play library for large language model inference across various hardware platforms. Power-Infer (and Power-Infer2) are designed for consumer-level GPUs with offloading support when loading weights for LLMs. ExecuTorch, MediaPipe, OpenLLM, and BentoLLM facilitate the deployment of open-source LLMs as local APIs or OpenAI-compatible API endpoints optimized for high throughput and cloud efficiency. However, none of them support sparse Mamba on mobile devices.”

Response to Weakness 3:

We have included the latency results for 30% sparsity in Table 3. With this level of sparsity, our framework achieves a 1.1x to 1.2x speedup on mobile devices compared to the dense model and a 3.5x to 4.7x speedup compared to llama.cpp, highlighting its efficiency in mobile deployment.