New Amazon paper shows KV-cache policy is not just an inference optimization; but it can change the training regime the model needs.

If your LLM will forget parts of its context at inference, train it to forget that way too: this paper shows matching fine-tuning to the KV-cache policy can prevent long-context failures.

Sparse attention lets long-context inference use a fixed-size KV cache by keeping only part of the model’s past context. But models are often fine-tuned with full attention, then asked to work with missing memory at inference.

That mismatch can break behavior. In 128k-token tests, models trained with full attention often produced long, nonsensical answers under sparse inference, while models fine-tuned with the same cache policy learned to answer and stop normally.

The method makes this policy-matched training practical for arbitrary cache policies, including computing gradients for a 4B model on a 40 GB A100.

– arxiv. org/abs/2608.19920

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