Training-Free Activation Sparsity in Large Language Models

We thank the reviewer for the kind review! Please find below our point-by-point response regarding your feedback:

W1: Although I understand that the paper focuses on the training-free scheme, providing a discussion and/or experimental results on fine-tuned sparse activation models could make the paper stronger and more interesting… If post-sparsification fine-tuning is prohibited due to the attributes of TEAL, this should be discussed in the paper.

Q1: Are there any limitations or difficulties caused by TEAL when it comes to training or fine-tuning the sparsified model?

We thank the reviewer for the interesting suggestion. TEAL lends itself well to fine-tuning. We fine-tune Llama-3-8B and use a LoRA rank of 32, meaning ~1% of parameters are trainable. We use a learning rate of 0.0002, and fine-tune on 30M tokens (CATS [1] uses around ~160M tokens. 100k tokens is a bit low). We use C4 for the fine-tuning dataset and WikiText for the evaluation dataset. We use greedily optimized sparsity levels.

With fine-tuning we find marginal improvements at lower sparsities (50-60%), and more significant improvements at higher sparsities (70-90%).

The standard deviation of the sparsity by TEAL is not reported. Since the sparsity level changes from input to input, providing more statistical information can help in understanding the proposed method better.

We find statistics on the model-wide sparsity of Llama-3-8B for decoded tokens on the GSM8K dataset. We set the sparsity level to 50% uniform sparsity.

The variance is nontrivial, but fairly small, and the per-token model-wide sparsities cluster tightly around the mean.

Does the Triton kernel need to be rewritten for different GPU architectures?

The kernel should support all triton-compatible GPUs out of the box, including NVIDIA GPUs with Compute Capability 8.0+, and AMD GPUs (ROCm 5.2+). The kernel would need to be rewritten for other accelerators --- though, the overall logic in the kernel is fairly architecture agnostic.

In Table 3, how was the inference speed tested? Was it tested on autoregressive text generation? If so, what are the lengths of the given and the generated text?

We use autoregressive text generation, and use the standard GPT-Fast inference benchmarking setup, which uses ~5 input tokens and 200 max new tokens.

Does the sparsity level change based on the domain of the text (e.g., code, novel, scientific paper, etc.)?

We set Llama-3-8B to 50% uniform sparsity, and measure the average model-wide sparsity level on a variety of datasets representing a variety of domains: GSM8K for Math, GitHub Code for Code, Tiny Textbooks for Novels, and Scientific Papers for Scientific Papers.

Preliminary results imply that code-based domains have slightly lower sparsity levels, although more work would be needed to concretely verify this. Overall, all domains have fairly similar sparsity levels.

[1]: CATS: Contextually-Aware Thresholding for Sparsity in Large Language Models. https://arxiv.org/abs/2404.08763

We thank the reviewer for the kind review! Please find below our point-by-point response regarding your feedback:

While the main focus is on Llama-like LLMs, discussing the applicability of TEAL to other architectures, such as Mixture-of-Experts and Mamba, would enhance the paper's scope.

We thank the reviewer for the interesting suggestion. We are actively working on this, and our intuition is that TEAL will extend nicely to both Mixture-of-Experts and Mamba architectures.

For Mixture-of-Experts models, we can sparsify each expert independently. Concretely, for a SwiGLU-based MoE, each expert will have two hidden states (before $W_{gate}/W_{up}$, before $W_{down}$), each of which we sparsify. We can additionally target the model router if it is insensitive to sparsification.

For Mamba, we can sparsify each of the hidden states that come before a linear layer. We would need to do a more in depth analysis on if any layer is particularly sensitive to sparsification.

The paper lacks a detailed analysis of the memory footprint of the TEAL method, especially in long context scenarios.

TEAL incurs minimal memory overhead. In practice, we represent each distribution (corresponding to a hidden state) as a histogram with 10,000 bins, and represent each threshold as a float. For Llama-3-8B, roughly

