Learning how to Forget: Fine-tuning for Long-Context Sparse Attention
arXiv:2608.199202026-08-21
Teaching AI models to forget the right things when reading very long documents
Language models that handle very long contexts need to limit the memory (KV cache) that stores past tokens, which means deciding what to discard. This paper introduces a new fine-tuning method that lets a model co-adapt with whatever cache-eviction policy is used, running on a single Nvidia A100 GPU with 40 GB RAM, and it often beats models trained the standard way with exact attention across sequence-parallel GPUs. The authors also release KeysAndValues, an open source library with this method and a faster implementation of the H2O caching policy.
What they did
- Long-context language models must store past token information in a KV cache that grows with context length, quickly exceeding GPU memory for contexts of hundreds of thousands of tokens
- Sparse attention methods like H2O fix the cache to a limited size and evict less useful entries, but until now there was no efficient way to fine-tune a model to work well with such a policy
- The authors combine nested activation checkpointing, CPU offloading, and a delta-encoding trick for the cache buffers (via PyTorch autograd saved tensor hooks) so fine-tuning needs only about as much memory as inference itself, working with any eviction policy
- They also rebuild the H2O policy with a Triton-based FlashInfer kernel that returns summed attention weights efficiently, making it much faster than prior implementations
- On Helmet benchmark tasks, models fine-tuned with their method under sparse attention often outperformed models trained via sequence parallelism (the standard exact-attention approach distributing context across multiple GPUs)
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 |
Why it matters
Handling long documents, multi-turn conversations, or long chains of tool calls is expensive on hardware, and this work shows a way to adapt models to memory-saving caching schemes without needing many GPUs. The released open source library also gives researchers an easy way to try out new cache policies or fine-tuning ideas.
Terms in this paper
- KV cache · memory that stores keys and values from previously processed tokens so the model can attend to them
- sparse attention · a technique that only keeps or attends to a limited subset of past tokens to save compute and memory
- H2O (heavy-hitter oracle) · a cache policy that scores cache entries by past usefulness and evicts the least useful ones
- sequence parallelism · the standard approach of splitting a long context across multiple GPUs to compute exact attention
- activation checkpointing · a memory-saving technique that recomputes intermediate results during backpropagation instead of storing them all
Original 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.
Authors · Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter
Read on arXiv