Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers

TL;DR

REFORM is a novel inference framework that enhances long-context processing in transformers by efficiently compressing, gathering, and recomputing key-value caches, leading to significant performance improvements and resource savings.

Contribution

It introduces REFORM, a two-phase approach for long-context processing that combines incremental compression, cross-layer embeddings, and selective recomputation, outperforming existing methods.

Findings

Over 52% performance gain on RULER at 1M context length

30% reduction in inference time

5% decrease in peak memory usage

Abstract

As large language models increasingly gain popularity in real-world applications, processing extremely long contexts, often exceeding the model's pre-trained context limits, has emerged as a critical challenge. While existing approaches to efficient long-context processing show promise, recurrent compression-based methods struggle with information preservation, whereas random access approaches require substantial memory resources. We introduce REFORM, a novel inference framework that efficiently handles long contexts through a two-phase approach. First, it incrementally processes input chunks while maintaining a compressed KV cache, constructs cross-layer context embeddings, and utilizes early exit strategy for improved efficiency. Second, it identifies and gathers essential tokens via similarity matching and selectively recomputes the KV cache. Compared to baselines, REFORM achieves over 52% and 34% performance gains on RULER and BABILong respectively at 1M context length. It also outperforms baselines on Infinite-Bench, RepoEval, and MM-NIAH, demonstrating flexibility across diverse tasks and domains. Additionally, REFORM reduces inference time by 30% and peak memory usage by 5%, achieving both efficiency and superior performance.