Anchor-based Large Language Models (2024)

Abstract: Large language models (LLMs) predominantly employ decoder-only transformer architectures, necessitating the retention of keys/values information for historical tokens to provide contextual information and avoid redundant computation. However, the substantial size and parameter volume of these LLMs require massive GPU memory. This memory demand increases with the length of the input text, leading to an urgent need for more efficient methods of information storage and processing. This study introduces Anchor-based LLMs (AnLLMs), which utilize an innovative anchor-based self-attention network (AnSAN) and also an anchor-based inference strategy. This approach enables LLMs to compress sequence information into an anchor token, reducing the keys/values cache and enhancing inference efficiency. Experiments on question-answering benchmarks reveal that AnLLMs maintain similar accuracy levels while achieving up to 99% keys/values cache reduction and up to 3.5 times faster inference. Despite a minor compromise in accuracy, the substantial enhancements of AnLLMs employing the AnSAN technique in resource utilization and computational efficiency underscore their potential for practical LLM applications.

## introduction large language models (llms) primarily utilize decoder-only transformer architectures, which necessitate caching keys/values information for historical tokens during the auto-regressive inference to supply contextual information and avoid redundant computation @xcite @xcite . however, due to their immense size and high parameter count, a considerable amount of gpu memory is required for loading. furthermore, as the length of input text grows, storing keys/values caches requires more and more gpu memory, as evidenced in in-context learning, complex instructions, and extended conversations @xcite @xcite , which is not conducive to scenarios with limited computational resources. an alternative approach entails recalculating these extensive inputs, which, however, results in increased time overhead. therefore, this study aims to reduce the storage demand for keys/values caches during the inference phase of llms, improving the memory efficiency and, consequently, accelerating the inference speed as well. in a recent study, @xcite demonstrate that label words in prefix demonstrations can act as anchors during inference, providing an effective context compression approach for improving inference efficiency in in-context learning. however, in practical applications, not all prefix inputs or demonstrations contain label words suitable for compressing information, making the reliance on label words a less universal approach for text information compression. additionally, @xcite observe that llms tend to attend to only a few, yet consistent, prefix tokens during inference. however, the specific tokens utilized are often unpredictable and uncontrollable. these observations raise an intriguing question: do natural language texts contain anchor points that compress the overall semantic information of sequences? in this context, previous studies on sequence embeddings have shown that the hidden state of a special token in neural network models can encapsulate semantic information @xcite . furthermore, contemporary llms typically utilize the causal self-attention mechanism during both training and inference phases @xcite , attending on each prior token. this suggests that the final token in a sequence may be better suited to serve as a natural information compression point compared to other tokens, as they cannot observe future tokens. therefore, a methodical approach that identifies and exploits sequence anchor tokens in a dependable and controllable manner is essential for compressing sequence information, effectively reducing keys/values caches, and improving inference efficiency for llms. to this end, we propose novel anchor-based large language models (anllms), equipped with an innovative anchor-based self-attention network (ansan) and an anchor-based inference strategy. the ansan is designed to compel the models to compress sequence information into the anchor token (the last token in our implementation) during the training process, with the aid of anchor-based attention masks. during inference, the anchor-based inference strategy retains the keys/values caches of anchor tokens, which have aggregated the entire sequence information, and discards those of nonanchor tokens, thereby reducing memory demands. specifically, the anchor-based attention masks for ansan serve two objectives: 1) to ensure anchor tokens attend exclusively to tokens within the same sequence, preventing attention to other sequences, and 2) to direct non-anchor tokens' attention to previous sequence anchors, blocking the other nonanchor tokens from previous sequences. it is noteworthy that the technique of anchor-based attention bears similarities to the principles underlying sparse attention @xcite . however, unlike the existing research that employs sparse attention to extend the context length of llms @xcite , our method focuses on continually pre-training the model to compress sequence information into the anchor token. in our implementation, we utilize the publicly available redpajama datasets @xcite to continuously pre-train the open-source llama2 models @xcite , resulting in an-llms that incorporate our proposed anchor-based attention mechanism. experimental results on question answering benchmarks, as depicted in figure 1 , reveal that our method achieves up to a 99% reduction in keys/values caches and up to a 3.5-fold increase in inference acceleration ratios, while maintaining comparable accuracy to the original model. despite a minor decrease in accuracy (within 1.5%), these findings underscore the significant improvements in computational efficiency and memory utilization offered by our method. ## related work our research is inspired by the recent investigation into the understanding of in-context learning (icl) within llms by @xcite . in their study, the authors delve into the underlying mechanisms of icl, emphasizing the influence of label words in demonstration examples on information flow. they reveal that these label words serve as anchors, wherein semantic information converges into these anchors during inference, subsequently directing the llms' final predictions. motivated by their findings, our objective is to extend this feature to natural language modeling by guiding sequence information compression into manually designed anchor tokens, rather than solely relying on label words. this is crucial because natural language texts may not always contain an explicit label. the most relevant method to our approach in the existing literature is the learning to compress prompts with gist tokens @xcite . their approach centers around compressing task-specific prompts by fine-tuning the model using the proposed gist masking, thereby enforcing prompt compression. however, there are several crucial divergences between our study and theirs. unlike their focus on compressing a task prompt, our objective lies in training the llm to condense sequence information into the anchor tokens. consequently, our approach can be universally applied to a range of tasks without requiring task-specific training, a feature not shared by gist tokens, as the anchor tokens are seamlessly incorporated into the model's lan-guage modeling. furthermore, our anchor-based attention masks account for information compression within a sequence and information interaction between sequences, thus extending beyond the mere compression of task prompts. on the other hand, flashattention @xcite and pagedattention @xcite both present memory-efficient attention mechanisms for llms. while they focus on optimizing attention computation and subdividing attention processing, our proposed method offers a distinct approach that specifically targets the compression of sequence information into anchor tokens, making it orthogonal to these existing works. ## experimental results as evident from the results presented in ## analysis to further elucidate our method's insights, we conduct a natural language generation experiment with the german-to-english (de2en) translation task. we evaluate the models using comet-da @xcite , indicating translation quality, and the keys/values cache reduction c ⇓ metric, denoting memory efficiency as previously described. in line with previous findings, anllms accept a minor accuracy trade-off (about 3 comet-da points) for enhanced memory efficiency. all llms are finetuned on the alpaca dataset, combined with the newstest2017-2020 datasets, following @xcite . results are presented in ## conclusion llms have emerged as a significant research area in the field of artificial intelligence. however, despite their exceptional performance across various natural language tasks, the practical application of these models is limited by their significant memory overhead and time efficiency. implementing llms on resource-constrained devices, such as smartphones, poses a unique challenge. to address this issue, we propose anchor-based llms with the ansan technique. our experiments demonstrate that by sacrificing a marginal 1.5% in precision, our approach saves 99% of keys/values cache memory while simultaneously improving inference speed by up to 3.5 times. our methods' application in machine translation showcases their compatibility and flexibility, effectively enhancing memory efficiency for practical use. our novel approach is practical, straightforward, flexible, and compatible with existing methods, paving the way for further adoption of llms in real-world applications.