CATS: Context-Aware Thresholding for Sparsity in Large Language Models
Authors: Donghyun Lee
Jaeyong Lee
Genghan Zhang
Mo Tiwari
Azalia Mirhoseini
Venue: COLM (2024)
Keywords: efficient inference, sparsity, context-aware inference
TL;DR: We make LLMs faster and more efficient by inducing sparsity; our prescription for doing so applies to many LLMs.

Abstract: The dramatic improvements in Large Language Models (LLMs) come at the cost of increased computational resources for inference. Recent studies ameliorate the computational costs of LLMs by increasing their activation sparsity but suffer from significant performance degradation on downstream tasks.
In this work, we introduce a new framework for sparsifying the activations of LLMs and reducing inference costs, dubbed $\underline{C}$ontextually $\underline{A}$ware $\underline{T}$hresholding for $\underline{S}$parsity (CATS).
CATS is a relatively simple algorithm that is easy to implement and highly effective.
At the heart of our framework is a new non-linear activation function.
We demonstrate that CATS can be applied to various models, including Mistral-7B and Llama2-7B \& 13B, and outperforms existing sparsification techniques across multiple tasks.
More precisely, CATS-based models achieve downstream task performance within $\sim$ 99\% of their base models at activation sparsity levels of 50\%, even without any fine-tuning.
Moreover, with fine-tuning that targets only 1\% of the parameters, CATS-based models not only converge faster but also achieve better task performance than competing techniques.
Finally, we develop a custom GPU kernel for the efficient implementation of CATS that translates the activation sparsity of CATS to real wall-clock time speedups.
Our custom kernel implementation of CATS results in a $\sim$15\% improvement in wall-clock inference latency of token generation. We release our code, experiments, and datasets at https://github.com/ScalingIntelligence/CATS.

Published as a conference paper at COLM 2024
CATS: Contextually-Aware Thresholding for Sparsity in Large Language Models
Je-Yong Lee!
Donghyun Lee Department of Computer Science University College London donghyun. lee 21 fucucl.... uk
Mathematical Institute
Oxford University je-yong lee@xorc.ox..cc
Genghan Zhang, Mo Tiwari, and Azalia Mirhoseini Department of Computer Science Stanford University {zgh23, motiwari, azalia @staeford..addeeeoooooo edu
Abstract
The dramatic improvements in Large Language Models (LLMs) come the cost of increased computational resources for inference. Recent studies ame- liorate the computational costs of LLMs by increasing their activation spar- sity but suffer from significant performance degradation on downstream tasks. In this work, we introduce new framework for sparsifying the acti- vations of LLMs and reducing inference costs, dubbed Contextually Aware Thresholding for Sparsity (CATS). CATS is relatively simple algorithm that is easy to implement and highly effective. At the heart of our frame- work a new non-linear activation function. We demonstrate that CATS be applied to various models, including Mistral-7B and Llama2-7B & 13B, and outperforms existing sparsification techniques across multiple tasks. More precisely, CATS-based models achieve downstream task perfor- mance within 999 of their base models at 50% activation sparsity, even without fine-tuning. Moreover, with fine-tuning that targets only 1% of the parameters, CATS-based models not only converge faster but also achieve better task performance than competing techniques. Finally, we develop a custom GPU kernel for efficient implementation of CATS that translates the activation sparsity of CATS real wall-clock time speedups. Our custom kernel implementation of CATS results in ~15% improvement in wall- clock inference latency of token generation. We release our code, experi- ments, and datasets at https // htwpa//oooooallanneellllieeeg
Introduction
LLMs have demonstrated remarkable success across a variety of fields (Devlin al. 2018 Brown et al. 2020: Achiam et al. 2023 Brohan et al. 2023), however, this progress comes with significant computational costs. The training of GPT3 is estimated to have consumed over 3,000,000 GPU-hours and emitted three thousand times the CO2 equivalent of round-trip flight from San Francisco to New York (Patterson et al.) 2021). Furthermore, inference costs often eclipse training costs for models that serve trillions of queries. As such, there is significant interest in reducing the inference costs of LLMs while preserving task performance.
Various techniques have been proposed to mitigate LLM inference costs. These approaches often based on quantization (Frantar al. 2022 Dettmers et al. 2022), pruning (Mal et al. 2023 Sun et al. 2023), and other forms of weight sparsification Frantar & Alistarh
(2023
Mixture Experts (MoE) techniques have emerged as particularly promising and
These authors contributed equally co-first authors.

---

Published as a conference paper at COLM 2024
4.04
CV0 c.0 180 K05 LO oL0 kil was com
(a) Llama2 Layer O.
(b) Llama2 Layer 15.
(c) Llama2 Layer 31.
Daleca Treeshat
021-
600
oLa mis Bio c2t0 Acsee Fer Arseaic
6000
30
(d) Mistral 7B Layer O.
(e) Mistral 7B Layer 15.
(f) Mistral 7B Layer 31.
Figure 1: Histograms of post-MLP activations of different layers in different models. Subfig- ures (a), (b), and (c) correspond to Layers 0, 15, and 31 in Llama2-7B, respectively. Subfigures (d), (e), and (f) correspond to Layers 0, 15, and 31 in Mistral 7B, respectively. The absolute threshold indicates 50% sparsity, where values smaller than the threshold are considered negligible in our technique and thus zeroed out.
are employed by current sfate-of-the-art LLMs (Shazeer et al. 2017 Lepikhin et al. 2020 Fedus et al., 2022c Jiang et al.) 2024).
MoE techniques activate only subset of parameters at each inference stage, thereby re- ducing memory and computational requirements compared to using the entire model. Prevailing implementaioons of MoE techniques introduce many multi-layer perceptrons (MLPs; the "experts") and dynamically select which experts to multiply with the hidden vector. This selection is performed by à "router"- small neural network trained to deter- mine the appropriate experts to activate based on the input (Lewis et al., 2021 Rajbhandari et al. 2020).
Concurrently, recent work has observed that activations in the MLP blocks of LLMs are sparse (Liu et al. al. 2023b Mirzadeh et al., 2023). This implies that only few rows (or columns) of the corresponding weight matrices are required for the forward pass. Intuitively, if we could predict a priori which elements of the weight matrices were unnecessary via an oracle, we could obviate their respective computations, This is thematically similar to MoE approaches: the activated neurons of the weight matrices can be viewed as activated "experts" and the oracle can be seen as the "router."
We observe that the activation patterns of common LLMs suggest a path to such an oracle. Tigure[ssowe a histogram of the post-MLP activations for Layers 0, 15, and 31 for Llama-7B and Mistral-7B on a sample of 500 data points from the RefinedWeb dataset (Penedo et al., 2023). Many of the activations are concentrated about 0; if these activations could be made exactly 0, the corresponding weights of the MLP blocks could be made unnecessary during inference. It is this observation that motivates our study.
In this work, we make the following contributions:
1. We draw connection between the MoE framework and multiplication performed by dense matrices in the MLP blocks of LLMs.
2. We introduce a new sparsification procedure based on a novel activation function, dubbed CATS (for Contextually Aware Thresholding for Sparsity), motivated by
2

---

Published as a conference paper at COLM 2024
empirical evaluation of activation distributions (Figure 1 Crucially, CATS allows for a controllable level of sparsity.
3. We demonstrate that, without any fine-tuning, CATS can be applied to various models, including Mistral-7B and Llama2-7B & 13B, and achieves comparable downstream task performance even at sparsity levels as high as 50%.
4. We demonstrate that, with parameter-efficient fine-tuning, CATS outperforms an existing sparsification technique in downstream task performance at the same sparsity level and number fine-tuning steps.
5. We provide a custom GPU kernel implementation that exploits the sparsity of CATS and achieves a 155% improvement in wall-clock inference latency of token generation over the dense models.
Related Work
Significant recent work focuses reducing the inference costs of LLMs. Approaches that utilize mixture-of-experts or activation sparsity are most similar to our work.
Mixture-of- offExxerts (MoE) techniques induce effective sparsity in LLMs by determining which subset of subnetworks (the "experts") to activate during the inference pass, often via trained "router" subnetwork. This is popular line of work with significant research interest (Shazeer et al., 2017 Hazimeh et al. 2021 1: Zhou et al. 2022: Lewis et al.) 2021 Roller et al. 2021 Zuo al. Komatsuzkk
al. 2022; Rajbhandari et al.) 2022: /Zhang et al.) 2022a b Fedus et al. 2022al Zoph et al. 2022
al. 2022: Lou al. |
Kudugunta et al.) 2021 Fedus al. 2022c Lepikhin al. 2000 Du et al.
2022; Fedus al.
2022b:JJaag et al. 2024). For a review of Moe models, we reter the reader to (Fedus et al.
2022a
Activation Sparsity: Activations are known to be sparse in LLMs that utilize ReLU non- linearities in their MLP blocks (Li et al. 2022]; however, the reasons for this are not well- understood Hoefler et al. 22021) Nonetheless, activation sparsity induced by ReLU non- linearities has been explored to reduce memory usage and inference time (Jaszczur al., 2021 Liu et al. 2023b; Szatkowski et al. 2023). Recent work in this area has framed the rows of weight matrices in MLP layers as experts, similar to our work, and/ or deploys a small neural network to predict which activations will be non-zero to reduce inference costs (Zhang et al. 2024 Liu et al.) 2023b) in these ReLU-based models.
Crucially, however, recent state-of-the-art LLMs such as Mistral-7B [Jiang et al. 2023) Llama2-7B (Touvron et al. 2023), and Gemma (Team et al. 2024)) employ MLP'block based on more complex nonlinearities that do not inherently induce sparsity Mirzadeh et al. (2023). As such, most of the work on ReLU-based activation sparsity is inapplicable to these models. To the best of our knowledge, ReLUfication the only work that attempts bridge this gap (Mizadde et al. 2023). ReLUfication replaces the SiLU and GeLU activation functions in LLMs with ReLU to induce sparsity. ReLUfication is the primary baseline against which we compare CATS. In contrast with ReLUfication, CATS contains controllable level of sparsity. Furthermore, in Section 5 we demonstrate that CATS demonstrates significantly better downstream task performance and fine-tuning efficiency than ReL.Ufcctiio.
We note that Zhang et al. (2024) concurrent to our work. In contrast with their work, however, our work is not empirical evaluation of existing activation functions. Rather, we propose a new framework for sparsifying LLMs. Our framework utilizes a novel activation function and enables controllable sparsity. We validate the performance of CATS in extensive evaluations and provide a custom GPÚ kernel that translates CATS' sparsity to real wall-clock time gains in Section 5
We discuss additional research areas on LLM efficiency, such as quantization, structure pruning, knowledge distillation, and hardware-aware optimization in Appendix
3

