💬 NLP

Learning how to Forget: Fine-tuning for Long-Context Sparse Attention

This paper introduces a fine-tuning method for transformer models that enables efficient long-context inference with sparse attention on moderate hardware by allowing models to co-adapt with various KV cache policies, often outperforming exact attention approaches.

Original authors: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Published 2026-08-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Modern artificial intelligence systems that generate human-like text rely on a mechanism that acts as a short-term memory, allowing them to remember what was said earlier in a conversation or a long document. This memory is stored in a digital buffer that grows larger with every new word the system processes. For these systems to function well over long stretches of text, this memory needs to be vast, but the computer hardware required to hold it is expensive and limited. When the memory buffer fills up, the system must decide which old pieces of information to discard to make room for the new. If it throws away the wrong information, the system loses its ability to reason or answer questions accurately. This creates a difficult trade-off: keeping the memory small saves money and allows the system to run on standard equipment, but it risks losing the context needed to be smart.

Researchers have long tried to solve this by teaching the system to be selective about what it keeps, a process known as sparse attention. However, a new study reveals a critical flaw in how these systems have been trained up until now. Most existing methods train the AI using a perfect, unlimited memory and then try to force it to operate with a limited one later. The researchers found that this approach fails because the AI never learned how to function under the specific constraints it would actually face. By training the model to forget intentionally and adapt to a fixed memory size from the very beginning, the team demonstrated that the system could perform significantly better than those trained with unlimited resources, even when running on a single, moderately powerful computer chip.

The team, led by scientists at Amazon Web Services and the University of Amsterdam, developed a new way to fine-tune these large language models. Instead of using massive supercomputers to simulate perfect memory, they taught the models to co-adapt with a specific memory management policy. Imagine a librarian who is trained to organize books in a library with infinite shelves, only to be told later to work in a tiny room with a single shelf. The librarian trained in the big library would likely struggle to prioritize what to keep in the small room. In contrast, the method proposed in this paper trains the librarian directly in the small room, teaching them exactly which books to keep and which to discard based on the rules of that specific space. This allows the model to learn the rhythm of its own limitations, rather than trying to unlearn the habits of having too much space.

The researchers tested this approach on a model with four billion parameters, a size that is substantial but manageable. They ran their experiments on a single graphics card with 40 gigabytes of memory, a setup that is affordable for many organizations compared to the clusters of dozens of cards required by previous methods. They compared their new training technique against the standard method, which uses a technique called sequence parallelism to split the memory load across multiple expensive devices. The results showed that the models trained with the new method often outperformed the standard ones, particularly when the task required the system to generate specific, concise answers rather than long, rambling text. In several tests involving complex questions and data extraction, the standard method produced outputs that were far too long and filled with random, nonsensical numbers, while the new method learned to stop at the right time and provide the correct single value.

A key part of this success was improving the "heavy-hitter oracle," a popular strategy for deciding which information to keep. This strategy works by tracking which pieces of information the model pays the most attention to over time. The researchers found that the original version of this strategy was slow and inefficient. They rewrote the underlying code to work much faster, allowing the system to calculate these importance scores without slowing down the entire process. This optimization meant that the system could make smart decisions about what to forget in real-time, without needing the massive computational power that usually accompanies such tasks. The team also released a new open-source software library to make these techniques available to others, aiming to lower the barrier for anyone wanting to build long-context AI systems without needing a fortune in hardware.

The study highlights that the way a model is trained is just as important as the hardware it runs on. When the researchers forced the model to train with the exact same memory constraints it would face during use, it learned to navigate those constraints effectively. In one specific test involving JSON data, the standard method failed completely, unable to find the correct data points, while the new method succeeded in identifying them about half the time. This suggests that the ability to handle long contexts is not just a matter of having more memory, but of teaching the system how to manage the memory it has. The findings indicate that for many applications, the most effective path forward is not to build bigger computers, but to teach the software to be more efficient with the resources it already possesses.

