CATS: Context-Aware Thresholding for Sparsity in Large Language Models

Authors: je-yong  lee
donghyun  lee
genghan  zhang
mo  tiwari
azalia  mirhoseini
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-20 21:33:25

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.

Main paper: 
 introduction: 
 LLMs have demonstrated remarkable success across a variety of fields (Devlin et 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 GPT-3 is estimated to have consumed over 3,000,000 GPU-hours and emitted three thousand times the CO 2 equivalent of a 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 are often based on quantization (Frantar et al., 2022; Dettmers et al., 2022) , pruning (Ma et al., 2023; Sun et al., 2023) , and other forms of weight sparsification Frantar & Alistarh (2023) . Mixture of Experts (MoE) techniques have emerged as particularly promising and are employed by current state-of-the-art LLMs (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2022c; Jiang et al., 2024) .MoE techniques activate only a subset of parameters at each inference stage, thereby reducing memory and computational requirements compared to using the entire model. Prevailing implementations 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 a "router"-a small neural network trained to determine 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., 2023b; Mirzadeh et al., 2023) . This implies that only a 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. Figure 1 shows 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 a 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 an 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 of fine-tuning steps.5. We provide a custom GPU kernel implementation that exploits the sparsity of CATS and achieves a ∼15% improvement in wall-clock inference latency of token generation over the dense models. 