---

Published as a conference paper at COLM 2024
3 Background
Motivation: As described in Section 1 MoE models selectively activate expert subnetworks via trained router. Crucially, we may view the rows (or columns) of MLP matrices as experts in an MoE model. To identify the layers most likely to benefit from this MoE perspective (where many activations can be zeroed), we examine the activations of different layers in LLMs. Figure demonstrates that activations of the Gated-MLP layers tend to concentrate around zero across different LLMs. This behavior suggests that many neurons of MLP layers minimally affect the output in future operations.
We now describe the components of LLMs that our work aims to accelerate: the Gated-MLP blocks. They are commonly used in LLMs, including in the Llama2 family of models, Mistral-7B, and Gemma. A Gated-MLP block consists of several fully-connected layers and performs the following computation:
Gated-MLP Blocks:
Gated-MLP(x) (SiLU gate) (xWup))
(1)
where x E RRbxd, Wgate, Wup E RRdxm, Wdown E Rxd * indicates element-wise multiplica- tion, and
SiLU(x) = x sigmoid(x) 1+
(2)
Crucially, the operation SiLU((Wgate) can be viewed as the router in an MoE model. Through this lens, the columns of Wup and the rows of Wdown are the experts. If (x) always binary, i.e., or 0, it would turn on/ on/off off elements the remaining computation (multiplication by WapWaowal When SiLU((x) is not binary, it can be viewed as "soft" router that weighs the experts by different amounts.
Method: Contextually-Aware Thresholding for Sparsification (CATS)
We now describe CATS, a framework to accelerate the Gated-MLP blocks of LLMs. The CATS framework proposes new, simple activation function and exploits the sparsity induced by this activation. In Section 5, we apply CATS to Mistral-7B and Llama2-7B and show that CATS-based models still exhibit significant activation sparsity, even when fine-tuned.
4.1 Stage 1: Determining Cutoff Thresholdd
We assume we are given a desired sparsity level k (e.g., 70%) as input. For each Gated-MLP block in the LLM, compute the activations over random subset of the training data, limited to only 500 data points. We then compute the cutoff threshold as the kth percentile of the resulting values. Concretely, the cutoff threshold f
==
2 k)
(3)
where F represents the empirical cumulative distribution function of the absolute values of the activations for the given MLP block.
Figure shows histograms of the absolute values of activations for different MLP blocks
in various models on the RefinedWeb dataset [Penedo et al. 2023). A sparsity level of 70% corresponds to threshold of approximately 0.15; different sparsity levels correspond to different thresholds. We note that these thresholds are chosen and fixed before any further fine-tuning.

---

Published as a conference paper at COLM 2024
4.2 Stage 2: Sparsifying Gate-MLP Blocks
Given the cutoff threshold i 2 0 corresponding to the input sparsity level k, we wrap the SiLU(x) activations in each MLP block with the CATS activation. The CATS operation, denoted CATS, i), is defined as:
Xi. 00
x 2 ! if x <
(X)S =
(4)
Here, is the sparsification threshold and xj is the j-th element of the vector respectively. This results in new activation (SIUU)):
SiiLU(xWgate
(xWgaet 2 (iLLU gate) <
CATS (SiLLU gatte))
(5)
0
Intuitively, the resulting model zeros out activations that are likely to be close to 0 because their corresponding inputs were small. This procedure produces a trained model with sparse activations, whose performance can then be evaluated. We empirically validate that this procedure results in a model with an approximate sparsity level of k even after fine-tuning, detailed in Appendix
4.3 Custom Kernel Design
The previous subsections describe the procedure for sparsifying LLM's activations, obviating unnecessary computations, and reducing the required number of floating point operations (FLOPs) in each MLP block. We now translate this reduction in FLOPs to reduction in actual wall-clock latency and increase in generation throughput via a custom GPU kernel.
We focus on reducing the latency of the MLP blocks by reducing memory accesses, as the MLP blocks known to be memory-bound during inference when the batch size is small (Kim et al. 2023 As shown in Line 4 of the Custom GPU Kernel [ we first fuse the element-wise multiplication of v[Mask] into each tiling of x Wup [Mask]. Here, represents the hidden vector after applying the SiLU ac- tivation function, and Mask denotes a binary mask identifying elements of with large absolute values. This fusion saves memory oper- ations that would be necessary for storing and loading several times. We then directly use Mask to control which parts of the weight matrices Wup and Wdown to load, instead of using the compressed indices directly as in Zhang et al. (2023) This further improves the kernel speed because it avoids expensive synchronization operations. In Section 5.3 demonstrate how custom GPU kernel effectively reduces the inference latency of CATS-based models as sparsity increases.
Custom GPU Kernel 1 MLP using CATS
1: Input: threshold ! > 0, hidden layer x, weights Wgater Wdownr and Wup
2: < SiLU gate)
3: Mask t- if v 2 f else 0
4: x1 + (xWup Mask] * v ([Mask)
5: y + x, Wdown Mask]
Experiments
In this section, we describe the experiments with which we assess the performance of CATS. We first describe the experimental details that are common to all experimental settings. We then describe experiments on downstream task performance. Finally, we measure CATS effect on wall-clock time inference when implemented with the custom GPU kernel from Section 4 We find that CATS-based models outperform their ReLUfication versions in downstream task performance, with or without fine- sime-ttning,,,, and can exploit their sparsity for wall-clock inference time speedups over the base models.
5

---

Published as a conference paper at COLM 2024
We first describe the experimental setup, including base models, CATS-based models, metrics, datasets, and computational environment.
Base Models: We apply CATS to both Mistral-7B and Llama2-7B base models to verify it is generally applicable to different LLMs. We measure the performance each CATS model against the original base model. We also compare the performance to of the CATS-based models to the base model transformed by ReLUfication from Mirzadeh et al. (2023).
CATS-based Models: For given base model, we train three CATS-based variants that exhibit different sparsity levels in the MLP blocks: 50, 70%, and 90% activation sparsity. We call these models CATS-50%, CATS-70%, and CATS-90%, respectively, where the base models are clear from context.
Metrics: We compare models using several metrics. In the first set of experiments, we compare each model's accuracy on downstream tasks. In the second set of experiments, we compare each model's wall-clock time inference latency.
For the downstream task performance experiments, we use the OpenBookQA, ARC Easy, Winogrande, HellaSwag, ARC Challenge, PIQA, BoolQ, and SCI-Q datasets from the Eleuther AI Evaluation Harness (Gao et al. 2023) as in Mirzadeh et al. (2023) for ease of comparison; these tasks were originally chosen measure various abilities of the models across various domains, such as reading comprehension and reasoning. For the latency and generation task experiments, we assess the wall-clock inference time and perplexity score on the Ref test dataset (Penedo et al. 2023).
Datasets:
Computational Environment: All experiments were run on single machine with 8 L40S GPUs. Latency experiments were run single L40S GPU each 7B base model was able to fit in a single GPU RAM when performing inference in brain float 16 (BF16) or floating point 16/32 (FP16/33 /2) precision. We used DeepSpeed (Rasley et al. 2020 with BF16 precision to manage the high memory overhead during training. We also employed Low-Rank Adaptation (LoRA) (Hu et al., 2021) and targeted 1% of the parameters (Query and Key in attention modules, Wgaate and dowwn) in the fine-tuning experiments. During inference, we used the transf ormers v4. 36. 2 HuggingFace library, PyTorch v2.1. and CUDA v12.1. We used Triton v2.1 for our GPU kernels. All experiments were run in FP32 precision; changing this to FP16 did not materially affect results. All of our code, including a one-line script to set up an environment and reproduce all of our results, is available the supplementary material.
5.1 Downstream Task Performance
We now compare the downstream task performance of the CATS-based models to the baseline models in several settings and draw several conclusions.
CATS-based models perform comparably to the base models and outperform ReLUfi- cation in zero-shot accuracy without any fine-tuning: We first compare the performance of CATS-based models to the baseline models without any fine-tuning. In this setting, the CATS prescription is applied directly the base models, i.e., the activation functions are simply replaced in the MLP blocks and no fine-tuning is performed. Table | shows our results across 8 different benchmark tasks. CATS-based models demonstrate performance comparable to the unchanged, out-of-the-box base models, even at high sparsity levels. In particular, CATS 50% demonstrates performance comparable the base model. CATS sig- nificantly outperforms ReLUficiation in downstream task performance at the same sparsity level (90%).
CATS-based models perform more comparably to the base model and increasingly out- perform ReLUfication as the model size increases in zero-shot accuracy: Similar the previous evaluation, we assess the performance of various CATS configurations and baseline Llama-13B model across eight downstream tasks without any fine-tuning. Notably, the performance degradation of CATS-50% decreased from 1.46% at 7B to 0.65% at 13B. Conversely, the performance gap between CATS-50% and ReLUfication increased from 29.08% at 7B to 34.55% at 13B. We present the detailed results in Appendix

---

Published as a conference paper at COLM 2024
Model
Dataset
WG ace?
PIQA
SciQ
QA
HS
BoolQ
Are-E
Are-C
Avg
acct
acct
acct
acc?
ace
acct
accf
acc?
Mistral-7B
0.7419
0.8069
0.959
0.3260
0.6128
0.8370
0.8085
0.5034
0.6994
50%
0.7245
0,948
0.3200
0.6097
0,8193
0.5043
CATS 70%
0.7190
0.8003
0.929
0.292
0.6057
0,8028
0.7492
0.4693
0.6709
CATS 90%
0.5627
0.6001
0.422
0.212
0.3359
0.7086
0.3754
0.2773
0.4388
ReL.Ufiatiin
0.5043
0.5092
0.236
0:142
0.2580
0,4208
0,2723
0.2415
0.3230
Llama2-7B
0.6906
0.7807
0.94
0.314
0.5715
0.7774
0.7630
0.4343
0.6589
CATS 50%
0.6748
0.7693
0.927
0.322
0.5711
0.7263
0.7441
0.4121
0.6433
70%
0.294
0.5500
0.6590
0.7008
0.3805
0.6143
CATS 90%
0.5738
0.6627
0.611
0.212
0.3848
0.6284
0,4566
0.2816
0.4764
Rel.Ufication
0.4893
0.5408
0.2570
0.154
0.2586
0.6003
0.2795
0.2406
0.3525
Table 1: Downstream task performance of base models, CATS-based models at varying levels of sparsity, and ReL.fication across different benchmark datasets. The CATS versions of the base models demonstrate comparan comparable performance to original models at high sparsity levels, e.g., 50%. At higher sparsity levels, CATS still outperforms ReLUfication.
0.904
BAR0 3RELT BOO EEGL DAME GEAL 0ARO 20R1 1003 0.051 DARM
25g 10m
0,440
39641
8554
06505
DANI 36255
onees
(1008 0dd0 p011 3433 p010 1339
6.00
13502
5602
6.305
65399
0.903
33561 3500
3545
0.695
CATS 75%
CATS 70%
CAT5 92%
CXT5 60%
keliffrorica
33552
Protur Stape Eaarch
Proto 9 Sheps
Figure 2: Downstream task performance of the base model, CATS models with different sparsity levels, and ReL.Ufication across varying numbers of fine-tuning steps on the Re- Web dataset applied to Mistral-7B (left) and Llama2-7B (right). The CATS models exhibit faster convergence and greater fine-tuning efficiency than the ReLUfication variants. Furthermore, CATS-50% and CATS-70% demonstrate comparable performance to the base models without any fine-tuning (0 fine-tuning steps).
CATS-based models perform comparably to the base models and outperform ReLUfica- tion in zero-shot accuracy with "general" fine-tuning: In this setting, CATS is applied the base models Llama-7B and Mistral-7B. All models are then fine-tuned using LoRA (Hu et al., 2021), targeting only 1% of the parameters, on the RefinedWeb dataset (Penedo dataset[Penedo et al. (2023). Their downstream performance is subsequently measured across 8 evaluation datasets.
We emphasize that the dataset upon which the models are fine-tuned is different from the evaluation datasets in this setting. Figure demonstrates our results. We note several key observations:
1. CATS-based models still exhibit sparsity after fine-tuning (see Appendix C
2. CATS-50% demonstrates performance comparable to the base models, even without any fine-tuning. This is in contrast with ReLUficiation, which demonstrates poor performance without fine-tuning.
3. CATS-50%, CATS-70%, and CATS-90% all display better task performance than ReL.Ufccttin when controlling for the number of fine-tuning steps. In particular, even with very few fine-tuning steps, the CATS-based models achieve comparable performance to the base models.
4. CATS-based models, even with sparsity levels high 70%, achieve performance comparable to the base models within 500 steps of fine-tuning, whereas ReLUfica- tion does not.
7

---

Published as a conference paper at COLM 2024
CATS-based models perform comparably to the base models and outperform ReLUfica- tion with task-specific fine-tuning but without "general" fine-tuning: In this setting, the CATS prescription is applied to Mistral-7B. All variants are then fine-tuned for 10 epochs the training data and evaluated on test dataset for the Cola, SST2, and BoolQ datasets. Table 2 demonstrates results. Our observations are similar to those for "general" fine-tuning:
1. CATS-based models still exhibit sparsity after fine-tuning (see Appendix
2. CATS-50% demonstrates performance comparable to the base models. This is in contrast with ReLUficiation, which demonstrates a significant performance degradation.
3. CATS-50%, CATS-70%, and CATS-90% all display better task performance than ReLUfication.
Dataset/Sparsity
Base Model
0.5
0.7
0.9
ReL.Ufication
Cola
0.8667
0.8658 (-0.10%)
0.8552 (-1 (-1.3%%)
0.8303 (-4.21%)
0.6922 (-20.13%)
SST2
0.9644
0.9656 (+0.12%)
0.9702 (-0.60%)
0.9427 (-2.25%)
0.7856 (-18.55%)
BoolQ
0.8905
0.8862 (-0.48%)
0.8807 (-1.00%
0.7920 (-11.06%)
0.6624 (-25.61%)
Average 0.9072
0.9059 (-0.13%)
0.9020 (-0.52%)
0.8550 (-5.22%)
0.7134 (-19.38%)
Table 2: Downstream task performance of Mistral-7B and its CATS-based and ReLUfication variants across three different benchmark datasets. Top accuracies are marked in bold and second-highest in underline. Relative performance degradation is given in parentheses. CATS-50% demonstrates performance within 0.5% of the base model, whereas ReLUfication demonstrates a significant performance drop.
5.2 Generation task performance
We now evaluate the generation task performance the CATS-based models using the RefinedWeb dataset. Without any fine-tuning, CATS-50% exhibits only slight increase in perplexity scores. The difference in perplexity scores between the base models and CATS-50% further diminishes after 1300 fine-tuning steps.
Specifically, for Llama-7B, the perplexity increases by just 1.06% when comparing the base model to CATS-50% without fine-tuning, and this difference narrows to 0.60% with minimal fine- tonituning. Similarly, for Mistral-7B, the initial perplexity increase 1.21% is reduced
mere 0.45% after fine-tuning.
Model
Base Model
CATS 50%
CATS 50%*
Rel.Ufication
Rel.Ufication"
Llama-7B 2.18
2.203 (+1.06)
2.193 (+0600%
3.104 (+42.38%)
2.617 (+20.04%)
Mistral-7B 2.234
2.261 (+1.21%)
2.244 (+0.45%)
2.801 (+25.34%)
2.562 (+14.68%)
Table 3: Perplexity scores for different models and configurations. The best score is marked in bold and the second-best is underlined. Percentage differences relative to the base model are shown in parentheses. Models marked with are fine-tuned for 1300 steps.
5.3 Wall-clock Time Speedups for Inference
Activation sparsity in model is not sufficient to directly achieve wall-clock time inference speedups (Frantar & Alistarh 2023). In this subsection, we demonstrate that our custom GPU kernel translates the activation sparsity induced by CATS into tangible wall-clock time improvements.
CATS-based models translate their activation sparsity to wall-clock time speedups:
Figure shows the wall-clock inference time of the dense model compared to CATS imple- mented via the custom GPU kernel described in Section 4.3 for various sparsity levels of
8

---

Published as a conference paper at COLM 2024
0.9
1.1
1.0
0.8
0.9
0.7
Latency
0.8
Latency(ms) (ms)
0.6
0.7
0.5
0.6
0.5
0.4
Dense
Dense
0.3
Optimal
Optimal
0.3
CATS-wiit-CCustom-Krrnd
CATS-wiith-Custom-
0.2
0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8
0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8
Sparsity
Sparsity
(a) Mistral-7B
(b) Llama2-7B
Figure Latency of the original Mistral-7B MLP block (left, "Dense"), Llama-7B MLP block (right, "Dense"), and their CATS-based variants at different sparsity levels, compared to "Optimal." " Our custom GPU kernel improves the latency the CATS-based variants and achieves performance close to "Optimal" for most reasonable sparsity levels.
Generation Throughput (Tokens/s)
Generation Throughput (Tokensís)
25
20
15
5
Dense
Dense
CATSSwwith-Custom-keee
WATSSwith-CCsstm--ane
256
512
1024
2048
256
512
1024
2048
Generation Length (Tokens)
Generation Length (Tokens)
(a) Mistral-7B
(b) Llama2-7B
Figure 4: Throughput of Mistral-7B (left, "Dense") and Llama2-7B (right, "Dense") and CATS-50% with the custom GPU kernel. CATS-50% demonstrates significantly higher throughput.
CATS. We evaluate the latency of a single MLP block and the throughput of the genera- tion stage of the end-to-end inference. Mistral-7B contains 32 MLPs with m 14336 and d = 4096, and Llama2-7B contains 32 MLPs with m = 11008 and d 4096 (m and d are defined after Equation
with the dense with Matense m ("Dense") and the dense ML.P with Moptimal m * Sparsity ("Optimal"), the latter of which a proxy for the best wall-clock time we could hope to achieve. At 50% and 70% sparsity, the sparse kernel achieves approximately 40% and 70% speedup, respectively, over the original dense MLP. Latency measurements were obtained by conducting 20 rounds of warmups, repeating the kernel 80 times, and computing the geometric mean of the latency across each round. The comparison with Dense demonstrates that sparse kernel consistently outperforms the original ML.P. The comparison with
In Figures 3a and 3b we compare our method C"CATS-with-CCutoo-Koem

---

Published as a conference paper at COLM 2024
Optimal shows that our sparse kernel closely approaches the Optimal performance at low sparsity levels. However, as the sparsity level increases, the gap between our performance and Optimal widens, as expected. We note that our sparse kernel performs the same number of memory accesses as the Optimal; however, due to differences in access patterns, the methods result in different wall-clock time measurements. Optimal can perform worse than our sparse kernel when Moptinul does not match the shapes optimized by GPU libraries (Tillet & Cox 2017). Conversely, our sparse kernel can underperform compared to Optimal when the overhead of operations on zero values outweighs the benefits of reduced memory access.
In Figures 4a and 4b we compare dense models with CAISS Kerne (50% sparsity) on the throughput of the generation stage. The generation stage (or "decoding" stage) is known to be memory-bound (Kim al. 2023), which suggests CATS can improve inference througput. We test the generation throughput at batch size of 1 and beam width of 1, and record the latency from the first generated token to the last token. The throughput is calculated by the generated length divided by latency. The final throughput is averaged (geometric mean) over 50 samples from the RefinedWeb test dataset. CATS can accelerate the generation stage by 18% for Llama2-7B and ~21% for Mistral-7B at 50% sparsity.
Though we only test on Huggingface (Wolf et al. 2020), our methodology is orthogonal to the framework and thus can be used in other LLM serving systems such DeepSpeed (Rajbhandari et al. 2022) and TensorRT-LLM NNvidia 2024).
6 Discussion and Conclusion
We presented CATS, a novel framework for inducing and exploiting activation sparsity in L.Mss. At the heart of our framework the CATS activation, given in Equation that induces a controllable level of activation sparsity in LLMs. We also provide. custom GPU kernel implementation that exploits CATS's sparsity to achieve real wall-clock time gains in inference latency.
CATS-based models demonstrate downstream task performance comparable to unmodified base models and better than baseline models with no fine-tuning, even at sparsity levels as high as 50%. CATS-based models also exhibit better behavior than ReL Uficatiin at similar levels of fine-tuning, and often achieve performance comparable to the base model at high levels of sparsity, both with general and spec-fpec fine-tuning.
Limitations and Future Work: Our work leaves several opportunities for future work. Most importantly, our empirical evaluations of CATS were restricted to the Mistral-7B and Llama2-7B base models. While we suspect CATS would also apply to other, larger models, we leave a precise empirical study to future studies. Future work may also investigate how to apply techniques similar to CATS to other MLP architectures beyond GLPed-MLP, or to attention layers but without task performance degradation. It may be possible, for example, use recent techniques accelerate attention layers (such those from Zhang al. (022aa) and Voita et al. (019) in conjunction with CATS.
10

---

Published as a conference paper at COLM 2024
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. Xiv preprint arXiv:2303.08774 2023.
Haoli Bai, Wei Zhang, Lu Hou, Lifeng Shang, Jing Jin, Xin Jiang, Qun Liu, Michael Lyu, and Irwin King. Binarybert: Pushing the limit of bert quantization. arXiv preprint arXio:2012.15701, 2020.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choromanski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, et al. Rt-2: Vision-language-action models transfer web knowledge to robotic control. arXi preprint ar arXit:2307.15818, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Tri Dao. FlashAttention-2:: Faster attention with better parallelism and work partitioning. 2023.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances Neural Information Processing Systems, 2022.
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, 0000033883033222 2022.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized Iims, 2023.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre- training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 2018.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of- pp. 5547-5569. PMLR, 2022.
In International Conference on Machine Learning,
William Fedus, Jeff Dean, and Barret Zoph. A review of sparse expert models in deep learning. Xiv preprint ar arXii:2209.0166: 2022a.
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-39, 2022b.
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-39, 2022c. ISSN 1533-7928. URL httt:://jmm hwwpjill...or///eeeeeeeeeeeeeee. wwtgg/apprs//v2//// 0998. html
Elias Frantar and Dan Alistarh. Sparsegpt:
Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pp. 10323-10337. PMLR, 2023.
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate
post-training quantization for generative pre-trained transformers. ar arXio preprint arXiv:2210.17332,, 2022.
Daniel Y. Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré. Flashfftconv: Efficient convolutions for long sequences with tensor cores, 2023a.
11

---

Published as a conference paper at COLM 2024
Yao Fu, Hao Peng, Litu Ou, Ashish Sabharwal, and Tushar Khot. Specializing smaller language models towards multi-step reasoning, 2023b.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. Megablocks: Efficient sparse training with mixture-of-experts Proceedings of Machine Learning and Systems, 5, 2023.
Leo Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noachh, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL hwtp///www...ggrrrcoo//////////0
Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdhery, Maheswaran Sathiamoorthy, Yihua Chen, Rahul Mazumder, Lichan Hong, and Ed Chi. Dselect-k: Differentiable selection in the mixture of experts with applications to multi-task learning. Advances in Neural Information Processing Systems, $4:29335-293411 2021.
Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning:
Pruning and growth for efficient inference and training in neural networks. Journal of Machine Learning Research, 22(241):1-124, 2021.
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. ar preprint ar Xir::2066009655 2021.
Sebastian Jaszczur, Aakanksha Chowdhery, Afroz Mohiuddin, Lukasz Kaiser, Wojciech Gajewski, Henryk Michalewski, and Jonni Kanerva. Sparse is enough in scaling trans- formers. Advances in Neural Information Processing Systems, 34:9895-9907, 2021.
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arv preprint ar Kio:2401.04088 2024.
Sehoon Kim, Coleman Hooper, Thanakul Wattanawong, Minwoo Kang, Ruohan Yan, Hasan Genc, Grace Dinh, Qijing Huang, Kurt Keutzer, Michael W Mahoney, et al. Full stack optimization of transformer inference: survey. ar Xiv preprint ar arXin:2302. 14017, 2023.
Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of-experts from dense checkpoints, ar Xiv preprint ar xixi:22212.05055, 2022.
Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, Minh- Thang Luong, and Orhan Firat. Beyond distillation: Task-level mixture-of-experts for efficient inference. Xiv preprint arXiv:2110.03742, 2021.
Eldar Kurtic, Daniel Campos, Tuan Nguyen, Elias Frantar, Mark Kurtz, Benjamin Fineran, Michael Goin, and Dan Alistarh. The optimal bert surgeon: Scalable and accurate second- order pruning for large language models. Xiv preprint ar xrXiv:2203.07259, 2022.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arv preprint ar XrXiv:2006.166 16688, 2020.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning, pp. 2265-6274 PMLR, 2021.
12

---

Published as a conference paper at COLM 2024
Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank
Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, et al. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. arXiv preprint ar arXio:2210.06313, 2022.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. Xrv preprint ar Xiw::230000188000 2023a.
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning, pp. 22137-22176. PMLR, 2023b.
Yuxuan Lou, Fuzhao Xue, Zangwei Zheng, and Yang You. Cross-token modeling with conditional computation. ar preprint arXiv:2109. 02000, 2021.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning of large language models. Xiv preprint drXio:2305.11627, 2023.
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. arviv preprint arXiv:2310.0455 2023.
Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Multimodal contrastive learning with limoe: the language-image mixture of experts. Advances in Neural Information Processing S 35:9564-957,, 2022.
Nvidia. Tensorrt-llm, 2024. URL UtLpapteee///vadaaa.......o/e//r
Haojie Pan, Chengyu Wang, Minghui Qiu, Yichang Zhang, Yaliang Li, and Jun Huang. Meta-kd: A meta knowledge distillation framework for language model compression across domains. arXiv preprint Xiv:2012.012666, 2020.
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 preprint Xiv:2110 103500, 2021.
Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay, The refinedweb dataset for falcon IIm: outperforming curated corpora with web data, and web data only. aiv preprint Xrx::2306.00111 2023.
Markus N. Rabe and Charles Staats. Self-attention does not need (n22) memory, 2022.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Netwworking, Storage and Analysis, pp. 1-16. TEEE, 2020.
Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Advancing mixture-
Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed of-experts inference and training to power next-generation ai scale. In International Conference on Machine Learning, pp. 18332-18346. PMLR, 2022.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. pp. 3505-3500 2020.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, et al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34:17555-17566, 2021.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture- of-experts layer. aixiv preprint ar Xiv::1700 2017.
13

---

Published as a conference paper at COLM 2024
Mingjie Sun, Zhuang Liu, Anna Bair, and Zico Kolter. A simple and effective pruning approach for large language models. arxio preprint arXiv:2306.11695, 2023.
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. ar preprint ar arXiv:1908.09355 09355, 2019.
Siqi Sun, Zhe Gan, Yu Cheng, Yuwei Fang, Shuohang Wang, and Jingjing Liu. Contrastive distillation intermediate representations for language model compression. arXiv preprint arXi::2009 2020.
Filip Szatkowski, Bartosz Wójcik, Mikolaj Piórczyríski, and Kamil Adamczewski. Sadmoe: Exploiting activation sparsity with dynamic-k gating. arXiv e-prints, pp. arXiv-2310, 2023.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv preprint Stir:2403.08295,, 2024.
Philippe Tillet and David Cox. Input-aware auto-tuning of compute-bound hpc kernels. In Proceedings of the international conference for high performance computing, networking, storage and analysis, pp. 1-12, 2017.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brían Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, lliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi- head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. Xiv preprint ar xixx:1905.005 09410, 2019.
Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. arv preprint ar 11X1::191. 2019.
Peter West, Chandra Bhagavatula, Jack Hessel, Jena D. Hwang, Liwei Jiang, Ronan Le Bras, Ximing Lu, Sean Welleck, and Yejin Choi. Symbolic knowledge distillation: from general language models to commonsense models, 2022.
Samuel Williams, Andrew Waterman, and David Patterson. Roofline: an insightful visual performance model for multicore architectures. Communicatioos of the ACM, 52(4):65-76, 2009.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: State- of-the-art natural language processing. In Proceedings of the 2020 conference on empirical methods natural language processing: system demonstrations, pp. 38 38-45, 2020.
Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models. arxiv preprint ar arXiv:2204.00400 2022.
Zhongming Yu, Guohao Dai, Shang Yang, Genghan Zhang, Hengrui Zhang, Feiwen Zhu, June Yang, Jishen Zhao, and Yu Wang. Hypergeft A framework enabling efficient fusion for hypergraph neural network on gpus. Proceedings of Machine Learning and Systems, 5, 2023.
14

---

Published as a conference paper at COLM 2024
Ofir Zafrir, Ariel Larey, Guy Boudoukh, Haihao Shen, and Moshe Wasserblat. Prune once for all: Sparse pre-trained language models. arXiv preprint ar Xxio:2111.007544 2021.
Genghan Zhang, Yuetong Zhao, Yanting Tao, Zhongming Yu, Guohao Dai, Sitao Huang, Yuan Wen, Pavlos Petoumenos, and Yu Wang. Sgap: towards efficient sparse tensor algebra compilation for gpu. CCF Transactions on High Performance Computing, 5(2):210- 227, 2023.
Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. arXiv preprint #/Xiv:2210005144, 2022a.
Zhengyan Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. MoEfica- tion: Transformer ferwaforward Layers are Mixtures of Experts. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Findings of the Association for Compula- tional Linguistics: ACL 2022, pp. 877-890. Association for Computational Linguistics, org/ 2022
2022b. doi: 10.18653/ 10.18653/v11202222 http:///1/2222..nnnneeeeee
URL
Badiings-acl.71
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 Ilms. arXiv preprint arXie:2402.03884, 2024.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing, Advances in Neural Information Processing Systems, 35:7103-7114, 2022.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. Xiv preprint ar Xrix:2202..8800, 2022.
Simiao Zuo, Xiaodong Liu, Jian Jiao, Young Jin Kim, Hany Hassan, Ruofei Zhang, Tuo Zhao, and Jianfeng Gao. Taming sparsely activated transformer with stochastic experts. ar preprint ar arXio:2110.04266 2021.
A Additional Related Work
In this appendix, we discuss additional veins of related work.
Hardware-Aware Optimization that relies on customizing the algorithm implementation for the underlying hardware can result in significant performance speedup Dao et al. (2022);
Fu et al. (2023a), especially for sparse kernels Gale et al. (2023) Yu et al. (2023). Recent hardware-aware methods LLMs have shown to be highly effective in lowering the cost of
attention operation Rabe & Staats 2022 Dao 2023); Liu et al. (2023a). Similar to attention operation, MLP is also memory-bounded
bounded on highly parallel machines like GPU Kim et al.. (2023). The sparsity has the potential to expedite MLP because can increase the arithmetic intensity. Based on the Roofline analysis Williams et al. (2009), higher arithmetic intensity means shorter wall-clock time for memory- -bounded operations. In this work, we focus on leveraging the sparsity to reduce the memory transfers associated with the MLP weights. We do so by designing algorithmic optimizations that adaptively induce sparsity and implementing hardware-aware optimizations that translate the achieved nominal sparsity into actual wall-clock time speedup.
Structural pruning techniques induce sparsity by setting certain weights to zero so their corresponding activations need not computed Wang et al. (2019); Kurtic et al. (2022); Xia et al. ) Zafrir et al. (2021); Ma et al. (2023). However, applying such techniques naively may not result in actual wall-clock time speedups if the resulting sparsity pattern does not lower the number of General Matrix Multiplication (GEMM) calls. Furthermore, the pruning pattern is determined at the model level and is not adaptive to the inputs, which may result in degradation in task performance.
15

---

Published as a conference paper at COLM 2024
Quantization and Knowledge Distillation from larger models to smaller models are other popular forms of LLM inference optimization [ai et al. (2020); Frantar et al. (2022) Dettmers et al. (2023); Sun et al. (2019) P et al. (2020); West et al. (2022 Fu et al. (2023b). These methods often reduce the memory and computational complexity at the cost of performance degradation require extensive finetuning. Our work can be applied to quantized or distilled models as well, although the achieved sparsity level on these models may differ.
B
Accelerating Attention Layers
B.1
Method
In this section, discuss how apply CATS to reduce the inference costs of attention layers inside Transformer blocks. The basic operations of a Transformer block can be written as:
(Attentttttoonn
(6)
where is the hidden vector right before the i-th layer and where have excluded operations like batch normalization, positional embedding, residual connections, etc. for simplicity. (For more details the variants of Attention layers and those used in our
models, we refer the reader to Touvron et al. (2023) and Jingg et al. (202)).)
The new equation for i-th transformer layer, where we wrap the previous layer with CATS activations, becomes:
(CAT(S (AATTS (Atenntion(CAATS, (x)))
(7)
where f and f,2 are the sparsification thresholds for the CATS operations applied before the MLP and attention layers, respectively, in the i-th transformer layer.
We verify that this operation results in sparse activations in Appendix
B.2 Experimental Results
CATS can also be applied to accelerate the attention blocks of LLMs: We also apply CATS to accelerate the computation of attention layers. Our approach inspired by "Stage 2" of ReLUficiation [Mirzadeh et al. 2023).
Due to space constraints, we only measure the performance CATS-50% applied the base Mistral-7B model and measure zero-shot task performance. We fine-tune both models for 2000 fine-tuning batches of 16 examples each. Stage 2 CATS, which appplies CATS to both the MLP and Attention blocks, demonstrates average downstream task performance of 66.84% across the 8 different evaluation tasks from Section 5 whereas the base Mistral-7B model demonstrates an average task performance of 69.94%. In contrast, the original CATS, applied only to the MLP layers, demonstrates an average task performance of 69.21% %.
Our results demonstrate that CATS also be applied to the attention layers of LLMs, albeit with a slight (4.3% relative) performance degradation. Future work may investigate how to apply CATS in way that better preserves the performance of the model.
C Target sparsity vs. actual sparsity
Figure 5
demonstrates the the sparsity of each layer of Mistral-7B and Llama2-7B after CATS has been applied and fine-tuning has been performed on the RefinedWeb dataset.
The average sparsity of each model (dashed lines) is roughly equal to the target sparsities (indicated by the legend).
Table 4 demonstrates the average layer sparsity of each model after task-specific fine-tuning on the datasets used for this experimental setting in Section The observed sparsity levels are approximately equal to the target sparsity levels.
16

---

Published as a conference paper at COLM 2024
Jurrup Sandy
029 hand.
(a) Sparsity of Mistral-7B.
(b) Sparsity of Ilamma-7B.
Figure DAsS-based models still exhibit sparsity after general fine-tuning on the Refined- Web dataset.
Dataset/Sparsity
0.5
0.7
0.9
Cola
49.629
68.926
87.6
BoolQ
49.196
68. 444
87.571
SS12
48.727
68.738
87.882
Average
49.184
68.703
87.684
Table 4: CATS-based models' final sparsity after specific fine-tuning each task. They continue to exhibit sparsity after task-specific fine-tuning.
Future work might focus on enforcing a minimum sparsity layer-wise, ie., by zeroing out at least enough neurons to enforce the desired sparsity level for each layer. Such work could investigate the tradeoffs between sparsity, latency, and downstream task performance.
D Details on Custom GPU Kernel Design
The previous subsections describe the procedure by which we sparsify the activations of an LLM, obviate some computations, and reduce the required number of FLOPs. Though significant recent work has focused on FLOPs as proxy for inference cost, other work has demonstrated that reducing FLOPs is not sufficient to reduce real wall-clock inference latency Liu et al. (002bb). However, predictable sparsity patterns can be exploited to reduce floating point operations (FLOPs) during inference. We now translate the reduction in FLOPs to an actual wall-clock latency reduction via several custom GPU kernel optimization techniques.
The operations of the Gated-MLP with the CATS activation functions are:
CATS(SILU
(8)
Mask [0/> (elementwise)
(9)
= (o' (x
(10)
where v', Wlp and Woown are V, Wup, and Waown masked by Mask (for the matrices Wup and Wdownr the entire column j is 0 Mask, = 0, i.e., the mask is broadcast across columns).
Mask is sparse, then Equation 110) performs two sparse matrix multiplications. In fact, only coordinates (respectively, rows) of a (respectively, Wup and Wdown) corresponding to nonzero coordinates Mask need to be loaded into memory. Since the MLP layer at inference time is known to be memory-bound Kim et al. (2023), the latency be reduced
17

---

Published as a conference paper at COLM 2024
if the memory access is reduced. We exploit these observations to translate the reduction in FLOPs to a real wall-clock time reduction in inference.
Custom GPU Kernel 2 MLP using CATS
Custom GPU Kernel 3 MLP using CATS without atomic operations
1: Input: threshold > hidden layer x,
weights Wgate, down/ and Wup
1: Input: threshold > 0, hidden layer
2: v f CATS (SiLU(xWgatt)
weights Wgate, Wdown/ and Wup
3: Mask t 1 |)] 2 f else O
v < CATS (xWate)
4: idcs indices where Mask
Mask 1 p 2 f else
5: <- (xWupp (idcs] o[idcs])
4: f (xWup Mask] v o[Maak)
6: y < X1 Wdown [dccs
5: y < X1 down [ask]
describe Equations and contain several optimizations.
Algorithms and
8 in lower-level pseudocode. Algorithms
Optimization 1:
We fuse the element-wise multiplication of vi idcc] into each tiling of xWup ides) shown in Line of Algorithm 2
We use an efficient algorithm from Deja Vu Vu et al. (2023b) to compute x1 (idss without the element-wise multiplication by a[idcs]. In this manner, we fuse several operations and save the memory operations for storing and loading x1 several times.
The atomic operations in Line 4 of2 however, introduce extra overhead. Line 4 compresses a one-hot mask to a compressed coordinate array and requires atomically appending to the idcs. GPUs, however, cannot efficiently perform such atomic operations because of their massively parallel nature.
We therefore introduce another optimization in Algorithm to reduce the memory loading incurred by the atomic operations. In Algorithm we directly use Mask to control which parts of weight matrices to load, instead of the condensed idcs. Algorithm| has more operations than Algorithm| because it directly assigns the unloaded elements to instead of squeezing out the zero values before computation. Algorithm 3 does not skip the zero operations in a fine-grained way because the sparsity in this problem is not asymptotically high Zhang al. (2023), which means the operation reduction does not compensate for the performance loss caused by complex control logic. Figure 6 the ablation experiment results
Optimization 2:
1.4
1.0
1.5
1.1
1.0
Latenofon
6.7
0.8
0.7
8.5
0,5
Dans
Optinal
0.4
Optinal
Desse
0,4
Cuth Kome!
CATS
0,2
Baseire
0.2
0.0
0.1
0.2
0.5
0.4
0.5
0.7
0.8
0.0
0.5
0.2
0.4 Sparsity
0.5
0.6
0.8
Sparsity
(a) CATS of Mistral-7B's MLP.
(b) CATS of Llama2-7B's MLP.
Figure 6: Ablation study on kernel optimizations.
18

---

Published as a conference paper at COLM 2024
PyTorch
34
2 import triton. lragon. as tl
35 V
3 V SILU(X e k.gate)
El.1oad nask-flag, other=0.0)
4 Nack torch. abs >- threshold
37 aсc (.to(t1..1oe:32)
5
38 X_1 X_1 rn
6 Triton 1: (x W_up [Nask]] [Mask]
39 tl.stere(X_1, acc. float16), maskern
7 # Input
40
8
41
=_ M_doun [Nask]]
9 V: [X]
42 Input
W_up: DM.00 Stored column-naaoo
43 X_1:
Nack: [M]
44 W_down: [M,3]
Dutput
45 Mask: [M]
X_1: [M]
46
# Tunable paranetera
47 Y: [K]
BLOCK_M 14,3,16,32)
paramsters
16 BLOCK_N N (256,512,1024) Configuration
49 BLOCK_M = 316,32,64,
50 BLOCK
grid (X BLOCK_N.)
51 Configuration
Kernel [grad]
52 grid
// BLOCK_M, // BLOCK_X)
atart_n ti..prgrrn_.de
53 Karnel [grii]
21 xn start_n BLOCK H t1. ar.arange BLOCX_M)
54 start_a tl. pa-prrgrannidOO id(o)
tra.gea BLOCK_N)
55 start_n t1.prooram_
23 Nask Maax
56 start_. BLOCK_M + tl. arange BLOCK_N)
flag (Naek >
57 In start__ BLOCK_N tl. (o, BLOCK_N)
25 4_up W_up (rn(: None]
rn :]]
58 Mask Nask
X rm
59 flag tl. (Mank)
t1..eeroa
(BLOCK dtype=t1.-loat33)
60 W_down W__doun (rm None]
Dione, :])
28 i_naek [lon(; None]
61 X_1 X_1 + xn
range(U, o, -BLOCK]N):
s1. mack-flagl; Neme], other=0.0)
30
t1.1oad(V_up mask-i_nask, other-0.0)
63 [1..oed(X_1)
x #1.1oad(I)
64 acc "tl.eun(s x[;,None]] t).to(t1.ffoss
(u :])one,;), totttl.flsstll
65 Y Y
W_up += BLOCK_K 33
66 El.atonic_ndd(V, acc.to(tl.fleattee)
Figure 7: Triton pseudo-code for Algorithm
E Additional Experiments
We conducted an additional experiment to test whether this trend holds for larger models. We evaluated the Llama 13B model on the same downstream tasks and observed that the performance degradation becomes even more minimal. This leads to the tentative conclusion that CATS becomes more effective for larger models.
Additionally, we compared the performance of CATS with Wanda Sun et al. 20033,, a method that prunes weights per-output basis based on the product of weight magnitudes and input activation norms. As Wanda achieves unstructured sparsity of weights, it is important to note that this approach may not lead wall-clock time improvements in practice.
Model
WG
PIQA
ScQ
QA
HS
BoolQ
Axc-E
Arc-C
Avg
Llama2-13B
0.7230
0.7905
0.9460
0.3520
0.6006
0.8055
0.7946
0.4838
0.6870
CATS 50%
0.7111
0.7884
0.9460
0.3540
0.6037
0.7706
0.7870
0.4812
0.6805
CATS 70%
0.6946
0.7862
0.9310
0.3360
0.5987
0.7419
0.7546
0.4497
0.6616
CATS 90%
0.5596
0.6828
0.5820
0.2080
0.3960
0.6416
0.4726
0,3029
0.4807
ReL.Ufication
0.5522
0.2648
0.4355
0.2757
0.2483
0.3350
Wanda
0.7079
0.7873
0.9450
0.3200
0.5710
0.8122
0.7605
0.4300
0.6667
Table 5: Zero-shot performance of 13B-based models, CATS-based models for varying levels of sparsity, ReLUfication, and Wanda 50% across different benchmark datasets. CATS 50% significantly outperforms other techniques, which suggests CATS scales well with model
size.
19


Reviews:
[{"review_id": "mGYvCuetSO", "reviewer": "Reviewer_pqMf", "summary": "This paper proposes a method to sparsify Gate-MLP Blocks of LLM, which is dubbed Contextually Aware Thresholding for Sparsity (CATS). This work aims to accelerate the Gated-MLP blocks that utilize the SiLU activation function. The resulting model zeros out activations when $|x W_{gate}|$ is smaller than the cutoff threshold.\n\nWithout any fine-tuning, CATS can be applied to various base models, including Mistral-7B and Llama2-7B, and achieves comparable performance on downstream tasks.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewers,\n\nThank you for your thorough review of our submission. Your insights have been invaluable in refining our research. In response to your suggestions, we have conducted additional experiments to further evaluate the effectiveness of CATS. Our key findings are as follows:\n1. CATS 50% preserves quality in generative tasks including perplexity and arithmetic reasoning.\n2. CATS 50% outperforms other baselines across both Llama 7B and Llama 13B models on eight zero-shot tasks.\n3. CATS demonstrates promising scalability, becoming more robust as model size increases.\n\nWe also share an anonymized reproduction script [here](https://shorturl.at/AGmLJ).\n\nBelow, we provide a detailed, point-by-point response to your comments.\n\n---\n\nDear Reviewer pqMf,\n\n[Response to R1]\nWe have conducted extensive comparison experiments to understand the behaviour of CATS and Wanda 50% across different model sizes and various evaluation tasks.\n\nFirst, we compared the performance of different techniques on the 8 evaluation tasks covered in the paper. On Llama 7B, CATS 50% outperformed Wanda in 5 out of 8 tasks (PIQA, OpenBookQA, HellaSwag, Arc-E, and Arc-C), with average scores of 0.6433 and 0.6366, respectively. On Llama 13B, CATS 50% outperformed Wanda in 7 tasks (WinoGrande, PIQA, SciQ, OpenBookQA, HellaSwag, Arc-E, and Arc-C), with average scores of 0.6850 and 0.6667. Notably, CATS 70% achieved an average score of 0.6616, demonstrating that CATS scales well with model size, becoming more robust as the model grows larger.\n\nAdditionally, CATS 50% and 70% outperformed Wanda on the gsm8k (arithmetic reasoning) task. Specifically, the scores were: Llama 7B original: 12.59, CATS 50%: 12.28, CATS 70%: 9.10, CATS 90%: 1.21, and Wanda: 4.24. CATS 50% performs similarly to the original model even without any fine-tuning. This suggests that CATS preserves quality on a generative task, while Wanda significantly loses it.\n\n[Response to Q1]\nThank you for this suggestion. We will clarify this phrasing in the revised manuscript.\n\n[Response to Q2]\nOur approach is a more fine-grained version of the conventional MoE. We use activations as a router to decide which rows of the weight matrix to activate. This provides a new perspective on framing a Gated-MLP block using SwiGLU, offering an opportunity to speed up computations.\n\n[Response to Q3]\nThank you for the detailed review. We will correct this typographical error in the revised version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. CATS can be applied to various base models without any fine-tuning and achieves comparable performance in downstream tasks.\n2. When fine-tuned, CATS surpasses existing state-of-the-art sparsification techniques in downstream task performance, considering the same sparsity level and number of fine-tuning steps.\n3. The proposed method demonstrates good hardware affinity, resulting in approximately a 15% improvement in wall-clock inference latency for token generation compared to dense base models.\n4. The proposed method is characterized by its simplicity and effectiveness.", "reasons_to_reject": "1. It lacks extensive baselines, like Wanda.", "questions_to_authors": "1. \"achieve downstream task performance within 1-2% of their base models \" => \"achieve downstream task performance within 1-2% reduction of their base models \".\n2. I am not sure that the first contribution is solid, which draws connection between the MoE framework and multiplication performed\nby dense matrices in the MLP blocks of LLMs.\n3. \"$W_{up}$\" => \"$W_{down}$\" at the line under Equation (1).", "extra_scores": {}}, {"review_id": "vq88iXrv2S", "reviewer": "Reviewer_kwq1", "summary": "This work introduce sparsity in gated MLP blocks by thresholding the gates. The idea is to determine the threshold values for pruning in each block simply judged by the empirical CDF of absolute activations on a sample of training data. Empirical results on LLMs show that the sparsity does not significant drop the prediction accuracies on various tasks with faster inference speed thanks to the simple trick in kernel design in implementation.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewers,\n\nThank you for your thorough review of our submission. Your insights have been invaluable in refining our research. In response to your suggestions, we have conducted additional experiments to further evaluate the effectiveness of CATS. Our key findings are as follows:\n1. CATS 50% preserves quality in generative tasks including perplexity and arithmetic reasoning.\n2. CATS 50% outperforms other baselines across both Llama 7B and Llama 13B models on eight zero-shot tasks.\n3. CATS demonstrates promising scalability, becoming more robust as model size increases.\n\nWe also share an anonymized reproduction script [here](https://shorturl.at/AGmLJ).\n\nBelow, we provide a detailed, point-by-point response to your comments.\n\n---\n\nDear Reviewer kwq1,\n\nThank you for your detailed review of our submission. We appreciate your feedback and have incorporated your suggestions. \n\n[Response to R1]\nThe way CATS computes the threshold is task-agnostic. We sample a small subset of only 100 examples from a pretraining dataset (RefinedWeb) to compute the thresholds. This is a one-time initialization step. For example, the threshold for CATS, evaluated on eight zero-shot tasks, was computed only once, without prior access to any datasets other than RefinedWeb. This initialization ensures that CATS can effectively operate in zero-shot settings without requiring task-specific tuning. Thank you for allowing us to clarify this important contribution. We will elaborate on this in the revised paper.\n\n[Response to Q1]\nThe efficiency gain is not due to the masking itself, but rather from reducing I/O overhead by only loading and computing on a sparse subset of the weight matrix's rows and columns. This reduction in I/O operations and computations is made possible by our custom kernel, which efficiently handles sparse operations. The result is a significant reduction in latency during inference, leading to faster performance without compromising accuracy.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The proposed method is very simple yet performing quite well on standard benchmarks in terms of the trade-off of prediction accuracies and latencies.\n- The approach sounds reasonable in that it is based on simple statistics on training data.", "reasons_to_reject": "- The threshold has to be computed for each task, and thus, it might not be effective for zero-shot setting in which there exists no data for tuning.", "questions_to_authors": "It is a rather surprise that the simple masking in Section 4.3 yields speed up in inference, but is it mainly because it has the large number of values in a block?", "extra_scores": {}}, {"review_id": "mrjXl7HGhT", "reviewer": "Reviewer_hNTR", "summary": "This paper introduces CATS (Contextually-Aware Thresholding for Sparsity), a new method designed to reduce the inference costs LLMs by increasing their activation sparsity. The authors develop a new nonlinear activation function which is applied to various base models including Mistral-7B and Llama2-7B. They demonstrate that CATS-based models can achieve downstream task performance within 1-2% of their base models any fine-tuning, even at activation sparsity levels as high as 50%. A custom GPU kernel implementation translates the activation of sparsity of CATS to real wall-clock time speedup.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewers,\n\nThank you for your thorough review of our submission. Your insights have been invaluable in refining our research. In response to your suggestions, we have conducted additional experiments to further evaluate the effectiveness of CATS. Our key findings are as follows:\n1. CATS 50% preserves quality in generative tasks including perplexity and arithmetic reasoning.\n2. CATS 50% outperforms other baselines across both Llama 7B and Llama 13B models on eight zero-shot tasks.\n3. CATS demonstrates promising scalability, becoming more robust as model size increases.\n\nWe also share an anonymized reproduction script [here](https://shorturl.at/AGmLJ).\n\nBelow, we provide a detailed, point-by-point response to your comments.\n\n---\n\nDear Reviewer hNTR,\n\n[Response to R1]\nWe recognize this limitation that arose due to limited computational budgets. We have conducted additional experiments to understand the effectiveness of CATS and baseline techniques on Llama2 13B. Following the suggestion of a reviewer pqMf, we included Wanda (50% sparsity) that prunes weights using activations into our baselines.\n\nThe evaluation of eight zero-shot tasks (WinoGrande, PIQA, SciQ, OpenBookQA, HellaSwag, BoolQ, Arc-E, and Arc-C) demonstrates that CATS performs comparably to the original model without any finetuning. CATS 50% surpasses Wanda in 7 tasks out of 8 tasks except BoolQ and outperforms ReLUfication in all tasks. The average scores of Llama2 13B, CATS 50%, CATS 70%, ReLUfication, and Wanda are 0.6870, 0.6805, 0.6616, 0.3350, and 0.6667, respectively. This hints that CATS is likely to scale well with model size, showing a tendency to become more robust as the model grows larger. We will try to further investigate the scaling property of CATS by incorporating a 70B-based model and MoE models in the paper. \n\n[Response to Q1]\nWe ran an additional zero-shot experiment on an instruction-tuned Llama2 7B model to demonstrate the effectiveness of CATS when no fine-tuning data. The eight tasks discussed above are formatted for instruction tuning. We applied the same CATS technique as before. Surprisingly, the performance gap between the base model (average score of 0.6569) and the CATS 50% model (average score of 0.6520) is smaller than that for a non-instruction-tuned model. Additionally, CATS 50% even slightly surpasses the base model in three tasks (OpenBookQA, HellaSwag, and BoolQ) without any fine-tuning. We emphasize that we did not fine-tune CATS on any datasets.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper addresses an important and challenge in deploying LLMs, their significant inference costs.\n2. Extensive experiments have been conducted with different models, demonstrating the general applicability of CATS.\n3. The authors provide an efficient GPU kernel implementation that translates the model's theoretical advantages into practical speed-ups.", "reasons_to_reject": "The empirical assessments of CATS were limited to 7B base models, leaving its applicability to larger models or, more critically, instruction-tuned models, an open question. Furthermore, exploring its effectiveness on MoE models would be valuable, given their rising popularity.", "questions_to_authors": "How to adapt a instruction-tuned model? It would sometimes be challenging if the finetuning data is not available.", "extra_scores": {}}, {"review_id": "NtR53rYdZK", "reviewer": "Reviewer_UDEW", "summary": "This paper proposes a novel technique called CATS that is motivated by the fact that LLM FFN activations usually have a large portion of close-to-zero values. CATS uses some data to calibrate for a threshold that will lead to a desired level of sparsity, and then at inference time zero-out the activations with magnitudes below threshold. This can be implemented as an efficient GPU kernel to bring sparsity into actual speed ups. When CATS are applied to pretrained LLMs such as Llama2 7B and Mistral 7B, we get little loss of performance without further training at 50% sparsity, and with general/task-specific finetuning, CATS can recover some lost performance at even higher sparsity and outperform ReLUfication baselines.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer UDEW,\n\n[Response to R1]\nWe conducted additional studies to evaluate CATS's performance on generative tasks, including perplexity and arithmetic reasoning. Our findings show that CATS performs comparably to the original model while outperforming other techniques. For perplexity, the original Llama 7B model scores 2.18. CATS 50% scores 2.203 without fine-tuning, improving to 2.193 with minimal fine-tuning (1300 steps). Similarly, the original Mistral 7B model scores 2.234, while CATS 50% scores 2.261 (no fine-tuning) and 2.244 (minimal fine-tuning). CATS also preserves arithmetic reasoning skills without any fine-tuning. On the gsm8k task, the Llama 7B model scores 12.59, CATS 50% scores 12.28 and CATS 70% scores 9.10 with no fine-tuning. We will include these evaluations and more extensive analyses on generative tasks.\n\n[Response to R2]\nThank you for your suggestion. For Custom GPU Kernel 1, lines 2 and 3 are implemented in PyTorch. Line 4 uses a Triton kernel to distribute the m dimension of W_up across thread blocks, processing tiles of x * W_up along the d dimension and fusing the element-wise multiplication. Line 5 uses another Triton kernel to distribute the m and d dimensions of W_down across thread blocks, processing one tile of x_1 * W_down, and atomically accumulating the result. The full code is in `flash_gemv/flash_gemv/kernels.py` inside https://shorturl.at/AGmLJ. We will include the pseudo-code for Triton kernels in the appendix.\n\n[Response to R3]\nYour understanding is correct; the current kernel implementation works for batch size = 1. While batch size = 1 is common in real-world and real-time applications, it is extendable to batch size > 1 by manually unrolling the batch dimension or by aggregating the sparse masks in a batch and performing GEMM instead of GEMV for each tile of x and W. We will clarify this in the paper.\n\n[Response to Q1]\nThe intended explanation was that CATS performs well when directly fine-tuned on a downstream dataset without prior fine-tuning on the RefinedWeb dataset. We appreciate your suggestion and will revise the statement to remove confusion.\n\n[Response to Q2]\nWe will correct the typographical error where the m x d matrix should be W_down.\n\n[Response to Q3]\nYour understanding is correct. Our algorithm targets the generation phase when the batch size is 1. For larger batch sizes, we obtain a d-dimensional mask by applying a logical \"AND\" operation to the d-dimensional masks of each batch of data.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "CATS is well-motivated from prior work's empirical observations, tackles the important problem of LLM inference, and authors empirically demonstrate it provides a decent speed up with little loss of quality.\n\nThe authors also offer a kernel implementation of CATS as an artifact, which can benefit practical applications and further research.", "reasons_to_reject": "1. The model quality evaluations are constrained to classification (posed as log-likelihood evaluation) tasks. Given LLM inference cost is primarily in generation, as a method to speed up inference CATS is not demonstrated to preserve quality on generative tasks.\n2. The kernel description in the paper is too simplified to convey clearly how it makes threshold-based unstructured sparsity fast. Given the kernel is implemented in triton, I suggest the authors specify the pseudo-code in a triton-like structure with launch grids, program ids and thread block level pseudo code. Current pseudo-code is too high-level, while it serves the purpose to demonstrate the concept, it doesn’t really inform how the actual kernel could be implemented.\n3. This point is a continuation of #2 - given the current description of the kernel, my best guess of how the kernel works seems to suggest that it is restricted to batch size=1 (more elaboration for this in question section below), but this limitation is not mentioned anywhere. (I would take this point back if this is a misunderstanding)", "questions_to_authors": "Page 7 paragraph about task-specific finetuning - the boldfaced statement says “…outperforms ReLUification in zero-shot accuracy…”. Given this is task-specific finetuned model, it does not make sense to say “zero-shot accuracy”.\n\nPage 4 after eqn 1, the m x d matrix should be W_down but seems to be typo’ed to W_up\n\nAppendix D - tensor shapes seems to not match unless batch size is 1. We know x is (b x d), therefore xW_gate is (b x m), therefore v is also (b x m) and mask is (b x m). In such a case, how could a (b x m) mask be applied to mask W_up, a (d x m) matrix (or W_down, m x d)? If b = 1 and we squeeze out the batch dim, it would then make sense that mask is just an m dimensional vector and therefore broadcastable to W_up and W_down, but I'm not sure if I'm having the correct understanding here.", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper designs a sparsification method named CATS to reduce inference costs. The technique is simple but effective. All reviews are positive.
Decision: Accept