Learning how to Forget: Fine-tuning for Long-Context Sparse Attention
arXiv:2608.199202026-08-21
긴 문서를 읽는 AI가 '무엇을 잊을지'까지 함께 배우게 만들다
AI 언어모델이 아주 긴 문서를 처리하려면 그동안 읽은 내용을 저장하는 캐시 공간을 아껴 써야 하는데, 이 연구는 그 캐시를 어떻게 비울지 결정하는 정책과 모델이 함께 학습하도록 만드는 새 훈련법을 제안한다. 값비싼 여러 대의 GPU 없이 A100 GPU 한 대로도 훈련이 가능하며, 여러 실험에서 기존의 '정확한 어텐션'으로 훈련한 모델보다 오히려 성능이 좋았다. 연구팀은 이 방법과 개선된 H2O 캐시 정책을 담은 KeysAndValues라는 오픈소스 라이브러리도 함께 공개했다.
무엇을 했나
- 초거대 문맥(수십만 토큰)을 처리하는 언어모델은 이전에 읽은 토큰 정보를 담는 'KV 캐시'가 문맥 길이에 비례해 커져 GPU 메모리를 감당하기 어려워지는 문제가 있다
- 이를 해결하려고 캐시 크기를 고정하고 덜 중요한 정보를 지우는 '희소 어텐션' 기법들(H2O 등)이 나왔지만, 지금까지는 모델을 이런 캐시 정책에 맞춰 다시 훈련시키는 효율적인 방법이 없었다
- 연구팀은 캐시 항목 사이의 관계를 압축해서 저장하는 방식과 체크포인팅, CPU로 임시 데이터를 옮기는 기법을 결합해, 어떤 캐시 정책이든 A100 GPU 한 대 수준의 자원으로 모델을 미세조정할 수 있게 했다
- 가장 널리 쓰이는 캐시 정책인 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 캐시 · 모델이 이전에 처리한 토큰들의 정보(키, 값)를 저장해두는 메모리 공간
- 희소 어텐션 · 모든 과거 토큰을 다 보지 않고 일부만 골라 계산량과 메모리를 줄이는 기법
- 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에서 원문 보기