Prefixing Attention Sinks can Mitigate Activation Outliers
for Large Language Model Quantization

The fact that there exists usually large entries in LLM activations, or outliers, has been reported by multiple works. Kovaleva et al. (2021) and Timkey and van Schijndel (2021) report the existence of outliers in large transformer-based language models (e.g., BERT), and find that they appear mostly in a small number of channels and layers. Bondarenko et al. (2021) make a similar observation in the context of quantization, and finds that quantizing the activations lead to a large degradation in generation quality; the work also reports that semantically meaningless tokens can have higher tendencies to have outliers. Dettmers et al. (2022) confirm the same finding while quantizing GPT-scale models, and studies how the model scale affects the prevalence of outliers over tokens and layers. More recently, Sun et al. (2024) investigates a similar phenomenon in newer LLM variants and confirms that certain tokens are more likely to suffer from outliers.

A line of work proposes to mitigate outliers in LLM activation quantization by applying different scaling factors or precision to each channel. Bondarenko et al. (2021) splits activation channels into several groups and perform quantization on each group. LLM.int8() (Dettmers et al., 2022) applies higher precision (e.g., FP16) to a small number of outlier-prone channels, while quantizing the other channels to lower bits (e.g., INT8). These works, however, are difficult to be implemented in conventional hardwares, as they requires scaling along the contracting dimension of matrix multiplication.

Another line of work proposes to quantize the activations per-token to reduce the impact of outliers with better hardware acceleration. Many of these works adopt reparameterization of weights to mitigate the outliers further. ZeroQuant (Yao et al., 2022) applies per-tensor quantization and knowledge distillation to achieve reasonable INT8 quantization performance. SmoothQuant (Xiao et al., 2023), Outlier Suppression+ (Wei et al., 2023), and OmniQuant (Shao et al., 2024) migrates the activation magnitudes to the weights to normalize the scales of the activations. More recently, QuaRot (Ashkboos et al., 2024) rotates the activations so that the outlier magnitudes are distributed over multiple axes in the reparametrized space. While these methods are effective, per-token quantization are typically slower than per-tensor quantization at the same quantization precision as it requires larger scale which has a size of the number of tokens.

Notably, Xiao et al. (2023) also provides two options for per-tensor activation quantization: one with dynamic quantization range, and another with static range. While these options tend to be faster than per-token (with static range being the fastest), their generation quality is much lower than per-token, especially on recent models such as LLaMA-3 (Touvron et al., 2023).

Recent works report an intriguing phenomenon in large transformers, termed attention sink. Xiao et al. (2024) find that a small number of semantically meaningless tokens, usually at the beginning of the sequence, tend to receive unusually large attention. Darcet et al. (2024) make a similar observation for vision transformers, and show that training ViTs with additional meaningless tokens can help make the attention structures more semantically meaningful. Bondarenko et al. (2023) hypothesize that the sink tokens may be the root cause of the activation outliers, and propose a new architecture that prevents the outliers from emerging when pretrained from scratch. Our work shares a similar intuition, but critically differs in that we mitigate outliers by fine-tuning the pretrained LLM. This means no modification to the network architecture is needed and does not need to train the model from scratch.

Modern language models, typically based on decoder-only architecture, are built as a sequence of transformer blocks which process a sequence of tokens to predict the next token (Vaswani et al., 2017). That is, at each decoding step, the transformer $f(\cdot)$ performs:

where $t_{1},\cdots,t_{n}$ are the preceding tokens used as context. LLM iteratively applies Eq. 1 autoregressively to generate text as a sequence of tokens.

As the context length grows, the computational cost to process all previous tokens also grows larger, slowing down the generation significantly. A popular solution is to cache and reuse the keys and values of the preceding tokens computed during the previous iteration. This trick relies on the fact that preceding tokens affect the outcome of the current token only through their keys and values:

where $q_{i},k_{i},v_{i}$ denotes the query, key, and value vectors for each token and $s_{i,j}$ denotes the attention score for the $i$-th token query on $j$-th token key. By storing and reusing the cached values (called KV cache), one only needs to process new tokens as:

where $t_{n+1:n+l}$ denotes $l$ tokens given at the current step and $k_{1:n},v_{1:n}$ are the keys and values of the preceding context of length $n$, computed during the previous iteration. During the prefill phase, $l$ may be the length of the prompt, and during the decoding phase, we can simply use $l=1$, processing only a single token at a time.

Quantization is an act of casting a high-precision tensor (typically FP) into a lower-precision tensor (typically INT), to save the memory to store and computation to process the tensor. In neural network quantization, a popular choice is the linear quantization, which performs

where $\mathbf{X}_{\mathrm{int}},\mathbf{X}_{\mathrm{fp}}$ denotes the quantized and original tensors, $z,s\in\mathbb{R}$ denote zero-point and scaling factor, and $\mathrm{round}(\cdot)$ denotes the rounding operation. The scaling factor is typically selected as

where $N$ denotes the number of bits for the integer format. The zero-point is determined as either $z=\min(\mathbf{X}_{\mathrm{fp}})$ oder $z=0$, for the asymmetric and symmetric quantization, respectively.

By quantizing both activation and weight matrices, one can avoid performing computation-heavy FP matrix multiplications. That is, for the case of symmetric quantization (for simplicity), we approximate:

where the right-hand side can be computed using an integer matrix multiplication, and a single multiplication of FP16/32 quantities (for scaling factors). The combined scaling factor need not be multiplied back to the matrix immediately, and can be used in the subsequent operations directly.

In many cases, the scaling factors $s_{\mathrm{W}},s_{\mathrm{X}}$ can be pre-computed based on the validation set statistics. This method, called static-range quantization, enables more acceleration than computing these values dynamically during the inference.

In LLMs, the activation $\mathbf{X}_{\mathrm{fp}}$ tends to have a very large entry (Dettmers et al., 2022; Sun et al., 2024). In such case, the magnitude of $\max(\mathbf{X}_{\mathrm{fp}})$ and $\min(\mathbf{X}_{\mathrm{fp}})$ will be very large, making the scaling factor $s_{\mathrm{X}}$ very large. This leads to a high sparsity in the tensor $\mathbf{X}_{\mathrm{int}}$, and a much degraded generation quality.

This problem can be alleviated in various ways: One can change the scaling factor dynamically over time (i.e., per-tensor dynamic quantization), or apply different scaling factors for each channel or token (i.e., per-channel/token quantization). As these methods require on-the-fly computations of scaling factors, the methods are typically slower.

The drawback of finer quantization granularity becomes more significant in the distributed setup, as it affect the communication cost between nodes.

To see this, consider the case of multiplying matrices with tensor parallelism, e.g., Megatron-LM (Shoeybi et al., 2019). Comparing with the per-tensor static quantization, per-tensor dynamic quantization requires an additional AllReduce operation over the nodes to aggregate the (high-precision) scaling factor. The overhead is even more significant for per-token dynamic quantization, as the number of scaling factors is multiplied by the number of tokens, increasing the cost of AllReduce.