Spark Transformer: Reactivating Sparsity in Transformer FFN and Attention
Authors: Chong You
Kan Wu
Zhipeng Jia
Lin Chen
Srinadh Bhojanapalli
Jiaxian Guo
Utku Evci
Jan Wassenberg
Praneeth Netrapalli
Jeremiah J. Willcock
Suvinay Subramanian
Felix Chern
Alek Andreev
Shreya Pathak
Felix X. Yu
Prateek Jain
David E Culler
Henry Levy
Sanjiv Kumar
Venue: NeurIPS 2025 poster ()
Keywords: LLM
activation sparsity
Primary area: deep_learning

Abstract: The discovery of the *lazy neuron phenomenon* (Li et al., 2022), where fewer than 10% of the feedforward networks (FFN) parameters in trained Transformers are activated per token, has spurred significant interests in *activation sparsity* for enhancing large model efficiency. While notable progress has been made in translating such sparsity to wall-time benefits across CPUs, GPUs, and TPUs, modern Transformers have moved away from the ReLU activation function crucial to this phenomenon. Existing efforts on re-introducing activation sparsity, e.g., by reverting to ReLU or applying top-k masking, often degrade model quality, increase parameter count, or complicate training. Sparse attention, the application of sparse activation to the attention mechanism, often face similar challenges.
    
This paper introduces the Spark Transformer, a novel architecture that achieves high activation sparsity in both FFN and the attention mechanism while maintaining model quality, parameter count, and standard training procedures. Our method realizes sparsity via top-$k$ masking for explicit control over sparsity level. Crucially, we introduce *statistical top-k*, a hardware-accelerator-friendly, linear-time approximate algorithm that avoids costly sorting and mitigates significant training slowdown from standard top-k operators. Furthermore, Spark Transformer reallocates existing FFN parameters and attention key embeddings to form a low-cost predictor for identifying activated entries. This design not only mitigates quality loss from enforced sparsity, but also enhances wall-time benefit. Pretrained with the Gemma-2 recipe, Spark Transformer demonstrates competitive performance on standard benchmarks while exhibiting significant sparsity: only 8\% of FFN neurons are activated, and each token attends to a maximum of 256 tokens. This translates to a 2.5x reduction in FLOPs, leading to decoding wall-time speedups of up to 1.79x on CPU and 1.40xon GPU.

