Sparse Learning for State Space Models on Mobile
Authors: Xuan Shen
Hangyu Zheng
Yifan Gong
Zhenglun Kong
Changdi Yang
Zheng Zhan
Yushu Wu
Xue Lin
Yanzhi Wang
Pu Zhao
Wei Niu
Venue: ICLR 2025 Poster (2025)
Keywords: Mamba
Pruning
Mobile
Primary area: applications to computer vision, audio, language, and other modalities

Abstract: Transformer models have been widely investigated in different domains by providing long-range dependency handling and global contextual awareness, driving the development of popular AI applications such as ChatGPT, Gemini, and Alexa.
State Space Models (SSMs) have emerged as strong contenders in the field of sequential modeling, challenging the dominance of Transformers. SSMs incorporate a selective mechanism that allows for dynamic parameter adjustment based on input data, enhancing their performance.
However, this mechanism also comes with increasing computational complexity and bandwidth demands, posing challenges for deployment on resource-constraint mobile devices.
To address these challenges without sacrificing the accuracy of the selective mechanism, we propose a sparse learning framework that integrates architecture-aware compiler optimizations. We introduce an end-to-end solution--$\mathbf{C}_4^n$ kernel sparsity, which prunes $n$ elements from every four contiguous weights, and develop a compiler-based acceleration solution to ensure execution efficiency for this sparsity on mobile devices.
Based on the kernel sparsity, our framework generates optimized sparse models targeting specific sparsity or latency requirements for various model sizes. We further leverage pruned weights to compensate for the remaining weights,  enhancing downstream task performance.
For practical hardware acceleration, we propose $\mathbf{C}_4^n$-specific optimizations combined with a layout transformation elimination strategy. 
This approach mitigates inefficiencies arising from fine-grained pruning in linear layers and improves performance across other operations. 
Experimental results demonstrate that our method achieves superior task performance compared to other semi-structured pruning methods and achieves up-to 7$\times$ speedup compared to llama.cpp framework on mobile devices.

Published as a conference paper at ICLR 2025
SPARSE LEARNING FOR STATE SPACE MODELS ON MOBILE
Xuan Shen Hangyu Z Yifan Gong' Zhenglun Kong²
Changdi Yang' Wei Niu
Zheng Zhan' Yushu Wu! Xue Lin'
Yanzhi Wang' Pu Zhao
NNottheastern University UUniversity of Georgia Haarvard University
{shen. www.://ooortassteee.ede.. (h285760, wniu]@uga.edu
ABSTRACT
Transformer models have been widely investigated in different domains by pro- viding long-range dependency handling and global contextual awareness, driv- ing the development of popular AI applications such as ChatGPT, Gemini, and Alexa. State Space Models (SSMs) have emerged as strong contenders in the field of sequential modeling. challenging the dominance of Transformers. SSMs in- corporate a selective mechanism that allows for dynamic parameter adjustment based on input data, enhancing their performance. However, this mechanism also comes with increasing computational complexity and bandwidth demands, posing challenges for deployment on resource-constraint mobile devices. To ad- dress these challenges without sacrificing the accuracy of the selective mecha- nism, we propose a sparse learning framework that integrates architecture-awarc compiler optimizations. We introduce end-to-end solution- CT kernel spar- sity, which prunes n elements from every four contiguous weights, and develop a compiler- bompiler-based acceleration solution to ensure execution efficiency for this spar- sity on mobile devices. Based on the kernel sparsity, our framework generates optimized sparse models targeting specific sparsity or latency requirements for various model sizes. We further leverage pruned weights to compensate for the remaining weights, enhancing downstream task performance. For practical hard- ware acceleration, we propose C7 -sesciiccc optimizations combined with a layout transformation elimination strategy. This approach mitigates inefficiencies aris- ing from fine-grained pruning in linear layers and improves performance across other operations. Experimental results demonstrate that our method achieves su- perior task performance compared to other semi-structured pruning methods and achieves up-to 7 7x speedup compared to llama.cpp framework on mobile devices.
INTRODUCTION
Recent research advancements have significantly heightened interest in State Space Models (SSMs). Building on the foundation of the Kalman filter model (kal 1960 SSMs have been further improved
to address long-range dependencies with parallel training. Works (Gu et al. 202 la 2022 Gupta
al. 2022) propose SSM-based models designed to process sequence data across a variety of tasks and modalities. Recent work, Mamba (Gu & Dao 2023b), integrates marying parameters into the SSM, enabling the model to selectively propagate or forget information. Additionally, Mamba introduces hardware-aware parallel algorithm designed accelerate training and inference. Com- pared to quadratic attention, which becomes prohibitively expensive with longer sequence lengths, Mamba's subquadratic-tim architecture more efficient and better suited for handling long se- quences. Mamba's ' exceptional scaling performance highlights its potential as an effective alterna- tive to the Transformer model /Vaswani al. 2017 for generative language modeling tasks.
To make advanced language processing accessible to more people and address privacy concerns, de- ploying Mamba models on mobile devices is promising strategy to improve the accessibility and usability of SSMs. A report indicates that by the end of 2020, nearly 6 billion smartphones were in
"qqual Contribution
Corresponding Author

---

Published as a conference paper at ICLR 2025
use worldwide, experiencing annual growth rate of 4% (Samsung), and demonstrating increasing
processing capabilities (Huynh et al. 2017| Xu et al., 2018
Chen et al. 2023; Zhan et al. 2021
Wu et al. 2022 Li et al. 2022 2023a Shen al.) 2024c) Nuu al. 2024a). Deploying Mamba
models on mobile devices ensures functionality in offline scenarios and reduces reliance on costly cloud services. However, the hardware- aware parallel algorithm in Mamba models is specifically optimized for GPUs, similar optimizations for mobile devices have yet to be explored. The com- putational complexity in Mamba poses significant challenge for resource-constraint mobile devices, where bandwidth is limited compared to desktop-level GPUs (Huynh et al. 2017 Lee et al. 2019). Moreover, the projection mechanism within Mamba demands high throughput memory. Unlike powerful server platforms, mobile devices typically have less memory size and are constrained by limited battery capacity, which restricts their computing performance and memory bandwidth. For instance, state-of-the-art framework llama.cpp (contributors] 2023a) takes over 2.5s to generate one token with Mamba-2.8B model on a high-end mobile device (Qualcomm 2017), highlighting the need for further optimization to make Mamba models viable for mobile deployment.
To tackle the mentioned challenges on mobiles, we propose a sparse learning framework that in- corporates architecture-avve aware compiler optimizations for the acceleration of Mamba models on mo- bile devices. Inspired by the modern hardware architecture of Single Instruction Multiple Data (SIMD) (Cypher & Sanz 1989 Mitra al. 2013 Khorasani et al. 2015) units, which are optimized to load and process four-element vectors in parallel. We focus on investigating the sparse patterns within every four elements improve the hardware efficiency with the support of our compiler optimization. We first introduce the CT kernels which prune n elements from every four contigu- ous weights. We further propose a sparse learning framework to thoroughly optimize the pruning strategy for these kernels, i.e., determine the value of n for each kernel (four-element vector) and the corresponding pruned n elements. To preserve task performance while applying sparsity and achieving significant acceleration, we profile the effectiveness loss (as an accuracy predictor), spar- sity loss and latency loss for each choice of n for each kernel. With the detailed profiling, given the specific sparsity latency goals, our sparse learning framework targets to learn mask to choose a value of n for each kernel, so that the accuracy loss is minimized while satisfying the sparsity/la- tency constraint. In detail, to render the sparse learning process differentiable, we define the pruning strategy (mask) through probabilities assigned to different kernels for every set of four consecutive weights throughout the entire model. Finally, introduce compensation algorithm the remaining weights by utilizing the pruned weights, optimizing overall model functionality. The compensation mainly based on the classic OBS update (Hassibi al., 1993 Singh & Alistarh, 2020 Frantar et al. 2021 Zhao et al. 2024) for the weight reconstruction, leveraging calibration with only 128 training samples. Regarding hardware acceleration, we propose a unique design for CT CT-specific optimizations that includes weight reordering and an efficient method for storing sparse weights. For other operators with intensive memory movement, we design a layout transformation climination to decrease the bandwidth demands without the need for data layout changes. These optimizations are crucial to mitigate the performance degradation associated with fine-grained CT pruning versus structured pruning or dense configurations. Experiments show that our algorithmic approaches can achieve better performance with the same sparsity on different scales of Mamba models compared to other cally, we reduce the perplexity from 212.9 to 28.97 and enhance the accuracy from 35.6% to 41.0% on Mamba-130M model compared to Wanda (Sun al. 2023) with 2:4 pattern. Our comprehen- sive ablation study demonstrates the effectiveness of our mixed kernel design and the compensation methods. We implement the sparse model with our proposed kernels on mobile devices and achieve a practical on-device speedup of up to 7 x compared to llama.cp. We summarize our contribution:
rectify
pruning methods with fixed sparsity patterns. Specifi-
1. We design a special kernel CF and with a set of comprehensive compiler optimizations, includ- ing CT specific optimizations and layout transformation elimination strategy on mobile devices.
2. We propose the sparsity- the optimal pruning strategy with the proposed kernels for Mamba models.
sparsity-oriented oriented and/or latency- sparse learning framework to explore
3. We propose the weight compensation algorithm for the rectification of the sparse model weights by calibrating with only 128 samples, thereby further enhancing the model effectiveness.
4. Experiments show that our framework can achieve better task performance than other semi- structure pruning methods and achieve pratical on-device speedup up to 7x compared to llama.cpp.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
2.1 STATE SPACE MODELS
The work (Gu et al. 2021a) a initially models long sequences using structured state spaces rather than Transformers (Vaswami et al. 2017) or Convolutional Neural Networks (CNNs), sparking interest in exploring state space models. The memory usage in Transformer increases with the context length, making it difficult to efficiently process long-context windows or multiple parallel batches without substantial hardware resources. Recently, the work (Fu al. (022) fills the performance gap between SSMs and Transformers in language modeling, and Mamba (Gu & Dao 2023a) introduces new Mamba structure as the general sequence model backbones. Mamba introduces an input- dependent selection mechanism into SSMs and benefits from linear scaling in sequence length, surpassing traditional Transformers across multiple model sizes. However, previous study Zhan ct al. (2024a) does not fully investigate the redundancy inherent in SSMs and hardware accelcration resource-constraint mobile devices, leaving this research area largely under-explored.
2.2 DNN INFERENCE ACCELERATION ON MOBILE
As machine learning applications (Li ct al. 20236 2024a Zhang ct al. 2022 Yang et al. 2023
on mobile devices continue to grow, there is a strong emphasis on optimizing frameworks of deep neural network (DNN) inference on mobile. Efforts such as works including (Han et al. 2016 Niu
et al. 2020 have primarily focused on accelerating traditional CNNs. General inference frameworks that support both server and mobile platforms for different neural networks, such as (Niu et al. 2021 2024b) offer advanced features including operator fusion, memory planning, shape inference, quantization, and tensor offloading. More recently, there has been trend towards working with large models like llama.cpp (contributors) 2023a), ex Lx LaMa (exllama contributors) 2023). ML.C- LLM (contributors) 2023b,, and fastL_.M (The fastllm contributors, 2023). Yet many of these efforts
either overlook model efficient techniques (Shen et al.] 2024d e b. 2023 2022
el al. 2024b
Zhan et al. 2024b) or fail support SSMs on mobile platforms.
Thruggput (Tokens/s)
3 MOTIVATION AND BACKGROUND
Higher Throughput with
Larger Model Size
To showcase the superior efficiency and mobile-friendliness of Mamba models over Transformers, we conduct a throughput com- parison between them. As shown in Figure by comparing the model size and throughput (lested on Oneplus 11 mobile phone) un- der the same configuration (such as the same batch size and input sequence length), Mamba models can achieve higher throughput with model size similar or even larger than the Transformer models from various L.L fam-
25
Mamba
10
OpenE
Qwea2.5
X
BL.OOM
2.2
24
2.6
2.8
3.2
3.1
3.6
Log Purameters (M)
Figure I: Throughput v.s. Log of Parameters (M)
ilies (Mehta et al. 2024 Yang et al.] 2024)
Le Scao et al. 2023 leading to a better trade-off between throughput and model size.
Besides, in practice, under the same memory usage, Mamba models typically can use a larger batch size, resulting in 4-5x higher inference throughput than Transformer of similar size. The reason is that, unlike Transformers, Mamba models don't 't require the KV cache @Gu & Dao 2023a), reducing memory usage and allowing for larger batch sizes with higher throughput.
Furthermore, because of Mamba's 's ability to selectively remember the relevant token while ignoring everything else in between, it can tion heads task OOlsson et al. 2022). ). it generalizes perfectly to million-length sequences, or 4000 longer than it saw during training, while other method goes beyond 2x (Gu & Dao 2023a).
extremely long context with length up 1M. On induc-
Due to the efficiency and optimizing Mamba models for superior on-mobile performance.
of Mamba models over Transformers, we focus on
3

