通过分块注意力机制,让大模型推理提速近4.5倍,同时保持98%以上准确率。
ChunkLLM: A Lightweight Pluggable Framework for Accelerating LLMs Inference
- 引入可插拔的QK适配器与分块适配器,动态识别关键文本块。
- 在12万词长文本上实现4.48倍加速,长文本性能保留98.64%。
- 训练时冻结主干网络,仅微调适配器,适合部署在资源受限设备。
基于Transformer的大模型在自然语言处理和计算机视觉中表现优异,但因自注意力机制随输入token数量呈二次增长,导致计算效率低下。近期研究提出基于块选择与压缩的方法缓解此问题,但存在语义不完整或训练-推理效率差等缺陷。为此,我们提出ChunkLLM,一种轻量级、可插拔的训练框架。具体包括两个组件:置于每层Transformer中的QK适配器(Q-Adapter和K-Adapter),用于特征压缩与块注意力获取;位于模型底层的分块适配器,利用上下文语义信息检测分块边界。训练阶段冻结主干参数,仅训练适配器。设计了注意力蒸馏方法以提升关键块的召回率。推理时仅当检测到分块边界才触发分块选择,从而加速。在多种长短文本基准数据集上的实验表明,ChunkLLM在短文本任务上表现相当,长文本任务性能保持98.64%,关键值缓存保留率达48.58%。尤其在处理12万词长文本时,相比原始Transformer最高提速4.48倍。
原文摘要 · Abstract (English)
Transformer-based large models excel in natural language processing and computer vision, but face severe computational inefficiencies due to the self-attention's quadratic complexity with input tokens. Recently, researchers have proposed a series of methods based on block selection and compression to alleviate this problem, but they either have issues with semantic incompleteness or poor training-inference efficiency. To comprehensively address these challenges, we propose ChunkLLM, a lightweight and pluggable training framework. Specifically, we introduce two components: QK Adapter (Q-Adapter and K-Adapter) and Chunk Adapter. The former is attached to each Transformer layer, serving dual purposes of feature compression and chunk attention acquisition. The latter operates at the bottommost layer of the model, functioning to detect chunk boundaries by leveraging contextual semantic information. During the training phase, the parameters of the backbone remain frozen, with only the QK Adapter and Chunk Adapter undergoing training. Notably, we design an attention distillation method for training the QK Adapter, which enhances the recall rate of key chunks. During the inference phase, chunk selection is triggered exclusively when the current token is detected as a chunk boundary, thereby accelerating model inference. Experimental evaluations are conducted on a diverse set of long-text and short-text benchmark datasets spanning multiple tasks. ChunkLLM not only attains comparable performance on short-text benchmarks but also maintains 98.64% of the performance on long-context benchmarks while preserving a 48.58% key-value cache retention rate. Particularly, ChunkLLM attains a maximum speedup of 4.48x in comparison to the vanilla Transformer in the processing of 120K long texts.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。