LLoCO: Learning Long Contexts Offline

TL;DR

LLoCO introduces a method for learning and compressing long contexts offline, enabling large language models to process up to 128k tokens efficiently and accurately, significantly reducing inference costs and improving speed.

Contribution

The paper presents LLoCO, a novel offline context learning approach that extends LLM context windows and enhances long document processing efficiency.

Findings

Extends context window from 4k to 128k tokens.

Achieves up to 7.62x speed-up during inference.

Reduces inference token usage by 30x.

Abstract

Processing long contexts remains a challenge for large language models (LLMs) due to the quadratic computational and memory overhead of the self-attention mechanism and the substantial KV cache sizes during generation. We propose LLoCO, a novel approach to address this problem by learning contexts offline through context compression and in-domain parameter-efficient finetuning with LoRA. Our method enables an LLM to create a concise representation of the original context and efficiently retrieve relevant information to answer questions accurately. Our approach extends the effective context window of a 4k token LLaMA2-7B model to handle up to 128k tokens. We evaluate our approach on several long-context question-answering datasets, demonstrating that LLoCO significantly outperforms in-context learning while using $30\times$ fewer tokens during inference. LLoCO achieves up to $7.62\times$ speed-up during inference and $11.52\times$ higher throughput during finetuning, substantially reduces the cost of long document question answering. This makes it a promising solution for efficient long context processing. Our code is publicly available on https://github.com/jeffreysijuntan/lloco.