[Paper Club] Writing in the Margins: Chunked Prefill KV Caching for Long Context Retrieval
Sep 19, 2024 · 53:53
Umar Jamil from Writer presents 'Writing in the Margins,' a paper on improving long-context retrieval in LLMs by leveraging chunked prefill of the KV cache. The method splits a long prompt into chunks, generates query-relevant annotations (margins) after each chunk, and appends them at the end before generating the final answer. This approach solves the 'lost in the middle' problem without fine-tuning, and unlike traditional RAG or separate summarization, it avoids re-prefilling the entire context, reducing cost from double to single prefilling. The technique is compatible with any transformer model and includes overlapping margin generation and classification within the same batch request. Benchmarks show consistent improvements across models, and the implementation is open-source on GitHub.