related work: 
 Significant recent work focuses on reducing the inference costs of LLMs. Approaches that utilize mixture-of-experts or activation sparsity are most similar to our work.Mixture-of-Experts (MoE) techniques induce effective sparsity in LLMs by determining which subset of subnetworks (the "experts") to activate during the inference pass, often via a trained "router" subnetwork. This is a popular line of work with significant research interest (Shazeer et al., 2017; Hazimeh et al., 2021; Zhou et al., 2022; Lewis et al., 2021; Roller et al., 2021; Zuo et al., 2021; Komatsuzaki et al., 2022; Lou et al., 2021; Mustafa et al., 2022; Rajbhandari et al., 2022; Zhang et al., 2022a; b; Fedus et al., 2022a; Zoph et al., 2022; Kudugunta et al., 2021; Fedus et al., 2022c; Lepikhin et al., 2020; Du et al., 2022; Fedus et al., 2022b; Jiang et al., 2024) . For a review of MoE models, we refer the reader to (Fedus et al., 2022a) .Activation Sparsity: Activations are known to be sparse in LLMs that utilize ReLU nonlinearities in their MLP blocks (Li et al., 2022) ; however, the reasons for this are not wellunderstood Hoefler et al. (2021) . Nonetheless, activation sparsity induced by ReLU nonlinearities has been explored to reduce memory usage and inference time (Jaszczur et 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 blocks 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 is the only work that attempts to bridge this gap (Mirzadeh 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 a controllable level of sparsity. Furthermore, in Section 5, we demonstrate that CATS demonstrates significantly better downstream task performance and fine-tuning efficiency than ReLUfication.We note that Zhang et al. (2024) is concurrent to our work. In contrast with their work, however, our work is not an 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 GPU 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 A. ) := x * sigmoid(x) = x 1 + e -x(2)Crucially, the operation SiLU(xW gate ) can be viewed as the router in an MoE model. Through this lens, the columns of W up and the rows of W down are the experts. If SiLU(x) is always binary, i.e., 1 or 0, it would turn on/off elements of the remaining computation (multiplication by W up W down ). When SiLU(x) is not binary, it can be viewed as a "soft" router that weighs the experts by different amounts.4 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 a 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. 
stage 1: determining cutoff threshold: 
 We assume we are given a desired sparsity level k (e.g., 70%) as input. For each Gated-MLP block in the LLM, we compute the activations over a 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 t isEQUATIONwhere F represents the empirical cumulative distribution function of the absolute values of the activations for the given MLP block.Figure 1 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 a 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.Given the cutoff threshold t ≥ 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 as CATS t (•), is defined as:CATS t (x j ) := x j , if |x j | ≥ t 0, if |x j | < t (4)Here, t is the sparsification threshold and x j is the j-th element of the vector x, respectively.This results in a new activation CATS t (SiLU(•)):CATS t (SiLU(xW gate )) = SiLU(xW gate ) |SiLU(xW gate )| ≥ t 0 |SiLU(xW gate )| < t (5)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, as detailed in Appendix C.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 a reduction in actual wall-clock latency and increase in generation throughput via a custom GPU kernel.Custom GPU Kernel 1 MLP using CATS 1: Input: threshold t > 0, hidden layer x, weights W gate , W down , andW up 2: v ← SiLU(xW gate ) 3: Mask ← 1 if |v| ≥ t else 0 4: x 1 ← (xW up [Mask] * v[Mask]) 5: y ← x 1 W down [Mask]We focus on reducing the latency of the MLP blocks by reducing memory accesses, as the MLP blocks are 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 1, we first fuse the element-wise multiplication of v[Mask] into each tiling of xW up [Mask] . Here, v represents the hidden vector after applying the SiLU activation function, and Mask denotes a binary mask identifying elements of v with large absolute values. This fusion saves memory operations that would be necessary for storing and loading x 1 several times. We then directly use Mask to control which parts of the weight matrices W up and W down 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, we demonstrate how our custom GPU kernel effectively reduces the inference latency of CATS-based models as sparsity increases. 
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-tuning, and can exploit their sparsity for wall-clock inference time speedups over the base models.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 as base models to verify it is generally applicable to different LLMs. We measure the performance of 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) .For a 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.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.Datasets: 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 to 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 RefinedWeb test dataset (Penedo et al., 2023) .Computational Environment: All experiments were run on a single machine with 8 L40S GPUs. Latency experiments were run on a single L40S GPU as 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/32) 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, W gate , and W down ) in the fine-tuning experiments. During inference, we used the transformers v4.36.2 HuggingFace library, PyTorch v2.1.2, and CUDA v12.1. We used Triton v2.1.0 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 in the supplementary material.We now compare the downstream task performance of the CATS-based models to the baseline models in several settings and draw several conclusions.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 to the base models, i.e., the activation functions are simply replaced in the MLP blocks and no fine-tuning is performed. Table 1 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, at CATS 50% demonstrates performance comparable to the base model. CATS significantly outperforms ReLUficiation in downstream task performance at the same sparsity level (90%).CATS-based models perform more comparably to the base model and increasingly outperform ReLUfication as the model size increases in zero-shot accuracy: Similar to the previous evaluation, we assess the performance of various CATS configurations and a 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 E. CATS-based models perform comparably to the base models and outperform ReLUfication 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 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 2 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 ReLUfication 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 as high as 70%, achieve performance comparable to the base models within 500 steps of fine-tuning, whereas ReLUfication does not.In this setting, the CATS prescription is applied to Mistral-7B. All variants are then fine-tuned for 10 epochs on the training data and evaluated on test dataset for the Cola, SST2, and BoolQ datasets. Table 2 demonstrates our results. Our observations are similar to those for "general" fine-tuning:1. CATS-based models still exhibit sparsity after fine-tuning (see Appendix C).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 CATS-50% demonstrates performance within 0.5% of the base model, whereas ReLUfication demonstrates a significant performance drop.We now evaluate the generation task performance of the CATS-based models using the RefinedWeb dataset. Without any fine-tuning, CATS-50% exhibits only a 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-tuning. Similarly, for Mistral-7B, the initial perplexity increase of 1.21% is reduced to a mere 0.45% after fine-tuning. Activation sparsity in a 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. 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 generation 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 1).In Figures 3a and 3b , we compare our method ("CATS-with-Custom-Kernel") with the dense MLP with m dense = m ("Dense") and the dense MLP with m optimal = m * Sparsity ("Optimal"), the latter of which is a proxy for the best wall-clock time we could hope to achieve. At 50% and 70% sparsity, the sparse kernel achieves approximately a 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 our sparse kernel consistently outperforms the original MLP. The comparison with 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 m optimal 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 CATS-with-Custom-Kernel (50% sparsity) on the throughput of the generation stage. The generation stage (or "decoding" stage) is known to be memory-bound (Kim et al., 2023) , which suggests CATS can improve inference througput. We test the generation throughput at a 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 as DeepSpeed (Rajbhandari et al., 2022) and TensorRT-LLM (Nvidia, 2024) . 
discussion and conclusion: 
 We presented CATS, a novel framework for inducing and exploiting activation sparsity in LLMs. At the heart of our framework is the CATS activation, given in Equation 5, that induces a controllable level of activation sparsity in LLMs. We also provide a 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 ReLUfication at similar levels of fine-tuning, and often achieve performance comparable to the base model at high levels of sparsity, both with general and task-specific fine-tuning. Fu et al. (2023b) . These methods often reduce the memory and computational complexity at the cost of performance degradation or 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. 
 
 Appendix: In this section, we discuss how we can apply CATS to reduce the inference costs of attention layers inside Transformer blocks. The basic operations of a Transformer block can be written as:EQUATION)where x is the hidden vector right before the i-th layer and where we have excluded operations like batch normalization, positional embedding, residual connections, etc. for simplicity. (For more details on the variants of Attention layers and those used in our models, we refer the reader to Touvron et al. (2023) and Jiang et al. (2023) .)The new equation for i-th transformer layer, where we wrap the previous layer with CATS activations, becomes:EQUATION) where t i,1 and t i,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 C.We also apply CATS to accelerate the computation of attention layers. Our approach is inspired by "Stage 2" of ReLUficiation (Mirzadeh et al., 2023) .Due to space constraints, we only measure the performance of CATS-50% applied to 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 an 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 can 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.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 3 datasets used for this experimental setting in Section 5. The observed sparsity levels are approximately equal to the target sparsity levels. 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 a proxy for inference cost, other work has demonstrated that reducing FLOPs is not sufficient to reduce real wall-clock inference latency Liu et al. (2023b) . 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:EQUATION)where v ′ , W ′ up , and W ′ down are v, W up , and W down masked by Mask (for the matrices W up and W down , the entire column j is 0 if Mask j = 0, i.e., the mask is broadcast across columns).If Mask is sparse, then Equation (10) performs two sparse matrix multiplications. In fact, only coordinates (respectively, rows) of v (respectively, W up and W down ) corresponding to nonzero coordinates of 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 can be reduced 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 Algorithms 2 and 3 describe Equations ( 8)-( 10) in lower-level pseudocode. Algorithms 2 and 3 contain several optimizations. The atomic operations in Line 4 of 2, 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.Optimization 2: We therefore introduce another optimization in Algorithm 3 to reduce the memory loading incurred by the atomic operations. In Algorithm 3, we directly use Mask to control which parts of weight matrices to load, instead of the condensed idcs. Algorithm 3 has more operations than Algorithm 2 because it directly assigns the unloaded elements to zero 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 et al. (2023) , which means the operation reduction does not compensate for the performance loss caused by complex control logic. Figure 6 These authors contributed equally as co-first authors. 


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

