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

2026-08-20

AI Summary

本文提出一种针对稀疏注意力模型的微调方法，适用于任何KV缓存策略，可在有限硬件预算下提升长上下文推理性能。

Abstract

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.

Problem

• sparse attention
• long-context inference
• KV cache

Innovation

• fine-tuning
• sparse attention
• KV cache policy
• hardware budget
• long-context inference