Spark Transformer: Reactivating Sparsity in Transformer FFN and Attention
Chong You'
Kan Wu
Zhipeng Jia
Lin Chen"
Srinadh Bhojanapalli
Jiaxian Guo
Utku Evci
Jan Wassenberg
Praneeth Netrapalli
Jeremiah J. Willcock
Suvinay Subramanian
Felix Chern
Alek Andreev
Shreya Pathak
Felix Yu
Prateek Jain
Henry M. Levy
David E. Culler
Sanjiv Kumar
Google
Abstract
The discovery the lazy neuron phenomenon the feedforward networks (FFN) parameters in trained Transformers are activated per token, has spurred significant interests activation sparsity for enhancing large model efficiency. While notable progress has been made in translating such sparsity to wall-time benefits across CPUs, GPUs, and TPUs, modern Transformers have moved away from the ReLU activation function erucial to this phenomenon. Existing efforts on re- re-introducing activation sparsity, e.g., by reverting to ReLU, applying top-k masking parameter count, complicate training. Sparse atlention, the application of sparse activation the attention mechanism, often face similar challenges.
54
where fewer than 10% of
sparse predictor, often degrade model quality, increase
This paper introduces the Spark Transformer, novel architecture that achieves high activation sparsity in both FFN and the attention mechanism while maintaining model quality, parameter count, and standard training procedures. Our method realizes sparsity via top-k masking for explicit control over sparsity level. Cru- cially, we introduce statistical top-k, a wwwaae--cceleraato-frieeeeee approximate algorithm that avoids costly sorting and mitigates significant training slowdown from standard top-k operators. Furthermore, Spark Transformer reallo- cates existing FFN parameters and attention key embeddings to form a low-cost predictor for identifying activated entries. This design not only mitigates quality loss from enforeed sparsity, but also enhances wall-time benefit. Pretrained with the Gemma-2 recipe, Spark Transformer demonstrates competitive performance standard benchmarks while exhibiting significant sparsity: only 8% of FFN neurons are activated, and each token attends to a maximum of 256 tokens. This translates 2.5x reduction in FLOPs, leading to decoding wall-time speedups of up to 1.79x on CPU and 1.40 on GPU.
linear-time
Introduction
The machine learning field has experienced rapid expansion in the scale of Transformer models 3 [7] [ significantly advancing the state-of-the-an in language understanding and generation. However, this pursuit of larger models is often constrained not by inherent limitations in model quality but by the soaring computational costs 76 68 associated with the number of parameters. This
"Equal contribution
New at xAI
Now at Anthropic
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(High)
Quality
(Low)
1.0
Gemma-2
sporstty (Saseline
LLow
72-
Gemmaz
Spark Tensfrrire
Activatton sparsity FFN
50
3.5
sparatyy FFN
25
(Relabie) FLORS
ReLUe coangge ..000
Effiliency
0.7
ReLU
Spark FFN (hop
(b) Downstream evaluation
Topk biz
Prefill speed (mshoken)
Decode speed (mahtken)
0.5
Spark Transformer mha mapppe
(1ight
Gemma-2 Spark Transformer
28
141
1.000
1.001
1.002
2. 10003
1.004
1.005
1.005
86
(Pelatiiel Loss
(a) FLOPs per token vs. quality (1/6 of full training)
(c) CPU inference wall time
Spark Transformer improves inference efficiency via activation sparsity in both FFN and attention, while being nearly quality neutral. (a) Comparison to prior work in terms of relative FLOPs per token 8k sequence length (y-axis) vs relative training loss (x-axis [ [m] We use standard Gemma-2 [Q as baseline, which has no activation sparsity. [+] Methods employing activation
Figure 1:
sparsity within the FFN layers only. Our Spark FFN achieves the most most favorable trade-off' compared to ReLU, ReLU², and Topk, which refer to standard Gemma-2 with activation function switched to ReLU (64), ReLU [92), and the composition of Topk and GELU, respectively. [A] Combining Spark FFN (with 8% activated parameters) with Spark Attention (with at most 256 attended tokens), our Spark Transformer achieves performance comparable to Gemma-2 while reducing FLOPs to 40%. (b) Evaluation standard downstream tasks confirms near-quality neutrality of Spark Transformer. (c) Prefill decode wall time demonstrate 1.86x/1. /xxp//.64xxx speedup resulting from FLOPs reduction. Results are obtained on a 4-Core CPUs for prompts of 4096 tokens. For prefill, the prompt is chunked into batches of 64 tokens, following a default setup of gegma. cpp (33.
challenge is compounded by the trend towards models capable of processing increasingly longer input sequences 722, where computational demands scale proportionally with input length.
Activation sparsity has emerged as prominent technique for mitigating the computational burdens associated with both large model sizes and long input sequences. For large models, activation sparsity reduces computational cost by activating only small fraction of the model's 's parameters for each input. This approach has attracted considerable attention due to the observed lazy-neuron phenomenon (54) where natural activation sparsity occurs in the feed-forward networks (FFNs) of traditional Transformer models like TS 71] and ViT [26). without explicit enforcement [91]. Subsequent work has successfully demonstrated wall-time benefits from activation sparsity multiple hardware
platforms, including CPU 90]. GPU
82 58. and TPU 89].
Despite the great success, fundamental challenge arises in the application of sparse activation to the latest and state-of of-e------hht-- models. That is, the inherent sparsity, which pivotal for obtaining efficiency, is largely absent due to the adoption of gated non-ReLU activation functions [19,, widely adopted in, e.g., Mistral [11, Gemma 30 LLAMA 277. This motivates the following question:
Can we re- re-introduce a high level of activation sparsity in recent Transformers,
without compromising their quality?
To answer this question, prior work has explored reverting to ReLU variants 64
or incorporating top-k thresholding [89 81 83.. However, (Challenge #1) these methods often lead to degradation in model quality (see Figure lap. Furthermore, (Challenge #2) top-k k thresholding based methods not only introduce non- dineeiffeeentaabiiity that compromises model quality, but also involve sorting to obtain maximally activated neurons, which is inefficient on ML accelerators, such as TPUs [43], with possibly 10x training slowdown (see Figure by. Finally, low-cost predictor is often introduced to identify activated parameters, critical for maximizing efficiency benefits [58] 90, 82 89. However, (Challenge #3) the inclusion of such predictor increases training pipeline complexity, incurs additional training costs and parameters, and exacerbates model quality loss.
92
A All models in Figure la were trained with standard Gemma-2 recipe for 1/6 of the full pretraining iterations.
All other results in this paper obtained for fully pretrained Spark Transformer.

---

Sparse attention, the application of activation sparsity to the attention mechanism, presents similar challenges. A common strategy top-k attention [35, which applies a top-k mask to the attention coefficients. This can be augmented with a low-cost predictor to further enhance efficiency [74 88 [50] Nevertheless, achieving both high sparsity and accurate prediction without resorting complex procedures, training slowdown from top-k, and compromising quality remains an open problem.
Contributions. This paper introduces the Spark Transformer, a novel architecture that achieves strong level of activation sparsity in both FFN and attention mechanism with minimal impact on quality, providing a positive answer to the question posed above.
Spark Transformer comprises Spark FFN and Spark Attention, both of which exploit the interpretation of FFNs and attention mechanisms key-value lookup tables (32 to provide unified framework for sparsity and a low-cost predictor (see Figure 2 Our low-cost predictor is constructed by repurposing subset of the dimensions of the query and key vectors to compute
importance score for each key-value pair (see Section (). This design addresses Challenge #3 by avoiding the introduction of extra parameters, enabling all model parameters be trained in a single stage.
Our key technical tool for introducing sparsity is a statistical top-k operator, which is applied to the predicted scores in Spark FFN and Spark Attention to select the activated keys. In particular, statistical top-k is a comaarco algorithm for approximate nearest neighbor search, which addresses the issue of high computation cost of standard top-k algorithms (i.e., Challenge #2). As explained in Section this is achieved by fitting a Gaussian distribution to the activation scores and estimating threshold that selects the top entries. Moreover, it enjoys continuously differentiability almost everywhere by the usage of a soft-thresholding operator, making the network more amenable to optimization
Finally, we demonstrate that Spark Transformer alleviates quality loss (i.e., Challenge #1) by perform- ing a full pretraining using the Gemma-2 recipe [30]. Comparison with prior work shows that Spark Transformer exhibits a more favorable trade-off between FLOPs reduction and quality measured by pretraining loss (Figure La). Further evaluation standard benchmarks confirms that Spark Transformer closely matches the performance of Gemma-2 (see Figure lb). despite exhibiting a high degree of sparsity: only 8% of FFN neurons are activated, and each token attends to at most 256 tokens. Leveraging this sparsity, we assess the model's inference efficiency on CPUs, demonstrating wall-time speedups of 1.86x for prefill and 1.64 for decode (see Figure lc). Furthermore, up to 1.4x wall-time speedup is obtained on NVIDIA T4 GPU (see Section 4 This enhanced efficiency broadens access to high-quality models for users with limited access to high-FLOP hardware, such as high-end GPUs and TPUs.
2
Spark Transformer
This section describes Spark FFN and Spark Attention, the two components of Spark Transformer.
2.1 Spark FFN
FFNs in a standard Transformer are two-layer multi-layer perceptrons that map an input token q I ddnat to an output
Standard K) V.v oooo((K 9) E Rdman
(1)
In above, {K, VJ C IRdeo are trainable model parameters, and o() is nonlinear activation function. We ignore the dependency on layer index to simplify the notations.
Each matrix multiplication in eq. ( has 2dmodel du FLOPs hence overall the computation cost 4ddmadel - dg. Previous work shows that when a is ReLU, the activation map a (TT q) is very sparse after model training. The sparsity can be used trivially to reduce the computation costs in the caleulation of its product with the second layer weight matrix V (54] reducing the overall FLOPPs count of FFN to 2dmodel (dar + k). where k < da is the number of nonzero entries in the activation. Note that the sparsity cannot be used to reduce the computation costs associated with K, which constitute half of the total FLOPS in FFN.
SStatistical top-k may also be applicable to Mfxture-of -xpeertss,, which face same non di--ifffreentiabiiity and inefficiency challenges in their routing mechanism. A recent paper [87] also used the soft- boft-thresholdng in MoE but still relies on costly sorting operator. Hence, the linear complexity of statistical top-k could address a key computational bottleneck, especially MoEs trend toward larger number of experts [88 36).
3

---

Input: q
Dense Matrual
Dense Mattmu
Input: q
Matmul
K
Statitical Top- y max(xx-(x)) 0)
K
100
Outout
Output
Dense Matmiu
Sparse Matrice
V
V
Figure 2: Architecture of Spark FFN and Spark Attention. (Left) Unified illustration of standard FFN (I.e., eq. (1) and standard Attention (i.e., eq. (6)) In the case of FFN, q E Rd is the input, K and V are the first and second layer weights, respectively, and o() is GELU. In the case of Attention, E R2 is the query, K and V are key and value matrices, respectively, and o() is softmax. (Right) Unified illustration Spark FFN (i.e., eq. (2)) and Spark Attention (i.e., eq. (7) In the case of Spark FFN. () is GELU and (² is identity. In the case of Spark Attention, di () is softmax and 072() is softplus. In both cases, Statistical- Top, (i.e.. eq. (IO)) is applied to introduce sparsity, which enables sparse matrix multiplication with K₂ and V that reduces FLOPs count.
In order to reduce FLOPs count in the first layer of FFN as well, we introduce Spark FFN as follows:
Spark-FFN(q q: K1. K2, V. krr) def V. (a (KT ( q(r::)
(2)
In above, K1 C RRd, K2 C (dunds ( x)x,s and V S Au xdr are trainable parameters, and the activation o() is taken to be GELU (37] following Gemma Top, is introduced for obtaining sparsity, with & being a hyper-parameter specifying the sparsity level. Concretely, Top;: preserves the largest k values in the activation vector, while setting other values O. In this paper, we use the efficient statistical top-k presented in Section which avoids sorting activation values. Finally, the input is split into q[:r] and q(r:]. which contain the first r and the rest of the dimensions, respectively, with It is introduced of the location of the nonzero entries, which allows us to obtain efficiency benefits in computing K. q(r: i] and the multiplication with V. This is discussed in detail below.
being Hlper-parametee
that the term KJ q(:r) serves low-rank predictor
FLOPs per Token. Direct implementation of the Spark-FFN without exploiting sparsity has the same number of FLOPs as the standard FFN in eq. I
2r du + (ddmoste) - r) dr + 2dhmodel du = 4dmmmtl X du,
(3)
where the three terms are from multiplication with K,, K2, and V, respectively. In Spark-FFN, one may first compute the term K, q(:r] as a low-rank predictor. After passing its output through Topg. which selects approximately the k most important entries, followed by the activation function a(, we obtain sparse output. Importantly, after obtaining the sparse output there no need to perform the full computation of the other two matrix multiplications eq. i.e. KJJ q[r: and the multiplication with V. Instead, one can perform sparse matrix multiplication with a drastically reduced FLOPs count:
(4) which is an increasing function of r. In other words, r controls the computation cost. We provide ablation study in Section C.4 to show that the best model quality obtained when dinad In this case, the total FLOP count of Spark FFN is approximately dmodel du + 3 dmodel k. which is a 4-times reduction from eq. 3 when l is small.
2r du + 2(dmmnte r) k + 2dmodal 1 = 2dn k) rr 4d k.
Relation to gated activation. Many of the most recent Transformers, including Gemma, use a variant of the standard FFN in eq. (I) where the activation function is replaced with gated one:
Gated-FFN (q; K1, K2, V) V [a(K]] q) の (Kg q))
In above, (K1, K2, V} C Ta X Note that when compared with the FFN eq. for quality studies, d' is usually taken to be 2/3 de to be iso -s--paraeeter count 777

---

Our Spark FFN in eq. (2) bears some resemblance to Gated FFN in that both have two linear maps in the first layer and one the second layer. The difference lies in that 1) Spark FFN adds a Top; to obtain sparsity, and 2) the input to the first layers of Spark FFN are obtained from splitting the dimensions of the input.
2.2 Spark Attention
In standard multi-head attention layer, an input @ E Rdasse is mapped to a query, key, and a value vector of dimension dam as (+) w(() E Rdm k(l) W E Rdm, v(+) w (i)
E Rdm
for each head Here, (Wl
(8) W (8) W() C Rdd ddmin are trainable weights.
Collecting all the key and value vectors in the context of x into and ( = [w] Una E Mun, attention conducts the following computation:
K() [ kin R x.
Standard- Aatiaaaddd-teetinnnggg K, V) def softmax (K") 4) TRddam
(6)
where we omit the dependency on for simplicity. The computation cost associated with eq. (6) is 4damm Tetx for each head. Finally, output from all heads are concatenated followed by a linear map to project to dmodel-
Note that eq. (6) has the same form as FFN in eq. ④Th except for the choice of nonlinearity. Hence, following similar strategy in obtaining Spark FFN, here we present Spark Attention as
Spark- Seark-Attention q: K, V, ,,, k) def V (a (Topp (K] g(:r)) @ 02 02 (22 alr:))
(7)
In above, K, RT and K₂ pddma-))xma contain the first r rows and the rest of the rows from K. respectively, and and 02 are nonlinear functions. Empirically, we find that taking T1 softmax and 02 softplus gives good results, Finally, Top, (o)) refers an operator that keeps the largest h values of the input while setting the rest to coo. This operator be implemented using statistical top-k, which we explain Section 3
FLOPs per Token. With a direct implementation the number of FLOPs in eq. 7 is given by 4damm Maxx which is the same as the FLOPs for eq. (6). However, by noting that the output of the softmax is expected to be sparse with approximately k nonzero entries, the computation costs associated with KJ q(: and in the multiplication with can be drastically reduced. In particular, we take dum then the FLOPs per token becomes
dmoode/aaxx + 3d min (Katt, Netx):
(8)
which is nearly 4x reduction when kam Tletx-
3 Statistical Top-k
This section introduces Statistical-Top Topp. efficient algorithm for implementing the Top; operators in Spark FFN (i.e. eq. (2)) and Spark Attention (i.e. eq. 77).
Recall that the soft soft-thresholding operator is defined for an arbitrary vector x Rd and a scalar threshold 0 E IR as
Goft-Threshol((z,.0 0) del
max(# - 0- 1, 0} RR
where 1 and O are d-dimensional vectors with all entries equal to and 0. respectively. The soft- thresholding operator shifts each entry of the left by @ and then thresholds the result zero
We define Statistical-Top as the following mapping from Rd to Rd: def Soft-Threshold(a.. (xx k)),
Statistical-Top..e (œ)
(10)
where
((x, k) det mean(#) + std(2r) Q(1 /'
k
(11)
In above, we define mcan(x) def E xn and std( def + compute the sample mean and standard deviation of the entries of the input æ. respectively. Q(-) is
Din (o: mcan()))² which
5

---

the quantile function (i.e., inverse of the cumulative distribution function) of the standard Gaussian distribution. In Spark Transformer, eq. (TO) is used as the Top;, operator eq. ( For the operator Topp (00) in eq. ( slight variant of eq. (10 is used where the entries below the threshold (ax, k) are set to oo instead of O.
cellistical-Top. operates by first computing threshold ((æ; k) such that approximately &o entries of æ exceed it, and then applying the soft-threshoding operator with this threshold to to obtain sparse output. We discuss these two components in the next two subsections.
3.1 Threshold Estimation
The threshold ((x, k) in eq. (10) is designed such that, if the entries of are drawn from a Gaussian distribution, approximately k out of the d entries will exceed this threshold. To understand this, let and a denote the mean and standard deviation the underlying Gaussian distribution. Its quantile function is given by μ + a Q(p) for E (0, 1). Consequently, due to the properties of quantile functions, we expect roughly p entries of x to exceed μ + a Q(1 p). practice, since μ and are unknown, they are replaced with the sample mean mcan(x) and the sample standard deviation std(2)) (z),, respectively.
The following theorem formalizes this argument.
Theorem 3.1. Let a E IRd be a vector with entries drawn i..dd from N(u, o²). For any 1 < K d-1, let ((x, k) be a scalar defined in eq. (10). Take any ö E 1) and assume d max 2, log 9). With a probability of at least 1 - 6, the number of entries of x that are greater than ((x, k),, i.e.
card ((i E [d] 24. > ((x,, k)}), satisfies
card (fi E [d] > ((x, b)}) kl
I 4
log
6
1 (++1 -2 log min { k
(2-1-3)
1: / 41)
(12)
d
Theorem
3.1 provides relative error bound between k and the true number of entries of æ that exceed k. This bound is maximized when k 1 k d 1. Consequently, the worst-case bound
log d & which vanishes as d increases. Notably, the error bound becomes O (V log
is O
when k = (d), demonstrating even faster convergence.
Computation cost. The computation of the threshold 6(x, k)) is highly efficient and resembles the operations used in LayerNorm layers, requiring only 2d FLOPs to compute the mean and standard deviation of the samples. This contrasts sharply with naive sorting-based based approach, which has Od d) complexity.
While the Gaussian quantile function Q(-) lacks closed-form solution, high-precision piecewise approximation algorithms with constant complexity are available standard software packages like SciPy (86), readily applicable to our needs.
3.2 Sparsification
Given the threshold ((x, k). a straightforward approach to obtain a sparse vector is to set all entries of æ below the threshold to zero, preserving the remaining values. This operator, sometimes referred to as hard thresholding [10], suffers from discontinuity, potentially hindering its suitability for gradient-descent-based training.
To address this, Statistical-Top. employs the noft-thresholding operator defined in eq. 9
g
This
operator first shrinks all entries of by the threshold ((æ, k) and then sets all entries below 0 to 0. Soft thresholding offers the advantages of being continuous and differentiable almost everywhere (except when entries of æ coincide with ((x,, k)).
For complete differentiability, one can utilize a smoothing function like the Huber loss 39 defined element-wiss on an input x as:
22
for |æ|| < 6, otherwise.
Huber(z; 8)
(13)
6 (|z] (
The continuous differentiability of the mapping a + Tlbeer(laaistiiall-oo....... is estab- lished below:

---

Percentage of nonzero
218
NUmE O nonz Ronzer
216
214
M
212
Step=0k
210
Step-OK
Step=120k
208
Scop-120k
Step-240K
Stop-240K
Step-360k
206
Step= 50UK
Step=480k
204
4top-430k
202
0
10
15
20
25
5
10
20
25
Layer index
Layer index
(a) Sparsity level in Spark FFN
(b) Sparsity level in Spark Attention
Figure 3: Sparsity in the intermediate activation of Spark FFN and Spark Attention across 26 layers at selected training steps. For FFN we report the percentage of nonzero entries out of da = 13824 entries. For Attention, we report the number of nonzero entries (i.e., attended tokens). Our hyyper-prramett choice to have 8% nonzeros Spark FFN and at most 256 nonzeros Spark Attention.
Theorem 3.2. For any ó O, the function 2. R defined as HuberrStaatssiica--Top. (x))
(14)
continuously differentiable.
Note that eq. Huber(I; 5)/6 find that using a non-zero ε does not improve model quality, and therefore we set & = O for simplicity.
14) converges to Blatistical-Top. as 5 - 0, which can be seen as a and Statistical-Top, (a:) always non-negative. In practice, however,
Finally, soft thresholding admits variational form (sec, c.g., (67
Soft-Thresholl(z, 0) arg min 0||=|| + =
(15)
=20
This formulation seeks vector z that minimizes both its squared & distance to the input ær and its & norm, with the threshold balancing these terms. Given the pparsity-promotion norm, soft thresholding effectively finds sparse approximation of the input æ. This variational form also connects Statistical- Top,, with other top-k algorithms in the literature; see Section D.3
nature of the &
4 Experiments
In this section, we present an experimental evaluation of Spark Transformer using the Gemma-2 2B model. Gemma-2 2B is decoder-only Transformer with 2 billion parameters, pretrained on 2 trillion tokens of primarily English text data (see (30 for details). To evaluate Spark Transformer, we train model by substituting the standard FFN and Attention in Gemma-2 2B with their Spark Transformer counterparts (Spark FFN and Spark Attention, respectively). This Spark Transformer model is trained using the same procedure and data as the Gemma-2 2B model.
Implementation details. Gemma-2 uses model dimension of dmodell = 2304. For FFN, Gemma-2 uses the Gated FFN in eq. (5) with d'r 9216. We replace it with Spark FFN in eq. (2) with dr = 13824 so that the parameter count keeps the same. In addition, we take k to be I 1106, which gives a sparsity level of 8%. and r 1024 = dmmodel/2 (due sharding constraints, r can only be multiple of 256). For Attention, Gemma-2 alternates between a global attention that have span of 8192 tokens, and local attention with a 4096 window size, both with dain 256. We replace both with Spark Attention in eq. (6) where for the latter we use the same 4096 window size. For hyper-parameters, we use = 256, i.e. each token attends to at most 256 tokens, and r 128 damn/2 2. Gemma-2 uses Rotary Position Embedding (84) which is applied to g and the columns of K eq. (6). For Spark Attention eq. we apply this position encoding to q(:r).
qir:] :]. the columns of K, and the columns of K2.
4.1 Quality
We evaluate Spark Transformer on a suite of benchmarks that used in the Gemma-2 -2 paper 330 and report the result in Figure [b] We observe that Spark Transformer matches the quality of Gemma-2 while having drastically reduced FLOP count per token.
The near-quality neutrality of Spark Transformer distinguishes it from related work on enforcing activation sparsity in FFN, which often lead to quality loss. To demonstrate this, pretrain variants
7

---

0
mask
O 8080 X
K 2
FFN output
Red: skipped memory access /compute
(a) Vector-Masked Matrix Multiplication
(b) Sparse Vector-Matrix Multiplication
Figure 4: Illustration of the matrix multiplication implementation using sparse activation. (a) Vector- Masked Matrix Multiplication takes a dense vector qir :]. a dense matrix K2 and a mask from statistical top-k on K g(: to compute u == (K2 q(r :)) :)))))))))) o mask. skips memory loading and compute associated with the masked columns. (b) Sparse Vector-Matie Multiplication takes a sparse activation vector u to compute weighted sum of rows in the dense matrix V. skips loading and computation of rows corresponding to O's in To optimize performance, implement Sparse Vector-Matrix Multiplication using tiling, which helps minimize cross-CPU core synchronization.
of Gemma-2 2 2B with 1) activation function switched to ReLU [92, 2) activation function switched to ReLU² [64]. and 3) Top-k thresholding applied before GELU. Due to high training cost, all models pretrained for 1/6 of the standard Gemma-2 training iterations. In Figure la report the training losses and FLOPs per token relative to those of the standard Gemma-2 models either suffer from large quality loss (i.e., ReLU and Topk), do not lead to sufficient FLOPSs reduction (i.e., ReLUU²). In contrast, our Spark FFN achieves less quality loss with more FLOPs reduction Finally, the combination of Spark FFN with Spark Attention introduces additional FLOPs reduction and quality benefits, enabling an almost neutral quality of Spark Transformer with large overall FLOPs reduction.
can be seen that these
4.2
Sparsity
To verify the effectiveness of statistical top-k, we report the level of sparsity measured in terms of percentage nonzeros in FFN and the number of nonzeros in Attention. At the beginning model training, we observe that statistical top-k produces close to 8% nonzeros in FFN (see Figure Ba). which aligns well with our hyper-parameter choice of using k/d# = 8% in Spark FFN. This is expected as the model parameters, particularly K in Spark FFN, are randomly initialized, hence the entries the activation maps are drawn from Gaussian distribution which in accordance with the assumption of statistical top-k. The Gaussian assumption is longer guaranteed after training, but we empirically observe to hold approximately (see Section C.I and statistical top-k k reliably produce sparsity level close to 8% until the end of training at 480k steps. Sparsity in attention is reported in Figure 36] which show that the number of attended tokens is below our hyper choice of 256 in Spark Attention throughout training. In particular, the numbers syper-parameter parameter much smaller because the results are from averaging over all tokens many which have a context length of less than 256. Finally, observe comparable levels of sparsity during evaluation (see Section
4.3 Inference Efficiency
We evaluate the efficiency benefits of Spark Transformer over standard Gemma-2 on both CPUs and GPUs. For CPU evaluation, we use gemma. cpp 33 the official C++ inference engine optimized for CPUs. For GPU evaluation, we use llama. cpp [3 supports running wide selection of [.LM models on GPUs. We modify both frameworks to support sparse matrix multiplication operators, which exploit sparsity in both FFN and Attention layers. Our implementation leverages vector SIMD operations
widely-used LLM inference engine which
[0) existing modern CPU, and customized CUDA kernel for GPU. In particular, our implementation reduces not only computational FLOPs but also memory bandwidth requirements, directly accelerating memory- bound scenarios; see Figure 4 for
"Notably, the only difference between Spark FFN and Topk, when both have 8% activated neurons, is in the introduction of the low-cost predictor. Hence, this predictor not only leads to FLOPs reduction but also (surprisingly?) improves model quality.

---

30
Cptinimation
Decoote speed rodemmo: F1
Spari: IFN optirn Spark: FiNeat A4to Optiiiiee
Soart PN Octinizze
FFN Oxtimized
Spart: Ffteedcm Opdinteod
Decoot apsed
3043
(100me))
wngth [n50w])
(a) 16 Core CPU VM.
(b) 4 Core CPU VM.
(c) NVIDIA T4 GPU.
Figure Spark Transformer decoding speedup from activation sparsity on various hardware platforms. We report decoding speed of Spark Transformer without handware optimization for sparse activation, with hardware optimization for sparsity Spark FFN only, and with hardware optimization for sparsity in both Spark FFN and Spark Attention. All experiments use a decode batch size of 1.
an illustration and Section B in Appendix for details. We show that Spark Transformer significantly improves the efficiency of transformer models, even in highly FLOP-constrained environments such as CPUs.
CPU results. Figure and and 5b report the decoding speed under varying prompt lengths on a 4-Core or 16-Core CPU. We see that with hardware optimization for sparse activation both Spark FFN and Spark Attention, a speedup that ranges from 1.35x to 1.79x can be achieved on 16-Core CPU depending on the prompt length. For short prompts (e.g., 256 tokens) and long prompts (e-g. 4096 tokens), the hardware optimization for Spark FFN and Spark Attention provide the most speedup, respectively.
Figure further highlights the efficiency of Spark Transformer in both prefill and decode phases. During the prefill, the prompt is usually chunked into batches since the process bounded by memory bandwidth. This may reduce the benefit of activation sparsity as different tokens in a chunk may activate different subsets of parameters (in FFN) and attend different subsets of tokens (in Attention). However, Figure Tc shows that Spark Transformer maintains strong performance with chunk size of tokens, following the default setup geema cpp A more detailed performance
analysis of batching/chunking is provided in Section C.33 In addition, Spark Transformer significantly
outperforms Gemma-2 during decoding (with batch size=1). 1).
GPU results. We also evaluate the efficiency gain from sparsity on low-profile GPUs. Figure 5c reports the decoding speed under varying prompt lengths on NVIDIA T4 GPU. Similar the CPU case, we see Spark Transformer achieves decode speedup ranging from 1.25x to 1.40x.
Step time relative without top-k
4.4 Training Efficiency
statistical top-k jox.lax.approx_mee cacpppppoo_xxxpp_X_eeaa
The introduction of top-k operator is expected to lead training slowdown due the extra computational operators. This slowdown can be prohibitively large if one uses the standard approximate top-l operator pro- vided in JAXX namely jax.laxx wat..xx.appooo_p....... approx_maxx k [14). Specifically, this JAX top-k operator is optimized to achieve TPU peak performance and has controllable recall target, which we vary on the x-axis in Figure 6 It can be observed that the JAX top-k leads to more than 10x slowdown even when operating on a small recall of 50%. In contrast, the slowdown from statisti- cal top-k is with a very small amount, demonstrating its efficiency. Finally, we do not provide the quality of models trained with JAX top-k since such models take very long time train.
an²
10'
0.50
0.55
0.60
0.65
0.70
0.75
0.90
0.05
0.50
Recall
Figure 6: Comparison of training slowdown from using our statistical top-k vs the stan- dard top-k (i.e.,
jaax (ic...jax.. lax. appraxppppaa...ee... _aaooo____________
[44 relative to not using any top-k.
Discussions
This paper introduces the Spark Transformer architecture to reduce the FL.OPs in both the FFN and Attention of Transformer models. Because the FFN and Attention components dominate the

---

computational cost in large Transformers with long contexts, the overall FLOP count for decoding token also drastically reduced (see TableC..T, leading notable computation efficiency benefits on appropriate hardwares. Specifically, the benefit is obtained by selectively activating only part of the model parameters and limiting the attended context for each input. This principle of sparse activation finds compelling parallel in neuroscience, where studies reveal sparse activity patterns in the brain as a key factor in its remarkable efficiency 7 48). While hardware limitations currently hinder the full exploitation of sparse activation, our demonstration of practical wall-time reduction of
Spark Transformer on CPU and GPU, together with prior evidence for related techniques [58] [2]
including TPU [89], highlight its potential. We hope that this work opens avenues for research into alternative hardware better suited for sparse computations, circumventing the hardware lottery (3
and potentially leading to greater efficiency gains in the future. Additional discussions related work for activation sparsity, including Mixture-of-Expert are provided in Section [
This work also opens promising avenues for future research on efficient inference of large language models, particularly in combining Spark Transformer's ' architectural efficiency with other leading optimization techniques. We briefly discuss two key synergies here.
Synergy with Speculative Decoding. Spark Transformer is highly complementary speculative decoding, As a target model, its faster inference directly accelerates the primary verification bot- tleneck. As a draft model, its near-quality neutrality and high speed make an ideal candidate for generating high-quality drafts, potentially leading to higher token acceptance rates and greater overall speedups. A full discussion is provided in Section D.5
Synergy with Quantization. We also hypothesize a strong synergy with quantization. The benefits expected be multiplicative, as Spark Transformer reduces the number of operations while quantization reduces their cost. More importantly, unlike standard pruning which preserves high- magnitude outliers, our statistical top-k operator uses soft-thresholding. This shrinks the dynamic range of activation distribution, which may reduce sensitivity to activation quantization. A detailed analysis of this mechanism is available in Appendix section D.6
Acknowledgments and Disclosure of Funding
The authors acknowledge helpful discussion with Zonglin Li (Anthropic), Daliang Li (Anthropic), Amir Yazdanbakhsh (Google), Zeid Samoail (Google), Abhishek Kumar (Amazon), Vlad Feinberg (Google), Veeru Sadhanala (Google) at various stages of this project.
References
[1] Bo Adler, Niket Agarwal, Ashwath Aithal, Dong H Anh, Pallab Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan Cohen, et al. Nemotron-4 340b technical report. arXiv preprint arXiv:2406.1 11.4061.7044, 1704, 2024.
[2] Joshua Ainslie, Tao Lei, Michiel de Jong. Santiago Ontanon, Siddhartha Brahma, Yury Zemlyan- skiy, David Uthus, Mandy Guo, James Lee-Thorp. Yi Tay, et al. Colt5: Faster long-range transformers with conditional computation. In Conference Empirical Methods in Natural Language Processing, 2023.
[3] Ebtesam Almazrouei, Hamza Alobeidli, Abdulaziz Alshamsi, Alessandro Cappelli, Ruxan- dra Cojocaru, Mérouane Debbah, Étienne Goffinet, Daniel Hesslow, Julien Launay, Quentin Malartic, et al. The falcon series of open language models. arXiv preprint 1rXiv:2311.16867 2023.
[4] Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, et al. Palm 2 technical report. arXiv preprini arXiv:2305. 10403, 2023.
[5] David Attwell and Simon B Laughlin. An energy budget for signaling in the grey matter of the brain. Journal of Cerebral Blood Flow Metabolism, 21(10):11331144 2001.
[6] Pranjal Awasthi, Nishanth Dikkala, Pritish Kamath, and Raghu Meka. Learning neural networks with sparse activations. In The Thirty Seventh Annual Conference on Learning Theory, pages 406-425. PMLR, 2024.
10

---

[7] Alison L Barth and James FA Poulet. Experimental evidence for sparse firing in the neocortex. Trends in neurosciences, 35(6):345-355, 2012.
[8] Amir Beck and Mare Teboulle. A fast iterative Soinkaage-tresholling algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1):183-202, 2009.
[9] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- document transformer. arXiv preprint urXiv:2004.05/55 2020.
[10] Thomas Blumensath and Mike E Davies. Iterative thresholding for sparse approximations. Journal of Fourier analysis and Applications, 14:629-654, 2008.
[11] Shijie Cao, Lingxiao Ma, Wencong Xiao, Chen Zhang, Yunxin Liu, Lintao Zhang, Lanshun Nie, and Zhi Yang. Seernet: Predicting convolutional neural network feature-map sparsity through low-bit quantization. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 11216-11225,200 2019.
[12] Beidi Chen, Tharun Medini, James Farwell, Charlic Tai, Anshumali Shrivastava, et al. Slide: In defense smart algorithms over hardware acceleration for large-scale deep learning systems. Proceedings of Machine Learning and Systems, 2:291-300 2020.
[13] Charlic Chen, Scbastian Borgcaud, Geoffrey Irving. lean-Baptist Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.013//8 2023.
[14] Felix Chern, Blake Hechtman, Andy Davis, Ruiqi Guo, David Majnemer, and Sanjiv Kumar. Tpu-knn: K nearest neighbor search at peak flop/s. Advances in Neural Information Processing Systems, 95:15489-15500 2022.
[15] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[16] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. Approximating two-layer feedforward networks for efficient transformers. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 674-692, 2023.
[17] Xavier Suau Cuadros, Luca Zappella, and Nicholas Apostoloff. Self-conditioning pre-trained language models, In International Conference on Machine Learning, pages 4455-4473. PMLR, 2022.
[18] Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, al. Deepseekmoe: Towards ultimate expert specialization in mixture-of- of-experts language models. arXiv preprint 2xXiv:2401.660 2024.
[19] Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pages 933-941 PMLR, 2017.
[20] Andrew Davis and Itamar Arel. Low-rank approximations for conditional feedforward compu- tation in deep neural networks. arXiv preprint arXiv:1312.446) 2013.
[21] Yichuan Deng, Zhao Song, and Chiwun Yang. Attention is naturally sparse with gaussian distributed input. arXiv preprint arXiv:2404. 026900 2024.
[22] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Gpt3, int8 (): 8-bit matrix multiplication for transformers at scale. Advances in neural information processing systems, 55:30318-233 1--3332, 2022.
[23] Tim Dettmers, Ruslan Svirschevski, Vage Egiazarian, Denis Kuznedelev, Elias Frantar, Saleh Ashkboos, Alexander Borzunov, Torsten Hoefler, and Dan Alistarh. Spqr: A sparse-quantized representation for near-lossless Imm weight compression. arXiv preprint arXiv:2306.03000 2023.
11

---

[24] Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307.02480) 2023.
[25] Harry Dong. Beidi Chen, and Yuejie Chi. Towards structured sparsity in transformers for efficient inference. In Workshop on Efficient Systems for Foundation Models@ ICML2023, 2023.
[26] Alexey Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.11929 2020.
[27] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 2024.
(28) Aryeh Dvoretzky, Jack Kiefer, and Jacob Wolfowitz. Asymptotic minimax character of the sam- ple distribution function and of the classical multinomial estimator. The Anmals of Mathematical Statistics, pages 642-669, 1956.
[29] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-3,, 2022.
[30] Gemma Team. Gemma 2: Improving open language models at a practical size. arXiv preprint arXiv:2488.00/18, 2024.
[31] Georgi Gerganov. llama.cpp.
cpp 2025.
[32] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 5484-5495, 2021.
[33] Google Gemma.cpp. Google/gemma.cpp: lightweight, standalone c++ inference engine for google's gemma models.
gemma cpp. 2025.
[34] Jinyang Guo, Jianyu Wu, Zining Wang, Jiaheng Liu, Ge Yang. Yifu Ding, Ruihao Gong, Haotong Qin, and Xianglong Liu. Compressing large language models by joint sparsification and quantization. In International Conference on Machine Learning, 2024.
[35] Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory-efficient transformers via top-k attention. arXiv preprint arXiv:2106.06889 2021.
[36] Xu Owen He. Mixture of million experts. arXiv preprint arXiv:2407.04153, 2024.
[37] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.084155 2016.
[38] Sara Hooker. The hardware lottery. Communications of the ACM, 64(12):88 65, 2021.
[39] Peter J Huber. Robust estimation of location parameter. In Breakthroughs in statistics: Methodology and distribution, pages 492-518. Springer, 1992.
[40] Sebastian Jaszezur, Aakanksha Chowdhery, Afroz Mohiuddin, Lukasz Kaiser, Wojciech Gajew- ski, Henryk Michalewski, and Jonni Kanerva. Sparse is enough in scaling transformers. Advances in Neural Information Processing Systems, 84:9895-9907,, 2021.
[41] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825. 2023.
[42] Huiqiang Jiang, YUCHENG LI, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context IIms via dynamic sparse attention. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
12

---

[43] Norm Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing. Brian Towles, Clifford Young, Xiang Zhou, Zongwei Zhou, and David A Patterson. Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. In Proceedings of the 50th Annual International Symposium on Computer Architecture, ISCA '23, New York, NY, USA, 2023. Association for Computing Machinery.
[44] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08300 2020.
[45] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.
[46] Guillaume Lample, Alexandre Sablayrolles, Mare' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. Advances in Neural Information Processing Systems, 32, 2019.
[47] Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selection. Annals of statistics, pages 1302-1338, 2000.
[48] J Quinn Lee, Matt Nielsen, Rebecca McHugh, Erik Morgan, Nancy S Hong, Robert J Sutherland, and Robert McDonald. Sparsity of population activity in the hippocampus is task-invariant across the trisynaptic circuit and dorsoventral axis. Behavioural Brain Research, 423: 13790, 2024.
[49] Je- Yong Lee, Donghyun Lee, Genghan Zhang, Mo Tiwari, and Azalia Mirhoseini. Cats: Contextually-aware thresholding for sparsity in large language models. arXiv preprint arXiv:2404.08763, 2024.
[50] Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. [InfiniGen]: Efficient generative inference of large language models with dynamic (KV) cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 172-172, 2024.
[51] Tao Lei, Junwen Bai, Siddhartha Brahma, Joshua Ainslie, Kenton Lee, Yanqi Zhou, Nan Du, Vincent Zhao, Yuexin Wu, Bo Li, et al. Conditional adapters: Parameter-efficinn transfer learning with fast inference. Advances in Neural Information Processing Systems, 36:8152- 8172, 2023.
[52] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with condi- tional computation and automatic sharding. In International Conference Learning Represen- tations, 2020.
[53] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference Machine Learning, pages 19274-19286. 2023.
[54] Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J Reddi, Ke Yc, Felix Chern, Felix Yu, Ruiqi Guo, et al. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. In The Eleventh International Conference on Learning Representations, 2022.
[55] Yujun Lin, Song Han, Huizi Mao, Yu Wang, and Bill Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training. International Conference on Learning Representations, 2018.
[56] James Liu, Pragaash Ponnusamy, Tianle Cai, Han Guo, Yoon Kim, and Ben Athiwaratkun. Training-free activation sparsity in large language models. arXiv preprint arXiv:2408. 14690, 2024.
13

---

[57] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrilllidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of impor- tance hypothesis for I1m kv cache compression at test time. Advances in Neural Information Processing Systems, 36, 2024.
[58] Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivas- tava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient Ilms at inference time. In International Conference on Machine Learning, pages 22137-22176. PMLR, 2023.
[59] Chao Lou, Zixia Jia, Zilong Zheng, and Kewei Tu. Sparser is faster and less is more: Efficient sparse attention for long-rang, transformers. arXiv preprint arXiv:2406.1677 2024.
[60] Jinqi Luo, Tianjiao Ding, Kwan Ho Ryan Chan, Darshan Thaker, Aditya Chattopadhyay, Chris Callison-Burch, and René Vidal. Pace: Parsimonious concept engineering for large language 2024.
models. arXiv preprini arXiv:2406.0453
[61] Yuqi Luo, Chenyang Song, Xu Han, Yingfa Chen, Chaojun Xiao, Zhiyuan Liu, and Maosong Sun. Sparsing law: Towards large language models with greater activation sparsity. arXiv preprint
2024.
[62] Ahmed M Abdelmoniem, Ahmed Elzanaty, Mohamed- Slim Alouini, and Marco Canini. An efficient statistical-basee based gradient compression technique for distributed training systems. Pro- ceedings of Machine Learning and Systems, 3:297-322, 2021.
[63] Pascal Massart. The tight constant in the dvoretzty-Kieffer-ollvvvii inequality. The annals of Probability, pages 1269-1283, 1990.
[64] Iman Mirzadeh, Keivan Alizadeh, Sachin Mehta, Carlo C Del Mundo, Oncel Tuzel, Golnoosh Samei, Mohammad Rastegari, and Mehrdad Farajtabar. Relu strikes back: Exploiting activation sparsity in large language models. arXiv preprint arXiv:2310.04564, 2023.
[65] Ramchandran Muthukumar and Jeremias Sulam. Sparsity- aware generalization theory for deep neural networks. In Annual Conference on Learning Theory, pages 5311-5342. PMLR, 2023.
[66] Mattco Pagliardini, Daniele Paliotta, Martin Jaggi, and François Fleuret. Faster causal attention over large sequences through sparse flash attention. arXiv preprint arXiv:2306.01 160, 2023.
[67] Neal Parikh. Stephen Boyd, et al. Proximal algorithms. Foundations and trends in Optimiza- tion, 13)::277-339, 2014.
[68] David Patterson, Joseph Gonzalez, Quoc Le, Chen Liang, Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean. Carbon emissions and large neural network training. arXiv preprini arXiv: 104. 10350, 2021.
[69] Ze Peng, Lei Qi, Yinghuan Shi, and Yang Gao. Theoretical explanation of activation sparsity through flat minima and adversarial robustness. arXiv preprint arXiv:2339.03004 2023.
[70] Zihan Qiu, Zeyu Huang, and Jie Fu. Emergent mixture-of-experts Can dense pre-trained transformers benefit from emergent modular structures? arXiv preprint arXiv:23 10908, 2023.
[71] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
[72] Machel Reid, Nikolay Savinoy, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean- baptiste Alayrac. Radu Soricut, Angeliki Lazaridou. Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, arXiv preprint 0535v 2024.
[73] Liliang Ren, Yang Liu, Shuohang Wang. Yichong Xu, Chenguang Zhu, and Cheng Xiang Zhai. Sparse modular activation for efficient sequence modeling. Advances Neural Information Processing Systems, 36:19799-19822. 2023.
14

---

[74] Luka Ribar, Ivan Chelombiev, Luke Mudlass-Gaalee Charlie Blake, Carlo Luschi, and Douglas Ofr. Sparq attention: Bandwidth-efficien Ilm inference. arXiv preprint arXiv:2312.04985. 2023.
[75] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
[76] Or Sharir, Barak Peleg, and Yoav Shoham. The cost of training nlp models: A concise overview. arXiv preprint arXi::2004.08900, 2020.
[77] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
[78] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gas gated mixture-of layer. arXiv preprint arXiv: 701.06538, 2017.
[79] Shaohuai Shi, Xiaowen Chu, Ka Chun Cheung, and Simon See. Understanding top-k sparsifica- tion in distributed deep learning. arXiv preprint arXiv: 2019.
[80] SIMD Wikipedia. Single instruction, multiple data - wikipedia - en. www.iidiiiii.org... https
//en. wen.wikipedia. a. org
2024.
[81] Chenyang Song, Xu Han, Zhengyan Zhang, Shengding Hu, Xiyu Shi, Kuai Li, Chen Chen, Zhiyuan Liu, Guangli Li. Tao Yang, et al. Prosparse: Introducing and enhancing intrinsic activation sparsity within large language models. arXiv preprint 241v:2402.13516 2024.
[82] Yixin Song, Zeyu Mi, Haotong Xie, and Haibo Chen. Powerinfer: Fast large language model serving with a consumer-grade gpu. arXiv preprint arXiv:2312.1244 2023.
[83] Yixin Song, Haotong Xie, Zhengyan Zhang, Bo Wen, Li Ma, Zeyu Mi, and Haibo Chen. Turbo sparse: Achieving Ilm sota performance with minimal activated parameters. arXiv preprint arXiv:2406.05955, 2024.
[84] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:12706) 2024.
[85] Filip Szatkowski, Bartosz Wojcik, Mikotaj Piórczyńskii, and Simone Scardapane. Exploiting activation sparsity with dense dynamic-k mixture-of-experts conversion. In Workshop on Efficient Systems for Foundation Models @/@ ICML2024, 2024.
[86] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weekesser, Jonathan Bright, Stéfan J. der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E, A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SeiPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272, 2020.
[87] Ziteng Wang, Jun Zhu, and Jianfci Chen. Remoe: Fully diffecrentiable mixture-of-experts with relu routing. 2025.
[88] Shuo Yang, Ying Sheng, Joseph E Gonzalez, Ion Stoica, and Lianmin Zheng. Post-training sparse attention with double sparsity. arXiv preprint arXiv:2408.07 07092, 2024.
[89] Varun Yerram, Chong You, Srinadh Bhojanapalli.. Sanjiv Kumar, Prateck Jain, Praneeth Netra- palli, et al. Hire: High recall approximate top-k estimation for efficient IIm inference. arXiv preprins arXiv:2402.00336) 2024.
[90] Zhanpeng Zeng. Michael Davies, Pranav Pulijala, Karthikeyan Sankaralingam, and Vikas Singh. I Lookupffn: making transformers compute-lite for cpu inference. International Conference on Machine Learning, pages 40707-40718, PMLR, 2023.
15

---

[91] Zhengyan Zhang. Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. Moefication: Transformer feed-forward layers are mixtures of experts. In Findings of the Association for Computational Linguistics: ACI. 2022, pages 877-890, 2022.
[92] Zhengyan Zhang, Yixin Song, Guanghui Yu, Xu Han, Yankai Lin, Chaojun Xiao, Chenyang Song, Zhiyuan Liu, Zeyu Mi, and Maosong Sun. Relu² wins: Discovering efficient activation functions for sparse 1lms. arXiv preprint (.XXi::242.033855 2024.
[93] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661-34710, 2023.
[94] Zhenyu (Allen) Zhang, Zechun Liu, Yuandong Tian, Harshit Khaitan, Zhangyang Wang, and Steven Li. R- R-sparse: Rank-awee aware activation sparsity for efficient Ilm inference. International Conference on Learning Representations, 2025.
[95] Haizhong Zheng, Xiaoyan Bai, Beidi Chen, Fan Lai, and Atul Prakash. Learn to be efficient: Build structured sparsity in large language models. arXiv preprint arXiv:2402.06116, 2024.
16

---

This Appendix is organized as follows. In Section A we provide proofs to theoretical results in the main paper. Section B contains implementation details of sparse matrix multiplication on CPU. Dection[[rroviceee additional discussions on comparing with related work on activation sparsity in FFN (in Section D.2) and on statistical top-k (in Section D.4)
details and additional results on experiments. Finally, Section
[ presents
Appendix A Proofs
A.1 Proof to Theorem 3.1
Proof. In this proof we write def mean(x) and s def st) for brevity.
We first establish the concentration bounds that the empirical mean and standard deviation, i.e., E and are close the mean and true standard deviation, i.e., μ and of the underlying Gaussian, respectively. Recall from the definition of the chi-squared distribution that (d - 1) ~ x (d 1).
Using the Laurent-Massart bound on the tail probability of the chi-squared distribution [47 Corollary of Lemma 1J. we have
Pr Pr
(((d
1) 02
8² (d (d-1) (d 1)t ( 2t s 2e-
for every > O. We set
log g. Then, with probability of at least 1 - 8/3, we have
x²
6 8
(d - 1)
1 < 2 (d - 1) log
+ 2 2log
02
8
which implies
s²
log
log 6
log 6
6 log
log
1
<
2
+ 2
< 4
8
4
8
8
x2
1
d
d
d
where the last inequality uses the assumption that > max (2, log g. By rearranging the terms, we get
log
log 6 &
log 3 S
log
-
8
S a
max
1 - 8/
0
< S s 0
1 1+ 8
1 8)
d
d
d
which simplifies to
log
S
80
(A.1)
eq. A. provides a concentration bound for s. We now proceed to deriving a bound for μ. Towards that, notice that 8-μ c/va ~ (O, 1). By using the Mill's inequality that upper bounds the tail probability
of a standard normal distribution (i.e., if Z ~ N(00 1) and > 0, then Pr([Z] > t) S /2/3 have
we
T
1
3 2lg 8
5/3
Pr
S 8/3.
g/vd
2 log 3
Therefore, with probability at least 1
5/3, we have
- H
3
2 log
0/ v/
z
which yields
log
1 - p| S a
(A.2)
d
17

---

Combining eq. A.1
and eq. (.2))
with probability at least 1 28/3, we have
& )
0 k)
(1r + oQ(1
(A.3)
+ 8 al
Q0
1
d'
(A.4)
log 20
log
Sa
80
Q(1
d
(A.5)
d
We define the empirical cumulative distribution function (ECDF) of 22, 3d
Pa(x) = Σ
Σ
(aigu)-
iE(d]
Then, the number of the entries of x that greater than 8(x, k) may be written as Σ (m)00()))))) = d( - Fu(0(zzA)))
card (fi E [d] r; > k)))
iE|d]
Let F denote the cumulative distribution function (CDF) of NN,, 0²). By the Dvoretzky-Kiefer. we have
Wolfowitz inequality (28
63
Pr sup Fa(u) - F(u)) >t) S 2e -2dt?
LuER
Taking V log & and 0(xx k),, we obtain
Pr [Fa(0(z;))) F(0(z, A))
6
5
log
(A.6)
2d
3
Applying the union bound on eq. probability least 1 - 5:
A.3 and eq. A.6). we obtain that the following holds with
(2(0(2,A)) -1 d
(A.7)
Ea(O(z,, k)) - F(H + oQ(1 //
(A.8)
Ea(O(z, k)) - F(((x, k)){ F(((22 - +(#+F(++++ oQ(1= /0
(A.9)
0(x,1 k) - (m+ QQ11 - -
log
(A.10)
2d
6
V2To
2 log à
log 6
R Q(1
log
+ +8
(A.11)
2d
2m
d
541
log 6
1+
Q(1 QQ- -
(A.12)
In the above expression, the first equality stems directly from the definitions of F(-) and Q(-), which gives
F(O(x, k)) F(u oQ -
(( b/ ))
k d
where & denotes the CDF of the standard normal distribution.
To simplify eq. A.7 we consider two cases:
If < d/2, by Mill's inequality, we have
(22 F1//2
- (2221g 2log 11/2
k/d
k
1 - u1 22og
à
2log x
d 2 log
/2 log
I
18

---

where the last inequality is because 2 log d
2 1. Therefore
k
S Φ((
log
d
k
which gives
k
k
Q(1 - d) 1
2
log
= V
-2
k
d'
If d/2, we have
d
d -
Q(1 d) =
Q(1
S
- log
d
d
Combining the two cases, we get
Q(1 2
k
{ k
k
log min
d
d)
Plugging this into eq. (A.7) we obtain
Pa(O(z,))) (1 ") s /
log
1 5)
1 log min
d
Recall card (Li [d] xt 0e,,, k))) d(1 fa)) We conclude that with probability
at least 1 - 6, we have
|card ((i E [d] x4 > 0(x, k))) - kl s 24/ 4 ( 1 2 log min
k
d
d
A.2 Proof to Theorem 3.2
Proof.
The Huber statistical top-k in eq. (14) may be written as Huber (ubber(taatistical-Top. (x): 5)/5
SaoberrSSft-Trrsholdd.e A(x, k))))5,
(A.13)
where ((x, k) is defined in eq. (IOp. This function is the (multivariate) composition of two functions, namely, ((x, k) and Huber(Soft-T) Thrrstt-Trrshol 0)). In particular, the former is continuously differentiable (i.e., C²) in æ, since it is simply a linear combination of sample mean and sample standard deviation both of which are C'l functions. To establish the theorem, we only need to show that Huber Huber(Soft- Sooft-Thresholiii 8)) is also C function in (x, 8).
By definition, MuuberSSoft-Threshii
0)) is defined entry-wise on æ as
dx 80 4
if > 0 + 6;
Huber(Sofi- 0)) = [
0)2
if - x < 0 + 8;
(A.14)
0,
From here it is easy check that Huber(Soft-Thrrehhod 0)) is continuous in (xx 0). Its gradient with respect to (xc, 0) is given by
(6, -5),
if 0 6; if 1 x 0 + 8;
ShaberrSoft-Thrreholoe 0))
(I 0,0 - x)
(A.15)
Axx 0)
(0,0),
if < 0,
which is also continuous. This concludes the proof.
19

---

Appendix B Implementation Details on Sparse Matrix Multiplications
We describe how we implement sparse matrix multiplications for Spark FFN and Attention in gemma. cpp for CPU and llama. cpp for GPU. We start by focusing on batch size of one for decoding before expanding our discussion to larger batch sizes and prefill.
With batch size of 1, both Spark FEN and Spark Attention utilize two types of sparse vector- matrix multiplication: vector-masked matrix multiplication and sparse vector-matrix multiplication (Figure 4.. Given a vector and matrix w, vector-masked matrix multiplication multiplies q with the non-masked columns of based masking vector m. Masked columns yield zero output. Sparse vector-man matoix multiplication, on the other hand, involves vector that contains many zeroes being multiplied by dense matrix.
In Spark FFN, we perform vector-masked matrix multiplication for K. q/r (Figure 4a The masking vector is generated from the output Statistical- Tops (K, qq::r]): Based on the mask, Spark FFN skips loading the masked columns of w from DRAM (in CPU setup) or HBM (in GPU setup) and the associated computations. On CPU, Spark FFN utilizes SIMD operations (as in the original Gemma implementation). To further enhance performance, Spark FFN utilizes software CPU prefetching (builtin feaiitinnppeeeeeee frefetch) to overlap loading from DRAM to the CPU cache with computations. On GPU, Spark FFN utilizes customized CUDA kernel.
The same masking vector also identifies the zero entries in the intermediate vector that multiplied by matrix V (Figure 4b). For this sparse vector- metor- multiplication, we store the matrix in row format. Each CPU thread (or GPU warp) processes tile of the matrix while skipping the loading and computation of the masked rows. Prefetching and SIMD operations are applied similarly CPU in this context.
Spark Attention utilizes these two types of sparse matrix multiplication operators to accelerate qkv computations for each head.
When extending to decoding with batch sizes greater than one or prefill, we continue to use individual masks skip computations while using union of masks from each vector within the batch to create unified masks for memory loading. With larger batches, Spark transformer is expected to save less memory loading (vs. original Gemma), unless there is significant overlap in top-k positions within the same batch. Nonetheless, the Spark transformer consistently reduces FLOP by skipping computations based individual masks within the batch.
20

---

Percentage of nonzero
10
200
Number of nonzero
200
204
200
200
10 Layer indexx
Layer index
(a) Spark FFN sparsity
(b) Spark Attention sparsity
Figure C.I: Sparsity in the intermediate activation of Spark FFN and Spark Attention during evaluation (see Figure 3 for results during training). For FFN, we use a simple prompt "test" and report the percentage of nonzero entries in generating the 5th, 10th, and 15th token. For Attention, we report the nuber of nonzero entries at the 512th, 1024th, and 2048th token during prefill.
Appendix C Additional Experimental Results and Details
C.1 Distribution of Inputs Statistical Top-k
The underlying assumption for statistical top-k is that the activation vector upon which it is applied to, namely, the pre-GELU activation Spark FFN and the pre-softmax activation in Spark Attention, can be modeled as being drawn from i..d. Gaussian distribution. Here provide empirical evaluation on the distribution of these activation vectors for Spark Transformer. Results for Spark FFN and Spark Attention are provided in Figure C.4 and Figure C.5 that the distribution holds close proximity to a Gaussian, hence justifying the use of statistical top-k.
respectively. The results show
C.2 Sparsity Level During Evaluation
Complementing Figure 3 which reports sparsity level during pretraining, here we report the sparsity level during evaluation to confirm that statistical top-k produces the same level of sparsity during test time. The results are presented in Figure C.I for some arbitrarily selected tokens. For Attention, in particular, we select tokens at the positions 512, 1024, and 2048 which are all above our choice of A 256 for Spark Attention.
C.3 Batching Analysis
PdBl
Figure C.2: Spark Transformer vs. Gemma-2 Prefill Token/Sec with Varying Chunk Size. We use prompt length of 4096 tokens on a 16 core CPU VM.
Figure C.2 provides the performance comparison between Spark Gemma 2 and Gemma 2, measured in prefill throughput (tokens/sec) across varying chunk sizes. We use 4096-token prompt on a 16-core CPU VM. A similar trend is expected during the decoding phase with varying batch sizes.
Our analysis shows that Spark Transformer provides the highest speedup batch size 1, and again at large batch sizes (c.g. ~8), where the compute FLOP becomes the primary bottleneck.
21

---

For Gemma-2, 2, as seen in the figure, increasing batch/chunk size leads to a significant improvement in prefill throughput until the batch size reaches 8. This improvement occurs because batching reduces memory access by reusing weights across multiple tokens in the CPU cache. Once the computation becomes the bottleneck (i.e. batch = 8), further batching provides diminishing returns.
In contrast, Spark Transformer behaves differently. When the batch size increases from 1 to 2, we observe minimal throughput change. This is due to the lack of overlap in top-k positions between the tokens, resulting from the high sparsity. However, as the batch size increases beyond 4, Spark Transformer starts benefiting from weight reuse, similar to Gemma-2 2. Spark Transformer continues to show improvements in throughput until the batch size reaches approximately 64, where it eventually becomes FLOP- much later than Gemma-2 due to the reduced FLOP requirements.
Overall, Spark Gemma demonstrates the most significant gains in two scenarios: when the batch size is 1, a common setting for desktop mobile devices decoding, and when the batch size is large enough that FLOP becomes the dominant bottleneck.
C.4 Effect of and k in Spark FFN.
Spark FFN comes with two hyper-parameter namely which controls the rank hence FLOP count of the low-cost predictor, and & which controls sparsity of activation hence the FLOP count. In Figure C.3 we provide an ablation study on the effect of these two hyper-parameters by reporting the training loss curves in the first 25,000 training steps (which is around 5% of full training). From Figure C.3al the best choice of r is 1024 which is nearly half of dmwstel = 2304 (due to model sharding constraint, r' cannot be taken to be exactly a half of dmmodel). From Figure C.3b we see that the model quality is insensitive to choices of i that gives (5%, 10%) sparsity, but there is quality loss if we go sparser, e.g. 3% nonzeros.
2.05
2.05
de
2.00
2.00
dem
dy 10%
200
202
Sceps
Sceps
(a) Efffect r.
(b) Effect of k.
Figure C.3: Effect of hyper-parameters r and & in Spark FFN on training loss. A Gaussian filter of 200 is applied to smooth the loss curves. Models are trained with 1/20 of standard Gemma-2 training iterations.
C.5 Additional Ablation Studies
In this section, we provide ablation studies for understand- ing the effect of the individual components of Spark Trans- former. Towards that, we plot the training loss curves for Gemma-2 and Spark Transformer, see Figure C.6. Here, we restrict the first 80k training steps out of the 500k to- tal steps since it is costly to fully train all ablation models, and that 80k steps is sufficient for seeing the trend. We can see that Spark Transformer slightly lags behind Gemma- 2. However, as demonstrated in Figure [bb, that small difference in training loss does not lead to substantial difference in evaluation quality.
2.02
ne Cemmrad Adaen
3:43
Gemmma-- 103-
Loss
3.44
1.95
1:43
40000
Steps
Figure C.6: Ablation study in terms of training loss in the first 80k training steps (out of 500k total steps).
In our ablation studies below, we add a single component time Gemma-2 and report the quality impact.
Spark FFN vs Spark Attention. To understand the effect of Spark FFN vs Spark Attention, we conduct experiment where only attention is switched from a standard one to Spark Attention,
22

---

23
15
04
Col
Oli
Coldl
07
40
85
92
-20
-4
(a) Layer 0
Layer 8
(c) Layer 16
(d) Layer 24
17
43
03
12
(c) Layer 0
(1) Layer 8
(g) Layer 16
(h) Layer 24
Distribution of the entries of the input activation to statistical top-k in Spark FFN (see Figure C.5 for result of Spark Attention). The two rows correspond to activation at two positions
Figure C.4:
O and 1000 of an input, and the columns correspond to activation at four different depth levels {0, 8, 16, 24} of the 26-layer pretrained Spark Transformer. The input is the first 1000 tokens of httpp:://pallgh
com/ articletisesss hhtl prepended with the BOS token. We compare the empirical distribution (Empirical) with the Gaussian distribution whose mean and standard deviation (std) are computed the sample mean and std of the input (Fitted). We see that the Gaussian closely approximates the empirical distribution. We also compare the cutoff value estimated from the Gaussian, i.e., 0(x, k) used eq. (10) with k/d 5% (Cutoff for fitted), with the cutoff value for obtaining 8% nonzeros on the empirical distribution (Cutoff for empirical). It can be seen that these two cutoff values are close.
the first essay from
whereas the FFN remains the standard one. The result illustrated as Gemma-2 + Spark Attention in Figure C.6 It can be seen that Spark Attention provides a minor quality gain over Gemma-2. In comparing Gemma-2 Spurk Attention with Spark Transformer, this also shows that further adding Spark FFN slightly hurts model quality. As noted above, such small difference does not lead to substantial quality impact on the evaluation tasks. Hence, we conclude here that none of Spark FFN and Spark Attention has significant quality impact.
Sparsity enforcing vs Low-cost predictor. Sparsity enforcing via statistical top-k and low-cost activation predictor are two relatively independent components of Spark Transformer. This means that, upon the standard Gated FFN (see eq. (5) that is used in Gemma-2, which we rewrite here for convenience:
Gated-FFN(q: K1, K2, V) V. (a (KT q) @ (KJ 4))
(C.1)
we may choose only apply statistical top-k for enforcing sparsity, i.e.,
Took-Gated-FFNN(:: K1, K2, V)= V (a (Statistical- Top; (KJ T q))
(K₂ q))
(C.2)
Note that applying sparsifying function on the input to the nonlinear function o() as in eq. C.2 common choice in the literature of sparse activations, e.g.. 64
49 the main difference
between these works lies in the specific sparsity enforcing technique, see Table D. for summary. In addition to the sparsifying function, Spark FFN also has another architectural change for the purpose of introducing a low-cost predictor. Here, we rewrite Spark FFN for ease of comparison with eq. CC2
K, V, r) def V (a((Statistiicll (K) gltr)) ((:)))
(C.3)
Analogous to FFN. we may also only add statistical top-k to attention without the low-cost predictor, i.e., by switching from standard Attention in eq. (6 to the following:
Topk= (q; K, V) def V softmax (Statistical- Topp (-oe)
(K (K (KTq))
(C.4)
Here, we aim to understand the effect of introducing statistical top-k without the low-cost predictor. Towards that, we conduct an experiment where FFN and Attention in Gemma-2 are replaced with
eq. (C.2) and eq. (C.4). respectively. The result is illustrated as Gemma-2 + Top-k in Figure C.6
23

---

CripP OJ
Firce
6.24
6.3
(a) Layer 0
(b) Layer 8
(c) Layer 16
(d) Layer 24
0.21
6.15
0100
094
-3
-0
(c) Layer 0
Layer 8
(g) Layer 16
(h) Layer 24
Figure C.5: Distribution of the entries of the input activation to statistical top- in Spark Attention (see Figure C.4 for result of Spark FFN). The two rows correspond to activation for two different
attention heads, and the columns correspond to activation four different depth levels {0,8, 16, 24) of the 26-layer pretrained Spark Transformer. Model input is the first 1000 tokens of the first essay from fattps::/ppt //pp:l h/alpp::/ppallrraaamm www coomarrticles. html prepended with the BOS token, and we examine activation of the last token (i.e.. inner product between the query embedding of the 1001st token and all 1001 key embeddings). We compare the empirical distribution (Empirical) with the Gaussian distribution whose mean and standard deviation (std) are computed as the sample mean and std of the input (Fitted). We see that the Gaussian closely approximates the empirical distribution. We also compare the cutoff value estimated from the Gaussian, i.c., ((x, k) used in eq. (10) with k = 256 (Cutoff for fitted), with the cutoff value for obtaining top 256 entries on the empirical distribution (Cutoff for empirical). It can be seen that these two cutoff values are close.
Table C.I: FLOPs per token comparison: Spark Transformer vs. standard Transformer. In standard Transformer with model dimension dmodel. we assume multi-head attention where the sum of head dimensions equals dmodel. and an FFN with non-gated activation and width dr. Here, Nexx represents the context length for the target token. The computational cost is primarily determined by the FFN (assuming du >> dimuuei, which is typical) and the attention dot product (assuming a long context length). Spark Transformers introduce sparsity parameters, kn and kain- to reduce FLOPs. Setting kn 8% x dp and kaun 256 achieves a 3.2x FLOPs reduction in the FFN, a 4x reduction in the attention dot product, and a 2.5x reduction overall (assuming nex 8k) for Gemma-2B.
FLOPs per Toker
Operation
Standard Transformer
Spark Transformer (Ours)
FFN
Aducdada
dmase de 3dlodsska
Attention dot product
dandeilaax + 3dacdal min Kaams Nexx }
Attention lincar projection
8d
can be seen that the training loss becomes notably larger and the gap compared to Gemma-2 is further increasing with more training steps. This result demonstrates that while the low-cost predictor is introduced to Spark Transformer for reducing the cost in predicting the nonzero entries, it also helps in bridging the gap from the introduction of statistical top-k. In other words, Transformer with low-rank predictors in FFN and Attention is more amenable to activation sparsification without quality loss.
C.6 Computing FLOPs per Token
To understand the FLOPs saving reported in Figure Ta we provide a comparison between FLOPs of the major components of a Transformer, including FFN, Attention dot product, and Attention projections, with that of Spark Transformer. The results are presented in Table C.T
24

---

Appendix D
Additional Discussions
D.1 Additional Related Work
In the following, we review few lines of work closely related to ours.
Mixture of Experts (MoEs) may be considered as a particular case of sparsely activated models which group the neurons in FFN and activate all neutrons in selected groups [78 52). Neuron grouping has the benefit of being better suited for training accelerators compared to unstructured activation sparsity. However, training of MoEs incurs extra complexities in algorithmic design and requires special hardware support [29] Moreover, the structured nature of sparsity limits the model's flcxibility and expressiveness, and recent work advocates the use of a larger number of smaller experts On the other hand, the discovery of the naturally emerging unstructured activation sparsity
18 36.
has motivated the perspective of naturally emerging experts 191 25 16
70
85 95].
Sparse activation is common approach to improve the efficiency of large models and many tech- niques for low-cost activation prediction have been developed over the years, such as low-rank factorization (20). quantization (1I. product keys [46), hashing (12). etc. With the popularity of modern Transformer models, these techniques become natural choices 40) 90, 58 82] for reducing
their high computation costs. In particular. lot of the excitement comes from the discovery that the
activations in FFNs are naturally sparse [91] 54 61] and hence efficiency with activation sparsity is obtained without quality toll.
Our work falls into the category of the latest work in this direction that aims to bring the benefits to the latest generation large language models that do not have natural sparsity. Early attempts [4 69 92] seek to bring back sparsity by switching back to ReLU variants, but it usually incurs a quality loss. The quality gap may be largely bridged by more careful tuning. but the activation becomes less sparse (e.g. 25% nonzeros in LLAMA 7B (81). Top-k has become a more popular choice for obtaining sparsity recently [ and is able maintain neutral quality while offering strong sparsity, but only in selected layers [89 Moreover, such methods require finetuning to bring sparsity and also obtain predictor. Without doing finetuning, [49 56 94 obtained most 50% nonzeros under neutral quality. In contrast to these works, our work not only obtains 8% nonzeros in activation of all FFN layers, but also a predictor, all with a single-stage training. We provide summary of comparison to these methods in Section D.2
Finally, the usefulness of activation sparsity goes beyond efficiency. For example, theoretical studies show its benefits for model generalizability and learnability [65] Moreover, activated neurons may be associated with semantic concepts, which offers understanding of the working mechanism and enables manipulating the output of Transformer models [7,60]
Sparse attention broadly refers to the approach of attending to selected subset of tokens in the context as a means of reducing computation cost [21]
42 73). Works on sparse attention include those that use handcrafted attention patterns [15] 9 2 24 which feature simplicity, and learned
attention patterns [45] 75 (66] which feature better modeling capacity. However, learning attention
patterns often involve learning, e.g., hash table or k-means centers, which significantly complicates modeling. Closely related to our Spark Attention is the top-k attention [35, which obtains data- adaptive attention simply from top-k thresholding. Our work improves upon top-k attention by introducing a low cost predictor which enables an increased computational benefits from sparsity. Finally, KV pruning approaches drop selected tokens permanently as decoding proceeds 993 57]. and cannot achieve as high compression ratio as sparse attention based approaches.
D.2 Comparison with Related Work on Activation Sparsity in FFN
In Table D.I we provide summary of recent work on enabling FFN activation sparsity in the latest
LLMs, including ReLUification [64), ProSpare [8], HiRE 89), and CATS 49 Please see Section for discussion of these methods.
We can see that our Spark Transformer leads to a FLOPs reduction of -72%, which more than all the other methods. This comes at the cost of only a -0.9% quality loss, which is lower than most of the other methods except HiRE, In particular, HiRE is on par with ours in terms of quality loss but achieves less FLOPs reduction.
25

---

Table D.1: Comparison with related work on enforcing activation sparsity in FFN of LLMs. Spark Transformer has the largest FLOPs reduction with one of the smallest quality loss.
Main Techniques
Main Results
Enforce
Predict
Base model
Training cos
Sparsity (#zeros)
Quality
FFN FLOPs
sparsity
support
ReLUificatior
ReLU
None
OPT 1.3B
+0%
93%
-2%
-62%
ReLU
None
Falcon 7B
+2%
94%
-2.5%
-62%
Llama 7B
+3%
62%
-1.9%
-42%
ProSparse
ReLU +
None
Llama2 7B
+1.8%
88%
-1.1%
-59%
|| 1l
Llama2 13B
+6.7%
88%
-1.4%
-59%
2-layer FFN
Llaraa2 7B
NA
75%
NA
NA
Llama2 13B
NA
78%
NA
NA
HiRE
Group tops commonpath
Low-rank
PALM2 1B
+100%
80%
-0.8%
-60%
3
quantization
CATS
Thresholding
None
Mistral 7B
+0%
50%
-1.5%
-33%
Llama2 7B
50%
-2.4%
-33%
Spark
Statistical-
Partial
Gemma 2B
+0%
92%
-0.9%
-72%
Transformer
topk
dimensions
Comparing prefraining, finetuning, and "zero-shot" approaches. Existing approaches enforcing activation sparsity in FFNs can be categorised into three groups depending on when the sparsity is enforced during the training process.
Fretraining-based approaches are those whose activation sparsity is enforced from the very beginning the model pretraining, usually with certain architectural changes upon an established model. Hence, such methods require pretraining model from scratch. For example, ReLUification (44 is the method of switching the activation function in FFN from commonly used one, e.g. GELU, back to ReLU. When tested on OPT 1.3B, [64] applied this modification and pretrained version of OPT 1.3B with ReLU from scratch (Table D.1 contains a summary of relevant results for it).
Finetuning-basee approaches refer those that take existing pretrained model without activa- tion sparsity, and perform additional training steps usually with certain architectural changes to induce sparsity. For example, ReLUification discussed above as a pretraining based approach can also be used as finetuning-bass model ReLU and performing additional training steps. In [64, this approach is applied to Falcon 7B and Llama 7B, for which the results are summarized in D.I Other methods falling into this category, which we have summarized in Table D.1 include ProSparse 81 and HiRE (8)
approach by switching the activation function of pretrained
"Zero-shot" approaches. This refers to methods that enforces activation sparsity upon a pretrained model, without any additional training. CATS 49 is method of this kind.
Our Spark Transformer belongs to the category of pretraining-basee method, that is, it requires pre-training a model from scratch. It may not be used directly as a finetuning the drastic difference between the architecture of standard FFN with gated activation, i.e. eq. (5). and our Spark-FFi This appears to make Spark Transformer less favorable than finetuning based approaches, as the latter requires usually small fraction of the pretraining steps in the finetuning process, see Table [D.]] for summary. However, the long-term serving costs such models usually dominate overall expenditures, in which case the additional pre-training costs are amortized by the significant efficiency benefits realized from higher level of sparsity and FLOPSs reduction during model inference.
approach, due
26

---

D.3 Comparison of Statistical Top-k with Related Top-k Operators
The variational form in eq. 15) reveals connections of Statistical-Tope with other top-k algorithms in the literature. Specifically, [5]] defines soft top-k as
arg min
0 H(z) (z. x), s.t. =1 k. 0 < S 1,
(D.1)
where (2) is the entropy function. Another work [59] defines the Sparse operator
arg min - (=) (z, s.t. = = k, s 1
(D.2)
where Ha (x) (2) is the generalized Gini entropy.
Statistical-Topy in the form of eq. [1). as well as eq. (.1) and eq. (.2). can all be interpreted as finding an output that is close to the input subject to a sparsifying regularization. Their major difference lies in the choice of the sparse regularization. That is, soft top-k and SparseK uses entropy and Gini entropy, respectively, whereas Statistical-Top Topp uses 4 The choice of & makes Statistical- Topp,, superior in that it has a closed form solution provided by soft-thresholding, which only requires d FLOPs. In contrast, soft top-k and SparseK both do not have closed form solutions and require an iterative algorithm with FLOP count dependent the number of iterations. In addition, there no guarantee that soft top-k and SparseK can obtain (approximately) k nonzero entries as output.
Finally, we mention that ideas similar to statistical top-k have been used distributed training [55]. However, we are the first to introduce, adapt, and verify its effectiveness for activation sparsity. Additional discussions are provided in Section D.4
for the problem of
D.4 Additional Discussions on Statistical Top-
Novelty upon existing work. We note that ideas similar to our statistical top-k have appeared in the literature. In particular, (79) introduced the idea of fitting a Gaussian distribution to the entries of an input vector and estimating a threshold from quantile functions. Then, [62] extended the approach to additional distributions. In both cases, the method is used for solving the problem of distributed training. Here, we summarize our contribution upon these works:
We are the first to use statistical top-k for enforeing activation sparsity Transformers. Improving Transformer efficiency via activation sparsity has become very popular research topic (see Section 5). but may have been suffering from a lack of efficient top-k algorithms for enforcing sparsity. Hence, the introduction of statistical top-k may facilitate the development of this area.
Synergizing statistical top-k into Transformers is nontrivial. Since the method of statistical top-k is based on fitting a statistical distribution to the activation vector, there is the need to understand the distributions of different activations in order to determine which particular activation vector is suited for the application of statistical top-k and the associated choice of the distribution. In our case, we decide that statistical top-k should be applied to the activation before the nonlinear function (for FFN) and before softmax (for Attention) since entries of this vector provably follow a
'Please refer to Section 2.1 and Section 2.2 for the calculation of FLOPs for FFN and Attention, respectively. We omit non-lcading- order terms (c.g., those arising from embedding, normalization, and nonlincar layers) and exclude the number of layers as common multiplier.
'Total training cost relative to the base model. For finetuning based approaches, such as ReL Uification (on Falcon and Llama) and ProSparse, the total training cost includes both the pretraining cost and the finetuning cost.
'Quality loss relative to the base model. Here the numbers are based on the results reported in their respective papers. Note that different set of evaluation benchmarks is used in each paper. For ReL.Uficatiin, this set contains ARC-Easy, HellaSwag, Lambada (for OPT) and Arc-E, Arc-C, Hellaswag. BoolQ, PIQA, LAMBADA, TriviaQA, WinoGrande, SciQ (for Falcon and Llama). For ProSparse, this set contains Hvma MBPP, PIQA, SIQA. HellaSwag. WinoGrande, COPA, BoolQ, LAMBADA. and TyDiQA. For HiRE, this set contains WMT14/WMT22, SuperGLUE, Multiple QA datasets, and multiple discriminative tasks datasets. For CATS, this set contains OpenBookQA, ARC Easy, Winogrande. HellaSwag, ARC Challenge. PIQA. BoolQ, and SCI-Q. For Spark Transformer, the datasets are those reported in Figure Ressulss reported here are for the stage training their paper. Il "Specifically, -80% on odd layers only, and -60% average.
27

---

Gaussian distribution at random initialization. We also verify empirically that statistical top-k is
still reliable even after initialization.
We extended statistical top-k from using the hard thre-hhieshl operator with the estimated statisti- cal threshold to the noft-thresholing operator. This leads to a continuous optimization landscape that may have facilitated the optimization. Empirically, we found this choice provide quality benefits for Spark Transformer.
We provide the first theoretical justification for the correctness of statistical top-k, see Theorem 3.1
We reveal the conceptual connection between statistical top-k and several related top-k operators in the literature, Section D33 Such connections may motivate the development of more powerful top-k algorithms in the future.
Handling cases when the activation is sharded. The training of modern large Transformer models usually requires sharding certain model weights and activations across multiple computation devices, due to physical limitations on each device's memory. In particular, if sharding is used for the activations upon which the statistical top-k is applied to, i.e., the pre-GELLU activation in FFN and the pre-softmax activation in attention, extra care needs to be taken so that statistical top-k is applied correctly. While this has not been the case for our experiment on Gemma-2 2B, here we discuss possible solutions if this case arises, e.g. when training a larger Spark Transformer for which sharding relevant activations may become necessary.
Assume that an activation vector of length N is sharded over m devices, and we are interested in finding approximately the top-k entries of N with the largest value. There are two ways of applying statistical top-k for this purpose.
Global statistical top-k. Here we require each device to compute the mean and variance for entries on itself, then communicate them to all other devices. In this case, each device receives m 1 mean and variance values, which can be combined with mean and variance on its own to obtain global mean and global variance. Then, the rest of the steps in statistical top-k can be carried out on individual devices. In this method, the output is exactly the same as if applying statistical top-k without activation sharding. In terms of cost, there is extra computation coming from aggregating mean and variance from all devices, but the cost is very low it requires only ((k) FLOPs. The method also introduces a communication cost, but the cost is again small as each device only needs to send/ / receive 2k- floating point numbers.
Local statistical top-k. Here simply apply statistical top-k' to entries its device with k' = k/ m. The method is sub-optimal in the sense that it does not necessarily produce the same output applying the global statistical top-k. However, it has the benefit of not adding any computation and communication cost.
In cases where k < N, the global statistical top-k above is cheap enough and hence could be natural choice.
D.5 Synergy with Speculative Decoding
Speculative decoding is prominent technique for accelcrating the inference of large autoregressive models [53] [33. It employs a smaller, faster "draft model" to generate a block of candidate tokens, which are then verified in single, parallel forward pass by the larger, more accurate "target model". We view Spark Transformer as a highly complementary and synergistic technology, as its architecture enhances efficiency in both roles within the speculative decoding framework.
Spark Transformer as the target model. The primary computational bottleneck in speculative decoding is the parallel verification step, where the target model processes multiple tokens simultane- ously. Spark Transformer is ideally suited to accelerate this step.
As the target model, Spark Transformer directly reduces the latency of this bottleneck, leveraging the same wall-time speedups demonstrated in singlc-token decoding. While verifying multiple tokens in parallel (e.g., k = 2 to k = 4) necessarily activates more neurons than single-token pass, significant sparsity is still maintained. This for two primary reasons:
Small Draft Size: The number of speculative tokens is typically small, meaning the union of activated neurons remains small fraction of the total.
28

---

Aggregated Sparsity: Consecutive tokens in natural language sequence often share a substantial portion of their activated neurons. This phenomenon, termed aggregated sparsity [64], is an extension of the lazy neuron phenomenon 54 and the observed persistence of activation patterns for related inputs.
Because the union of activated FFN neurons and attended tokens remains sparse, Spark Transformer effectively reduces the computational cost of the most expensive part of the speculative decoding loop.
Spark Transformer as the Draft Model. An effective draft model must satisfy two criteria: must be significantly faster than the target model, and it must be accurate (i.e., its predictions must frequently align with the target model's). High accuracy is crucial as it dictates the token acceptance rate, which the primary driver of the overall speedup.
Spark Transformer is a strong candidate for a high-performance draft model. Our results demonstrate that it achieves rear-quality neutrality with its dense counterpart while operating at fraction of the inference cost. Using a Spark Transformer as the draft model - either a smaller version or even the same model with lower sparsity budget could lead much higher acceptance rate than standard distilled model of equivalent speed. This high fidelity, combined with low latency, makes it an ideal driver for maximizing the efficiency of the speculative decoding process.
D.6 Synergy with Quantization Techniques
We view Spark Transformer and quantization as two orthogonal and highly synergistic optimization axes. Quantization, such as conversion to FP8 INT8. reduces the computational cost and memory size of each individual operation and parameter. Spark Transformer, contrast, reduces the total number of operations and memory accesses by dynamically pruning FFN activations and attention tokens. Critically, our sparse implementation also reduces the memory bandwidth bottleneck by skipping the load of masked weights from HBM. The combined benefits are therefore expected be multiplicative, leading to significant compounded gains in inference throughput and energy efficiency.
A key question is whether model that is already sparsely activated is more sensitive to the precision loss from quantization. We hypothesize that, for activation quantization in particular, Spark Trans- former may reduce the model's sensitivity and be less prone to accuracy degradation compared to dense model. This is critical, as recent work has shown that standard sparsity and quantization can be in contradiction: magnitude-based sparsity preserves large-valus value outliers that are detrimental to guantization-awarr scaling 334] This hypothesis stems directly from the design of our statistical top-k operator.
The statistical top-k operator in Spark Transformer employs soft- theshhhresholdig.. This function outputs max( ((x, k), 0), effectively subtracting the learned threshold (xr, k) from all activated neurons. This operation inherently compresses the dynamic range of the activation tensor. A primary challenge in I.IM activation quantization is the presence of maggenma "outlier" values, which skew the quantization scaling factor and lead to significant precision loss for the more common, smaller values [22). By ppre-conditioning" the distribution and shrinking these outliers, our operator makes the activation tensor significantly more amenable to accurate low-bit representation.
Dynamic Range Compression:
Zero-Point Stability: The 92% of FFN neurons that in the top-k are set to This large volume of true zeros is perfectly and losslessly represented in any quantization format, avoiding the "near-zero" noise that can plague dense models.
For weight quantization, we hypothesize that the sensitivity would be broadly similar to the dense Gemma-2 counterpart, as Spark Transformer reallocates and reuses the full set of parameters rather than pruning them.
While a full empirical study of this interaction is a highly promising direction for future work, hypothesize that Spark Transformer is not only compatible with quantization but may actively facilitate it. This positions our work alongside other recent efforts to create novel quant-- representations (23), offering robust path to further acceleration."
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction clearly state the Spark Transformer's 's ability to achieve high activation sparsity in FFN and attention, maintain competitive performance, preserve parameter count, and use standard training, all supported by results.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: The paper discusses limitations such hardware not being fully optimized for sparsity, the need for full pretraining rather than finetuning, and specific batching scenarios where benefits are less pronounced
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "LLiitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
30

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Theorem 3.1 and 3.2 are stated with assumptions and their proofs are provided in Appendix A.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
The paper details the base model (Gemma-2 2B), hyper-parameters, training procedures, and inference engines (gemma.cpp. Ilama.cpp) used.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
31

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification:
Training data and code for Gemma-2 are proprictary.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. ec/ cublic/guidess htapli//uuides//deeeuubiissess
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Detailed information model dimensions, hyper-parameter (k, for FFN and Attention), activation functions, pretraining recipe, and inference setups (CPU/GPU, prompt lengths, batch sizes) are provided.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Extreme high pretraining cost prevents multiple independent trials.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
32

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: The paper specifies CPU core counts (4-Core, 16-Core CPUs) and GPU type (NVIDIA T4 GPU) used for inference evaluations, and notes that training was done using the Gemma-2 recipe.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: The paper does not mention any ethical concerns deviations from standard research practices in machine learning.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The paper highlights positive impacts, such as broadening access to high- quality models for users with limited access to high-FL.OP hardware. No significant negative socictal impacts are discussed.
33

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymenn of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
The paper focuses architectural improvements and efficiency, not releasing a new model or dataset that poses high risks for misuse.
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
The paper explicitly cites and credits Gemma-2 for the base model, and references gemma.cpp and llama.cpp for inference engines.
Justification:
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
34

---

The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thcode com/ comddatasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper introduces a new architecture ("Spark Transformer") but does not mention releasing standalone asset with accompanying documentation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The research does not involve crowdsourcing or human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
The research does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
35

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMS as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https://neurips: httt:p////nurripss.cc//offeeeeee Cc//Confereeccs///22////... for what should or should not be described.
36


Reviews:
[{"review_id": "mXY5IrVPUM", "reviewer": "Reviewer_hAy9", "summary": "This paper introduces the Spark Transformer, a novel architecture designed to re-enable activation sparsity in both feedforward networks (FFNs) and attention mechanisms of Transformer models. The authors propose a statistical top-k operator—an efficient, differentiable approximation to top-k masking that avoids expensive sorting. They also design a unified predictor framework using reallocated parameters to estimate activation importance without additional overhead. The Spark Transformer achieves high sparsity (e.g., 8% FFN activation, ≤256 attended tokens), significantly reducing FLOPs while maintaining near-baseline quality. Extensive experiments on Gemma-2 show 1.4x decoding speedups on GPU, demonstrating the practicality and efficiency of their approach.", "questions": "No other questions.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n1. The Spark Transformer introduces a unified architecture for both FFN and attention sparsity using reallocated model parameters as low-cost predictors, avoiding extra modules or parameters and maintaining training simplicity.\n2. The proposed statistical top-k operator provides an efficient, differentiable, and hardware-friendly alternative to traditional top-k methods, enabling sparse activation without significant training overhead.\n3. The model achieves significant FLOPs and wall-time reductions (up to 1.79× on CPU, 1.4× on GPU) while maintaining competitive accuracy on standard downstream benchmarks compared to dense baselines like Gemma-2 2B.\n\n**Weaknesses**:\n1. The experiments are conducted solely on the Gemma-2 2B model, leaving it unclear how well the proposed method scales to larger models or generalizes to more complex scenarios such as instruction-tuned, or multimodal settings.\n2. The paper introduces statistical top-k and shows that it outperforms standard top-k in empirical results. However, the paper lacks an in-depth analysis of why statistical top-k performs better. Is the improvement due to better differentiability? This is particularly important since the non-differentiability of sparse top-k has long been a bottleneck in training sparse MoE models, with many works proposing improved gating mechanisms. If statistical top-k indeed provides both performance gains and differentiability, it would be valuable to explore its potential applicability to MoE training, which the paper does not address.", "was_revised": true, "final_justification": "The methodology and results are clearly presented, and the paper adds value to ongoing research in the field. I will therefore keep my score as Accept.", "rebuttal": "We thank Reviewer hAy9 for their positive assessment and for recognizing our work as \"technically solid\" with \"high impact.\" We have found the comments very helpful in strengthening the paper and have revised the manuscript accordingly.\n\n# **Generalizability and Scaling (Evaluation on Gemma-2 2B only)**\n\nWe agree with the reviewer that evaluating scalability to larger models and generalizability to instruction-tuning are valuable directions for future research. We have clarified this in the paper's conclusion.\n\nOur primary goal in this work was to provide a **rigorous and foundational study** of the Spark Transformer's effectiveness. To this end, we deliberately focused our significant computational resources on the Gemma-2 2B model for two key reasons:\n\n- **Controlled and In-Depth Evaluation**: Concentrating on a single, state-of-the-art open model allowed for a thorough and controlled study. This approach is a common and necessary practice in resource-intensive architecture research before undertaking even costlier scaling experiments.\n\n- **Strong, Representative Evidence**: Demonstrating significant FLOP reduction and wall-time speedup on a strong baseline like Gemma-2 provides **compelling evidence** for our approach's viability, establishing the groundwork for future scaling law analysis.\n\n\n# **Analysis of Statistical Top-k and Applicability to MoE**\n\nThis is an excellent and insightful point. As detailed in the paper, the key advantages of our method are twofold:\n\n- **Advantages over Standard Top-k**:\n  - Efficiency: As shown in Figure 6, its linear-time complexity avoids costly sorting operations, yielding a >10x reduction in training overhead compared to optimized approximate top-k operators.\n  - Differentiability: Our use of soft-thresholding (Theorem 3.2) provides a smooth gradient signal, which we hypothesize leads to the superior convergence and final model quality we observed.\n\n- **Application to MoE Models**: This is a fantastic connection, and we thank the reviewer for suggesting it. Motivated by this comment, we investigated its applicability and found it to be a highly promising direction.\n  - A recent paper [a] uses a similar soft-thresholding technique for MoE routing but still relies on a costly sorting operator. Our method's ability to find the threshold in linear time could therefore address a key computational bottleneck, especially as MoEs trend toward a larger number of experts [17, 32].\n  - We believe this is an exciting avenue for future work and, to reflect the value of this insight, we will add a new discussion on this topic to the revised manuscript.\n\n[a] REMOE: ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing, ICLR 2025"}, {"review_id": "jSyZgSUrlc", "reviewer": "Reviewer_4xrP", "summary": "The paper modifies the transformer architecture into a new Spark Transformer model, leveraging sparsity in the model activations through predictors that decide which modules in the model are activated during inference. Sparsity is introduced in the training process, without any additional post-training modifications to the model. To efficiently identify sparse parts of the model, the authors introduce a novel statistical top-k operator. Spark Transformer version of Gemma-2B performs close to the original model using 40% of its FLOPS and provides 1.8 inference speedup on CPU and 1.4 inference speedup on GPU.", "questions": "Why is the idea only evaluated for a single, relatively small LLM? The authors justification for not providing the training code (“Training data and code for Gemma-2 are proprietary.”) indicates that the authors should have access to sufficient computational resources. If the authors cannot provide the code, the paper should at least contain strong experimental evidence that allows the reader to adequately compare the proposed method to the alternatives and judge its impact of wider set of models. However, all the experiments are contained to a single architecture (Gemma 2B) and the authors do not perform sufficient experiments to compare existing alternatives to their method even in this single setting.", "limitations": "Yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "# Strengths\n\n1. The paper introduces a method to directly integrate activation sparsity into the model via jointly trained sparsity predictors, which avoids the need to apply any post-training modifications and leads to an elegant solution.\n\n2. The proposed top-k operator is a novel and efficient alternative to standard top-k selection.\n\n3. Unlike many prior works that focus solely on FFN sparsity, this paper applies activation sparsity to attention mechanisms as well.\n\n4. The method supports efficient batched inference, as shown in Appendix C.3, which is not always given in sparse activation papers.\n\n5. Writing is very clear and easy to follow.\n\n# Weaknesses\n\n1. The empirical results are limited to a single model (Gemma-2B), and there are no results that would provide evidence that the method generalizes across architectures, model sizes, or tasks.\n\n2. The paper mentions related works on activation sparsity and provides a comparison between their approach and existing alternatives in Table D.1. However, this table contains only main results scraped from the other papers, not experimental results reproduced by the authors. Moreover, these results are not really compatible, as they refer to different models. Therefore, the paper does not provide sufficient results to convincingly position the method with respect to the other existing alternatives. Additionally, while wall-clock speedups for Spark Transformer are mentioned in the paper, the table comparing Spark with other methods includes only FLOPS, which do not directly translate to the inference efficiency.\n\n3. Paper emphasizes CPU speedups and does not provide results for common GPUs like A100, which makes it hard to judge the practical impact of the method for the most common deep learning accelerators.\n\n4. The paper gives a fairly detailed description of the architectural modifications, including hyperparameters, training setup, and implementation details. However, the authors do not share the code, which greatly affects the reproducibility of their work.\n\n# Summary\n\nWhile the paper proposes an interesting and potentially impactful idea, the lack of experimental evaluation makes it impossible to judge the impact of the proposed method in practice and judge whether it provides a significant efficiency improvements compared to already existing alternatives.", "was_revised": true, "final_justification": "After the discussion phase, I am confident that the paper’s technical contributions are strong and that it provides a clear practical and technical impact. I have chosen to keep my score unchanged to reflect certain evaluation concerns (as mentioned in more detail during the discussion), but I am now convinced that the positive aspects of the work outweigh the negatives and justify acceptance.", "rebuttal": "We thank Reviewer 4xrP for their feedback. We appreciate the acknowledgment of our work's strengths, including the elegant end-to-end training, the novelty of the statistical top-k operator, and the application of sparsity to both FFN and Attention. We address the noted weaknesses below.\n\n**Empirical Evaluation using Gemma-2B only**\n\nWe thank the reviewer for this suggestion. The decision to focus our study on Gemma-2B was driven by the substantial computational cost of pre-training. This approach allowed us to conduct a thorough and controlled investigation on a strong, representative open-weight model, a common practice in LLM pre-training research where resources are a key constraint.\n\nWhile larger models like Gemma-2 9B are an important future direction, their pre-training cost was beyond the scope of this project. By concentrating our resources, we were able to provide what is, to our knowledge, one of the first complete pre-training demonstrations of activation sparsity's feasibility in a modern LLM. We believe this foundational evidence is a valuable contribution and provides a strong empirical basis for future explorations on larger-scale models.\n\n\n**Code availability**\n\nWhile proprietary constraints on the full training data and scripts prevent a complete code release, we are committed to ensuring our work is reproducible. To this end, we have provided the following key components to the community:\n- A detailed architectural blueprint of our Spark Transformer (Section 2, Figure 2), enabling straightforward re-implementation.\n- Complete hyperparameter specifications for all FFN and Attention layers (Section 4.1), allowing for a faithful replication of our results.\n- Implementation details and pseudo-code for our custom sparse computation kernels (Appendix B), designed for direct integration into open-source frameworks like gemma.cpp.\n\n\n**Comparison with Alternatives (Table D.1) and FLOPS vs. Wall-Clock Time**\n\nThank you for these points. Regarding comparisons with prior methods, our results in Figure 1(a) show that Spark achieves a greater FLOP reduction than methods like ReLU, ReLU^2, and Top-k, while maintaining comparable model quality.\n\nOn the topic of evaluation metrics, we agree that direct wall-clock comparisons are valuable. However, fairly implementing and optimizing each baseline within our specific hardware and software framework (gemma.cpp) would be a significant engineering effort beyond the scope of this work. Therefore, to ensure a standardized and hardware-agnostic comparison against a wide range of published methods, we adopted FLOPs as our primary metric. This is an established practice in the activation sparsity literature (e.g., \"ReLU strikes back\"), allowing for more direct and fair comparisons with prior art.\n\n\n**GPU Evaluation on T4 vs. A100** \n\nWe thank the reviewer for the opportunity to clarify our hardware selection. Our choice of the NVIDIA T4 was strategic, aimed at demonstrating the practical benefits of the Spark Transformer beyond high-end data center accelerators. While prior work has validated activation sparsity on premium hardware like A100s and TPUs, our goal was to prove its effectiveness on more widely accessible GPUs.\n\nOur strong results on the T4, together with our CPU performance, show a clear path toward making powerful LLMs more efficient on commodity and resource-constrained hardware. Therefore, our findings do not simply replicate prior work but instead complement it, by confirming that these efficiency gains are achievable across a broader and more common class of hardware."}, {"review_id": "2ir49psQMU", "reviewer": "Reviewer_PzjZ", "summary": "1  Summary\nThe paper proposes Spark Transformer, an architecture that re‑introduces and exploits activation sparsity in both feed‑forward networks (FFN) and multi‑head attention. The key technical pieces are\n* Spark FFN – splits the input,  locate the k most relevant neurons via the \"statistical top‑k operator\", and performs the remaining matmuls sparsely.\n* Spark Attention – mirrors the same idea for query–key scores, limiting each token to at most 256 attended tokens.\n* Statistical top‑k – a linear‑time, differentiable approximation of top‑k based on Gaussian quantiles coupled with soft‑thresholding; analysis (Thm 3.1/3.2) bounds the expected error and shows differentiability.\n* An implementation on Gemma‑2 (2 B) shows ≈ 8 % active FFN neurons and ≤ 256 attention edges, yielding 2.5× theoretical FLOP reduction and 1.86× CPU / 1.25–1.40× GPU decode speed‑ups with negligible quality loss on standard LLM benchmarks.", "questions": "1. GPU speed‑up baseline – Could you provide absolute tokens/s for Spark versus a dense Gemma‑2 implementation on the same T4, so the community can reproduce the 1.25–1.40× figure? (Fig. 5’s GPU bars compare Spark without optimization  vs. Spark with optimization, not Spark vs. dense Gemma‑2)\n2. Huber versus ReLU – Have you tried dropping Huber (δ → 0) or replacing it with the ReLU‑style hard threshold? How does this affect convergence and final quality?\n3. Predictor overhead comparison – When integrating a Deja‑Vu‑style low‑rank predictor into a standard FFN (no splitting), what is the FLOP/latency trade‑off relative to Spark FFN?\n4. Attention sparsity generality – Did you evaluate Spark Attention on tasks requiring long‑range reasoning (e.g., PG‑19, LongBench) to ensure the 256‑edge cap does not hurt quality?\n5. Training stability – Were any gradient‑scaling or learning‑rate adjustments needed because of statistical top‑k, or does the recipe match Gemma‑2 exactly?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n1. Novel formulation utilizing activation sparsity which does not requires  Deja-vu-style \"predictor\".\n2. Linear‑time top‑k backed by theory and ablation showing <1 % training slow‑down versus >10× for JAX approx‑max‑k.\n3. End‑to‑end pre‑training at scale (2 T tokens) with near‑neutral accuracy on 8 diverse downstream tasks.\t\n4. Evidence of wall‑time wins on commodity 16‑core CPUs and a low‑profile NVIDIA T4 GPU.\n\n## Weaknesses\n1. GPU baseline ambiguity – Fig. 5 compares “optimized vs. un‑optimized” Spark kernels; the paper does not show a direct comparison against a dense Gemma‑2 baseline, making the reported 1.25–1.40× speed‑up hard to interpret.\n2. Regularisation ablation missing – Sec. 3.2 advocates Huber regularisation for smoothing statistical top‑k, but no experiment contrasts Huber with the common ReLU (or no) regulariser to isolate its benefit.\n3. Comparison to predictor‑based sparse methods – the paper qualitatively cites Deja‑Vu, PowerInfer, etc., yet does not provide a quantitative comparison on speed‑up vs. additional parameters / training cost, making it hard to gauge relative efficiency.", "was_revised": true, "final_justification": "Thank you for the response and additional results.\nNow all of my concerns have been resolved.\nThe idea of utilizing activation sparsity w/o prediction is very interesting and practical.\nI keep my original rating.", "rebuttal": "We thank Reviewer PzjZ for their thorough review and for recognizing the novelty of our formulation, the strength of our theoretical backing, and the evidence of wall-time wins. We address the specific technical questions below.\n\n# **Baseline Ambiguity in Figure 5**\n\nThis is an excellent point regarding clarity. To provide a direct comparison, we have conducted new decoding speed evaluations. We benchmarked our optimized Spark Gemma-2 against the original dense Gemma-2 (on CPU, with ongoing GPU experiments), and also measured the overhead of our architecture before sparsity optimizations were applied.\n\nThe results confirm a significant net speedup. The slight overhead from our operators is far outweighed by the gains from sparsity.\n\n| Platform           |  Dense Gemma-2 (Baseline) | Spark (No optimization) | Spark (Full optimization) |\n| -------------          | ---------------------------------------  |  -----------------------------    | --------------------------------- |\n| CPU (4-core)   | 1.00x                                         | 0.87x - 0.96x                  |  **1.42x - 1.64x**            |\n| CPU (16-core) | 1.00x                                       | 0.92x - 0.97x                  |  **1.35x - 1.79x**            |\n*Table. Decoding speed (tokens / sec) relative to the baseline*\n\n# **Huber Regularization Ablation**\n\nThank you for this question. As briefly mentioned (L209), we empirically investigated this. We found that a non-zero hyperparameter δ for the Huber loss provided no noticeable improvement in model quality or training stability over setting δ=0. Therefore, for simplicity and to avoid introducing an unnecessary hyperparameter, we proceeded with δ=0 in all experiments, which effectively reduces to the standard soft-thresholding operator.\n\n# **Comparison to Predictor-Based Methods like Deja-Vu**\n\nThis is a crucial point that distinguishes our work. Our approach targets a fundamentally different challenge than methods like Deja-Vu.\n\n- **Deja-Vu** was designed for older architectures (e.g., OPT) with naturally sparse activations. Its predictor identifies this pre-existing sparsity. \n- **Our work** targets modern LLMs with non-sparse gated activations (e.g., SwiGLU). Here, no natural sparsity exists. Our statistical operator is part of an architecture that induces structured sparsity from scratch while minimizing quality loss.\n\n# **Attention Sparsity on Long-Range Reasoning Tasks**\n\nThis is an important consideration. Our work focuses on architectural innovation within the context of the base model's capabilities. Since the Gemma-2 models we build upon are pre-trained for an 8k context length, they do not generalize well beyond that window. Benchmarking on long-context tasks would therefore not yield meaningful results for either the dense baseline or our Spark Transformer.\n\nWe agree that evaluating a version of Spark Transformer pre-trained for a long-context regime is a vital piece of future work, but it is beyond the scope of the current study.\n\n# **Training Stability and Recipe Adjustments**\n\nOne of the practical advantages of our method is its stability. The Spark components integrated seamlessly with the standard Gemma-2 training recipe, **requiring no adjustments** to the learning rate, optimizer, or gradient scaling. This highlights the robustness and ease of adoption of our approach."}, {"review_id": "Tdc3M3zRDK", "reviewer": "Reviewer_1Jfq", "summary": "The proposed Spark Transformer consists of two parts: Spark FFN and Spark Attention. The authors dedicate/repurpose a portion of hidden states to estimate the importance of weights and KV vectors against the given hidden state (e.g., K and V vectors). Further, Spark Transformer identifies top-K elements in an approximate manner, not requiring the costly process of sorting. On the CPU and low-profile GPU machines, the proposed Spark Transformer shows notable speedup ranging from 1.25x to 1.79x.", "questions": "* Can you further articulate the interaction of Spark Transformer and speculative decoding?\n* Can you further quantify or qualitatively discuss the effectiveness of Spark Transformer on conventional serving environments on the cloud?\n* Can you discuss the interplay between Spark Transformer and quantization techniques? (e.g., sensitivity of its accuracy while using different data types.)\n\nNot questions, but minor comments:\n*  Section vs. section: be consistent.\n* “an” NVIDIA GPU", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n+ The core ideas of embedding sparsity information within hidden states and avoiding the cost of sorting by approximately estimating top-k elements are well taken.\n+ The paper is well written and structured organizationally.\n+ Presenting results for CPUs and GPUs is another plus.\n\nWeaknesses:\n- Spark Transformer seems to target low-profile devices, where the degree of batching is not likely high (e.g., mobile devices). In this case, speculative decoding is gaining popularity, where the effective batch size of a ‘target’ model  increases. This paper does not even mention speculative decoding though.\n- It is not clear if the proposed idea can be highly effective to the cloud environments, which use more powerful GPUs/TPUs and where the batch size is neither extremely high (violating goodput in servers) or low (running without speculative decoding in mobile devices), thus, it would be limited by the memory bandwidth of fetching weights from main memory (e.g., HBM). In these scenarios, saving computation might not lead to any considerable improvement in overall performance and energy efficiency as the serving system is memory-bound and communication dominates computation in energy consumption. More discussion in these deployment scenarios will further strengthen the paper.", "was_revised": true, "final_justification": "Interaction with the authors and the reviews of the fellow reviewers help me keep my initial positive score unchanged.", "rebuttal": "We thank Reviewer 1Jfq for the excellent review and the insightful, forward-looking questions about our work's practical deployment. We are encouraged that the reviewer recognized the novelty of our core ideas and the quality of our presentation.\n\n# **Interaction with Speculative Decoding**\n\nThank you for raising this highly relevant point. We view Spark Transformer and speculative decoding as **highly complementary and synergistic technologies**. Our method enhances both roles within the speculative decoding framework.\n\n- As the Target Model: Spark Transformer directly accelerates the costly verification step. While verifying multiple tokens at once means more neurons are activated compared to single-token decoding, significant sparsity is still maintained for two reasons: 1) The number of speculative tokens is typically small (e.g., 2-4), and 2) Consecutive tokens often share activated neurons (a phenomenon termed \"Aggregated Sparsity\"). This allows Spark Transformer to **reduce the latency of the primary bottleneck** in the speculative decoding loop.\n- As the Draft Model: A great draft model is both fast and accurate. Because Spark Transformer nearly matches the quality of its dense counterpart at a fraction of the inference cost, it serves as an **ideal high-performance draft model**, potentially leading to higher token acceptance rates and greater overall speedups.\n\nWe will add a detailed discussion of this synergy to the revised manuscript.\n\n# **Effectiveness in Memory-Bound Cloud Environments**\n\nThis is a critical point for practical deployment. While LLM inference can indeed be memory-bound, our approach is specifically designed to provide substantial benefits in this regime by **reducing both computation and memory traffic**.\n\nOur innovation is not just a theoretical FLOP reduction. As detailed in Appendix B and illustrated in Figure 4, our custom sparse kernels are implemented to **avoid loading the weights corresponding to zeroed-out activations from High-Bandwidth Memory (HBM)**. Therefore, our activation sparsity translates directly into **reduced memory I/O**.\n\nThis means our reported ~2.5x FLOP reduction is coupled with a significant reduction in memory access, directly addressing the memory bandwidth bottleneck. This dual benefit also leads to **substantial energy savings**, a vital metric for large-scale cloud deployments. We will clarify this crucial point more explicitly in the revised paper.\n\n# **Interplay with Quantization**\n\nThis is another excellent suggestion. We consider activation sparsity and weight quantization to be **orthogonal optimization techniques** that can be combined for even greater efficiency.\n\nWhile we have not yet performed these experiments, we see no fundamental conflicts. Standard methods like post-training quantization (PTQ) or quantization-aware training (QAT) should be directly applicable to Spark Transformer. We hypothesize that because our model learns activation importance scores during training, it may even prove robust to the precision loss from quantization. Investigating this powerful combination is a promising avenue for future work, and we will add it to our discussion section."}]

Author rebuttal:

Meta-review: The paper proposes the Spark Transformer, a new architecture that reintroduces activation sparsity into both feed-forward networks (FFNs) and multi-head attention of Transformer models. The key idea is to selectively activate only the most important neurons or attention edges during inference, guided by lightweight predictors that repurpose existing parameters. Before the rebuttal, this paper receives one negative comment. However, after rebuttal and discussion, the major concerns of all the reviewers are addressed, especially about the novelty. Thus I also recommend to accept this paper.
Decision: Accept (poster)