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

Abstract: Despite recent advances in LLM quantization, activation quantization remains to be challenging due to the activation outliers. Conventional remedies, e.g., mixing precisions for different channels, introduce extra overhead and reduce the speedup. In this work, we develop a simple yet effective strategy to facilitate per-tensor activation quantization by preventing the generation of problematic tokens. Precisely, we propose a method to find a set of key-value cache, coined _CushionCache_, which mitigates outliers in subsequent tokens when inserted as a prefix. CushionCache works in two steps: First, we greedily search for a prompt token sequence that minimizes the maximum activation values in subsequent tokens. Then, we further tune the token cache to regularize the activations of subsequent tokens to be more quantization-friendly. The proposed method successfully addresses activation outliers of LLMs, providing a substantial performance boost for per-tensor activation quantization methods. We thoroughly evaluate our method over a wide range of models and benchmarks and find that it significantly surpasses the established baseline of per-tensor W8A8 quantization and can be seamlessly integrated with the recent activation quantization method.

## introduction tremendous capabilities of large language models (llms) come with a tremendous computational cost. modern language models often have over hundreds of billions of parameters, requiring significant memory and computation for prediction and training. for instance, opt-175b @xcite , one of the most popular open-sourced language models, requires at least 350gb of memory and the order of 10 18 floating point operations to generate a new token 1 @xcite . quantization is an effective strategy to reduce the computational cost of llms. recent works demonstrate that the precision of llm weight parameters can be greatly reduced by post-training quantization (ptq), with minimal degradations in its generation quality. for example, @xcite shows that one can quantize the weights of the llama3-70b to 4 bits, with less than 0.5%p drop in its zero-shot prediction accuracy. roughly, the reduced precision translates into 4× increase in the generation throughput, and similar reduction in memory requirements @xcite . llm activations, however, remain challenging to be quantized. the key obstacle is the activation outlier, i.e., a small number of activations that are substantially larger than others @xcite @xcite . such outliers elongate the quantization range and flattens out most non-outlier activations, leading to large performance losses even at w8a8 quantization. to address this issue, recent works propose to mitigate outliers based on various relaxations of the stringent static, per-tensor quantization. one line of work applies quantization separately to each channel depending on the outlier proneness @xcite . these methods, however, are difficult to be implemented on conventional hardwares. another line of work reparameterizes the activations and weights in a way that the impact of outliers are amortized @xcite . these algorithms focuses on attaining high generative quality by adopting per-token or dynamic quantization, leaving the most hardware-friendly option-static per-tensor quantization-less explored. to fill this gap, we take a novel approach for mitigating activation outliers in llms. in particular, we focus on answering the following key question: can we find a good prefix that mitigates the activation outliers in the subsequent tokens on a pretrained llm? our answer is positive; we develop a very simple yet effective method, coined cushioncache, to discover a prefix foot_0 which reduces the outlier in the following tokens processed by the given llm. by inserting this prefix, one then can quantize the activations of the llm with much smaller quantization error, leading to an improved generation quality. to design our method, we draw inspirations from a recent observation that the outliers may originate from attention sinks @xcite -the "no-operation" tokens that receive much attention from other tokens @xcite . by adding sink-like tokens as a prefix, one may be able to separate out outlier activations as well, rendering the subsequent tokens outlier-free. in a nutshell, our method works in two steps. 1. greedy initialization. we search for a sequence of sink-like prompt tokens in a greedy manner, so that the activations of the subsequent tokens are less prone to outliers (section 4.1). 2. quantization-aware prefix tuning. we train the greedily initialized prefix further to minimize the combined loss of the prediction loss and quantization error (section 4.2). our experiments demonstrate that the proposed cushioncache is highly effective in making llms more quantizable. the technique is versatile, consistently improving the quantized performance of llms under various scenarios, from per-token to per-tensor static quantization. the method can also be seamlessly combined with existing quantization algorithms to further boost their performances. to summarize, we contribute the following. ## related work outliers in llms. the fact that there exists usually large entries in llm activations, or outliers, has been reported by multiple works. @xcite and @xcite report the existence of outliers in large transformerbased language models (e.g., bert), and find that they appear mostly in a small number of channels and layers. @xcite 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. per-channel activation quantization. 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 outlierprone 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. per-token, with reparameterization. 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 @xcite applies per-tensor quantization and knowledge distillation to achieve reasonable int8 quantization performance. smoothquant @xcite , outlier suppression+ @xcite , and om-niquant @xcite migrates the activation magnitudes to the weights to normalize the scales of the activations. more recently, quarot @xcite 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. per-tensor quantization. notably, @xcite 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 llama3 @xcite . attention sinks and outliers. recent works report an intriguing phenomenon in large transform-ers, termed attention sink. @xcite 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. @xcite 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. ## preliminaries key-value cache. 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 @xcite . that is, at each decoding step, the transformer f (•) performs: 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: quantization. quantization is an act of casting a high-precision tensor (typically fp) into a lowerprecision tensor (typically int), to save the memory to store and computation to process the tensor. activation quantization with static range. 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: outliers and complications. in llms, the activation x fp tends to have a very large entry @xcite . in such case, the magnitude of max(x fp ) and min(x fp ) will be very large, making the scaling factor s x very large. this leads to a high sparsity in the tensor x 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. ## method we now describe cushioncache, an algorithm to find a prefix which can mitigate activation outliers in the subsequent tokens, thereby alleviating the quality degradation from activation quantization. cushioncache aims to find a set of prefix that minimizes the quantization error of the activations. more concretely, let x i denote activation of a transformer block for the input token t i . our goal is to minimize the squared difference between the original and the quantized activations, i.e., we solve the minimization (eq. 7) with a strategy based on prefix tuning. this is done in two steps: initializing prefixes based on greedily searched prompts (section 4.1), and quantization-aware prefix tuning (section 4.2). ## analysis we now conduct a brief sanity check. in particular, we ask the following questions. ## conclusion in this paper, we present cushioncache, a novel approach for mitigating activation outliers in llms to improve activation quantization performance. through extensive experiments, we demonstrate that cushioncache consistently enhances the per-formance of per-tensor activation quantization. our analysis shows that cushioncache effectively reduces the magnitude of activation outliers and redirects attention sinks, leading to more uniform and quantization-friendly activations. in contrast with other approaches to faciliate activation quantization, cushioncache is the first-up to our knowledge-to fundamentally alter the activation distribution itself without extensive training, making activations easier to quantize.