Learning how to Forget: Fine-tuning for Long-Context Sparse Attention
arXiv:2608.199202026-08-21
让处理超长文本的AI模型学会该忘记什么
处理超长文本的语言模型需要用一个容量有限的缓存来保存之前读过的内容,这就要决定该扔掉哪些信息。这篇论文提出了一种新的微调方法,让模型能与所用的缓存淘汰策略协同适应,只需一块40GB显存的Nvidia A100 GPU就能运行,而且在实验中往往比用传统精确注意力(序列并行)训练的模型表现更好。作者还开源了KeysAndValues库,里面包含这套方法以及更高效的H2O缓存策略实现。
他们做了什么
- 处理数十万token长文本的语言模型需要用KV缓存保存之前token的信息,而这个缓存会随文本长度线性增长,很快超出GPU显存承受范围
- H2O等稀疏注意力方法把缓存大小固定住,淘汰掉不太有用的旧信息,但此前一直缺少高效的方法让模型针对这种缓存策略重新训练
- 作者把嵌套式激活检查点、CPU数据卸载,以及利用缓存前后关联做差量编码(借助PyTorch自动求导的saved tensor hooks机制)结合起来,使微调所需资源接近推理所需资源,并且适用于任意缓存策略
- 他们还用Triton代码重写了H2O策略,搭配FlashInfer计算核心高效返回注意力权重之和,大幅提升了速度
- 在Helmet基准测试中,用该方法训练并配合稀疏注意力推理的模型,往往优于用序列并行(把长文本拆分到多块GPU上做精确计算的标准方法)训练的模型
Table 1: Results for long-context inference with 5 KV cache policies and chunk sizes 2048=2k,1024=1k (rows). The first row exact is for exact inference (sequence parallelism). We show SubEM values on test splits for different Helmet datasets nq, trivia_qa, hotpot_qa, pop_qa, limiting sequence lengths to 64k or 128k tokens. Columns us are for models trained using our novel method with the same cache policy in place, columns sp are for models trained with sequence parallelism. | 64k datasets | 128k datasets |
|---|
| nq | tri_qa | hot_qa | pop_qa | nq | tri_qa | hot_qa | pop_qa |
| us | sp | us | sp | us | sp | us | sp | us | sp | us | sp | us | sp | us | sp |
| exact | - | 50.7 | - | 79.8 | - | 60.0 | - | 62.7 | - | 50.7 | - | 68.7 | - | 46.3 | - | 57.0 |
| lr2k | 33.5 | 57.2 | 75.3 | 57.5 | 53.3 | 62.7 | 43.7 | 60.5 | 26.0 | 33.0 | 50.8 | 52.3 | 31.0 | 46.3 | 34.0 | 25.0 |
| slr2k | 47.3 | 56.5 | 74.5 | 60.8 | 50.0 | 67.3 | 44.0 | 56.7 | 26.0 | 33.7 | 61.2 | 51.8 | 34.0 | 42.0 | 37.7 | 22.2 |
| h2o2k | 47.2 | 70.8 | 78.0 | 72.0 | 53.0 | 68.7 | 57.5 | 44.7 | 24.2 | 40.7 | 47.8 | 63.7 | 19.3 | 26.0 | 53.3 | 49.8 |
| h2o2kno | 47.8 | 68.2 | 63.2 | 54.5 | 58.3 | 70.0 | 53.0 | 39.8 | 43.5 | 51.3 | 66.7 | 55.3 | 37.3 | 51.0 | 50.2 | 25.2 |
| h2o2kor | 49.5 | 73.3 | 66.3 | 65.7 | 57.3 | 68.7 | 62.2 | 45.5 | 45.3 | 58.8 | 71.2 | 71.0 | 36.7 | 44.3 | 50.2 | 33.3 |
| lr1k | 59.7 | 57.0 | 73.0 | 60.7 | 47.7 | 65.3 | 41.7 | 59.3 | 23.5 | 32.7 | 59.5 | 50.0 | 28.7 | 45.0 | 34.5 | 25.2 |
| slr1k | 37.8 | 57.0 | 59.8 | 59.7 | 52.7 | 65.0 | 46.8 | 57.0 | 29.3 | 36.2 | 58.2 | 49.7 | 33.7 | 44.7 | 34.3 | 21.2 |
| h2o1k | 62.0 | 72.8 | 79.7 | 72.0 | 55.0 | 68.0 | 59.3 | 45.8 | 23.2 | 41.3 | 51.7 | 58.7 | 25.3 | 24.3 | 51.0 | 50.5 |
| h2o1kno | 47.5 | 71.3 | 62.3 | 59.3 | 61.7 | 73.0 | 51.0 | 43.5 | 42.7 | 53.3 | 70.3 | 58.2 | 42.3 | 54.7 | 44.1 | 26.8 |
| h2o1kor | 49.0 | 72.2 | 75.7 | 68.7 | 57.0 | 66.7 | 55.0 | 47.3 | 44.3 | 60.5 | 72.0 | 74.8 | 31.0 | 41.3 | 51.0 | 30.8 |
Table 2: Results for 6 additional Helmet datasets not featured in Table 1 (context width 128k). Inference under 3 KV cache policies (chunk size 1024=1k), exact uses sequence parallelism (column). trn denotes model checkpoint being used: us uses our novel method with the same cache policy in place, sp is using sequence parallelism, no is the base checkpoint Qwen3-4B-Instruct-2507 (no fine-tuning). Note that metrics are different, depending on the dataset (see Table 3).| dataset | trn | exact | slr1k | h2o1kno | h2o1kor |
|---|
| trec_coarse | us | - | 96.0 | 96.4 | 96.2 |
| sp | 97.8 | 30.0 | 23.2 | 77.6 |
| no | - | 28.2 | 19.8 | 36.0 |
| nlu | us | - | 90.0 | 87.4 | 79.8 |
| sp | 90.2 | 28.6 | 32.8 | 74.0 |
| no | - | 24.8 | 30.0 | 21.2 |
| clinc150 | us | - | 97.4 | 96.8 | 94.0 |
| sp | 97.6 | 64.2 | 61.6 | 68.0 |
| no | - | 62.6 | 54.0 | 34.8 |
| inf_qa | us | - | 26.6 | 32.2 | 36.8 |
| sp | 40.8 | 2.2 | 2.2 | 3.3 |
| no | - | 2.5 | 2.9 | 3.4 |
| inf_mc | us | - | 40.0 | 42.0 | 54.0 |
| sp | 66.0 | 25.0 | 29.0 | 39.0 |
| no | - | 36.0 | 41.0 | 40.0 |
| json_kv | us | - | 49.0 | 50.0 | 3.0 |
| sp | 100.0 | 0.0 | 0.0 | 1.0 |
| no | - | 0.0 | 0.0 | 0.0 |
Table 3: Overview of the 10 Helmet tasks. Dev and Eval denote the number of instances in the training and evaluation partitions, respectively, at a single context-length setting.| Category | ID | Source | Metric | Dev | Eval |
|---|
| RAG | nq | Natural Questions | SubEM | 893 | 600 |
| trivia_qa | TriviaQA | SubEM | 876 | 600 |
| pop_qa | PopQA | SubEM | 192 | 600 |
| hotpot_qa | HotpotQA | SubEM | 787 | 300 |
| Many-shot ICL | trec_coarse | TREC | Accuracy | 1000 | 500 |
| nlu | SNIPS NLU | Accuracy | 2094 | 500 |
| clinc150 | CLINC150 | Accuracy | 2600 | 500 |
| Long-doc QA | inf_qa | InfiniteBench QA | ROUGE-F1 | 251 | 100 |
| inf_mc | InfiniteBench MC | Accuracy | 129 | 100 |
| Synthetic Recall | json_kv | JSON-KV | SubEM | 500 | 100 |
Table 4: Results for long-context inference with 5 KV cache policies and chunk sizes 2048=2k,1024=1k (rows). Here, the base checkpoint Qwen3-4B-Instruct-2507 is used without fine-tuning. The first row exact is for exact inference (sequence parallelism). We show 𝚜𝚞𝚋_𝚎𝚡𝚊𝚌𝚝_𝚖𝚊𝚝𝚌𝚑 values on test splits for different Helmet datasets nq, trivia_qa, hotpot_qa, pop_qa, limiting sequence lengths to 64k or 128k tokens. | 64k datasets | 128k datasets |
|---|
| nq | tri_qa | hot_qa | pop_qa | nq | tri_qa | hot_qa | pop_qa |
| exact | - | - | - | - | - | - | - | - |
| lr2k | 47.5 | 80.2 | 57.7 | 53.8 | 35.3 | 70.5 | 37.0 | 40.2 |
| slr2k | 45.3 | 80.2 | 55.3 | 52.0 | 37.5 | 69.7 | 34.7 | 34.8 |
| h2o2k | 42.8 | 70.8 | 44.7 | 59.0 | 20.5 | 63.2 | 12.3 | 27.3 |
| h2o2kno | 46.0 | 82.2 | 52.7 | 59.3 | 42.2 | 78.5 | 34.3 | 39.3 |
| h2o2kor | 44.2 | 74.7 | 47.7 | 61.0 | 43.7 | 79.8 | 25.0 | 43.3 |
| lr1k | 47.5 | 80.2 | 59.7 | 54.0 | 35.8 | 70.2 | 37.7 | 37.5 |
| slr1k | 49.0 | 79.0 | 55.0 | 54.3 | 37.3 | 67.3 | 32.0 | 35.7 |
| h2o1k | 43.2 | 71.7 | 41.7 | 59.8 | 21.0 | 62.8 | 13.0 | 27.0 |
| h2o1kno | 47.0 | 81.2 | 50.0 | 56.7 | 40.0 | 79.8 | 32.7 | 42.3 |
| h2o1kor | 47.5 | 76.3 | 46.7 | 58.2 | 41.8 | 79.8 | 23.7 | 39.3 |
Table 5: Results for long-context inference with setups not covered in the main text. We show 𝚜𝚞𝚋_𝚎𝚡𝚊𝚌𝚝_𝚖𝚊𝚝𝚌𝚑 values on test splits for different Helmet datasets nq, trivia_qa, hotpot_qa, pop_qa, limiting sequence lengths to 64k or 128k tokens. slr128, h2o128, h2o128no, h2o128or use chunk size S=128. qh2o2k and qh2o2kno are variants of Q-Hitter [81]. | 64k datasets |
|---|
| nq | tri_qa | hot_qa | pop_qa |
| slr128 | 38.8 | 66.2 | 51.7 | 41.0 |
| h2o128 | 44.5 | 65.3 | 55.7 | 54.3 |
| h2o128no | 49.0 | 74.3 | 55.0 | 56.3 |
| h2o128or | 46.8 | 67.3 | 48.3 | 52.3 |
| qh2o2k | 37.5 | 64.0 | 40.3 | 53.7 |
| qh2o2kno | 40.2 | 65.5 | 46.3 | 53.5 |
Table 6: Running time figures for training update step, for Helmet 128k datasets (columns), 5 KV cache policies and chunk sizes 2048=2k,1024=1k (rows). Batch size 8, running on 4 devices. The step from 2k to 1k is 11% to 13% more expensive for 𝚕𝚛,𝚜𝚕𝚛, 12% to 14% more expensive for 𝚑𝟸𝚘 variants. The step from 𝚕𝚛,𝚜𝚕𝚛 to 𝚑𝟸𝚘 variants is 2% to 3% more expensive for 2k, 3% to 4% more expensive for 1k. | nq | tri_qa | hot_qa | pop_qa |
|---|
| exact | 258.38 (15.14) | 266.05 (11.31) | 262.53 (8.18) | 236.74 (21.44) |
| lr2k | 326.07 (21.86) | 333.28 (16.56) | 330.76 (22.49) | 312.29 (27.71) |
| slr2k | 323.84 (21.19) | 333.03 (16.19) | 330.29 (22.55) | 310.73 (24.46) |
| h2o2k | 330.83 (21.98) | 344.83 (16.70) | 336.79 (23.13) | 316.02 (27.66) |
| h2o2kno | 331.48 (22.10) | 341.71 (16.65) | 337.77 (23.12) | 317.60 (25.16) |
| h2o2kor | 331.83 (22.25) | 344.85 (16.89) | 338.46 (23.25) | 316.94 (25.34) |
| lr1k | 364.10 (23.28) | 374.80 (18.62) | 375.38 (26.38) | 343.96 (27.99) |
| slr1k | 362.84 (23.20) | 370.55 (18.66) | 371.18 (25.77) | 344.77 (27.65) |
| h2o1k | 378.85 (24.88) | 385.19 (20.21) | 382.83 (26.94) | 359.51 (32.00) |
| h2o1kno | 378.61 (24.99) | 385.29 (19.32) | 388.77 (27.69) | 357.95 (29.02) |
| h2o1kor | 378.65 (24.79) | 385.64 (19.46) | 382.52 (26.91) | 358.99 (31.74) |
Table 7: Token length statistics of generated samples for 10 Helmet datasets (of context width 128k) and 3 cache logics. trn denotes model checkpoint being used: us uses our novel method with the same cache policy in place, sp is using sequence parallelism, no is the base checkpoint Qwen3-4B-Instruct-2507 (no fine-tuning). R is based on the ratio of output length to target length (in tokens), p128 (in percent) is the fraction of outputs of maximal size 128 (means, and stddevs over all test set samples). | trn | slr1k | h2o1kno | h2o1kor |
|---|
| | R | p128 | R | p128 | R | p128 |
| nq | us | 1.1±0.8 | 0.0±0.0 | 1.1±1.0 | 0.0±0.0 | 1.1±2.7 | 0.2±4.1 |
| sp | 35.5±24.1 | 99.5±7.1 | 35.5±21.6 | 100.0±0.0 | 36.2±22.1 | 99.3±8.1 |
| no | 35.2±22.0 | 97.7±15.1 | 35.8±22.8 | 99.2±9.1 | 34.9±22.0 | 97.8±14.6 |
| trivia_qa | us | 1.2±0.8 | 0.0±0.0 | 1.3±1.0 | 0.0±0.0 | 1.1±0.7 | 0.0±0.0 |
| sp | 35.4±24.1 | 97.3±16.1 | 38.2±25.2 | 99.2±9.1 | 42.0±24.4 | 96.7±18.0 |
| no | 43.4±28.5 | 87.8±32.7 | 46.5±28.9 | 96.0±19.6 | 48.2±30.4 | 95.7±20.4 |
| hotpot_qa | us | 1.0±0.5 | 0.0±0.0 | 1.0±0.6 | 0.0±0.0 | 1.3±1.4 | 1.0±9.9 |
| sp | 39.2±31.8 | 93.0±25.5 | 40.9±31.4 | 98.3±12.8 | 41.4±31.1 | 99.0±9.9 |
| no | 41.0±30.9 | 91.7±27.6 | 41.3±31.1 | 98.0±14.0 | 41.4±31.1 | 98.7±11.5 |
| pop_qa | us | 1.1±0.6 | 0.0±0.0 | 1.1±0.5 | 0.0±0.0 | 1.0±0.4 | 0.0±0.0 |
| sp | 53.8±24.9 | 98.8±10.7 | 54.4±24.7 | 99.5±7.1 | 55.0±24.7 | 98.7±11.5 |
| no | 58.9±29.3 | 93.5±24.7 | 60.0±30.9 | 98.8±10.7 | 58.2±28.9 | 96.8±17.5 |
| trec_coarse | us | 1.0±0.0 | 0.0±0.0 | 1.0±0.0 | 0.0±0.0 | 1.0±0.0 | 0.0±0.0 |
| sp | 127.6±6.1 | 99.6±6.3 | 127.3±8.4 | 99.0±9.9 | 128.0±0.1 | 99.6±6.3 |
| no | 128.0±0.0 | 100.0±0.0 | 128.0±0.0 | 100.0±0.0 | 128.0±0.1 | 99.4±7.7 |
| nlu | us | 1.0±0.1 | 0.0±0.0 | 1.0±0.1 | 0.0±0.0 | 1.0±0.2 | 0.0±0.0 |
| sp | 70.4±21.1 | 70.6±45.6 | 69.0±21.7 | 45.6±49.8 | 14.6±18.8 | 8.4±27.7 |
| no | 71.7±20.8 | 100.0±0.0 | 71.7±20.8 | 100.0±0.0 | 71.7±20.8 | 99.8±4.5 |
| clinc150 | us | 1.0±0.1 | 0.0±0.0 | 1.0±0.1 | 0.0±0.0 | 1.0±0.1 | 0.0±0.0 |
| sp | 30.2±28.6 | 43.8±49.6 | 32.6±31.0 | 49.8±50.0 | 57.6±32.2 | 87.6±33.0 |
| no | 59.9±19.7 | 100.0±0.0 | 59.9±19.7 | 100.0±0.0 | 60.1±20.5 | 100.0±0.0 |
| inf_qa | us | 1.1±0.8 | 0.0±0.0 | 1.2±0.8 | 0.0±0.0 | 1.2±0.8 | 0.0±0.0 |
| sp | 45.7±32.4 | 100.0±0.0 | 45.7±32.4 | 99.0±9.9 | 45.8±32.3 | 97.0±17.1 |
| no | 45.5±32.3 | 95.0±21.8 | 45.7±32.4 | 99.0±9.9 | 45.7±32.4 | 96.0±19.6 |
| inf_mc | us | 1.0±0.0 | 0.0±0.0 | 2.3±12.6 | 1.0±9.9 | 1.0±0.0 | 0.0±0.0 |
| sp | 128.0±0.0 | 100.0±0.0 | 128.0±0.0 | 100.0±0.0 | 127.9±1.2 | 98.0±14.0 |
| no | 128.0±0.0 | 100.0±0.0 | 128.0±0.0 | 100.0±0.0 | 128.3±3.6 | 97.0±17.1 |
| json_kv | us | 1.1±0.1 | 0.0±0.0 | 1.0±0.0 | 0.0±0.0 | 3.7±1.0 | 85.0±35.7 |
| sp | 4.1±0.3 | 100.0±0.0 | 4.1±0.3 | 100.0±0.0 | 4.1±0.3 | 98.0±14.0 |
Table 8: Token length statistics of generated samples for 10 Helmet datasets (of context width 128k) for training and inference with exact attention (sequence parallelism).| nq | tri_qa | hot_qa | pop_qa | trec_c |
|---|
| R | p128 | R | p128 | R | p128 | R | p128 | R | p128 |
| 1.1±1.1 | 0.0±0.0 | 1.0±0.6 | 0.0±0.0 | 1.1±0.6 | 0.0±0.0 | 1.0±0.4 | 0.0±0.0 | 1.0±0.0 | 0.0±0.0 |
| nlu | clc150 | inf_qa | inf_mc | json_kv |
| R | p128 | R | p128 | R | p128 | R | p128 | R | p128 |
| 1.0±0.1 | 0.0±0.0 | 1.0±0.0 | 0.0±0.0 | 1.2±0.9 | 0.0±0.0 | 1.0±0.0 | 0.0±0.0 | 1.0±0.0 | 0.0±0.0 |
为什么重要
处理长文档、多轮对话或调用大量工具的AI任务对硬件成本要求很高,这项工作让模型能够适配省内存的缓存方案,而不必依赖昂贵的多GPU配置。开源代码库也方便研究者尝试新的缓存策略或微调思路。
本文术语
- KV缓存 · 保存此前处理过的token的键和值信息,供模型后续注意力计算使用的内存空间
- 稀疏注意力 · 只保留或关注部分历史token,从而节省计算量和内存的技术
- H2O(heavy-hitter oracle) · 根据历史使用程度给缓存条目打分,淘汰最不重要条目的缓存策略
- 序列并行 · 把长文本拆分到多块GPU上进行精确注意力计算的标准训练方式
- 激活检查点 · 反向传播时不保存全部中间结果、而是按需重新计算以节省内存的技术
论文原文摘要（英文）
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.
作者 · Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter
在 arXiv 阅读最新论文
- LoRA-GA$^2$: Low Rank Adaptation with Multi-step Gradient Adaptive Alignment在正式微调前先偷看几步训练的梯度,让LoRA的初始化更聪明arXiv:2608.198002026-08-21
- Remember, Verify, or Ask? Cross-Family Evaluation of Memory Commitment in LLM AgentsAI助手在该向你提问的时候,却更愿意自己去核实事实arXiv:2608.195642026-08-21
- Robust Incomplete Multimodal Sentiment Analysis via Iterative Proxy Correction文本信息缺失或损坏时,这个AI不靠一次性猜测,而是反复修正猜测结果,从而更准确地判断情绪arXiv:2608.199712026-08-21
- Generating Diverse Personas for User Simulators to Test Interview Dialogue Systems要测试访谈式对话系统需要大量不同性格的虚拟用户,这项研究用大语言模型自动生成这些虚拟用户人设arXiv:2608.195492026-08-21
- Rethinking Patch Based Multivariate Time Series Forecasting with Semantic Structured Partitioning别再机械切分时间序列,按语义把它切成有意义的块arXiv:2608.199662026-08-21
- Reliable Financial Named Entity Recognition under Domain ShiftAI在正式文件里学到的自信,一到推特上就变得不可信arXiv:2608.195582026-08-21
- Bringing analytic rigor to agentic AI for science: The Brain Researcher platform for neuroimaging data analysis让AI分析脑影像数据时,把“为什么这个结论可信”也一并记录下来arXiv:2608.199022026-08-21
- GenMatch: An End-to-End Generative Matching Framework for Micro-View Order-Dispatching in Ride-Hailing滴滴把打车派单从预测-计算-匹配三段式流程改成一次生成完成,线上效果提升明显arXiv:2608.197512026-08-21
全部论文 →