Q-Sparse: All Large Language Models can be
Fully Sparsely-Activated

The Q-Sparse architecture is based on the Transformer architecture [27, 26] with modifications to enable sparsity in the activations.

Top-K Sparsity

The Transformer architecture uses nn.Linear to perform the projection in both attention and feed-forward layers, which can be written as:

where $\mathbf{X}\in\mathbb{R}^{N\times D}$ is the input tensor, $\mathbf{W}\in\mathbb{R}^{M\times D}$ is the weight tensor, and $\mathbf{Y}\in\mathbb{R}^{N\times M}$ is the output tensor. The nn.Linear operation is equivalent to the matrix multiplication operation.

We introduce a top-K sparsity function on top of the matrix multiplication operation. The top-K sparsity function is defined as:

where $\mathbf{M}\in\mathbb{R}^{N\times D}$ is the mask tensor that indicates the top-K activations in the input tensor $\mathbf{X}$ in terms of the absolute values, $\odot$ is the element-wise multiplication operation, and $\text{Top}_{k}$ is the function that selects the top-K elements in the tensors.

To reduce the interval around zero, we re-scale the tensor by its $L_{2}$ norm after performing the top-K sparsity function.

Quantized Top-K Sparsity

Recent works [28] have shown that quantization can be used to reduce the memory footprint and computational cost of LLMs without the loss of performance. We introduce a quantized version of the top-K sparsity function. The quantized top-K sparsity function is defined as:

where $\text{Q}(\cdot)$ is the quantization function that quantizes the input tensor $\mathbf{X}$ to a 8-bit representation:

where $\epsilon$ is a small constant to avoid division by zero, and $\gamma$ is the maximum absolute value in the input tensor $\mathbf{X}$.

Q-Sparse can be used with both full-precision and quantized LLMs. Specifically, the quantized version of Q-Sparse is compatible with 1-bit LLMs, such as BitNet b1.58 [28]. When using Q-Sparse with 1-bit LLMs, the quantization function is performed on the weight tensor $\mathbf{W}$:

where $\text{Q}_{w}(\cdot)$ is the quantization function that quantizes the weight tensor $\mathbf{W}$ to a 1.58-bit representation:

where $\alpha$ is the mean absolute value in the weight tensor $\mathbf{W}$:

Squared ReLU

To further improve the sparsity of the activations, we use the squared ReLU function [24] for the feed-forward layers. The squared ReLU function is defined as $\text{ReLU}(\mathbf{X})^{2}$.

Following the LLaMA architecture, we use the gated linear unit (GLU) for the feed-forward layers. The squared ReLU function is applied with the GLU function into a ReLU²GLU function. The ReLU²GLU function is defined as:

Most of the existing works [18] on training sparsely-activated models use the vanilla back-propagation algorithm to compute the gradient through the sparsity function:

where $\mathbf{M}$ is the mask tensor that indicates the top-K activations in the input tensor $\mathbf{X}$, and $\odot$ is the element-wise multiplication operation.

The vanilla back-propagation algorithm has a limitation. It zero-outs the gradients of the non-activated elements, which can lead to the vanishing gradient problem, especially when the sparsity ratio is high. In this work, we propose to use the straight-through estimator [2] to back-propagate the gradients through the sparsity function. In this way, the gradients are passed through the sparsity function without being zeroed-out. The straight-through estimator is defined as:

We visualize the average $l2$ norm of each projection’s gradient across different layers for dense model, Q-Sparse with and without STE. We adopt top-K as 50% for Q-Sparse. Without STE, the gradient is much smaller at the bottom layers, while STE can preserve the magnitude of the gradients. As shown in Figure 2, STE estimator significantly eases the issue of gradient vanishing, especially at the bottom of the layers. We present more visualizations for each components in the Appendix A.

Q-Sparse can be used in different settings, including training-from-scratch, continue-training, and finetuning. In the continue-train and finetuning settings, we use the same architecture and training procedure as in the training-from-scratch setting. The only difference is that we initialize the model with the pre-trained weights and continue training with the sparsity function enabled.

For the pre-trained models that do not have the squared ReLU function in the feed-forward layers, we apply the top-K sparsity function after the activated function (e.g., SiLU) in the feed-forward layers. It can improve the sparsity of the activations without changing the model architecture.

To determine the form of the scaling law of sparse-activated LLMs, we begin with a series of scaling experiments. In the experiments, we train a series of language models with Q-Sparse of various scales, ranging from 300M to 7B. The models are trained on the Redpajama dataset [3]. We use the Sentencepiece tokenizer from LLaMA to preprocess data. Besides Q-Sparse, we also train the dense baselines with the same datasets and settings. More details can be found in the Appendix B.

