让模型学会适应稀疏注意力策略，实现长文本推理的高效微调。
Learning how to Forget: Fine-tuning for Long-Context Sparse Attention
- 通过新微调方法使模型与稀疏注意力缓存策略协同优化。
- 仅需单张A100 GPU即可完成训练，性能优于精确注意力模型。
- 开源库KeysAndValues支持高效实现，适合长文本任务研究者。
此前许多工作通过稀疏注意力机制进行键值缓存选择与压缩，以在不增加硬件负担的情况下实现Transformer语言模型的长上下文推理。本文提出一种适用于任意键值缓存策略的新型微调方法，可在中等硬件预算（如单张40GB显存的Nvidia A100 GPU）下运行，并使模型与缓存策略共同适应，通常表现优于使用精确注意力（序列并行）训练的模型。此外，本文还实现了H2O稀疏注意力（实验中表现最佳的策略）的高效版本，配备专用缩放点积注意力核。KeysAndValues（https://github.com/awslabs/keys_values）是为长上下文推理与微调设计的新开源库，提供易用且高性能的代码支持所有所述方法。
原文摘要 · Abstract (English)
A lot of prior work addressed key-value (KV) cache selection and compression by sparse attention to enable long-context inference for transformer language models without excessive hardware budgets. We provide a new method for fine-tuning models with sparse attention. It works for any KV cache policy, runs on a moderate hardware budget (e.g., a single Nvidia A100 GPU with 40 GB RAM), and allows the model to co-adapt with the policy, often outperforming models trained with exact attention (sequence parallelism). We also provide an efficient implementation of H2O sparse attention (the leading policy in our experiments) with dedicated scaled dot product attention kernel support. KeysAndValues (https://github.com/awslabs/keys_values), a new open source library for long-context inference and fine-tuning, provides easy-to-use and performant code for all methods discussed here.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。