$num\textunderscore layers * num\textunderscore hidden \textunderscore states \textunderscore per \textunderscore layer * (num \textunderscore bins + 1) * bytes \textunderscore per \textunderscore value = 32 * 4 * 10001 * 4 = 4.88 \text{ MB}$

is used, which is 0.061% of the model weights’ footprint. Furthermore, distributions are stored on system RAM (as opposed to GPU DRAM), as the static thresholds can be configured offline. TEAL is just as performant in long context scenarios, as its memory footprint does not scale with sequence length.

In Figure 3 (right), why is the theoretical Latency higher than TEAL at small sparsity levels?

TEAL is faster than torch.matmul as the latter is not the most optimized for skinny matrix-multiplications, which past work also finds [1]. “Theoretical Optimal” is a fit that assumes that torch.matmul scales linearly with the sparsity level, so at 0% sparsity the two are equal and are both worse than TEAL. We will clarify this in the revision. We note that we compare against an improved baseline for the end-to-end results in Table 3, which uses TorchInductor to automatically generate faster matrix-multiplication kernels.

[1]: FlashDecoding++: Faster Large Language Model Inference on GPUs. https://arxiv.org/abs/2311.01282

We thank the reviewer for the kind review! Please find below our point-by-point response regarding your feedback:

Although TEAL is adapted for compatibility with newer activation functions like SwiGLU, the core concept of magnitude-based pruning is similar to prior work, which may limit its perceived innovation and contribution.

While magnitude-based pruning on activations is simple, yet effective, we believe there is significant contribution in 1) applying this to all hidden states of the model, not just the intermediate state of the MLP, 2) our greedy block-wise optimization method, which exploits the differing importance of layers in the model, and 3) our kernel, which significantly outperforms the DejaVu kernel in our setting.

The paper does not mention whether the code is available for reproducibility.

Indeed, we overlooked mentioning code availability in the paper. The code is currently available as supplementary material, and this will be reflected in the final manuscript.

1. Calibration Set vs. Evaluation Set for Pruning Decisions:

We note our choice to use a calibration set with fixed thresholds was inspired by CATS [1], which uses a similar method. We find the perplexity of Llama-3-8B on WikiText at 50% model-wide sparsity. We use the block-wise greedy optimized sparsity levels. We consider two datasets for calibration, C4 and FineWeb.

We observe little variation between different calibration sets. Each hidden state only has one degree of freedom (one threshold per hidden state), meaning it’s hard to overfit to a given calibration set.

2. Pruning Effectiveness in Highly Optimized Newer Models:

We note that all model compression methods are less effective on overtrained models (eg. weight quantization [2,3]), not just TEAL. For example, in 2:4 weight sparsity, MaskLLM [4] achieves 6.72 PPL vs. baseline 5.12 PPL for Llama-2-7B (App. C). However, it only achieves 8.50 PPL vs. baseline 5.76 PPL for Llama-3-8B (App. D). Accordingly, it seems to us that for more overtrained LLMs (at a fixed model size), TEAL would remain viable when applied more moderately, similar to methods like weight quantization.

We note that TEAL becomes more effective as model size increases. For example, both Llama-2-70B and Llama-3-70B exhibit substantially less degradation than Llama-2-7B and Llama-3-8B. Q-Sparse [5] finds a similar phenomenon in the pretraining setting, finding that the gap between sparsely activated models and dense baselines decreases as model size increases (Section 3.5). Thus, TEAL becomes more applicable as models increase in size over time.

[1]: CATS: Contextually-Aware Thresholding for Sparsity in Large Language Models. https://arxiv.org/abs/2404.08763

[2]: An Empirical Study of LLaMA3 Quantization: From LLMs to MLLMs. https://arxiv.org/abs/2404.14047

[3]: Scaling Laws for Precision. https://arxiv.org/abs/2411.04330

[4]: MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models. https://arxiv.org/abs/2409.17481

[5]: Q-Sparse: All Large Language Models can be Fully Sparsely-Activated. https://arxiv.org/abs/2407.10969

