# Fully-attentive and interpretable: vision and video vision transformers for
pain detection

Giacomo Fiorentini
Utrecht University
<EMAIL_ADDRESS>
&Itir Onal Ertugrul
Utrecht University
<EMAIL_ADDRESS>
&Albert Ali Salah
Utrecht University
<EMAIL_ADDRESS>

###### Abstract

Automatic detection of facial indicators of pain has many useful applications
in the healthcare domain. Vision transformers are a top-performing
architecture in computer vision, with little research on their use for pain
detection. In this paper, we propose the first fully-attentive automated pain
detection pipeline that achieves state-of-the-art performance on binary pain
detection from facial expressions. The model is trained on the UNBC-McMaster
dataset, after faces are 3D-registered and rotated to the canonical frontal
view. In our experiments we identify important areas of the hyperparameter
space and their interaction with vision and video vision transformers,
obtaining three noteworthy models. We analyze the attention maps of one of our
models, finding reasonable interpretations for its predictions. We also
evaluate Mixup, an augmentation technique, and Sharpness-Aware Minimization,
an optimizer, with no success. Our presented models, ViT-1 (F1 score 0.55
$\pm$ 0.15), ViViT-1 (F1 score 0.55 $\pm$ 0.13), and ViViT-2 (F1 score 0.49
$\pm$ 0.04), all outperform earlier works, showing the potential of vision
transformers for pain detection. Code is available at
https://github.com/IPDTFE/ViT-McMaster

##  1 Introduction

The International Association for the Study of Pain defines pain as “An
unpleasant sensory and emotional experience associated with, or resembling
that associated with, actual or potential tissue damage" [1]. In Europe, one
adult in five suffers from moderate to severe chronic pain, with major
consequences for their lives and well-being. Their ability to sleep, walk, do
chores, have sexual relations, live independently, and function normally feels
limited or restricted [2]. Pain is a major healthcare problem that medical
care needs to overcome.

Pain is a ubiquitous problem for hospital care as well, with a great deal of
research dedicated to pain analysis, quantification and understanding. To
quantify pain, visual analogue scales (VAS) [3] and similar metrics are
usually employed due to their convenience and simplicity. To measure pain with
VAS, the patient has to point at its pain level on a horizontal scale ranging
from absence to maximum pain. Unfortunately, this technique has the drawback
of being subjective and easily influenced, therefore leaving much to be
desired as the gold standard of pain assessment.

Furthermore, under many circumstances patients are unable to report their pain
levels, such as due to their mental and physical condition, making self-
reporting techniques unreliable and widely inapplicable [4, 5, 6]. In order to
overcome the limitations of VAS and individual checks, automation alongside
new metrics have to be employed. In this regard, facial expressions can be an
important means of communication for the emotional state of a person,
including their pain levels [7].

Facial expressions play an important role in communicating pain. The facial
action coding system (FACS) [7] is a framework based on the anatomy of the
facial muscles, and divides facial expression into 34 atomic components
defined as action units (AU) with scores ranging from A to E depending on
their intensity. While by itself this system contains no apparent information
on the pain levels of the subject, the Prkachin and Solomon Pain Intensity
(PSPI) score [8] identifies six AUs, grouped into four actions, that contain
most of the information on pain.

These actions are brow-lowering (AU4), orbital tightening (AU6 and AU7),
levator tightening (AU9 and AU10), and eye closure (AU43) [8]. The PSPI score
is computed by taking the highest intensity AU component of each action and
summing the numerical equivalent of their intensities (ranging from 1 to 5).
As AU43 (eye closure) has only one possible intensity value, PSPI is therefore
a 16-point pain scale.

$\displaystyle PSPI=AU4+max(AU6,AU7)+max(AU9,AU10)+AU43$

While the PSPI metric does not rely on self-reporting, eliminating one of the
aforementioned limitations, FACS coding requires an average training time of
three months, with each trained expert taking on average over two hours to
code a single minute of video [9]. In order to overcome this challenging
drawback, automation is needed to predict the PSPI scores directly [10].
Desirable properties for such automated pain detection models are
spatiotemporal reasoning [11], robustness to occlusion and changes in the
environment [12, 13], explainability [14], and accuracy [15]. Transformer
models meet many of these requirements, making them good candidates for pain
assessment pipelines.

Figure 1: Examples frames from the UNBC-McMaster dataset and their PSPI score
labels.

Transformers have achieved state-of-the-art performances in multiple tasks,
but only a few studies have researched their performance for pain detection
[16]. The possibility to analyze spatiotemporal relations through video
transformers [17], to extract attention maps and generate interpretations
intrinsic to the model [18] [19], the ability to fine-tune these models on
smaller datasets with good results [20], and their state-of-the-art
performance on other computer vision tasks [21] are promising for their
application towards pain assessment.

In this study, we evaluate the performance of vision transformers (ViT) and
video vision transformers (ViViT) for automated pain assessment from facial
features using a fully-attentive pipeline. Training is carried out on the
UNBC-McMaster dataset using PSPI labels, shown in Figure 1, under a variety of
configurations, pinpointing regions of interest in the hyperparameter space.
We extract attention maps and evaluate them, finding plausible interpretations
for the prediction of the model. Techniques that have been shown to boost
transformer performance are evaluated and adapted to the task, attempting to
maximize model performance for binary pain detection. We achieve state-of-the-
art performance on the task of pain detection for the F1 score metric,
demonstrating the potential of transformer models for automated pain
assessment, and building foundations for future transformer research on this
task.

The contributions of this paper are:

  * •

We propose the first fully-attentive pipeline for automated pain assessment
and achieve state-of-the-art performance on binary pain detection.

  * •

We identify regions of interest in the transformer hyperparameter space.

  * •

We compare the performance of vision and video vision transformers.

  * •

We visualize attention maps for the video vision transformer and show that
pain-specific facial regions are attended.

  * •

We show the impact of Mixup, an augmentation technique, and SAM, an optimizer,
on model performance.

##  2 Related work

#### Video Transformers

After their success in neural machine translation [22], transformers have been
used as standard in several NLP tasks. Yet, their application to vision-
related tasks is relatively new. Dostovitskiy et al. have proposed vision
transformers (ViT) and have shown that ViT outperforms CNN once it is trained
on very large databases [23]. Recently, video vision transformers (ViViT) have
been proposed to model spatiotemporal information and have been shown to
achieve state-of-the-art performance on activity recognition in several
settings. ViViT has outperformed earlier approaches that model spatiotemporal
information [24, 25] and other temporal extensions of ViT [17].

#### Automated pain detection

Recent work has shown that automated pain detection from facial expressions is
a feasible goal. Earlier works have focused on conventional machine learning
approaches such as Support Vector Machines [26, 27] and k-Nearest Neighbor
(kNN) [28] to detect pain using a number of features extracted from face
images. More recent ones have used deep learning approaches in which spatial
information is learned from the face images using convolutional neural
networks (CNN) [29, 30].

Modeling temporal information has been shown to be crucial as a static
approach based on Relevance Vector Regression [10] could not distinguish
between eye blinks and eye closures, which are pivotal for pain intensity
estimation. Recurrent convolutional neural networks (RCNN) [31] and a
combination of CNNs with a long short-term memory (LSTM) networks [32] have
been used to model spatiotemporal relationship among successive frames. They
have shown superior performance compared to the static approaches. Inspired by
these findings, we compare the performance of ViT and ViViT on automated pain
detection.

#### Transformers for pain detection

