Two are better than one: Context window extension with multi-grained self-injection

TL;DR

This paper introduces SharedLLM, a novel multi-grained context extension method for large language models that uses a dual-model architecture and a tree-structured data retrieval system to efficiently incorporate broader context without extensive retraining.

Contribution

SharedLLM presents a new architecture with a compressor and decoder LLMs and a tree-based data structure for efficient multi-grained context extension in LLMs.

Findings

Enables broader context understanding without retraining the entire model.

Reduces computational costs compared to continual pre-training on long data.

Efficient retrieval of multi-grained context information.

Abstract

The limited context window of contemporary large language models (LLMs) remains a huge barrier to their broader application across various domains. While continual pre-training on long-context data is a straightforward and effective solution, it incurs substantial costs in terms of data acquisition and computational resources. To alleviate this issue, we propose SharedLLM, a novel approach grounded in the design philosophy of multi-grained context compression and query-aware information retrieval. SharedLLM is composed of two short-context LLMs such as LLaMA-2, termed upper model and lower model. The lower model functions as a compressor while the upper model acts as a decoder. The upper model receives compressed, multi-grained context information from the lower model and performs context-aware modeling on the running text. Information transfer between the compressor and decoder occurs only at the lowest layers to refrain from long forward paths in the lower model and redundant cross-attention modules in the upper model. Based on this architecture, we introduce a specialized tree-style data structure to efficiently encode, store and retrieve multi-grained contextual information for text chunks. This structure, combined with a search algorithm, enables rapid encoding and retrieval of relevant information from various levels of the tree based on the input query. This entire process, wherein the sender and receiver are derived from the same LLM layer, is referred to as self-injection.