---

Published as a conference paper at ICLR 2025
4 PRELIMINARY
State Space Models (SSMs) are sequential models that can map 1-dimensional function or se- quence ztt) E IR to the output sequence y(t) E R through hidden state h(t) E IR as follows,
(1) Ah(() + Br((),
(1)
v(1) Ch((),
where N denotes the representation number, A E RNXN is evolution parameter, B E RNXI and Ce IR 1x are projection parameters.
The Mamha model (Gu & Dao 2023b) represents the discrete version of the continuous system for SSMs and incorporates timescale parameter A to facilitate the transformation of continuous parameters with the zero-order hold (ZOH) as follows,
A exp(AA),
(2)
B (AA) (ex)(AA) -)-- -BB.
After getting the discretized A and B, the discretization of Equation [ can be rewritten as follows,
h; = Ah + B
(3)
= Ch+-
At last, the Mamba model computes the output through global convolution as follows,
K (CB, CAB, ...,, C..
(4)
= x K
where denotes the output sequence, L denotes the length of the input sequence x and K E R- denotes one structured convolutional kernel.
METHODOLOGY
In this section, start by exploring the design philosophy for sparse patterns mobile devices. Next, we present a sparse learning framework aimed at optimizing the model's sparse structure through effective loss, sparsity loss, and latency loss for each layer. We then introduce compen- sation method that leverages pruned weights to further optimize the remaining weights. Finally, we illustrate of comprehensive compiler-enabled optimizations for proposed kernels.
5.1 SPARSE KERNEL DESIGN
Rationality of our sparse kernels We split the weights into multiple non-overlapping groups and each group has adjacent weights. Our kernel is designed as CT. which removes n elements from every group with four adjacent weights. This approach is inspired by the architecture of modern hardware's 's Single Instruction Multiple Data (SIMD) units that process groups of four elements at once (Cypher & Sanz 1989 Mitra al. 2013 Khorasani al. 2015p. Utilizing SIMD's ability to handle vectors of four elements in parallel boosts computational efficiency and performance (Chen) al., 2018, Alibaba, 2020p. By tailoring pruning kernels to fit the SIMD architecture,
hance the inference computations to fully leverage hardware capabilities, leading to higher efficiency with faster speed. Details are further explained in Section 5.6
Table 1: Latency profiling for different kernels.
Latency profiling In our sparse kernels with CT, can be different values leading to vari- ous sparsity and latency. To get a deeper under- standing for our C and n, we collect latency data from various kernels for our next sparse learning process. synthesized model with random weights is utilized for profiling on a mobile CPU since the weight values barely affect latency. The representative profiling results are shown in Table I We can establish the latency function Fr using Table This function then used the latency loss, detailed in Section
Kernl
C²
C
C
C
C
Sparisty 0%
25%
50%
75%
100%
Latency (ms) 37.14
29.66
22.71
19.05
5.74
4

---

Published as a conference paper at ICLR 2025
5.2 EFFECTIVENESS LOSs
Next we perform sparse learning for the SSM model to apply our sparse kernels for the whole model. To mitigate the accuracy degradation, it is essential to develop an accuracy predictor for our sparse learning, given that evaluating model performance on the full testing dataset is resource-intensinn and not conducive to model generalization. To address this problem. we introduce an effectiveness loss Lg based on the weight removal error, as the accuracy predictor. Considering the computational constraints, a global analysis of weight removal error presents significant challenges. Consequently, a layer-wise investigation emerges as a viable approach under these limitations.
Given the input X E P B for one layer with weight W E RR Pout x Pin, where B denotes the batch size, L is the sequence length, and Din/Dout denote the input/output dimensions, to mitigate the performance loss, the difference of outputs before and after pruning is minimized as follows,
M. L= WX (W @ M)
(S)
where 11.112 denotes the l2 norm, M E R Poui Di, is the sparse mask indicating the pruned locations, C denotes the element-wise multiplication, and W denotes the optimized weights. To make the problem tractable, we assume that the sparse mask is given and fixed during the optimization, and can have the following solution with detailed proof in Appendix A
Theorem 5.1. The optimal solution to Problem
with a fixed M can be obtained by the following,
W E e¹ (M, "M. "M!
x (2XX*)
(6)
and the minimal loss can be expressed as
T W. (2XX") -MM
(7)
The sparse locations are distributed in total of k rows in W, and their row indices can be denoted by eq. E. {0, 1] Dous x1 = I,...,k, k, where eq. is a one-hot vector with the qf element as I and all others as O. There are kū elements pruned in the df row and their indices can be represented by pp.-q, E {0, P) .j I,..., ki, which are also one-hot vectors similar to By: M, E g Xki, where the jth column of M, is [Ma].s = =pp:-40> Vj.
Remark 5.2. 2XX is not full rank with difficulties for the inversion (2XX*) )-1 the dampening technique is adopted to compute (2XX² +>I)= -1 instead, with y as the dampening ratio.
Remark 5.3. Although M is fixed during the optimization, the optimal M can be obtained by min- imizing the loss in Equation 7 i.e., each mask corresponds to loss in Equation 7 and the mask with the minimal loss is the optimal one. But it typically incurs unaffordable complexity find the optimal mask by comparing the losses of all masks.
Kernel-wise effectiveness loss Based on the optimal loss, we can obtain the effectiveness loss to estimate the weight removal error for each kemne.. For example, for C3 which has 2 zero elements among 4 elements, there are 6 combinations to select elements from 4, and we compute the loss for each combination following Equation 7 Then we sort the losses and find out the minimal loss as the effectiveness loss with the corresponding two pruned elements. We can perform the same process for other cases such as C with 3 zeros among 4 weights. In this way, the effectiveness loss for all cases (C), Cl. C₁, C. C+) of each kernel can be obtained. Note that once the effectiveness loss is determined, the corresponding pruning locations or weights for each kernel are also obtained as other combinations lead to larger weight removal error.
Effectiveness loss of the layer For the weight subset wij RI44 with = 1,..., Dout and j 1.... [Di./4] /4] where [-] denotes rounding up to the nearest integer, the weight effectiveness EW E P (PP.. D.../4] according to the removal strategy NW E {0, 1,, 3, 3, Dowt x[D../4] can be defined as follows,
(8) where (z k) denotes the function to generate the effectiveness loss for each kernel with & pruned elements according to Equation 7 and [NW denotes the number of pruned elements waj- Thus, we deliver the effectiveness loss LE by accumulating weight removal errors of all kernels:
E W Jk, = (w.,, IN (N 1ed),
LE = Σ E EE 1...
(9)
2.)
5

---

Published as a conference paper at ICLR 2025
5.3 SPARSITY AND LATENCY LOss
In our proposed sparse learning framework, the optimization of the sparse structure can be driven by accuracy or latency considerations. Thus, apart from the effectiveness loss, include the additional sparsity loss Ls or the latency loss Lr, to learn the sparse mask.
We first deliver the discrete sparsity SW E Doux x [D4/4]] of the subset weight Wa,j with the removal strategy NW as follows,
[S = [N /4
(10)
Subsequently, we define the sparsity loss Ls with the target sparsity ratio S, follows,
St avg(SW), ), if S, avg(SW) otherwise
Ls(sw, Sb) =
(11)
where avg(-) is the average function. We denote the discrete latency TW E Po of the sparse model using (-) defined in Section 5.1 with the look up table follows,
[T lag FrN
(12)
Then, the latency loss LT can be defined with the target latency Tr as follows,
Eis TW Tr, if Σi, TW T₃
CT Tí)
(13)
0. otherwise
5.4 SPARSE LEARNING
To optimize the removal strategy [NW].s for each sparse kernel Wa, one layer, we define the trainable probability mask as MW E (DDosx(Dee//)) with its subset m W E RRxx for each kernel. Each kernel has its unique effectiveness loss, sparsity loss and latency loss, corresponding ent pruned number, i.e, 0,1,2,3, or 4 with a total of cases. The probability mask m W is designed 1.j learn the number of pruned weights Way (choosing one from the cases), and is randomly initialized with the Gaussian distribution for the probability of removing 0.1,2.3, or 4 weights. We minimize the sum of the weight effectiveness and sparsity (or latency) losses with the sparse kernel selection denoted by the trainable mask MW with softmax operation to learn the appropriate spar- sity configuration for each kernel. The algorithm for sparse learning presented in Algorithm D featuring as the effectiveness loss ratio and Ti the target loss ratio. The 'max_index' retrieves the index the maximum value along a specified dimension. The output N Noutput generated by the algorithm represents the optimal sparse structure using the proposed kernel. Note that although only find the number pruned elements in each kernel for the sparse learning, the pruned locations corresponding to each choice of pruned number is already determined as discussed in Section 5.2 Thus, once W is learned, the pruned weights each kernel are also obtained immediately.
differ-
Algorithm 1: Sparse Learning
Input: W E R. Pin x Dout Tes Ft, S, or Tú, Create trainable mask MW E { Initialize all m C MW with Gaussian distribution.
[N.] Dave x [D4,/44] u 0,1,,33,
[E.j..; (w..j (N.kk)). = 0,1,2,3,4, Vi
S, N N,.4. 1, 1,2,3,
T. ===F(( [N u.kk), - 3, Vi,j
for e in [1, do
M' softmax(M dim=-1)
re (s(uum(( M'
Eg Ej Eg, Eg Ea), dim=-1)
+r, Ls (s(uu M' @
So. S, S2, Sz. S.), -im==1)),
or
+r₂ Cr(sum M' O (o. TI, T2, Ts. Tal, dim=-1) 1;)
Backward L
Update MW with SGD
Decay learning rate
end
Output:
N Noutput max maxiindex(ssftmmxxXMM dim=-1)

---

Published as a conference paper at ICLR 2025
5.5 COMPENSATION
After learning the sparse mask for each kernel in previous steps, we adjust the unpruned weights to offsst the accuracy loss from pruning. Specifically, we update the unpruned weights according to Equation 6 Since Cg. is a one-hot vector, Cx. x only has non-zero values in the qhth row with all zeros for all other rows. Thus, in Equation C each term with the index i in the sum just computes the ofth row in the outputs and the computation of the ofth row does not affect the a row, Vss + Specifically, based on Theorem 5.1 we have the following.
Jes. = [W] eº W.. ( (2xx") "M. M (2XX*)
(14)
Thus, can perform optimal adjustments row by row. For each row, Equation 14 incorporates the pruned locations across all kernels in the same row and considers their connections, as shown by
[M] (2Xx")-IM)
5.6 HARDWARE IMPLEMENTATION
For Mamba models utilizing C pruning, our compiler aims for maximum hardware efficiency during inference by implementing advanced optimizations and code generation. The optimization strategy consists of two primary approaches: (1) General optimizations, which includes operator fusion, static memory planning, and parameter tuning for Mamba models.
(2) C. -peecciif opti- mizations aimed at enhancing performance uniquely for the C Our compiler takes a model (in the form of computational graph) as input and generates source code output. Specifically, it pro- duces C++/Assembly code for mobile CPUs and OpenCL code for mobile GPUs. Due to the space limitations, our general compiler optimizations details in Appendix B.I
CT -seeciiiic Optimizations The irregular data access and computation patterns in pruned models significantly contribute to the inefficient execution of sparse DNNs (Sun et al.) 2023 Frantar & Alis-
tarh 2023 Our compiler tackles this inefficiency through two strategies: (1) reordering the weights offline to remove computational uncertainty and balance load disparities in sparse computation. (2) developing an efficient sparse weight storage method that utilizes the kernel information further reduce extra storage needs.
Figure 2 (a) and (b) illustrate examples of matrix multiplication with and without reordering. In Figure (a), without reordering, the process encounters typical challenges of sparse matrix mul- tiplication: control-flow divergence among threads, load imbalance across threads, and irregular memory access patterns. In Figure 2 (b), we initially group rows with similar numbers of non- zero elements in the weights to minimize load imbalance since different threads in a warp access contiguous rows. This approach reduces the number of idle threads caused by load imbalances. Within each row, we further reorder the same pattern of non-zero elements together to minimize the control-flow divergence. In Figure 2 (a) and (b), prior to reordering, the nested loop for computation induces if-else divergence, which is not hardware- frienddly for parallel computing units, particularly in mobile GPUs. However, after reordering, employ regular for loop that can perform identical computations for identical patterns contiguously, thereby eliminating branch divergence.
After reordering the matrix, store the model in compact format using our -peecciii format, named CT Compact as shown in Figure (c). Unlike CSR (Buluç et al., 2009) which only elim- inates zero weights, C Compact achieves higher compression ratio through hierarchical index structure that removes redundant column indices resulting from pruning with CT. This method effi- ciently conserves the limited memory bandwidth of mobile devices. The primary advantage of C Compact over CSR is its ability to store column indices more compactly. This efficiency comes from recognizing that multiple rows may share identical column indices after applying C kernels for pruning. To accomplish this, it utilizes three arrays: Ro Idx, Rc Idx, and K Idx, to represent the row and column index for each kernel after reordering, and pattern style, respectively. As we don't ' need to store the column/row information for each non-zero elements, this compact storage saves significant memory usage ranging from 25% to 75% (the number depends on the pruning ratio in the sparse weights) compared with the CSR format.
Layout Transformation Optimization Our previous optimizations are dedicated for the sparse model, targeting computational intensive operators, e.g., Cony, MatMul. However, Mamba's 's archi- tecture (computational graph), which relies heavily on layout transformation operators for tensor reshaping and transposing, demands high bandwidth memory (Gu & Dao 2023a). This issue is par-

---

Published as a conference paper at ICLR 2025
in to matre sto ccumple natii b.w.
(a) Before Reorder
in antris A.l For't blodi mantic keeee
(b) After Reorder
C2
4
2
3
2
comppas aocredng pataen (r is atse (riC2) olse (FCC))
2
3
N. mon Jog Ro sogeggeogg
(nng to K_Mppar
C
0
4
0
0
Cormppate Nodia
C7
2
2
0
For i [8] tok Miir4)
K M_I
(c) Sparse Storage
Re Ids:
Weight reorder and sparse storage. Numbers in (a) and (b) represent the pattern styles. Three extra indexes will be utilized during execution. Ro_Ida represents the original row index after reordering, while Rc_/dr denotes the original column index for each reordered kernel. And K Idw represents the stride information for kernels with identical patterns.
Figure 2:
ticularly critical on mobile devices, where bandwidth is limited compared to desktop GPUs (Huynh er al. 2017: Lee et al. 2019). leading to significant overhead. To further improve the performance, we propose layout transformation elimination strategy to fully eliminate the layout transformation operators while maintain the same accuracy. The key insight is that requiring the producer to create a layout based on the consumer's reduction dimension results in relatively low additional overhead compared to other options. First, compiler extract the operator information (e.g.. operator types, input feature sizes) and select the optimal layouts for each individual operator. Then, we eliminate the layout transformation operator in the computational graph (i.e., Transpose and Reshape). Fi- nally, we align the data layout produced by one operator to the needs of the subsequent operator during computation, and the layout transformation elimination technique minimizes the overhead associated with explicit data transformations, The detailed formalized algorithm for layout trans- formation elimination is provided in Appendix [.2 This approach significantly benefits mobile devices with limited bandwidth by enhancing data locality and reducing memory access costs, as reflected Table (our dense vs. llama.cpp).
EXPERIMENTS
6.1 EXPERIMENT SETUP
Sparse learning recipe We use Mamba models to test the effectiveness of our method. Our approach covers a variety of Mamba models, with parameters ranging from 130M to 2.8B. The selective SSM architecture in Mamba primarily uses linear projections for both input and output in each block, with a significantly smaller number of SSM parameters (projections for A, B, C, and A). We apply sparsity all these parameters. For sparse learning, we train over 1000 steps using the SGD optimizer at starting learning rate of 0.1, decaying it by 0.1 every 200 steps. The original model weights are frozen and sparse learning for Mamba-2.8B takes 50 mins on a A6000 GPU. We generate optimal sparse mask for each layer with ppursity-orienter sparse learning independently. The results generated with latency-orienter sparse learning are included in Appendix C
We prune the model in one-shot and evaluate the task performance on multiple common sense reasoning datasets including LAMBADA (Paperno et al. 2016), HellaSwag (Zellers et al. 2019) PIQA (Bisk et al. 2020), Are- easy (Clark et al. 2018), Arc-challenge (Clark et al. 2018), and WinoGrade (Sakaguchi et al. 2021) Perplexity on LAMBADA dataset and average accuracy on all mentioned datasets are provided.
Testing bed The latency evaluations conducted on Oneplus 11 mobile device equipped with Snapdragon Gen 2 SoC, featuring an octa-core Kryo CPU and Qualcomm Adreno 740 GPU with 16GB memory. Tests utilize all threads mobile CPUs and all pipelines on mobile GPUs. We report the performance comparison against llama.cpp (contributors 2023a), as is the only open- source framework support Mamba mobile devices yet. The GPU runs use 16-bit floating points while the CPU runs use 32-bit floating points. The data types used in the evaluation are aligned with llama.cpp's 's support for Mamba. Each experiment is repeated 50 times. Due small variance, only average results are reported. The batch size is set to for all models unless otherwise specified.
6.2 MAIN RESULTS
We show our main results of zero-shot performance with Mamba models with 130M 370M, 790M, 1.4B, and 2.8B in Table 2 including SparseGPT (Frantar & Alistarh| 2023) and Wanda (Sun et al., 2023). with 2:4 and 4:8 patterns which can accelerate the sparse networks Mishra al. 2021). As observed, our proposed
We compare method against other semi-structure pruning methods,
8

---

Published as a conference paper at ICLR 2025
Table 2: Main results of zero-shot performance with all scales of Mamba models. We compare
against semi-structuee pruning methods, including SparseGPT and Wanda, with 2:4 pattern.
Sparsity Ratio
LAMBDA
HellaSwag
PIQA
Arc-E
Are-C
WinoGrade
Avg.
Method
PPL L
Ace 1
Ace 1
Acc
Acc f
Acc1
Ace 1
Ace T
Mamba- 130M
16.07
44.3
35.3
64.5
24.3
51.9
44.7
SparseGPT 2:4
50%
69.80
26.1
29.8
58.8
37.5
22.7
52.4
37.9
Wanda 2:4
50%
212.91
15.4
29.4
56.8
38.7
21.1
52.5
35.6
Ours
30%
17.47
42,1
34,5
64,0
47.1
23.8
53.6
44,2
Ours Ours
50%
28.97
35.2
32.2
60.8
41.6
24.2
51.9
41.0
70%
57.29
27.7
29.9
58.9
36.4
22.9
50.1
37.8
Mamba-870M
8.14
55.6
46.5
69.5
55.1
28.0
55.3
51.6
SparseGPT 2:4
50%
27.93
35.8
34.7
61.5
40.6
23.1
51.8
41.2
Wanda 2:4
50%
82.52
22.6
31.9
60.1
40.3
22.4
51.7
38.1
Ours
30%
8.55
54.6
44.9
68.8
52.7
27.5
55.3
50.6
Ours Ours
50%
12.33
47.9
40.2
64.7
47.6
26.5
54.3
46.9
70%
22.09
38.9
35.2
61.4
40.5
23.7
51.6
41.9
Mamba-7900 790M
6.02
62.7
55.1
72.1
61.2
29.5
56.1
57.1
SpurseGPT 2:4
50%
13.69
40.1
64.0
24.7
55.3
Wanda
50%
43.76
28.5
36.9
62.8
43.3
22.6
55.3
41.5
Ours
30%
60.7
53.6
71.8
58.7
28.4
56.2
54.9
48.0
68.9
26.3
51.1
70%
11.85
48.9
41.8
64.7
44.5
25.8
55.4
46.9
Mamba-1.40 1.4B
64.9
74.2
32.8
61.5
59.7
SparseGPT 2:4
50%
8.87
54.3
44.5
66.5
49.1
24.3
$4.5
48.9
Wanda 2:4
50%
32.72
31.6
38.2
63.9
46.8
22.8
53.7
42.8
Ours
30%
5,08
65.1
58.2
73.2
64.0
32.4
60.7
58.9
Ours
50%
5.65
62.5
52.7
70.7
58.6
29.0
59.0
55.4
Ours
60%
7,30
57.5
46.0
68.3
51.6
26.3
$7.4
51.2
Ours Ours
70%
17.40
43.2
31.7
62.5
43.4
19.5
55.1
42.6
75%
19.65
42.0
35.7
61.1
41.2
22.9
54.4
42.9
Mamba-2.88
4.23
69.2
66.1
75.2
69.7
36.3
63.5
63.3
SparseGPT 2:4
50%
5.11
65.6
52.1
70.0
56.0
27.6
59.8
55.2
Wanda 2:4
50%
10.49
50.0
48.0
65.8
54.9
26.3
56.2
50.2
Ours
30%
4.18
69.4
65.2
75.6
69.2
36.4
62.6
63.1
Ours
50%
4.26
68.9
60.2
72.6
65.2
31.5
61.1
59.9
Ours 60%
4.72
66.7
54.0
71.1
57.3
28.4
59.4
56.1
Ours 70%
7,51
58.8
43.3
64.6
46.6
25.2
58.3
49.5
Ours
75%
15.86
46.6
36.2
61.0
40.7
22.8
56.2
43.9
pruning method shows a consistently significant improvement over SparseGPT and Wanda at the same sparsity ratio terms of the average accuracy across the six datasets. Taking Mamba-1.49 as example, our method with a sparsity ratio of 50% increases the average accuracy by at least compared with all baseline methods. Besides, compared to other methods with fixed 2:4 or 4:8 patterns, our method can achieve flexible sparsity ratios, such as 30% sparsity ratio with an accuracy comparable to the dense model. Full results are provided in Appendix D In summary, the results achieved by method indicate the effectiveness of our pruning technique in maintaining performance across different model scales and tasks. We also provide the comparison to transformer- based models in Table A3 and Table A4 and achieve superior task performance in terms of perplexity and average accuracy compared to transformer-based models under the same model size.
6.3 LATENCY RESULTS
We conduct the latency measurements of our method under different model scales on both the CPU and GPU of the Snapdragon 8 Gen 2 mobile platform to show the actual acceleration performance. The generation latency results are shown in Table 3 Our results demonstrate that, compared with the dense model, our method achieves approximately 1.1x to 1.2x acceleration at 30% sparsity, 1.3x to 1.5x acceleration at 50% sparsity and nearly 1.6x to 1.7 acceleration 70% and 75% sparsity, on both mobile CPU and GPU. To demonstrate the performance gains from general optimizations and layout transformation elimination, our compiler (dense) achieves up-to 4.5x speedup compared to laamaccpp (dense since llama.cpp does not yet support sparse models) mobile CPU. It is also

---

Published as a conference paper at ICLR 2025
Table 3: Latency results of Mamba with different scales Parplexity and 64 sequence length. SPD stands for speedup over lpp (red) and dense baseline (blue).
200
175
wio Compensation
w/ Compensation
150
Mobile CPU
Mobile GPU
Mamba
Method
Sparsity
Tokenss
SPD
Token/s
llama.cpp
100
130M
ours
0%
11.2
3.2x/1.0x
35.6
1,0x
ours
30%
12.1
3.5x/L.1x
40.7
1.2x
50%
149
4.3x/l13x
50.4
1.5x
ours
70%
18.1
5.2x/166x
54.8
1.6x
25
laama.cpp
1.5
1.0%
Size
130M
370M
790M
1.4B
2.7B
370M
ours 30%
1P9
6.1
6.9
4.6x/1.1x
4.1x7l.0x
21.5
19.6
1.1x
1,0x
Table 4:
Ablation analysis with or
ours
7.9
5,5x/1,3x
without weight compensation.
ours
70%
9.8
6.5x/1.6x
31.3
1.6x
lama.cpp
0%
0.7
1.0%
790M
ours
4.5x7L0x
Perplexity
Token's
ours 30%
ours
50%
4.3
3,3
4.7x/1.1x
6.3x/114
16.6
14.1
1.3x
1.1x
16
PPLToken's
Uniform
70%
5.0
7.0x/166x
20.4
1.6x
Mix
1.2
12
laaaacpp
0%
0.6
1.0%
10
1.4B
ours
2.1
3.5x/LLxx
10.1
1,0x
30%
2.5
4.1x/1.2x
10.9
1.1%
ours 50%
2.9
4,8x/14x
13.1
1.3x
0.4
ours
3.3
5.8x1166x
17.2
0.2
llama.cpp
0%
0.4
1.0x
2.8B
ours
0%
1.9
4.0%/T.0x
7.7
10x
C7 Sparsityy. 25%
/ 750
75755%
30%
2.2
4.7xf1.2x
ours 50%
2.6
5,5x/114x
11.5
1.5x
Table 5: Ablation analysis of uniform and mixed kernels strategies.
ours 75%
3.0
6.5x/L66x
13.0
1.7x
worth pointing out that our compiler works efficiently on mobile GPU, achieving approximately 3x to 5x speedup compared our CPU version. GPU results are not included for llama.cpp as not yet supported. Meanwhile, we present other hardware results including the peak memory during inference and energy consumption in Table A5 and Table A6 respectively. We observe that our sparse models require only a slight increase in memory while achieving notable energy savings and significant inference acceleration, demonstrating the efficiency of our method. Additionally, the results tested on another edge device, Xiaomi 6, which equipped with a Snapdragon 835, featuring octa-core CPU and an Adreno GPU with 6GB of memory, are included in Table A7
6.4 ABLATION STUDY
We conduct the ablation study in terms of weight compensation, and the results are shown in Fig- We experiment with two settings: with without weight compensation, for each scale of the Mamba model on the LAMBDA dataset. The consistent better task performance with our weight compensation technique demonstrates its effectiveness.
4
To verify the effectiveness of sparse learning, we also employ uniform kernel strategy for the ablation as shown Figure 5 The perplexity results evaluated with Mamba-2.8D model on the LAMBADA dataset, and the latency results are tested on mobile CPU. Compared with the uniform kernel strategy, our mixed kernel strategy with sparse learning achieves better task performance while maintaining similar latency performance, especially when the sparsity exceeds 25%.
CONCLUSION
In this paper, we propose a sparse learning framework with special compiler kernels for the accel- eration of Mamba models mobile devices. We introduce the C kernel designed to prune elements from every four contiguous weights. Then, we propose the sparse learning framework to explore the optimal pruning strategy with kernels for the weights in Mamba models. Besides, propose the weight compensation method to rectify the weights in sparse models with the pruned weights. Finally, we implement the sparse strategy with kernels mobile devices with the CT- specific optimizations and the layout transformation elimination strategy. Experiments show the effectiveness of method for the acceleration of Mamba models mobile devices.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
We sincerely appreciate the contributions of everyone involved in this work, with special thanks to the anonymous reviewers for their valuable feedback. This research was partially supported by the National Science Foundation (NSF) under awards CCF-2428108 and OAC-2403090. Any enrors and opinions are not those the NSF and are attributable solely to author(s).
REFERENCES
A new approach to linear filtering and prediction problems. 1960.
Alibaba. Mnn: universal and efficient inference engine. settp::/gitthu..com/aa MNN
2020.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical com- monsense in natural language. In Proceedings of the AAAI conference on arlificial intelligence, volume 34, pp. 7432-7439, 2020.
Aydin Buluç, Jeremy T Fineman, Matteo Frigo, John R Gilbert, and Charles E Leiserson. Parallel sparse matrix-vecte vector and paarix-transpose-verr vector multiplication using compressed sparse blocks. In Proceedings of the twenty-first annual symposium Parallelism in algorithms and architec- tures, pp. 233-244, 2009.
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. TVM: An automated end-to-end optimizing compiler for deep learning. In OSDI, pp. 578-594, 2018.
Yu-Hui Chen, Raman Sarokin, Juhyun Lee, Jiuqiang Tang, Chuo-Ling Chang, Andrei Kulik, and Matthias Grundmann. Speed is all you need: On-device acceleration of large diffusion models via gpu-aware optimizations. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pp. 4650-4654, 2023.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint
2018.
LLaMA.CPP contributors. LLaMA. http://pphaaaee.aaooaaooooarrrra cpp 2023a. http:/uwww..mmmmmmmmaaaaaalll---
com http.//a....c.ggggg.///aaaaaaaoa
MLC contributors. MLC-LLM. http:://tttbb..ommmmm/
2023b. URL
https::/
Robert Cypher and Jorge LC Sanz. Simd architectures and algorithms for image processing and computer vision. IEEE Transactions Acoustics, Speech, and Signal Processing, 37(1)):2158- 2174, 1989.
exllama contributors. exllama. httpp:://ggthh...com/bbbbbbootee https:///ithhb. wtt cwm/turbooerrp/ex........oo lama
lama 2023. URL
Elias Frantar and Dan Alistarh. SparseGPT: Massive language models can be accurately pruned in one-shot. arXiv preprint arXiv:2301.00774, 2023.
Elias Frantar, Eldar Kurtic, and Dan Alistarh. M-fac: Efficient matrix-free approximations of recond-order information. Advances in Neural Information Processing Systems, 35, 2021.
Daniel Y Fu, Tri Dao. Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212. /4052, 2022.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023a.
11

---

Published as a conference paper at ICLR 2025
Albert Gu and Tri Dao. Mamba: Linear-ting sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023b.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2/11.0003 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Com- bining recurrent, convolutional, and continuous-tine models with linear state space layers. Ad- vances in neural information processing systems, 34:572-585, 2021b.
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35:35971- 35983, 2022.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces, Advances in Neural Information Processing Systems, 35:22982-22994, 2022.
Seungyeop Han, Haichen Shen, Matthai Philipose, Sharad Agarwal, Alec Wolman, and Arvind Kr- ishnamurthy. Mcdnn: An approximation- based execution framework for deep stream processing under resource constraints. In Proceedings of the 14th Annual International Conference Mo- bile Systems, Applications, and Services (MobiSys), pp. 123-136. ACM, 2016.
B. Hassibi, D.G. Stork, and G.J. Wolff. Optimal brain surgeon and general network pruning. In IEEE International Conference on Neural Networks, pp. 293-299 vol.1, 1993. doi: 10.1109/ ICNN. 1993,29857..
Loc N Huynh, Youngki Lee, and Rajesh Krishna Balan. Deepmon: Mobile gpu- gpu-based deep learning framework for continuous vision applications. In Proceedings of the 15th Annual International Conference on Mobile Systems, Applications, and Services (MobiSys), pp. 82-95, ACM, 2017.
doi: 14.1345/308 333.3081366
Farzad Khorasani, Rajiv Gupta, and Laxmi N Bhuyan. Scalable simd-efficient graph processing on gpus. In 2015 International Conference on Parallel Architecture and Compilation (PACT), pp- 39-50. IEEE, 2015.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b- paramcter open-access multilingual language model. 2023.
Royson Lee, Stylianos I Venieris, Lukasz Dudziak, Sourav Bhattacharya, and Nicholas D Lane. Mobisr: Efficient on-device resor-resolution through heterogeneous mobile processors. In The 25th annual international conference on mobile computing and networking, pp. 1-16, 2019.
Yanyu Li, Pu Zhao, Geng Yuan, Xue Lin, Yanzhi Wang, and Xin Chen. Pruning-as-search: Efficient neural architecture search via channel pruning and structural reparameterization. arXiv preprint arXiv:2206.00 01 198, 2022.
Yanyu Li, Changdi Yang, Pu Zhao, et al. Towards real-time segmentation the edge. AMAF23/AAA
2023a. ISBN 978-1-57735-88044 doi: 10. 2..1609/aai..3372225223. URL httts:://daat /ttp:p:://doi..r//11..1 10..0.......00.0.
Yize Li, Pu Zhao, Xue Lin, Bhavya Kailkhura, and Ryan Goldhahn. Less is more: Data pruning for faster adversarial training. arXiv preprint arXiv:2302. 2023b..
Yize Li, Yihua Zhang, Sijia Liu, and Xue Lin. Pruning then reweighting: Towards data-efficient training of diffusion models. arXiv preprint arXiv:2409. 19128, 2024a.
Yize Li, Pu Zhao, Ruyi Ding, Tong Zhou, Yunsi Fei, Xiaolin Xu, and Xue Lin. Neural architec- ture search for adversarial robustness via learnable pruning. In Frontiers in High Performance Computing, 2024b.
Sachin Mehta, Mohammad Hossein Sekhavat, Qingqing Cao, Maxwell Horton, Yanzi Jin, Chenfan Sun, Iman Mirzadeh, Mahyar Najibi, Dmitry Belenko, Peter Zatloukal, and Mohammad Raste- gari. OpenELM: An Efficient Language Model Family with Open Training and Inference Frame- work. arXiv.org, April 2024. URL http:::/p//// http:s://///vvvv. wwwpvi.orrgaaeee abs 2404 1461 9v1
12

---

Published as a conference paper at ICLR 2025
Asit Mishra, Jorge Albericio Latorre, Jeff Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu, and Paulius Micikevicius. Accelerating sparse deep neural networks. arXiv preprint w0Xiv:2/0400837,, 2021.
Gaurav Mitra, Beau Johnston, Alistair P Rendell, Eric McCreath, and Jun Zhou. Use of simd vector operations to accelerate application code performance on low-powered arm and intel platforms. In 2013 IEEE International Symposium Parallel & Distributed Processing, Workshops and Phd Forum, pp. 1107-1116. IEEE, 2013.
Wei Niu, Xiaolong Ma, Sheng Lin, Shihao Wang, Xuehai Qian, Xue Lin, Yanzhi Wang, and Bin Ren. Patdnn: Achieving real-time dnn execution on mobile devices with pattern-based weight pruning. In Proceedings of the Twenty-Fift International Conference on Architectural Support for Programming Languages and Operating Systems, pp. 907-922, 2020.
Wei Niu, Jiexiong Guan, Yanzhi Wang, Gagan Agrawal, and Bin Ren. Dnnfusion: accelerating deep neural networks exccution with advanced operator fusion. Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation, pp- 883-898, 2021.
Wei Niu, Gagan Agrawal, and Bin Ren. Sod2: Statically optimizing dynamic deep neural network execution. In Proceedings of the 29th ACM International Conference Architectural Support for Programming Languages and Operating Systems, Volume 1, ASPLOS '24, pp. 386 4-00, New York, NY, USA, 2024a. Association for Computing Machinery. ISBN 9798400703720. 10. 10.11455/617233.36444666 UKL_fttes///01...oo0.0...111//1.
doi: htt:::ts:://// wttp:ioorgg.... 11t.://....11533661 3117/31113333 362486553626888
Wei Niu, Md Musfiqur Rahman Sanim, Zhihao Shu, Jiexiong Guan, Xipeng Shen, Miao Yin, Gagan Agrawal, and Bin Ren. Smartmem: Layout transformation elimination and adaptation for efficient dnn exccution on mobile. In Proceedings of the 29th ACM International Conference on Archi- tectural Support for Programming Languages and Operating Systems, Volume 3, pp. 916-931, 2024b.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv: 11895, 2022.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv: 2016
Qualcomm. Snapdragon 845.
https://en wik www.kk/pdaaa.g./aaaa..
Qualcomm Snapdragon systems on chipsh 2017.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver-
sarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-108 106, 2021.
Samsung.
https://ansights..smm httpm/ssunnncc/m/212112/ www. 2tmm/202100e
Xuan Shen, Zhenglun Kong, Minghai Qin, Peiyan Dong, Geng Yuan, Xin Meng, Hao Tang, Xiao- long Ma, and Yanzhi Wang. Data level lottery ticket hypothesis for vision transformers. arXiv preprint arXiv:2211.01483 2022.
Xuan Shen, Yaohua Wang, Ming Lin, Yilun Huang, Hao Tang, Xiuyu Sun, and Yanzhi Wang. Deep- mad: Mathematical architecture design for deep convolutional neural network. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6163-6173, 2023.
Xuan Shen, Peiyan Dong, Lei Lu, Zhenglun Kong, Zhengang Li, Ming Lin, Chao Wu, and Yanzhi Wang. Agile-quant: Activation Activation-guided guided quantization for faster inference of Ilms the edge. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 18944-18951, 2024a.
13

---

Published as a conference paper at ICLR 2025
Xuan Shen, Zhaoyang Han, Lei Lu, Zhenglun Kong, Peiyan Dong, Zhengang Li, Yanyue Xie, Chao Wu, Miriam Leeser, Pu Zhao, Xue Lin, and Yanzhi Wang. Hotaq: Hardware oriented token adaptive quantization for large language models. IEEE Transactions Computer-Aided Design of Integrated Circuits and Systems, pp. 1-1, 2024b. doi: 10. 109/TCCD.220444447777
Xuan Shen, Zhao Song, Yufa Zhou, Bo Chen, Yanyu Li, Yifan Gong, Kai Zhang, Hao Tan, Jason Kuen, Henghui Ding, et al. Lazydit: Lazy learning for the acceleration of diffusion transformers. arXiv preprint arXiv:2412. /2444, 2024c.
Xuan Shen, Zhao Song, Yufa Zhou, Bo Chen, Jing Liu, Ruiyi Zhang, Ryan A Rossi, Hao Tan, Tong Yu, Xiang Chen, et al. Numerical pruning for efficient autoregressive models. arXiv preprint arXiv:2412.12441 2024d.
Xuan Shen, Pu Zhao, Yifan Gong, Zhenglun Kong, Zheng Zhan, Yushu Wu, Ming Lin, Chao Wu, Xue Lin, and Yanzhi Wang. Search for efficient large language models. In NeurIPS, 2024e.
Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural network compression. H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 18098- 18109. Curran As- sociates, Inc., 2020. URL elff 3903ff www.00ff1113a336o3
pdf
Mingjie Sun, Zhuang Liu, Anna Bair, and J. Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint arXiv:2306. 2023.
The fastllm contributors. fastllm. /ttp 2023. URL
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Yushu Wu, Yifan Gong, Pu Zhao, et al. Compiler-aware neural architecture search for on-mobile real-time super-resolution. In ECCV, pp. 92-111. Springer, 2022.
Mengwei Xu, Mengze Zhu, Yunxin Liu, Felix Xiaozhu Lin, and Xuanzhe Liu. Deepcache: Princi- pled cache for mobile deep vision. In Proceedings of the 24th Annual International Conference on Mobile Computing and Networking, MobiCom '18, pp. 129-144, New York, NY, USA, 2018. ACM, Association for Computing Machinery. ISBN 978 9781450359030.
Yang, Baosong Yang, Binyuan Hui, Bo Zheng. Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li. Dayiheng Liu, Fei Huang, et al. Qwen2 technical report. arXiv preprint arXiv:2407.1067/ 2024.
Changdi Yang, Pu Zhao, Yanyu Li, et al. Pruning parameterization with bi-level optimization for efficient semantic segmentation on the edge. In CVPR, pp. 15402-15412, 2023.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? arXiv preprint arXiv:/905.00 /905.07830 2019.
Zheng Zhan, Yifan Gong, Pu Zhao, Geng Yuan, et al. Achieving on-mobile real-time super- resolution with neural architecture and pruning search. In ICCV, pp.
2021.
Zheng Zhan, Zhenglun Kong, Yifan Gong, Yushu Wu, Zichong Meng, Hangyu Zheng, Xuan Shen, Stratis loannidis, Wei Niu, Pu Zhao, and Yanzhi Wang. Exploring token pruning vision state space models. In The eiiht Annual Conference on Neural Information Processing Sys- tems, 2024a. URL http:pse//oopeeeeeeeeeeeeeeeeeee
Zheng Zhan, Yushu Wu, Yifan Gong, al. Fast and memory-effiiien video diffusion using streamlined inference. In NeurIPS, 2024b. URL WRLLattps::/ppeeee/eee... net Fit. NvXYOrkpil
14

---

Published as a conference paper at ICLR 2025
Yihua Zhang, Yuguang Yao, Parikshit Ram, Pu Zhao, Tianlong Chen, Mingyi Hong, Yanzhi Wang, and Sijia Liu. Advancing model pruning via bi-level optimization. Advances in Neural Informa- tion Processing Systems, 35:18309-193
2022.
Pu Zhao, Fei Sun, Xuan Shen, et al. Pruning foundation models for high accuracy without retrain- ing. In Findings of EMNLP 2024, pp. 9681-9694, ACL, November 2024, doi: 10. 1/2244 findings-emnlpe. URL httos: oatoss///cciaatnnooooooo..o//2/. ftneeeiiiiniiigg windinnsseinnnnnnnnnn...
566
15


Reviews:
[{"review_id": "0MuwM2l7M1", "reviewer": "Reviewer_vUZD", "summary": "This paper introduces a mobile-friendly solution for SSMs, targeted for on-device inference. Specifically, the authors propose a sparsification & pruning method of contiguous weights along with a reordering operator for efficient on-device execution. Last, they introduce a sample-efficient compensation algorithm that recovers any lost accuracy. Results showcase gains of from 3.2x to 7x without gradual accuracy degradation as a function of sparsity.", "questions": "### Evaluation\n\n* In the introduction, the authors quote that \"llama.cpp takes over 5.8s to generate a single token with Mamba-2.8B\", However, in table 3 of the evaluation, a similar setup is quoted at 0.4 tokens/sec = 2.5s per token. Could the authors clarify the source of this discrepancy?\n* Although the performance benefits seem impressive at first sight, the 7x gains quote is quite misleading, as it comes with significant accuracy degradation. For a sub-1% degradation, the speedup gain is not quoted on Table 3.\n* Why does Table 3 miss GPU execution for llama.cpp?\n* In Section 6.2, the average column seems to represent an non-weighted average across the datasets, thus not taking into consideration the size of each task. Is this correct?\n\n### Omissions / Extensions\n\n* Results have only been benchmarked on a single high-tier device with 16GB of memory. I am wondering how the proposed solutions work on lower-end devices that do not have these level of resources. At the same time, it would be very interesting to quantify the energy requirements of running such workloads on device.\n* A comparison with a similarly-sized Transformer-based LM would greatly enhance the evaluation and put the gains into perspective. Table A3 partly accomplished this, but there is no quantification of on-device performance.\n* Furthermore, it would be valuable to see how these gains compare with different compression methods, such as quantization for example.\n* Would there be any limitations of the method being applied to other modalities, such as vision (see Vision Mamba [c])\n* Since SSMs come with memory benefits, as quoted in §3, it would be important to highlight the peak memory consumption of the pruned models during inference.\n* It would be very insightful to visualise the sparsity per block in the resulting model to see if there is some kind or pattern in the pruning dynamics.\n\n### Questions\n\n* How does the compiler select the best layout per operator for different target devices?\n* What is the overhead during training of the models with the proposed method? Are there gains from pruning during training?\n\n[a] Xu, J., Li, Z., Chen, W., Wang, Q., Gao, X., Cai, Q., & Ling, Z. (2024). On-device language models: A comprehensive review. arXiv preprint arXiv:2409.00088.  \n[b] Liu, Z., Zhao, C., Iandola, F., Lai, C., Tian, Y., Fedorov, I., ... & Chandra, V. (2024). Mobilellm: Optimizing sub-billion parameter language models for on-device use cases. arXiv preprint arXiv:2402.14905.  \n[c] Zhu, L., Liao, B., Zhang, Q., Wang, X., Liu, W., & Wang, X. (2024). Vision mamba: Efficient visual representation learning with bidirectional state space model. Forty-First International Conference on Machine Learning (ICML).\n\n### Nitpicking\n\n* Table 1: sparsity typo\n* Numbering issue on remarks (i.e. missing remark 5.1)\n* Section 6.4: LAMBDA -> LAMBADA", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* The paper contributes both an algorithmic/architectural change and a system hardware component implementation for efficient on-device execution, across CPU and GPU backends, which is evaluated on device.\n* There has been great effort put in the evaluation and comparing against various baseline methods and models.\n* The proposed method, albeit involved, is straightforward and through the ablation we witness the importance of each component in the resulting accuracy.", "weaknesses": "* The proposed technique has only been applied to one SSM architecture (Mamba) and evaluated on a single high-tier device.\n* The on-device ML literature is quite old and there have been various contributions from 2018 onwards, also focusing on LLMs (see [a,b]).\n* The advertised gains are not quoted over the same accuracy threshold.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MBUz0IWfBL", "reviewer": "Reviewer_gxfe", "summary": "Recently, State Space Models (SSMs) are gaining attention in sequential modeling problems. However it comes with increasing computational complexity and bandwidth demands. This paper proposes a sparse learning framework with architecture-aware compiler optimizations. This proposal includes 1. optimized kernels, 2. sparsity or latency oriented learning framework that uses 1.\n\nThe paper is too abstract to provide sufficient insight to readers. It seems that the framework and the algorithm provided seems to be a reasonable contribution. hence the score of 5: marginally below the acceptance threshold.", "questions": "* Is there any performance impact associated with Remark 5.2.\n\n* How long does this \"sparse learning take\" it seems like the computation seems to be quite complicated and may prolong the \"learning\" by a big factor. Can you provide some data?\n\n* It seems that the work includes some compiler optimization that itself could be large enough to account for a separate paper. Can you provide more details to the infrastructure used? Is this some open-source work? Is this published anywhere?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "As the on-device AI is an important topic w.r.t privacy issues, it is an important direction to explore.\n\nThe paper seems to build a framework that reduces the computational intensity of the SSMs (which already is lower computational complexity than currently dominant Attention) seems to be a good direction.\n\nIt seems that the paper is claiming to have combined a number of compiler optimizations making it an end-to-end solution over simply exploring an algorithm-only or compiler-only or HW-only solution.", "weaknesses": "The paper is very difficult to understand. It seems that there are a lot of abstract explanations without concreteness that makes it difficult. Maybe adding some figures of what is really happening might help. For example, the paper states \"our kernel is designed as Cn4 , which removes n elements from every group with four adjacent weights.\" There may be multiple ways by which this could happen and the paper does not dive into details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8hX3Rf0Rjb", "reviewer": "Reviewer_6SJA", "summary": "The paper introduces a novel learning framework for state space models (SSMs) that emphasizes kernel sparsity to enhance performance on mobile devices. The approach optimizes for an ideal balance between sparsity, latency, and accuracy, enabling an efficient pruning strategy suited for mobile hardware. With a robust theoretical foundation, the framework accounts for sparsity levels and latency impact on accuracy. Further, the authors incorporate hardware efficiency through architecture-aware compiler optimizations, including weight reordering, sparse weight storage, and eliminating layout transformations, achieving state-of-the-art (SOTA) results in accuracy at comparable or higher sparsity levels.", "questions": "* Could additional metrics be included for deployment feasibility? Metrics like power consumption and memory consumption could provide more insights for real-world deployment on mobile devices.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Focus on Mobile Efficiency: The paper addresses a critical need for efficient execution of SSMs on mobile devices by optimizing for kernel sparsity without sacrificing accuracy, a valuable contribution to resource-constrained deep learning.\n* Strong Theoretical Foundation: The proposed framework’s theoretical backing adds credibility, ensuring that the optimizations in sparsity and latency are rigorously derived, rather than heuristically implemented.\n* Effective Hardware-Aware Optimization: Integrating architecture-aware compiler optimizations to handle sparse weight storage and reordering is a practical enhancement, boosting hardware efficiency while maintaining performance. This integration yields improved results over SOTA methods, showcasing either better accuracy for the same sparsity or higher sparsity for the same accuracy.* * Performance Presentation: The authors provide strong quantitative support,", "weaknesses": "* More discussion of latency trade-offs relative to different levels of accuracy and sparsity on the same model would be valuable for considering real-world deployment.\n* A minor improvement could be made by bolding or highlighting the best results in each column of Table 2 for easy reference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lMbOx45OpV", "reviewer": "Reviewer_iLUr", "summary": "The paper introduces a solution for accelerating inference using the Mamba SSM on mobile devices. The solution includes a framework for pruning the model’s weights and an assortment of optimizations to improve inference execution, such as weight reordering, operator fusion, and layout transformation elimination. Experiments using different sizes of the Mamba model show that the presented sparsification method achieves better accuracy than other semi-structure pruning methods, namely SparseGPT and Wanda. The authors’ solution outperforms llama.cpp when running inference with Mamba on the CPU of a Snapdragon 8 Gen2 SoC, and exhibits further speedup on the GPU.", "questions": "- Is your sparse kernel design and sparse learning approach unique, or is there prior related work?\n- Can your approach work for transformer-based models? If yes, have you done any accuracy comparisons?\n- Have you tested your approach with other SSMs? If yes, do you have any accuracy comparisons?\n- Why are your optimizations particularly good for mobile devices? Do you have any comparisons of the effects of your optimizations on mobile and desktop/server devices?\n- Can you comment on the novelty of your optimizations compared to related work?\n- Can you run your solution on non-mobile devices? If yes, do you have any results?\n- How can your approach be extended to NPUs?\n- You state in the paper that Mamba-370M achieves in your tests an average accuracy of 50.0%, while your 30%-sparse version achieves 50.6%. However, if one looks at the individual tests, there is no test where your version achieves higher accuracy than the original. Is there a typo?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "Overall, the sparse-learning framework for exploring pruning strategies for Mamba’s weights is promising. Furthermore, experimental results show an improvement in accuracy compared to other semi-structure pruning methods, SparseGPT and Wanda.", "weaknesses": "The paper presents an end-to-end solution, starting from a specific model, Mamba, pruning its weights, applying optimizations, and executing inference on a mobile SoC. Although such work must involve considerable engineering to implement the different stacks comprising the solution and gluing them together, it does not automatically produce any interesting results for the research-oriented community. Were there any particular challenges that had to be overcome, driving the development of specific innovations in the process? I do not see any such discussion in this submission. Although there can be innovations in the individual components, the full-stack presentation fails to sufficiently highlight them as, inevitably, we only get a rather shallow look into each one, both qualitatively and quantitatively. I am elaborating on that for every main component below.\n\nThe sparse-learning framework is very interesting, and the authors present improvements compared to other state-of-the-art pruning methods. However, the presentation is lacking. Are there other similar approaches to the Cn4 kernels? What is the related work here? SparseGPT and Wanda were initially developed for transformer-based models. Can the presented kernels tackle such models? If yes, how do they compare against the state of the art? If we are only looking at SSMs, why only Mamba? What about, e.g., S4 and H3?\n\nThe optimization workflow presented supposedly targets mobile devices. However, there is no explanation for why these optimizations are particularly good just for mobile devices and not all computational devices. The cited lack of “high throughput memory” (HBM?) on mobile devices is weak. All modern CPU and GPU architectures, both mobile and desktop/server, suffer from expensive data movement. Therefore, optimizations that reduce it are of general benefit. They may not offer equal benefit to all devices, but we cannot tell because there are no such comparisons in the paper. Maybe the optimizations take advantage of specific aspects of mobile architectures? There is no clear indication. There is a short discussion on SIMD units, but these are not fundamentally different between mobile and desktop/server devices.\n\nThere is a high-level description of the optimization workflow, but it is difficult to tell if any interesting innovations exist. Much prior work exists on optimizing sparse operations by reordering the non-zeros and introducing custom hierarchical sparse formats, such as the ParTI! Library (https://github.com/hpcgarage/ParTI). Other examples of prior work are DNNFusion (https://dl.acm.org/doi/10.1145/3453483.3454083) for operator fusion and SmartMem (https://dl.acm.org/doi/10.1145/3620666.3651384), which specifically addresses layout transformation elimination for mobile DNN execution. I am not saying that any of the above works are necessarily super relevant to the authors’ submission or that they need to be addressed, but if you are going to claim “a set of comprehensive compiler optimizations, including Cn4-specific optimizations and layout transformation elimination strategy on mobile devices” as a significant contribution, it will help to put your work into a better context.\n\nThe paper compares performance against llama.cpp. The authors provide an insight into why their solution is faster. Paraphrasing from their supplemental material, llama.cpp relies on a fixed pattern matching strategy to identify and fuse operation combinations, an approach that fails to recognize new combinations. Although llama.cpp is popular for executing transformer-based models, is it decent with SSMs? Aren’t there any better ways to execute Mamba and compare against them? If we are currently limited to llama.cpp because it is the only inference engine out there that currently supports Mamba and mobile, I would question if it is a “bad” baseline and how interesting the results are in the first place.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)