Several works have shown the success of using vision transformers for facial
expression recognition [33], and facial action unit detection [34]. However,
their application in automated pain assessment is very scarce. To the best of
our knowledge, the only existing work based on transformer technology for pain
intensity estimation is by Xu and Liu [16]. The pipeline presented in this
work focuses on end-to-end pain intensity estimation and includes both a CNN
and a transformer. Pain-related features are first identified and extracted
from the input images by a ResNet architecture with bottleneck attention
modules, then processed by a transformer model that predicts the pain
intensity. The successful performance of our model on a similar task, even
when only fine-tuning a pre-trained transformer, contradicts their finding
that a transformer alone does not work for pain assessment.

##  3 Methods

The base transformer model is pre-trained on ImageNet-21k [35] and fine-tuned
on ImageNet [36] at a resolution of 224x224 by HuggingFace [37] and shared
under the Apache License 2.0. It employs 16x16 patches, a CLS token for
classification purposes, and positional embeddings. It consists of 12
attention and 12 fully-connected layers, and employs 12 attention heads, for a
total of 86 million parameters. A classification head has been added to the
model, trained to interpret the CLS token and output a binary label
prediction. Finally, all fully connected layers are frozen, massively reducing
the training time required.

###  3.1 Dataset

The models are trained on the UNBC-McMaster dataset [38], one of the most
commonly used datasets for facial pain assessment. It consists of 48398 video
frames from 25 patients suffering from shoulder-related pain, captured as the
patients performed active and passive range-of-motion tests with each of their
limbs. The dataset is extremely imbalanced, with 82.7% of frames labeled with
a PSPI score of 0 and 10.9% having a score of 1 or 2 out of 16. For the
purposes of binary classification, we divide the dataset into two categories,
0 (no-pain) and 1 (pain), the latter category including images with a PSPI
score above 0. During training, the pain class is over-sampled to prevent
overfitting on the majority class.

The frames are divided into five folds, each containing samples from exactly
five patients. The splits are generated with the aim of maintaining a similar
number of pain samples across folds, achieved by pairing patients with the
fewest and most pain samples and shuffling four patients between folds to
further balance them, ensuring that each has a reasonable number of samples
for the minority class. Five-fold cross-validation guarantees that the models
learn to generalize painful features rather than overfitting on specific
patients.

The dataset is also processed for the purposes of video transformer training -
subsequent images are grouped in $2\times 2$ grids and labeled according to
the label of the last image of the 4-frame sequence. The use of multiple
subsequent frames aims to capture the dynamics of the facial expressions,
enabling the model to distinguish between the subject shutting their eyes due
to pain (AU43) and blinking [17], and other critical dynamics of facial pain.

###  3.2 3D registration

We perform 3D registration using PRNet [39], which gets a 2D face image as
input, performs 3D registration without requiring person-specific training,
and outputs a dense 3D mesh of the face. The result is achieved by regressing
the UV position map, a structure that records 3D coordinates of a complete
facial point cloud, from the input image. We then use Face3D tool [40] to
rasterize 2D image from frontalized 3D facial structure generated by PRNet as
shown in Figure 2a.

Figure 2: Data transformations through the pre-processing and transformer
pipeline

After this step, semantic correspondence is established across frames and
subjects. Consequently, visual words used in vision transformers are aligned
as given in Figure 2b.

###  3.3 Experiments

To determine the performance of vision and video transformers in automated
pain assessment, we conduct two sets of experiments. The first set of
experiments, using vision transformers, consists in tuning a single
hyperparameter and saving the best performing value to be used while tuning
the next parameter. The second set of experiments, identically structured, is
carried out using video transformers on $2\times 2$ image grids.

Due to the extreme imbalance of the labels, we evaluate the performance of the
model using the F1 score on the minority class. This way, we ensure that the
model prioritizes performance on the more difficult task of pain detection.
Furthermore, earlier studies carried out on this task used the F1 score
metric, making it possible to compare results.

We have tested 14 possible configurations, the first six seek the optimal
number of unfrozen attention layers for the transformer model, then four to
determine the optimal learning rate of the Adam optimizer, one to quantify the
effects of the Sharpness-Aware Minimization in combination with Adam, and
three for the impact of the Mixup augmentation [41] on the performance of the
transformer. All 14 configurations have been tested separately for the single-
image and the $2\times 2$ grid datasets, with the rationale that the use of
vision or video transformers is unlikely to be independent of each individual
hyperparameter.

First, all the fully-connected layers are kept frozen, leaving 12 attention
layers to be fine-tuned. However, while too few layers cannot be effectively
fine-tuned on a specific task, a higher number does not necessarily lead to a
better performance [42], necessitating the model to be evaluated with varying
amounts of unfrozen layers. Next, the learning rate of the Adam optimizer is
tuned, ahead of the introduction of the Sharpness-Aware Minimization (SAM)
optimizer.

Transformer models work best with large amounts of data, nevertheless, this
weakness might be mitigated with techniques such as SAM [43] and Mixup [41].
The SAM optimizer works in conjunction with the original optimizer, in our
case Adam, to prevent the model from converging to sharp local minima. While
it could potentially reduce overfitting on the small UNBC-McMaster dataset, it
also requires a second forward-backward pass, almost doubling the training
time required.

Mixup takes a percentage of the dataset and generates hybrid images by
blending frames with distinct labels. For example, a painful frame labeled
[1,0] is combined with a painless frame labeled [0,1], with a blending value
of 0.2. The resulting frame is labeled [0.8, 0.2], and consists of the sum of
the pixel intensity values from the painful and painless picture, the former
at 80% opacity, and the latter at 20%.

To better integrate Mixup with the pre-processed UNBC-McMaster dataset, one
further restriction is applied, allowing only images from the same patient to
be combined for the experiment. The intensity of Mixup can be adjusted through
its $\alpha$ parameter, causing images to be increasingly hybrid, and has been
configured according to previous research on Mixup and transformers [44].
Although these samples could allow for a more nuanced and linear function of
pain for the model to learn from, they might also result too noisy and
unnatural compared to other samples, further degrading the already limited
data available.

##  4 Results

Preliminary experiments have shown reasonable values for various parameters
such as learning rate (2E-04), batch size (16), and number of epochs (1).
Other important parameters for the initial training of the model are the drop-
out rate before the classification head (0.10), $\beta$ values (0.9, 0.999)
and $\epsilon$ (1e-08) of the Adam optimizer, weight decay (0), and the $\rho$
(0.05) of the SAM optimizer.

###  4.1 Number of unfrozen attention layers

The first step of the experimentation consists in identifying the optimal
number of unfrozen layers. The results can be seen in Figure 3. In total, 12
models are trained for the vision and video vision transformer with multiples
of two as the number of layers, from 2 to 12. For ViT, fine-tuning 12 layers
performs best (F1 score 0.47) while fine-tuning 6 layers achieves the second
best performance (F1 score 0.45). For ViViT, fine-tuning 6 layers performs
best (F1 score 0.55), while fine-tuning 12 layers achieves the second best
performance (F1 score 0.53).

Figure 3: Performance (F1 score) of ViT and ViViT models with different
numbers of unfrozen (fine-tuned) attention layers.

###  4.2 Learning Rate

For the second step, a large range of learning rates is tested to identify
regions of interest in the hyperparameter space. The initial learning rate of
0.0002 is both increased and decreased tenfold and a hundredfold. Performances
of the resulting models can be seen in Figure 4.

Figure 4: Performance (F1 score) of ViT and ViViT models with different
learning rates.