Reviews: review: paper_summary: paper_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.

Without any fine-tuning, CATS can be applied to various base models, including Mistral-7B and Llama2-7B, and achieves comparable performance on downstream tasks.
strength_weakness: strength_weakness: 1. It lacks extensive baselines, like Wanda.
questions: questions: 1. "achieve downstream task performance within 1-2% of their base models " => "achieve downstream task performance within 1-2% reduction of their base models ".
2. I am not sure that the first contribution is solid, which draws connection between the MoE framework and multiplication performed
by dense matrices in the MLP blocks of LLMs.
3. "$W_{up}$" => "$W_{down}$" at the line under Equation (1).
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: 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?
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: How to adapt a instruction-tuned model? It would sometimes be challenging if the finetuning data is not available.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
3. 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: questions: 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”.

Page 4 after eqn 1, the m x d matrix should be W_down but seems to be typo’ed to W_up

Appendix 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.

Comments: title: Thanks for the detailed response.
comment: Dear Reviewer hNTR,

Thank you for your comment. We would like to confirm that we used the official `meta-llama/Llama-2-7b-chat-hf`. Following your suggestion, we will investigate the effectiveness of CATS on the chat-based benchmarks and report the results in future revisions. Thank you again for your suggestion.
Thank you for your detailed response. The result does look great if that is the official `meta-llama/Llama-2-7b-chat-hf`. Additionally, it is confusing that the instruction-tuned model achieved a lower accuracy (0.6569) compared to the base model (0.6589). Would it be more appropriate for the authors to evaluate chat-based benchmarks for chat models, such as MT-Bench and AlpacaEval? (I am happy with current benchmark setup. This is just a suggestion for future revisions.)
Dear Reviewer UDEW, as the Author-Reviewer discussion period is coming to a close, we wanted to ask if you had any additional questions, comments, or feedback on our paper or rebuttal? We would be very happy to address them.
Dear Reviewer pqMf, as the Author-Reviewer discussion period is coming to a close, we wanted to ask if you had any additional questions, comments, or feedback on our paper or rebuttal? We would be very happy to address them.