用稀疏注意力让普通显卡跑百万级上下文
Exploiting Sparsity for Long Context Inference: Million Token Contexts on Commodity GPUs
- 每步只关注最相关的前k个词，减少计算量
- 仅用不到2%的输入词，性能保持95%以上
- 16GB显存就能跑百万词长文本，适合普通设备
当前对超长输入（数十万词）的推理需求日益增长，但传统Transformer模型在长上下文上的自注意力计算成本过高，难以在消费级硬件上运行。为此，我们提出一种可调节的稀疏化机制，在每个生成步骤中仅对最重要的前k个词进行注意力计算，显著降低前向传播开销。实验表明，该方法可在约16GB GPU内存下实现长达100万词的上下文推理。模型对键值稀疏化具有强鲁棒性：当仅关注不足2%的输入词时，仍能在RULER、AlpacaEval和Open LLM Leaderboard等基准测试中保持超过95%的原始性能。
原文摘要 · Abstract (English)
There is growing demand for performing inference with hundreds of thousands of input tokens on trained transformer models. Inference at this extreme scale demands significant computational resources, hindering the application of transformers at long contexts on commodity (i.e not data center scale) hardware. To address the inference time costs associated with running self-attention based transformer language models on long contexts and enable their adoption on widely available hardware, we propose a tunable mechanism that reduces the cost of the forward pass by attending to only the most relevant tokens at every generation step using a top-k selection mechanism. We showcase the efficiency gains afforded by our method by performing inference on context windows up to 1M tokens using approximately 16GB of GPU RAM. Our experiments reveal that models are capable of handling the sparsity induced by the reduced number of keys and values. By attending to less than 2% of input tokens, we achieve over 95% of model performance on common benchmarks (RULER, AlpacaEval, and Open LLM Leaderboard).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。