ViT performance peaks with a learning rate of 2E-05 (F1 score 0.55, model
ViT-1), followed by 2E-06 (F1 score 0.50). ViViT performs the best with a
learning rate of 2E-04 (F1 score 0.55, model ViViT-1), and obtains its second
best performance with a learning rate of 2E-06 (F1 score 0.49, model ViViT-2).
However, a peculiar trait emerges from the latter model, an extremely low
standard deviation across folds of the F1 score as visible in Table 1. The
models ViT-1 (al = 12, lr = 2E-05) and ViViT-1 (al = 6, lr = 2E-04) are the
best performing models of their type across all experiments, while ViViT-2 (al
= 6, lr = 2E-06) is the second best performing video vision transformer and
has a uniquely low standard deviation. Visible in Figure 5 is a comparison of
the best two ViViT models, showcasing the good performance across all folds
for ViViT-2 compared to ViViT-1.

Figure 5: Performance (F1 score) per fold of the two best performing ViViT
models.

###  4.3 Sharpness-Aware Minimization

The third step of experimentation introduces SAM to the model’s training,
however, this addition not only almost doubles the training time necessary but
also worsens the performance of ViViT (F1 score 0.40) and ViT (F1 score 0.50),
as shown in Figure 6.

Figure 6: Performance (F1 score) of ViT and ViViT models optimized with and
without the Sharpness-Aware Minimization (SAM).

###  4.4 Mixup

The fourth experimental step augments 20% of the dataset with the Mixup
technique, with three different $\alpha$ configurations. Mixup, even with the
additional restriction of combining images belonging to the same patient,
fails to contribute to the model’s performance even with its best parameter
($\alpha$ = 0.8) for the ViT model (F1 score 0.52) and ViViT model (F1 score
0.52), as shown in Figure 7.

Figure 7: Performance (F1 score) of ViT and ViViT models fine-tuned with Mixup
augmentation.

###  4.5 Comparisons with previous works

To our knowledge, recent work by Rudovic et al. [29] is the state-of-the-art
for automated binary pain detection on the F1 score metric. Their experimental
setup uses a CNN baseline (CDL) but focuses on federalized learning (PFDL),
achieving its best-performing model with this technique. As can be seen in
Table 1 our method not only achieves better performance on the F1 score metric
with our best performing ViT and ViViT models, but their method is also
outperformed by our ViViT-2 model, which trades off performance for more
consistent results across folds.

Table 1: Model results on the UNBC-McMaster dataset Model Name | F1 score | AUC
---|---|---
CDL [29] | 0.46 $\pm$ 0.18 | -
PFDL [29] | 0.47 $\pm$ 0.20 | -
SPTS + CAPP [38] | - | 0.84
SPTS + SAPP + CAPP [45] | - | 0.85
ViT-1 |  0.55 $\pm$ 0.15 | 0.88
ViViT-1 |  0.55 $\pm$ 0.13 | 0.86
ViViT-2 | 0.49 $\pm$ 0.04 | 0.76

The F1 score is affected by the skew in the labels but AUC is not [46]. Given
that our labels are highly imbalanced, we also report AUC values and compare
our results with the works that also report AUC. We compare our top-performing
models ViT-1 (AUC 0.88) and ViViT-1 (AUC 0.86) against SPTS + CAPP (AUC 0.84)
[38] and SPTS + SAPP + CAPP (AUC 0.85) [45], and find them to outperform
previous works despite not being optimized for this metric.

###  4.6 Qualitative Analysis

Previous works have shown that attention maps can be used to generate visual
interpretations for the predictions of vision transformers [19, 34]. To
demonstrate this feature we will perform a qualitative analysis of the
attention maps of ViViT-1 given the sample visible in Figure 8, whose pain
label is correctly predicted by the model.

Figure 8: The pain sample used to qualitatively analyze the performance of
ViViT-1. The first frame in the top-left has a PSPI score of 0, while the
remaining three have a PSPI score of 10.

As shown in Figure 9, the last attention layer of ViViT-1 has disentangled
representations across its attention heads. These representations partially
overlap with AU43 (eyes closed, head 0), AU4 (brow-lowering, head 1) AU6-7
(orbital tightening, head 2), while others capture a large area of the face
(head 3).

Figure 9: Attention maps of individual heads of the final attention layer.
Disentangled representations similar to AUs are present in heads 0-2, while
head 3 captures most of the face. The activations are thresholded between 0.7
(blue) and 1 (red).

In Figure 10a, the maximum value of the attention patches for the ViViT-1
model is shown, obtained with attention rollout [47]. Attention rollout is a
transformer technique that combines information from every attention layer,
capturing its flow through the model. In Figure 10b, we show instead the
combined maximum values of the heads of the last attention layer for ViViT-1.
While the strongest activations are found in the forehead and cheek area for
the final layer (b), the flow of information instead clearly originates from
the inner brow, lip corner and cheek area (a), which are areas of significance
according to facial pain assessment literature [7].

Figure 10: Attention maps obtained with max rollout (a) and maximum values of
the last layer (b). Frame (b) is thresholded between 0.7 (blue) and 1 (red).

The model is clearly capable of generating intrinsic and plausible
interpretations for its predictions. Facial regions having higher attention
weights in the attention maps are the ones that show a change in appearance
and shape during several actions that are observed during a painful
expression. It shows that the model effectively detects pain from actions of
relevant facial regions.

##  5 Discussion

#### Number of unfrozen layers

The pre-trained transformer model is overall successful across a large variety
of parameters for the task of pain detection, contrary to earlier findings on
the topic [16]. While the models perform consistently no matter the number of
layers, the region around 6 and 12 layers stands out as the better choice both
for ViT and ViViT, warranting a deeper investigation of similar parameters and
setting a precedent for future work. Despite sharing the top two
configurations, the vision and video vision transformers perform best with
different numbers of layers, distinguishing their configurations for the
following steps.

#### Learning rate

Comparably, learning rate proves to be a far more delicate parameter, with two
of the configurations achieving the worst performance overall across all
experiments for ViT and ViViT. Learning rates lower than 2E-03 are generally
high-performing, with ViT peaking around 2E-05, achieving the best model
performance across all configurations, and ViViT performing best with 2E-04, a
much larger learning rate.

Furthermore, ViViT scores its second best performance at the learning rate of
2E-06 with an extremely low standard deviation across folds, a sign of good
generalization. Consistency is a desirable trait for all models, and even more
so for delicate tasks in the medical field. Therefore, while achieving only
the second best performance according to the metric chosen for this
experiment, it possesses a desirable trait and indicates a second region of
interest for the tuning of this parameter.

#### Sharpness-Aware Minimization and Mixup

The SAM optimizer and Mixup augmentation fail to improve the model’s
performance across a variety of configurations for both ViT and ViViT. While
SAM decreases the standard deviation across folds for the ViViT model, it does
so by affecting all folds negatively rather than pushing their performance
towards an average. Mixup appears ineffective in generating meaningful samples
for the model to learn from despite the constraints applied, perhaps due to
the delicate nature of FACS and PSPI encoding.

#### ViT and ViViT

ViViT fails to outperform ViT despite the strong case for facial dynamics in
pain detection literature. While the configurations we use may be limiting the
performance of the model, we believe that ViViT models which compute spatio-
temporal attention separately, such as the other ViViT implementations
described in the original paper [17] might suit this task better. Separating
the temporal and spatial attention would benefit the model by allowing larger
sequence lengths while avoiding quadratically increased computational time.
Furthermore, spatially or temporally-local computed attention might track
better the delicate facial dynamics necessary for pain detection.

#### Societal Impact

To our knowledge, the negative impact of this research is limited to the
physical resources used to train and run the models. The presented models
should not be used in a clinical context, for which they are untested.

#### Hardware and Training Time

Training has been performed with a variety of GPUs on Google Colab and Kaggle.
Total training time for all 14 configurations is around 30 hours for ViT on
Kaggle and around 80 for ViViT on Google Colab.

