하나에서 쌓아 올리기: 컨텍스트 윈도우 확장을 위한 멀티 스케일 자가 주입

Stacked from One: Multi-Scale Self-Injection for Context Window Extension

최신 대규모 언어 모델(LLM)의 제한적인 컨텍스트 윈도우는 다양한 분야로의 활용을 가로막는 주요 병목 현상입니다. 긴 컨텍스트 데이터를 활용한 지속적인 사전 학습은 간단한 해결책이지만, 엄청난 데이터 획득 및 계산 비용을 초래합니다. 이러한 문제를 해결하기 위해, 본 논문에서는 멀티 그레인 컨텍스트 압축과 쿼리 기반 정보 획득을 기반으로 하는 새로운 프레임워크인~ extit{SharedLLM}을 제안합니다. SharedLLM은 두 개의 스태킹된 짧은 컨텍스트 LLM으로 구성됩니다. 하위 모델은 압축기로 작동하고, 상위 모델은 디코더로 작동합니다. 하위 모델은 긴 입력을 압축하여 컴팩트한 멀티 그레인 표현으로 변환하고, 이 표현은 상위 모델에서 컨텍스트 인지 처리를 위해 전달됩니다. 효율성을 극대화하기 위해, 정보 전달은 최하위 레이어에서만 이루어지며, 긴 순방향 연산과 중복된 크로스 어텐션 연산을 피합니다. 상위 및 하위 모델이 동일한 기본 LLM 레이어를 기반으로 파생되는 이 전체 프로세스를~ extit{자가 주입(self-injection)}이라고 합니다. 이 아키텍처를 지원하기 위해, 특수한 트리 기반 데이터 구조를 사용하여 컨텍스트 정보의 효율적인 인코딩 및 쿼리 기반 검색을 가능하게 합니다. extit{SharedLLM}은 8K 토큰의 시퀀스로 학습되었지만, 128K 토큰을 초과하는 입력에 대해서도 효과적으로 일반화됩니다. 광범위한 긴 컨텍스트 모델링 및 이해 벤치마크에서, extit{SharedLLM}은 강력한 기준 모델과 동등하거나 더 나은 성능을 달성하며, 효율성과 정확성 간의 최적의 균형을 유지합니다. 또한, 이러한 설계 선택은 extit{SharedLLM}의 메모리 사용량을 크게 줄이고, 상당한 추론 속도 향상($2 imes$ 스트리밍 방식 대비, $3 imes$ 인코더-디코더 아키텍처 대비)을 제공합니다.

The limited context window of contemporary large language models (LLMs) remains a primary bottleneck for their broader application across diverse domains. Although continual pre-training on long-context data offers a straightforward solution, it incurs prohibitive data acquisition and computational costs. To address this challenge, we propose~\modelname, a novel framework based on multi-grained context compression and query-aware information acquisition. SharedLLM comprises two stacked short-context LLMs: a lower model serving as a compressor and an upper model acting as a decoder. The lower model compresses long inputs into compact, multi-grained representations, which are then forwarded to the upper model for context-aware processing. To maximize efficiency, this information transfer occurs exclusively at the lowest layers, bypassing lengthy forward passes and redundant cross-attention operations. This entire process, wherein the upper and lower models are derived from the same underlying LLM layers, is termed~\textit{self-injection}. To support this architecture, a specialized tree-based data structure enables the efficient encoding and query-aware retrieval of contextual information. Despite being trained on sequences of only 8K tokens, \modelname~effectively generalizes to inputs exceeding 128K tokens. Across a comprehensive suite of long-context modeling and understanding benchmarks, \modelname~achieves performance superior or comparable to strong baselines, striking an optimal balance between efficiency and accuracy. Furthermore, these design choices allow \modelname~to substantially reduce the memory footprint and yield notable inference speedups ($2\times$ over streaming and $3\times$ over encoder-decoder architectures).