Technical Summary: Learning how to Forget: Fine-tuning for Long-Context Sparse Attention

1. Problem Statement

Modern large language models (LLMs) require the ability to process extremely long contexts for tasks such as multi-turn conversations, tool use, and chain-of-thought reasoning. While standard transformer implementations scale quadratically in compute and linearly in memory with context width, sparse attention methods have been proposed to achieve linear time and constant memory scaling. These methods utilize a fixed-size Key-Value (KV) cache where information is evicted based on specific policies (e.g., Heavy-Hitter Oracle, H2O) once the cache is full.

However, a significant gap exists between inference and training in this domain:

Training Constraints: Fine-tuning models with sparse attention policies is notoriously difficult on moderate hardware budgets. Standard automatic differentiation (autograd) requires storing the full computation graph, including KV cache buffers for every chunk of the sequence. For long contexts, this memory requirement exceeds available GPU memory, often necessitating expensive multi-GPU setups with sequence or context parallelism (SP/CP).
Co-adaptation Gap: Models trained with exact attention (using SP/CP) often perform significantly worse when deployed with sparse attention inference. This is because the model is never exposed to the specific memory constraints and eviction logic it will face during inference.
Implementation Inefficiency: Leading sparse attention policies like H2O are underutilized in state-of-the-art (SotA) inference libraries (e.g., vLLM) because current implementations are slow, lacking optimized kernels for the specific requirements of sparse attention (e.g., returning summed attention weights).

2. Methodology

The paper proposes a novel fine-tuning method that allows transformer models to co-adapt with arbitrary KV cache policies on moderate hardware (demonstrated on a single Nvidia A100 40GB GPU).

2.1 Nested Activation Checkpointing with CPU Offloading

To address the memory bottleneck of training on long sequences ( $N \gg N_C$ , where $N_C$ is the cache size), the authors employ a nested activation checkpointing strategy:

Outer Loop (Layers): The backward pass iterates over model layers.
Inner Loop (Cells): The sequence is divided into "cells" (groups of chunks). For each cell, a forward pass is run to store necessary KV cache buffers to CPU memory.
Replay Mechanism: Instead of differentiating through the KV cache policy (which is often non-differentiable or expensive), the forward pass logs all eviction decisions. During the backward pass, these decisions are replayed using "replay caches" that mimic the forward eviction logic without requiring gradient computation through the policy itself.

2.2 Delta Encoding for KV Cache Buffers

The most critical memory optimization involves how KV cache buffers are stored within the autograd graph.