##  6 Conclusion

In this paper, we have used vision and video vision transformers trained on
the UNBC-McMaster dataset for binary pain detection. We have achieved state-
of-the-art performance using the F1 score, identified regions of interest in
the transformer hyperparameter space, compared the performance of vision and
video transformers on this task, and obtained intrinsic plausible
interpretations for the performance of the model. Results show that pre-
trained transformers can be applied toward pain assessment with good results,
after a single epoch of training and on a small unbalanced dataset. Future
work could include different augmentation techniques, leave-one-patient-out
validation, longer sub-sequences, and more efficient transformer
architectures.

## References

  * [1] S. Raja, D. Carr, M. Cohen, N. Finnerup, H. Flor, S. Gibson, F. Keefe, J. Mogil, M. Ringkamp, K. Sluka _et al._ , “The revised international association for the study of pain definition of pain: concepts, challenges, and compromises. pain,” _press. doi_ , vol. 10, 2020.
  * [2] H. Breivik, B. Collett, V. Ventafridda, R. Cohen, and D. Gallacher, “Survey of chronic pain in europe: prevalence, impact on daily life, and treatment,” _European journal of pain_ , vol. 10, no. 4, pp. 287–333, 2006.
  * [3] M. Hayes, “Experimental developement of the graphics rating method,” _Physiol Bull_ , vol. 18, pp. 98–99, 1921.
  * [4] A. B. Ashraf, S. Lucey, J. F. Cohn, T. Chen, Z. Ambadar, K. M. Prkachin, and P. E. Solomon, “The painful face–pain expression recognition using active appearance models,” _Image and vision computing_ , vol. 27, no. 12, pp. 1788–1796, 2009.
  * [5] Z. Hammal and J. F. Cohn, “Automatic, objective, and efficient measurement of pain using automated face analysis,” in _Social and interpersonal dynamics in pain_. Springer, 2018, pp. 121–146.
  * [6] A. C. d. C. Williams, H. T. O. Davies, and Y. Chadury, “Simple pain rating scales hide complex idiosyncratic meanings,” _Pain_ , vol. 85, no. 3, pp. 457–463, 2000.
  * [7] P. Ekman and W. V. Friesen, “Facial action coding system,” _Environmental Psychology & Nonverbal Behavior_, 1978.
  * [8] K. M. Prkachin and P. E. Solomon, “The structure, reliability and validity of pain expression: Evidence from patients with shoulder pain,” _Pain_ , vol. 139, no. 2, pp. 267–274, 2008.
  * [9] E. A. Clark, J. Kessinger, S. E. Duncan, M. A. Bell, J. Lahne, D. L. Gallagher, and S. F. O’Keefe, “The facial action coding system for characterization of human affective response to consumer product-based stimuli: a systematic review,” _Frontiers in psychology_ , vol. 11, p. 920, 2020.
  * [10] S. Kaltwang, O. Rudovic, and M. Pantic, “Continuous pain intensity estimation from facial expressions,” in _International Symposium on Visual Computing_. Springer, 2012, pp. 368–377.
  * [11] Z. Ambadar, J. W. Schooler, and J. F. Cohn, “Deciphering the enigmatic face: The importance of facial dynamics in interpreting subtle facial expressions,” _Psychological science_ , vol. 16, no. 5, pp. 403–410, 2005\.
  * [12] S. Rezaei, A. Moturu, S. Zhao, K. M. Prkachin, T. Hadjistavropoulos, and B. Taati, “Unobtrusive pain monitoring in older adults with dementia using pairwise and contrastive training,” _IEEE Journal of Biomedical and Health Informatics_ , vol. 25, no. 5, pp. 1450–1462, 2020.
  * [13] Y. Sun, C. Shan, T. Tan, X. Long, A. Pourtaherian, S. Zinger _et al._ , “Video-based discomfort detection for infants,” _Machine Vision and Applications_ , vol. 30, no. 5, pp. 933–944, 2019.
  * [14] S. Tonekaboni, S. Joshi, M. D. McCradden, and A. Goldenberg, “What clinicians want: contextualizing explainable machine learning for clinical end use,” in _Machine learning for healthcare conference_. PMLR, 2019, pp. 359–380.
  * [15] S. Walter, S. Gruss, S. Frisch, J. Liter, L. Jerg-Bretzke, B. Zujalovic, and E. Barth, ““what about automated pain recognition for routine clinical use?” a survey of physicians and nursing staff on expectations, requirements, and acceptance,” _Frontiers in medicine_ , p. 990, 2020.
  * [16] H. Xu and M. Liu, “A deep attention transformer network for pain estimation with facial expression video,” in _Chinese Conference on Biometric Recognition_. Springer, 2021, pp. 112–119.
  * [17] A. Arnab, M. Dehghani, G. Heigold, C. Sun, M. Lučić, and C. Schmid, “Vivit: A video vision transformer,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021, pp. 6836–6846.
  * [18] X. Chen, C.-J. Hsieh, and B. Gong, “When vision transformers outperform resnets without pre-training or strong data augmentations,” _arXiv preprint arXiv:2106.01548_ , 2021.
  * [19] K. Han, A. Xiao, E. Wu, J. Guo, C. Xu, and Y. Wang, “Transformer in transformer,” _Advances in Neural Information Processing Systems_ , vol. 34, 2021.
  * [20] R. Strudel, R. Garcia, I. Laptev, and C. Schmid, “Segmenter: Transformer for semantic segmentation,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2021, pp. 7262–7272.
  * [21] H. Chen, Y. Wang, T. Guo, C. Xu, Y. Deng, Z. Liu, S. Ma, C. Xu, C. Xu, and W. Gao, “Pre-trained image processing transformer,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2021, pp. 12 299–12 310.
  * [22] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention is all you need,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [23] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly _et al._ , “An image is worth 16x16 words: Transformers for image recognition at scale,” _arXiv preprint arXiv:2010.11929_ , 2020.
  * [24] Q. Fan, C.-F. R. Chen, H. Kuehne, M. Pistoia, and D. Cox, “More is less: Learning efficient video representations by big-little network and depthwise temporal aggregation,” _Advances in Neural Information Processing Systems_ , vol. 32, 2019.
  * [25] B. Jiang, M. Wang, W. Gan, W. Wu, and J. Yan, “Stm: Spatiotemporal and motion encoding for action recognition,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2019, pp. 2000–2009.
  * [26] Z. Hammal and J. F. Cohn, “Automatic detection of pain intensity,” in _Proceedings of the 14th ACM international conference on Multimodal interaction_ , 2012, pp. 47–52.
  * [27] P. Werner, A. Al-Hamadi, K. Limbrecht-Ecklundt, S. Walter, S. Gruss, and H. C. Traue, “Automatic pain assessment with facial activity descriptors,” _IEEE Transactions on Affective Computing_ , vol. 8, no. 3, pp. 286–299, 2017\.
  * [28] Z. Zafar and N. A. Khan, “Pain intensity evaluation through facial action units,” in _2014 22nd International Conference on Pattern Recognition_. IEEE, 2014, pp. 4696–4701.
  * [29] O. Rudovic, N. Tobis, S. Kaltwang, B. Schuller, D. Rueckert, J. F. Cohn, and R. W. Picard, “Personalized federated deep learning for pain estimation from face images,” _arXiv preprint arXiv:2101.04800_ , 2021.
  * [30] F. Wang, X. Xiang, C. Liu, T. D. Tran, A. Reiter, G. D. Hager, H. Quon, J. Cheng, and A. L. Yuille, “Regularizing face verification nets for pain intensity regression,” in _2017 IEEE International Conference on Image Processing (ICIP)_. IEEE, 2017, pp. 1087–1091.
  * [31] J. Zhou, X. Hong, F. Su, and G. Zhao, “Recurrent convolutional neural network regression for continuous pain intensity estimation in video,” in _Proceedings of the IEEE conference on computer vision and pattern recognition workshops_ , 2016, pp. 84–92.
  * [32] P. Rodriguez, G. Cucurull, J. Gonzàlez, J. M. Gonfaus, K. Nasrollahi, T. B. Moeslund, and F. X. Roca, “Deep pain: Exploiting long short-term memory networks for facial expression classification,” _IEEE transactions on cybernetics_ , 2017.
  * [33] F. Ma, B. Sun, and S. Li, “Facial expression recognition with visual transformers and attentional selective fusion,” _IEEE Transactions on Affective Computing_ , 2021.
  * [34] C. Wang and Z. Wang, “Progressive multi-scale vision transformer for facial action unit detection,” _Frontiers in Neurorobotics_ , vol. 15, 2021.
  * [35] T. Ridnik, E. Ben-Baruch, A. Noy, and L. Zelnik-Manor, “Imagenet-21k pretraining for the masses,” _arXiv preprint arXiv:2104.10972_ , 2021.
  * [36] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, A. C. Berg, and L. Fei-Fei, “ImageNet Large Scale Visual Recognition Challenge,” _International Journal of Computer Vision (IJCV)_ , vol. 115, no. 3, pp. 211–252, 2015.
  * [37] T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Y. Jernite, J. Plu, C. Xu, T. L. Scao, S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush, “Transformers: State-of-the-art natural language processing,” in _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_. Online: Association for Computational Linguistics, Oct. 2020, pp. 38–45. [Online]. Available: https://www.aclweb.org/anthology/2020.emnlp-demos.6
  * [38] P. Lucey, J. F. Cohn, K. M. Prkachin, P. E. Solomon, and I. Matthews, “Painful data: The unbc-mcmaster shoulder pain expression archive database,” in _2011 IEEE International Conference on Automatic Face & Gesture Recognition (FG)_. IEEE, 2011, pp. 57–64.
  * [39] Y. Feng, F. Wu, X. Shao, Y. Wang, and X. Zhou, “Joint 3d face reconstruction and dense alignment with position map regression network,” in _Proceedings of the European conference on computer vision (ECCV)_ , 2018, pp. 534–551.
  * [40] YadiraF, “face3d,” https://github.com/YadiraF/face3d, 2018.
  * [41] H. Zhang, M. Cisse, Y. N. Dauphin, and D. Lopez-Paz, “mixup: Beyond empirical risk minimization,” _arXiv preprint arXiv:1710.09412_ , 2017.
  * [42] J. Lee, R. Tang, and J. Lin, “What would elsa do? freezing layers during transformer fine-tuning,” _arXiv preprint arXiv:1911.03090_ , 2019.
  * [43] P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur, “Sharpness-aware minimization for efficiently improving generalization,” _arXiv preprint arXiv:2010.01412_ , 2020.
  * [44] X. Chen, C.-J. Hsieh, and B. Gong, “When vision transformers outperform resnets without pre-training or strong data augmentations,” _arXiv preprint arXiv:2106.01548_ , 2021.
  * [45] P. Lucey, J. F. Cohn, I. Matthews, S. Lucey, S. Sridharan, J. Howlett, and K. M. Prkachin, “Automatically detecting pain in video through facial action units,” _IEEE Transactions on Systems, Man, and Cybernetics, Part B (Cybernetics)_ , vol. 41, no. 3, pp. 664–674, 2010.
  * [46] L. A. Jeni, J. F. Cohn, and F. De La Torre, “Facing imbalanced data–recommendations for the use of performance metrics,” in _2013 Humaine association conference on affective computing and intelligent interaction_. IEEE, 2013, pp. 245–251.
  * [47] S. Abnar and W. Zuidema, “Quantifying attention flow in transformers,” _arXiv preprint arXiv:2005.00928_ , 2020.