We thank the reviewer for the kind review! Please find below our point-by-point response regarding your feedback:

It would be nice to include formal definitions to various matrices…

We thank the reviewer for the suggestion, and will incorporate this in the next revision.

There are no comparisons with structured and semi-structured pruning…

We thank the reviewer for the interesting suggestion. Below we compare against MaskLLM [1], a state-of-the-art approach to semi-structured 2:4 weight sparsity that learns weight masks through differentiable relaxation. The authors have not released checkpoints yet, so we trained our own masks on Llama-3-8B, using 2B tokens from C4. We compare against TEAL with greedily optimized sparsities, and evaluate on WikiText.

We see that TEAL outperforms MaskLLM. We emphasize that TEAL is training-free and MaskLLM is fairly expensive (on Llama-3-8B, uses 2B tokens and requires ~8B frozen params, ~12B learnable params. This took us roughly 800 H100 hours). Interestingly, the combination of the two methods works well, meaning this isn’t an “either-or” situation.

We additionally compare speed on Llama-2-7B on 1xA6000 in single-batch decoding (using MaskLLM’s reported numbers in Table 6):

We note this may not be the most fair comparison as 2:4 weight sparsity is more performant in high-batch settings, and can additionally speed up prefill.

In particular, comparison with fine-tuning free ReLUfication seems to be a little bit unfair…

We initially included these results to emphasize that simply replacing SiLU with ReLU is insufficient with respect to accuracy in the training-free case. However, we agree that comparing to ReLUfication in the training-free setting doesn't properly represent the method as intended, so we will remove this comparison in the next revision.

Line 107: Which recent work?

This refers to CATS [2] -- we will make the reference more explicit in the next revision.

Could the authors please explicitly explain why activation sparsity error in figure 8 increases with the batch size? Does the algorithm nullify the activations independently across different samples of a batch or activations are zeroed according to some aggregate statistics?

Activation sparsity results in speedup by avoiding the transfer of unneeded weight channels from off-chip to on-chip memory during decoding. However, this is an issue in batched settings. If a given activation value is extremely large for one client, but relatively small for another client, the weight channel that corresponds to both clients needs to be loaded. This means that on a per-error basis, activation sparsity is less efficient in batched settings.

We propose to zero out activations according to an aggregate statistic. In particular, using the distribution of average activation values over the batch dimension. The resultant sparsification criterion is batch dependent, but is still entirely characterized by a threshold.

It is mentioned in the paper that the efficient low-level Triton kernel is implemented based on Deja vu kernel. Is it possible to apply that kernel directly to the task of sparsifying SwiGLU and GeGLU transformers, and if it is, how it fares against the proposed method in terms of quality?

This depends on the method being used. Using the DejaVu kernel with the original intended method (low rank predictor to get a boolean output sparsity mask, apply input sparsity on $W_{up}/W_{gate}$, apply output sparsity to $W_{down}$) results in substantial degradation as it was originally intended for the 95+% sparsity introduced by ReLU in the intermediate states of MLPs.

Using the DejaVu kernel with TEAL (create boolean sparsity mask from input activations, apply input sparsity to target matrix) is feasible and identical from an accuracy standpoint, but is slower than our kernel (as shown in Figure 3), as we make some additional optimizations (eg. fusing the sparsity mask creation).

The paper seems to imply but not to explicitly state that the result of (SiLU(xWgate ) ⊙ xWup is also sparsified. Are there any other places where activation sparsity is applied but not mentioned?

We will make this more explicit. Activation sparsity is applied to the four hidden states in the model, visualized in Figure 2: Before $W_q/W_k/W_v$, before $W_o$, before $W_{up}/W_{gate}$, and before $W_{down}$.

[1]: MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models. https://arxiv.org/abs/2409.17481

[2]: CATS: Contextually-Aware Thresholding for Sparsity in Large Language Models. https://arxiv.org/abs/2404.08763