Everything K-culture — comebacks to K-beauty, straight to your inboxGet it in your inbox›
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.
METAL MEDIA explanatory visual
Teaching AI models to forget the right things when reading very long documents
01Long-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
02Sparse 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
03The 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
04They also rebuild the H2O policy with a Triton-based FlashInfer kernel that returns summed attention weights efficiently, making it much faster than prior implementations
05On 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)
An explanatory diagram made by METAL MEDIA, not a figure supplied by the paper's authors.
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.
