Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

COUNTDOWN: Contextually Sparse Activation Filtering Out Unnecessary Weights in Down Projection

Jaewon Cheon, Pilsung Kang

TL;DR

This work tackles the inference bottleneck of large language models by redefining FFN sparsity as a global, linear combination over the down projection matrix. It introduces two CountDown methods: M-CountDown, which uses indirect indicators from the up and gate projections and is predictor-free, and D-CountDown, which uses the direct coefficient vector for fine-grained pruning. The approach achieves substantial acceleration, including up to 90% computation reduction with as little as 5.5% performance loss in ideal conditions, and up to 29.4% better preservation over CATS for practical settings, aided by specialized Triton kernels. The paper also contrasts SPIdeal (theoretical upper-bounds) with SPPrac (real-world deployment) considerations and provides extensive empirical validation across multiple LLMs and tasks, demonstrating robust gains in both task performance and throughput.

Abstract

The growing size of large language models has created significant computational inefficiencies. To address this challenge, sparse activation methods selectively deactivates non-essential parameters during inference, reducing computational costs in FFNN layers. While existing methods focus on non-linear gating mechanisms, we hypothesize that the sparsity of the FFNN layer lies globally in the form of a linear combination over its internal down projection matrix. Based on this insight, we propose two methods: M-COUNTDOWN, leveraging indirect coefficients, and D-COUNTDOWN, utilizing direct coefficients of the linear combination. Experimental results demonstrate that D-COUNTDOWN can omit 90% of computations with performance loss as low as 5.5% ideally, while M-COUNTDOWN provides a predictor-free solution with up to 29.4% better performance preservation compared to existing methods. Our specialized kernel implementations effectively realize these theoretical gains into substantial real-world acceleration.

COUNTDOWN: Contextually Sparse Activation Filtering Out Unnecessary Weights in Down Projection

TL;DR

This work tackles the inference bottleneck of large language models by redefining FFN sparsity as a global, linear combination over the down projection matrix. It introduces two CountDown methods: M-CountDown, which uses indirect indicators from the up and gate projections and is predictor-free, and D-CountDown, which uses the direct coefficient vector for fine-grained pruning. The approach achieves substantial acceleration, including up to 90% computation reduction with as little as 5.5% performance loss in ideal conditions, and up to 29.4% better preservation over CATS for practical settings, aided by specialized Triton kernels. The paper also contrasts SPIdeal (theoretical upper-bounds) with SPPrac (real-world deployment) considerations and provides extensive empirical validation across multiple LLMs and tasks, demonstrating robust gains in both task performance and throughput.

Abstract

The growing size of large language models has created significant computational inefficiencies. To address this challenge, sparse activation methods selectively deactivates non-essential parameters during inference, reducing computational costs in FFNN layers. While existing methods focus on non-linear gating mechanisms, we hypothesize that the sparsity of the FFNN layer lies globally in the form of a linear combination over its internal down projection matrix. Based on this insight, we propose two methods: M-COUNTDOWN, leveraging indirect coefficients, and D-COUNTDOWN, utilizing direct coefficients of the linear combination. Experimental results demonstrate that D-COUNTDOWN can omit 90% of computations with performance loss as low as 5.5% ideally, while M-COUNTDOWN provides a predictor-free solution with up to 29.4% better performance preservation compared to existing methods. Our specialized kernel implementations effectively realize these theoretical gains into substantial real-world acceleration.

Paper Structure

This paper contains 33 sections, 9 equations, 5 figures, 12 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. ReLU-based Sparse Activation
  4. Non-ReLU Sparse Activation
  5. Generalization of Sparse Activation
  6. Problem Formulation
  7. Limitation of Comparative Methodology
  8. Threshold Variants
  9. Implementation of Sparse Activation
  10. SPIdeal and SPPrac
  11. Constructing SPPrac for CountDown
  12. Hardware Aware Kernel Design
  13. Experiments
  14. Experimental Setup
  15. Downstream Task Performance
  16. ...and 18 more sections

Figures (5)

  • Figure 1: Comparison of sparsity determinations: our approach determines sparsity from the full FFN computation (turquoise box), whereas conventional methods like CATS Lee2024-zb rely solely on non-linear activations (red box).
  • Figure 2: CountDown Pipeline. Note that $h_{pre}=x\cdot W_{gate}$. Left (a): In M-CountDown, we determine which parameters to activate by binarizing densely computed $u$ with pre-calculated $\hat{\tau_\text{M}^k}$. Right (b): In D-CountDown, low-rank predictors $\bigl(\theta_\text{A}$, $\theta_\text{B}\bigr)$ determine which parameters to activate.
  • Figure 3: Kernel Speed for Llama-3.1-8B-Instruct. CATS, M-CountDown and D-CountDown show their respective SPPrac kernel speeds, Full and Optimal show Dense while $int({d_{inter}} \times k)$ instead of $d_{inter}$ for the Optimal. Results for other models are in \ref{['fig:apdx_kernel']}.
  • Figure 4: Tornado plots of CIF and CAF across whitening ratios. Bars to the right indicate the proportion of $\textup{CIF}^k(\cdot,alive)$, while those to the left indicate $\textup{CIF}^k(\cdot,dead)$.
  • Figure 5: All results for kernel speed.