A2SF: Accumulative Attention Scoring with Forgetting Factor for Token Pruning in Transformer Decoder

TL;DR

This paper introduces A2SF, a novel token pruning method for transformer decoders that uses a forgetting factor to fairly evaluate token importance, improving model accuracy and memory efficiency.

Contribution

A2SF is the first to incorporate a forgetting factor into accumulative attention scoring for decoder token pruning, addressing uneven token importance evaluation.

Findings

A2SF improves LLaMA 2 accuracy by up to 7.8% in 1-shot settings.

A2SF enhances OPT and LLaMA models' token selection effectiveness.

The method effectively reduces memory usage while maintaining or improving accuracy.

Abstract

Recently, large language models (LLM) based on transformers are facing memory bottleneck issues due to KV cache, especially in long sequence handling. Previous researches proposed KV cache compression techniques that identify insignificant tokens based on Accumulative Attention Scores and removes their items from KV cache, noting that only few tokens play an important role in attention operations. However, we have observed that the existing Accumulative Attention Score is not suitable for the transformer decoder structure. In the decoder model, the number of times the Attention Score accumulates varies depending on the order of token appearance due to the effect of masking, causing an uneven comparison between tokens. To solve this, we propose Accumulative Attention Score with Forgetting Factor (A2SF) technique, which introduces a Forgetting Factor in the Attention Score accumulation process. A2SF applies a penalty to the past Attention Score generated from old tokens by repeatedly multiplying the Forgetting Factor to the Attention Score over time. Therefore, older tokens receive a larger penalty, providing fairness among different ages of tokens. Through the fair comparison among tokens, we can more effectively select important tokens. We have verified the accuracy improvement through A2SF in the OPT and LLaMA models and A2SF improves the accuracy of LLaMA 2 by up to 7.8% and 5.1% on 1-shot and 0-shot.