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

TL;DR - This paper introduces a hardware-efficient method for fine-tuning language models with sparse attention, allowing models to co-adapt to arbitrary KV-cache policies for long-context inference. It can outperform models trained with exact attention while running on a single 40 GB Nvidia A100 GPU.

• Supports any KV-cache selection or compression policy.
• Provides an optimized H2O sparse-attention implementation with dedicated scaled dot-product attention kernels.
• Reduces the hardware requirements of long-context fine-tuning compared with exact-attention sequence parallelism.
• Releases the methods through KeysAndValues, an open-source library for long-context inference and fine-tuning.