The observed losses of the sparsely-activated models and the dense baselines are shown in Figure 3. We summarize the findings as below:

• •

  The performance of the sparsely-activated models scales with the model size and the sparsity ratio.

• •

  Given a fixed sparsity ratio $S$, the performance of the sparsely-activated models follows a power-law scaling law with regards to the model size $N$.

• •

  Given a fixed parameters $N$, the performance of the sparsely-activated models follows an exponential-law scaling law with regards to the sparsity ratio $S$.

• •

  As the parameters $N$ scales, the performance gap between the sparsely-activated models and the dense baselines decreases.

According to these findings, our main hypothesis is that the performance of the sparsely-activated models follows a combination of a power-law scaling law with regards to the model size $N$ and an exponential-law scaling law with regards to the sparsity ratio $S$.

With a fixed sparsity ratio $S$, the scaling law should follows [11]’s scaling law, which can be written as:

where $\alpha(S)$ is the scaling exponent, and the scaling factor $A(S)$ is a function of the sparsity ratio $S$. Given any model size $N$, the function $L(N,S)$ should follow the Lipschitz continuity with regards to the sparsity ratio $S$. Therefore, the scaling exponent $\alpha(S)$ should be a non-decreasing function. Given any model size $N$, the function $L(N,S)$ is increasing with the sparsity ratio $S$, so $\alpha(S)$ should be a non-increasing function. Above all, the scaling exponent $\alpha(S)$ should be a constant, and the scaling function can be written as:

According to the above finding, the performance of the sparsely-activated models follows an exponential-law scaling law with regards to the sparsity ratio $S$. Therefore, the scaling factor $A(S)$ should also follow an exponential law. Besides, given any model size $N$, the scaling function is increasing with the sparsity ratio $S$. Therefore, the scaling factor $A(S)$ should be a non-decreasing function. The scaling factor $A(S)$ can be written as:

where $B$ is the scaling factor for extremely sparse LLMs, $C$ is the scaling factor for dense LLMs, and $\beta$ is the scaling exponent of the scaling factor $A(S)$ with regards to the sparsity ratio $S$.

We fit the parameters of the scaling law to the observed losses of the sparsely-activated models. We use the L-BFGS algorithm [20] to minimize the Huber loss [9] between the predicted and observed log loss.

Nach [8], $\delta$ is set as $10^{-3}$. We select the best fit from a grid of initialisations around possible local optimas. $E$, $B$, $C$, $\alpha$ and $\beta$ are estimated as 1.86, 0.01, 1.89, 0.10 and 0.05, respectively.

Given the above scaling law, we can derive the performance of the sparsely-activated models and the dense baselines with the same model size $N$ and the same sparsity ratio $S$. The performance gap between the sparsely-activated models and the dense baselines decreases as the model size $N$ scales. The performance gap can be written as:

Since $\alpha$ is a constant that satisfies $\alpha>0$, the performance gap decreases as the model size $N$ scales. It means that given a large enough model size $N$, the performance of the sparsely-activated models can eventually match the performance of the dense baselines with the same model size.

The scaling law can also be transformed into a form that is dependent on the activated parameters $N_{a}$, which reflects the effective compute (i.e., FLOPs) of the model during inference:

where $N_{a}$ is the number of activated parameters in the model, which is equal to $N\times(1-S)$. Since $A(S)$ is an increasing function and $(1-S)^{\alpha}$ is a decreasing function, there exists a sparsity ratio $S^{*}>0$ that minimizes the loss of the sparsely-activated models. This leads to the inference-optimal scaling law of the sparsely-activated models:

It shows that the performance of the sparsely-activated models is better than the dense baselines with the same inference compute budget. We further solve the optimal sparsity ratio $S^{*}$, finding that $S^{*}\approx 45.58\%$. It means that a sparsely-activated model with a sparsity ratio of 45.58% (or $1.84N_{a}$ parameters) can achieve the best performance with the same inference budget $N_{a}$. We follow the same process to estimate the inference-optimal scaling law for 1.58-bit Q-Sparse models. We find that the optimal sparsity ratio is 61.25% (or $2.58N_{a}$ parameters). Figure 4 shows the inference-optimal scaling curves of the sparsely-activated models with full-precision and 1.58-bit weight. It shows that with the same performance, the sparsely-activated models can achieve a significant reduction in the number of activated parameters or FLOPs during inference.

The inference-optimal scaling law shows that the performance of the sparsely-activated models can be optimized by adjusting the sparsity ratio $S$. It can be used to guide the training of the sparsely-activated models and to optimize the performance of the models during inference.