FocusLLM: Precise Understanding of Long Context by Dynamic Condensing

TL;DR

FocusLLM addresses the challenge of extending fixed decoder context lengths to process very long sequences without prohibitive training costs. It introduces dynamic condensing and parallel decoding to extract and fuse chunk-level information, keeping the original decoder frozen while adding a small set of trainable projections. By partitioning input into chunks, injecting dynamic prompts, and aggregating candidate tokens across chunks, it reduces attention complexity to $O((L/n)^2)$ and achieves effective long-context understanding up to $400K$ tokens within a training budget of $0.5B$ tokens, outperforming diverse baselines on LongBench and ∞-Bench. The approach demonstrates strong language modeling and downstream task performance with favorable memory and inference-time characteristics, and code is available at GitHub.$

Abstract

Empowering LLMs with the ability to precisely understand long contexts is crucial for many downstream applications. However, handling long contexts with conventional transformer architecture requires substantial training and inference resources. Existing context condensing methods cannot accurately understand the full context, as there is a considerable amount of information loss in the condensing process. To address these issues, we present FocusLLM, a framework designed to extend the fixed context length of any decoder-only LLM, allowing the model to focus on relevant information from very long sequences. FocusLLM first divides long text input into chunks based on the model's original context length. It then employs the dynamic condensing process to distill crucial information from each chunk. Ultimately, through the novel parallel decoding mechanism, FocusLLM can integrate the extracted information into its local context. FocusLLM stands out for great training efficiency and versatility: trained with an 8K input length and with much less training cost than previous methods, FocusLLM exhibits superior performance across downstream tasks and maintains strong language modeling ability when handling extensive long texts, even up to 400K tokens. Our code is available at https://github.com/leezythu/FocusLLM.

Figures (8)

• Figure 1: A comparison between FocusLLM and previous context scaling methods on the passkey retrieval task, including CEPE, LongLLaMA and Activation Beacon. Our method extrapolates beyond the original context length of LLaMA, achieving 99% accuracy at a context length of 400K, with less training cost.
• Figure 2: One decoding step of the FocusLLM framework. A small fragment of the local context (denoted as the dynamic prompt) is appended to each chunk. The representations of the candidate tokens, obtained through dynamic condensing and parallel decoding, are then concatenated and integrated back into the local context.
• Figure 3: FocusLLM exhibits a more efficient growth pattern in memory usage compared to previous methods.
• Figure 4: Comparison of inference time. The time taken by FocusLLM is superior to previous methods.
• Figure 5: Perplexity under different chunk size with the total sequence length fixed as 8K on three datasets.