# Attention Score is not All You Need for Token Importance
Indicator in KV Cache Reduction: Value Also Matters

Zhiyu Guo, Hidetaka Kamigaito, Taro Watanabe
Nara Institute of Science and Technology
{guo.zhiyu.fy1, kamigaito.h<EMAIL_ADDRESS>

###### Abstract

Scaling the context size of large language models (LLMs) enables them to
perform various new tasks, e.g., book summarization. However, the memory cost
of the Key and Value (KV) cache in attention significantly limits the
practical applications of LLMs. Recent works have explored token pruning for
KV cache reduction in LLMs, relying solely on attention scores as a token
importance indicator. However, our investigation into value vector norms
revealed a notably non-uniform pattern questioning their reliance only on
attention scores. Inspired by this, we propose a new method: Value-Aware Token
Pruning (VATP) which uses both attention scores and the $\ell_{1}$ norm of
value vectors to evaluate token importance. Extensive experiments on
LLaMA2-7B-chat and Vicuna-v1.5-7B across 16 LongBench tasks demonstrate VATP’s
superior performance.

Attention Score is not All You Need for Token Importance
Indicator in KV Cache Reduction: Value Also Matters

Zhiyu Guo, Hidetaka Kamigaito, Taro Watanabe Nara Institute of Science and
Technology {guo.zhiyu.fy1, kamigaito.h<EMAIL_ADDRESS>

##  1 Introduction

Recent studies have focused on scaling the context sizes of Transformer-based
Vaswani et al. (2017) large language models (LLMs) in addition to scaling
data, compute, and model size. For example, the context size has increased
from 2048 tokens in GPT-3 Brown et al. (2020) and LlaMA1 Touvron et al.
(2023a) to 2 million tokens in Gemini 1.5 Pro Reid et al. (2024). Longer
context sizes enable LLMs to address tasks that extend beyond conventional
capabilities, such as book-length summarization Chang et al. (2024), SWE-agent
Yang et al. (2024), and many-shot in-context learning Agarwal et al. (2024).
However, the enormous inference costs of LLMs limit their applications.
Therefore, in addition to model weight compression Dettmers et al. (2022); Sun
et al. (2024b), enhancing the efficiency of long-context inference has become
increasingly important.

LLMs utilize an auto-regressive framework in which tokens are produced
sequentially. The generation of each token relies on the tokens generated
before it. During generation, the key and value tensors of previously
generated tokens, known as the KV cache, have to be preserved in memory
throughout the generation process for attention computation. The memory cost
of the KV cache scales linearly to the batch size and sequence length. This
prohibitive memory cost has become a critical bottleneck limiting the
applications of long-context LLMs.

One of the approaches for improving long-context inference efficiency is token
pruning, which has been extensively explored for BERT Goyal et al. (2020);
Zhao et al. (2022); Guan et al. (2022). However, these methods necessitate a
complicated fine-tuning process to restore optimal performance. Given the
extensive text corpora and the considerable size of LLMs, such fine-tuning
becomes exceptionally challenging and less preferred. Fortunately, recent
studies Zhang et al. (2023); Liu et al. (2023); Ge et al. (2024); Xiao et al.
(2024) have explored token pruning for KV cache reduction without the need for
fine-tuning, indicating that a significant number of tokens can be pruned with
minimal impact on performance during token generation. It is notable that
these studies unanimously chose to rely solely on the attention score as the
token importance indicator in LLMs. This choice is reasonable for LLMs, as
training additional token importance predictor Guan et al. (2022) is
computationally expensive. Nevertheless, before establishing the attention
score as the default choice for the token importance indicator, we pose a
timely question: Are there any essential elements that may have been
accidentally omitted when considering pivotal tokens for KV cache reduction?

(a) Layer 1-2

(b) Layer 3-31

(c) Layer 32

Figure 1: Typical attention map (logarithmic) and value vector norm patterns
in LLaMA2-7B-chat. Key observations include: (1) The $\ell_{1}$ norms are non-
uniformly distributed across tokens in all layers and heads. (2) In figure
(b), for most heads in layers 3-31, regardless of the input text, there are
two attention sink Xiao et al. (2024) tokens at the beginning of the text.
Contrary to their massive attention scores, their $\ell_{1}$ norms are close
to 0 (highlighted in red). (3) In some heads of the last layer, the second
attention sink token in figure (b) has a smaller attention score than other
tokens, while its $\ell_{1}$ norm is significantly larger than those of other
tokens.

Since the output of the attention mechanism is the result of the
multiplication of each token’s attention score with its corresponding value
vector, we investigated the value vectors of LLMs. We found the $\ell_{1}$
norm of each token is highly non-uniformly distributed, showing distinct
differences in magnitude. Previous study Xiao et al. (2024) identifies the
attention sink tokens with massive attention scores. We find, in contrast to
the attention score, the value vector norms of the attention sink tokens are
much smaller than other tokens. Such a phenomenon is similar to the finding in
small Transformer models Kobayashi et al. (2020). When considering each
token’s effects on the attention output, their value vector should also be
considered.

Building upon this observation, we introduce a new approach termed Value-Aware
Token Pruning (VATP). Unlike traditional methods that rely solely on attention
scores, VATP augments the attention score with the norm of the value vector,
providing a robust metric for evaluating token importance. Specifically, we
propose a novel token pruning metric, where the KV cache of each token is
assessed based on the product of its attention score and the $\ell_{1}$ norm
of the corresponding value vector. We conduct extensive experiments on the
LLaMA2-7B-chat and Vicuna-v1.5-7B models, evaluating VATP across 16 long-
context tasks from the LongBench Bai et al. (2023) benchmark. The results
demonstrate that VATP outperforms attention-only baselines across a wide
variety of tasks. Our research clearly reveals the critical, yet previously
overlooked, role of the value vector in KV cache reduction, questioning the
prevailing belief that attention score is all you need for determining token
importance in LLMs.

##  2 Related Work

Recent works investigate reducing the KV cache of unimportant tokens during
auto-regressive generation. H2O Zhang et al. (2023) dynamically retain a
combination of recent and heavy-hitter tokens in the KV cache, which are
identified based on the accumulated attention scores. Concurrent work, i.e.,
Scissorhands Liu et al. (2023), uses the attention score from a history window
as the token importance indicator. FastGen Ge et al. (2024) adjusts its
compression strategies to align with the attention structure of each head.
StreamLLM Xiao et al. (2024) keeps the attention sink tokens together with the
sliding window tokens to anchor the attention computation and stabilize the
model’s performance. Notably, these methods primarily focus on attention
scores as the basis for determining which tokens to prune, our work for the
first time explores the critical role of value vector in token pruning.

##  3 Value-Aware Token Pruning

This section introduces Value-Aware Token Pruning (VATP), starting from
observations and concluding with our algorithm.

###  3.1 Observations

The output of an attention head at step $t$ is defined as follows:

$\text{Attention}(Q,K,V)_{t}=\sum_{i\leq t}a_{i}^{t}\boldsymbol{v}_{i}$ (1)

where $a_{i}^{t}$ is the attention score of query token $t$ to token $i$, and
$\boldsymbol{v}_{i}$ is the value state of token $i$. The attention output for
the token $t$ is thus a weighted sum of the value states $\boldsymbol{v}_{i}$
of all preceding tokens $i\leq t$, where the weights are the corresponding
attention scores $a_{i}^{t}$. The goal of token pruning is to remove tokens
that have a minimal impact on the attention output. From Equation (1), each
token’s influence on the attention output is determined by both the attention
score $a_{i}^{t}$ and the value vector $\boldsymbol{v}_{i}$.

Here, we jointly analyze the attention maps and the corresponding value
vectors. In Figure 1, the $\ell_{1}$ norm of value vector exhibits a highly
non-uniform distribution across all layers and heads. Notably, the two
attention sink tokens111The first token is the starting word token, the second
token is often the token representing the first period (.) or newline token
(\n) in the text, there are corresponding to massive activations as discussed
in (Sun et al., 2024a).  often show a striking contrast between their
attention scores and value vector norms. This observation is similar to the
study of small Transformer models (Kobayashi et al., 2020).

###  3.2 Methodology

The above observation highlights the importance of considering both the
attention score and the value vector norm together. Such a dual consideration
provides a more comprehensive understanding of each token’s influence on the
attention output. Consequently, it becomes obvious to implement token pruning
strategies that take into account attention score and value vector norm
simultaneously.

#### Attention Score

H2O Zhang et al. (2023) uses the accumulated attention score as token
importance indicator. Specially, the token importance score for a given token
$k$ at decoding step $t$ is calculated as:

$S_{k}^{t}=\sum_{k\leq j\leq t}a_{k}^{j}$ (2)

Scissorhands Liu et al. (2023) use the attention score based on the history
window with size $w$.

$S_{k}^{t}=\sum_{\max(t-w,k)\leq j\leq t}a_{k}^{j}$ (3)

#### Value-aware Pruning Metric

Motivated by the success of LLM weight pruning work Wanda Sun et al. (2024b),
which evaluates model weight importance by the product of its magnitude and
the corresponding input feature norm, we propose a new metric to evaluate
token importance. For each token in an attention head, its importance is
evaluated by the product of its attention score $S_{k}^{t}$ and the
corresponding value vector norm. Specifically, the score for the token $k$ at
decoding step $t$ is defined by:

$I_{k}^{t}=S_{k}^{t}\cdot\left\|\boldsymbol{v}_{k}\right\|_{1}$ (4)

where $\left\|\boldsymbol{v}_{k}\right\|_{1}$ is the $\ell_{1}$ norm of of
token $k$’s value vector. The attention score $S_{k}^{t}$ can be either Eq.
(2) or Eq. (3). We empirically find $\ell_{1}$ norm performs better than
$\ell_{2}$ norm in Appendix B. The computation of VATP metric is
straightforward by jointly considering the attention score and value vector.

#### Attention Sink Tokens

In our metric, the importance scores of attention sink tokens are downgraded,
and they could be accidentally removed. While the value updates from those
tokens may be small, the attention distribution of the rest tokens will be
largely shifted after the removal, leading to deteriorated performance Xiao et
al. (2024). Thus we intentionally keep the first $\mathbf{F}$ tokens.

| Method | Single-Document QA | Multi-Document QA | Summarization | Few-shot Learning | Synthetic | Code
---|---|---|---|---|---|---|---
| 1-1 | 1-2 | 1-3 | 2-1 | 2-2 | 2-3 | 3-1 | 3-2 | 3-3 | 4-1 | 4-2 | 4-3 | 5-1 | 5-2 | 6-1 | 6-2
LlaMA2-7B-chat | All Budget | 19.12 | 20.99 | 37.55 | 30.55 | 27.44 | 8.31 | 27.77 | 20.67 | 24.39 | 58.33 | 86.22 | 39.14 | 3.89 | 9.67 | 59.88 | 48.61
StreamLLM | 15.4 | 18.6 | 25.96 | 28.19 | 23.59 | 7.08 | 23.87 | 19.97 | 22.52 | 56.67 | 86.45 | 38.72 | 3.87 | 2.62 | 58.55 | 48.28
H2O | 18.4 | 18.83 | 33.67 | 30.18 | 25.74 | 7.85 | 26.18 | 21.12 | 23.44 | 58.67 | 85.35 | 39.0 | 4.37 | 7.0 | 59.4 | 49.09
w/ VATP | 18.77 | 19.6 | 35.31 | 29.95 | 27.15 | 8.44 | 26.08 | 21.14 | 23.76 | 58.33 | 86.09 | 38.74 | 4.39 | 8.33 | 59.56 | 49.46
Scissorhands | 18.5 | 19.32 | 36.35 | 29.5 | 25.51 | 8.59 | 25.42 | 20.35 | 23.86 | 57.33 | 85.55 | 38.77 | 4.38 | 6.0 | 58.33 | 48.86
w/ VATP | 19.4 | 19.53 | 36.58 | 29.57 | 27.71 | 9.66 | 26.17 | 20.46 | 23.63 | 58.0 | 85.98 | 38.9 | 4.18 | 10.0 | 59.39 | 48.71
vicuna-v1.5-7B | All Budget | 18.67 | 23.39 | 39.25 | 27.48 | 19.62 | 8.09 | 30.84 | 22.85 | 24.7 | 64.33 | 86.53 | 39.69 | 4.33 | 13.0 | 50.15 | 36.52
StreamLLM | 16.97 | 21.55 | 26.01 | 23.79 | 16.94 | 5.83 | 26.6 | 21.94 | 22.48 | 62.67 | 86.4 | 39.54 | 2.0 | 11.33 | 49.56 | 37.79
H2O | 18.46 | 21.84 | 32.99 | 26.86 |  19.79 | 6.04 | 27.92 | 23.2 | 23.78 | 64.0 | 79.06 | 39.19 | 4.33 | 11.67 | 51.38 | 36.86
w/ VATP | 18.86 | 21.89 | 36.94 | 28.23 | 19.47 | 7.72 | 28.57 | 23.21 | 23.74 | 64.33 | 86.57 | 40.02 | 4.33 | 13.0 | 50.18 | 37.54
Scissorhands | 18.19 | 20.56 | 33.83 | 26.08 | 18.93 | 5.99 | 26.14 | 22.62 | 23.24 | 61.67 | 80.26 | 39.7 | 4.33 | 9.0 | 49.91 | 35.05
w/ VATP | 18.99 | 21.95 | 37.63 | 28.22 | 20.3 | 7.98 | 27.82 | 23.45 | 23.44 | 62.33 | 86.36 | 39.89 | 4.33 | 13.0 | 48.73 | 36.11

Table 1: Performance of different token pruning methods on LongBench at 50% KV
cache budget. To streamline the text, following Bai et al. (2023), we refer to
the dataset as ID (eg., 1-1 map to NarrativeQA, 2-2 map to 2WikiMultihopQA);
the mapping from ID to the dataset and evaluation metrics are available in
Table 3 of Appendix.

##  4 Experiments

###  4.1 Settings

#### Models

We use two open LLMs, LLaMA2-7B-chat Touvron et al. (2023b) and
Vicuna-v1.5-7B-16k Zheng et al. (2023). For LLaMA2-7B-chat, we set the max
sequence length as 4K. For Vicuna-v1.5-7B-16k, we set the max sequence length
as 8K due to GPU memory limitation. We conduct all experiments using one A6000
GPU.

#### Dataset

To extensively assess the effectiveness of our method in real-world scenarios,
we select all the English tasks in LongBench Bai et al. (2023) as our
evaluation benchmark. The Longbench benchmark consists of 16 English tasks,
each containing between 150 and 500 samples. This benchmark encompasses a
diverse array of long-text tasks, including question answering, text
summarization, few-shot learning, synthetic tasks, and code completion. The
detailed information about the dataset is in Table 3 in Appendix C. We use the
official task-specific prompts to evaluate task-wise performance of
instruction-tuned LLMs.

#### Baselines

We choose several token pruning works: StreamLLM Xiao et al. (2024), H2O Zhang
et al. (2023), Scissorhands Liu et al. (2023). The entire KV cache is used for
assessing the performance degradation. We incorporate VATP directly into H2O
and Scissorhands, resulting in two variants: “H2O w/ VATP” and “Scissorhands
w/ VATP”. Implementation details are in Appendix A.

###  4.2 Results

#### Main results

While individual task results may exhibit variability, the aggregate results
presented in Table 1 are more stable and reliable. For the LlaMA2-7B-chat
model, the VATP method surpasses H2O in 12 out of 16 tasks and outperforms
Scissorhands in 13 out of 16 tasks. Similarly, for the Vicuna-v1.5-7B-16k
model, VATP exceeds H2O in 12 out of 16 tasks and Scissorhands in 14 out of 16
tasks. Note that for the tasks where VATP does not surpass the baseline, its
performance is still very comparable to the baseline. However, in certain
tasks (e.g., 1-3 in Vicuna), VATP significantly outperforms the baseline. The
overall performance gains demonstrate the effectiveness of our proposed
method.

#### KV budget ratio variation.

In Figure 2, we test the performance on 2WikiMultihopQA at different KV cache
reduction ratio. Scissorhands w. VATP yields the best performance in nearly
all the KV budget ratios, in the high reduction ratio, it outperforms
Scissorhands significantly. The improvement of H2O w/VATP over H2O is mainly
at the less aggressive reduction ratio.

Figure 2: Performance on 2WikiMultihopQA of the LLaMA2-7B-chat with varying KV
Cache Ratios.

#### Inference Efficiency.

VATP maintains the inherent simplicity of baseline methods by introducing
negligible computation and memory overhead compared with H2O and Scissorhands,
since the size of the value vector norm is $\frac{1}{d_{\text{head}}}$ of the
KV cache size, ${d_{\text{head}}}=128$ for a 7B model. In practice, we observe
no noticeable latency or memory difference between H2O and VATP at the same KV
budget. For comprehensive benchmarks on inference efficiency of token pruning,
we defer the reader to H2O (Zhang et al., 2023).

##  5 Conclusion

This study addresses a critical yet previously overlooked aspect of token
pruning in LLMs — the value vectors. Motivated by the observed highly non-
uniform distribution of value vector norms, we propose a novel token pruning
method, Value-Aware Token Pruning (VATP), that jointly considers both the
attention score and the value vector norm. Our extensive experiments
demonstrate that VATP consistently outperforms traditional attention-only
methods across various long-context tasks.

## Limitations

Our work has the following limitations:

#### Flash attention support:

When integrating with H2O, we need to calculate the accumulated attention
score. However, the current implementation of Flash attention does not return
the attention matrix. Without integrating Flash attention, the memory cost of
prompt prefilling remains $O(n^{2})$. Although it’s unnecessary to store the
attention matrix for all layers simultaneously, handling extensive context
significantly increases the memory cost during prompt prefilling. We
anticipate this issue will be resolved with future improvements in H2O
implementations.

#### Compatibility with grouped-query attention:

Similar to Scissorhands and H2O, our method is currently not applicable to
grouped-query attention Ainslie et al. (2023). Exploring the use of token
pruning to further reduce the KV cache in this context represents a promising
direction for future research.

## References

  * Agarwal et al. (2024) Rishabh Agarwal, Avi Singh, Lei M Zhang, Bernd Bohnet, Stephanie Chan, Ankesh Anand, Zaheer Abbas, Azade Nova, John D Co-Reyes, Eric Chu, et al. 2024.  Many-shot in-context learning.  _arXiv preprint arXiv:2404.11018_.
  * Ainslie et al. (2023) Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. 2023.  GQA: Training generalized multi-query transformer models from multi-head checkpoints.  In _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pages 4895–4901, Singapore. Association for Computational Linguistics.
  * Bai et al. (2023) Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. 2023.  Longbench: A bilingual, multitask benchmark for long context understanding.  _arXiv preprint arXiv:2308.14508_.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901.
  * Chang et al. (2024) Yapei Chang, Kyle Lo, Tanya Goyal, and Mohit Iyyer. 2024.  Booookscore: A systematic exploration of book-length summarization in the era of LLMs.  In _The Twelfth International Conference on Learning Representations_.
  * Dettmers et al. (2022) Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. 2022.  GPT3.int8(): 8-bit matrix multiplication for transformers at scale.  In _Advances in Neural Information Processing Systems_.
  * Ge et al. (2024) Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. 2024.  Model tells you what to discard: Adaptive KV cache compression for LLMs.  In _The Twelfth International Conference on Learning Representations_.
  * Goyal et al. (2020) Saurabh Goyal, Anamitra Roy Choudhury, Saurabh Raje, Venkatesan Chakaravarthy, Yogish Sabharwal, and Ashish Verma. 2020.  Power-bert: Accelerating bert inference via progressive word-vector elimination.  In _International Conference on Machine Learning_ , pages 3690–3699. PMLR.
  * Guan et al. (2022) Yue Guan, Zhengyi Li, Jingwen Leng, Zhouhan Lin, and Minyi Guo. 2022.  Transkimmer: Transformer learns to layer-wise skim.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 7275–7286, Dublin, Ireland. Association for Computational Linguistics.
  * Kobayashi et al. (2020) Goro Kobayashi, Tatsuki Kuribayashi, Sho Yokoi, and Kentaro Inui. 2020.  Attention is not only a weight: Analyzing transformers with vector norms.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 7057–7075, Online. Association for Computational Linguistics.
  * Liu et al. (2023) Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava. 2023.  Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time.  In _Thirty-seventh Conference on Neural Information Processing Systems_.
  * Reid et al. (2024) Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean-baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. 2024.  Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.  _arXiv preprint arXiv:2403.05530_.
  * Sun et al. (2024a) Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu. 2024a.  Massive activations in large language models.  In _ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models_.
  * Sun et al. (2024b) Mingjie Sun, Zhuang Liu, Anna Bair, and J Zico Kolter. 2024b.  A simple and effective pruning approach for large language models.  In _The Twelfth International Conference on Learning Representations_.
  * Touvron et al. (2023a) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023a.  Llama: Open and efficient foundation language models.  _Preprint_ , arXiv:2302.13971.
  * Touvron et al. (2023b) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, and et al. 2023b.  Llama 2: Open foundation and fine-tuned chat models.  _Preprint_ , arXiv:2307.09288.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _Advances in neural information processing systems_ , 30.
  * Xiao et al. (2024) Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. 2024.  Efficient streaming language models with attention sinks.  In _The Twelfth International Conference on Learning Representations_.
  * Yang et al. (2024) John Yang, Carlos E. Jimenez, Alexander Wettig, Kilian Lieret, Shunyu Yao, Karthik Narasimhan, and Ofir Press. 2024.  Swe-agent: Agent-computer interfaces enable automated software engineering.  _Preprint_ , arXiv:2405.15793.
  * Zhang et al. (2023) Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Re, Clark Barrett, Zhangyang Wang, and Beidi Chen. 2023.  H2o: Heavy-hitter oracle for efficient generative inference of large language models.  In _Thirty-seventh Conference on Neural Information Processing Systems_.
  * Zhao et al. (2022) Jing Zhao, Yifan Wang, Junwei Bao, Youzheng Wu, and Xiaodong He. 2022.  Fine- and coarse-granularity hybrid self-attention for efficient BERT.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 4811–4820, Dublin, Ireland. Association for Computational Linguistics.
  * Zheng et al. (2023) Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric. P Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. 2023.  Judging llm-as-a-judge with mt-bench and chatbot arena.  _Preprint_ , arXiv:2306.05685.

##  Appendix A Implementation Details

For each input sequence, we set the KV cache budget to 50% of the input prompt
length. We assign uniform KV budgets across different heads and layers, as a
uniform strategy is more practical to achieve actual inference improvements in
hardware. FastGen Ge et al. (2024) uses a non-uniform strategy, thus we
haven’t chosen it as a baseline. For StreamLLM, we set the number of attention
sink tokens to 20 for LLaMA2-7B-chat and 40 for Vicuna-v1.5-7B-16k. For VATP,
we intentionally keep the first $\mathbf{F}$ tokens, where $\mathbf{F}=20$ for
LLaMA2-7B-chat and $\mathbf{F}=40$ for Vicuna-v1.5-7B-16k. Given a KV budget
of $k$ tokens, the number of tokens selected by importance score in H2O is
$\frac{k}{2}$, with a local window size also of $\frac{k}{2}$. In
Scissorhands, following Liu et al. (2023), the number of tokens selected by
importance score is $k-10$, with a local window size of 10 and a history
window size of $w=400$. When integrating with Scissorhands and H2O, the only
differences are the token importance score and intentionally keeping attention
sink tokens.

##  Appendix B Ablation Study

Table 2 shows the F1 scores for 4 QA tasks under different types of norm for
value vector: $\ell_{1}$, $\ell_{2}$, and $\ell_{\infty}$. Overall, the
$\ell_{1}$ norm achieves the highest average performance with an average F1
score of 28.00, indicating that $\ell_{1}$ norm performs better across the
evaluated tasks compared to $\ell_{2}$ and $\ell_{\infty}$ norms. Thus we use
$\ell_{1}$ norm in all the experiments.

Task |  $\ell_{1}$ Norm |  $\ell_{2}$ Norm |  $\ell_{\infty}$ Norm
---|---|---|---
Qasper | 19.60 | 18.47 | 18.76
MultifieldQA (en) | 35.31 | 35.48 | 35.11
HotpotQA | 29.95 | 30.30 | 30.20
2WikiMQA | 27.15 | 26.55 | 26.76
Average | 28.00 | 27.70 | 27.71

Table 2: F1 Scores under Different Norms for 4 QA Tasks

##  Appendix C Dataset Details

We select the English subset from Longbench Bai et al. (2023). Table 3 shows
the information of 16 tasks we use in the experiments.

ID | Dataset | Avg len | Metric | #data
---|---|---|---|---
1-1 | NarrativeQA | 18,409 | F1 | 200
1-2 | Qasper | 3,619 | F1 | 200
1-3 | MultiFieldQA-en | 4,559 | F1 | 150
2-1 | HotpotQA | 9,151 | F1 | 200
2-2 | 2WikiMultihopQA | 4,887 | F1 | 200
2-3 | MuSiQue | 11,214 | F1 | 200
3-1 | GovReport | 8,734 | Rouge-L | 200
3-2 | QMSum | 10,614 | Rouge-L | 200
3-3 | MultiNews | 2,113 | Rouge-L | 200
4-1 | TREC | 5,177 | Accuracy (CLS) | 200
4-2 | TriviaQA | 8,209 | F1 | 200
4-3 | SAMSum | 6,258 | Rouge-L | 200
5-1 | PassageCount | 11,141 | Accuracy (EM) | 200
5-2 | PassageRetrieval-en | 9,289 | Accuracy (EM) | 200
6-1 | LCC | 1,235 | Edit Sim | 500
6-2 | RepoBench-P | 4,206 | Edit Sim | 500
Table 3: The dataset statistics in LongBench include several key metrics.
’Source’ indicates where the context originates. The ’Avg len’ (average
length) is measured by the number of words for datasets in English (or code).
’Accuracy (CLS)’ represents classification accuracy, while ’Accuracy (EM)’
denotes exact match accuracy.