Linear Recurrence: The authors observe that the KV cache buffer for a new chunk ($keys'$) is a linear transformation of the previous buffer ($keys$), differing only by the values of the new tokens and the specific slots overwritten by the eviction policy.
Delta Storage: Instead of storing the full $keys$ and $values$ tensors for every chunk in the computation graph, the method stores only the delta (the difference) and the indices of the overwritten slots.
Autograd Saved Tensors Hooks: Using PyTorch's autograd_saved_tensors_hooks, the authors implement a "packing" mechanism. During the forward pass, full tensors are replaced by references to (index, delta) tuples. During the backward pass, the full tensors are reconstructed ("unpacked") from the deltas and subsequent buffers. This reduces the memory footprint of the autograd graph by a factor of $k$ (the number of chunks per cell), bringing it down to levels comparable to inference.

2.3 H2O Policy Improvements

The paper introduces specific improvements to the Heavy-Hitter Oracle (H2O) policy:

Batch Independence: The original H2O implementation aggregated scores across batch dimensions. The new method allows eviction decisions to be made independently for each batch position.
Normalization: A normalized H2O score is introduced to prevent bias toward tokens that have been in the cache longer.
Kernel Optimization: To enable efficient training and inference, the authors provide Triton code that returns summed attention weights alongside a FlashInfer Scaled Dot Product Attention (SDPA) kernel. This allows H2O to run efficiently without resorting to slow, naive SDPA implementations.

3. Key Contributions

Novel Fine-Tuning Method: A method to fine-tune transformer models with sparse attention and arbitrary KV cache policies on resources comparable to sparse attention inference (e.g., single A100). It combines nested activation checkpointing, CPU offloading, and delta encoding via autograd hooks.
H2O Implementation & Optimization: Methodological improvements to the H2O policy (batch independence, normalization) and a high-performance implementation using dedicated SDPA kernels that return summed attention weights, taking a significant step toward latencies competitive with SotA inference libraries.
Comprehensive Evaluation: Experiments on the Helmet benchmark suite (64k and 128k context windows) demonstrating that while performance varies by dataset, the proposed method allows models to co-adapt with KV cache policies, often outperforming models trained with sequence parallelism when evaluated under sparse attention inference.
KeysAndValues Library: The release of an open-source library (awslabs/keys_values) providing performant, easy-to-use code for long-context inference and fine-tuning with sparse attention.

4. Experimental Results

The authors evaluated a Qwen3-4B-Instruct model on the Helmet benchmark across 64k and 128k context lengths.

Comparison: Models were trained using:
- SP: Sequence Parallelism with exact attention (SotA baseline).
- US: The proposed method with the specific KV cache policy in place.
- NO: The base checkpoint with no fine-tuning.
Findings:
- Mixed Results on Table 1 Datasets: For the datasets in Table 1, results were mixed and inconclusive. The SP baseline (exact attention) achieved the best results for nq and hotpot qa. For trivia qa, the base checkpoint with no fine-tuning (no) performed best, indicating that fine-tuning did not help in this specific case. The proposed method (us) outperformed SP on pop qa.
- Strong Performance on Table 2 Datasets: For the datasets in Table 2, us strongly outperformed both sp and no.
- Failure Modes of SP: SP-trained models exhibited a failure mode where outputs were excessively long and contained random nonsense when subjected to sparse attention eviction, likely because they were never trained to handle the "forgetting" mechanism. In contrast, us learned to stop properly and usually output a single number or correct target.
- Policy Variance: Results varied significantly by dataset. While H2O variants performed well on several tasks, the lastrec (lr) policy outperformed H2O on datasets such as nq and hotpot qa. Additionally, for nq and hotpot qa, the SP baseline (exact attention) often achieved the best results, whereas for pop qa and other datasets in Table 2, the US method with sparse attention policies was superior.
- Efficiency: The H2O implementation with the new kernel support represents a major improvement in efficiency, taking a big step towards latencies competitive with SotA libraries, though the paper acknowledges that current sparse attention latency remains significantly higher than methods using sequence or context parallelism.

5. Significance and Claims

The paper claims that sparse attention has been underutilized in practice primarily due to latency issues and the difficulty of fine-tuning models to adapt to specific eviction policies.

Feasibility: The authors demonstrate that high-quality fine-tuning for long-context sparse attention is feasible on moderate hardware, removing the need for massive distributed training clusters solely for the purpose of adapting to a cache policy.
Necessity of Co-adaptation: The results strongly suggest that models must be trained under the same constraints (eviction policies) they will face during inference to avoid performance degradation.
Future Direction: The paper argues that by simplifying the KV cache structure and clarifying SDPA kernel requirements (e.g., returning summed weights), the community can better integrate sparse attention into SotA inference libraries.
Limitations: The authors acknowledge that while their method improves latency, the sequential nature of sparse attention decisions may still result in higher latency compared to RingAttention or Sequence Parallelism for very large hardware budgets. They plan to address this via kernel fusion and asynchronous CPU offloading in future work.

In summary, the paper provides a practical pathway to "learning how to forget," enabling models to specialize for long-context inference on accessible hardware. While performance relative to sequence parallelism varies by dataset, the method often outperforms exact-attention training in scenarios requiring adaptation to specific sparse attention policies.