Extensible Embedding: A Flexible Multipler For LLM's Context Length

TL;DR

This work tackles the challenge of extending LLM context windows without expensive fine-tuning by introducing Extensible Embedding, a plug-in module that converts long inputs into high-density extensible embeddings via a lightweight embedder and a down-sampling factor $k$, trained with a two-stream autoregressive scheme while keeping the downstream LLM fixed. The approach enables flexible context extension up to hundreds of thousands or even millions of tokens, with favorable memory and time characteristics and compatibility across fine-tuned derivatives. Extensive experiments on long-context language modeling and understanding demonstrate improved perplexity and QA/summary performance while maintaining efficiency and minimal training cost. This work provides a practical pathway to scalable long-context reasoning in real-world AI deployments.

Abstract

Large language models (LLMs) call for extension of context to handle many critical applications. However, the existing approaches are prone to expensive costs and inferior quality of context extension. In this work, we propose Extensible Embedding, which realizes high-quality extension of LLM's context with strong flexibility and cost-effectiveness. Extensible embedding stand as an enhancement of typical token embedding, which represents the information for an extensible scope of context instead of a single token. By leveraging such compact input units of higher information density, the LLM can access to a vast scope of context even with a small context window. Extensible embedding is systematically optimized in architecture and training method, which leads to multiple advantages. 1) High flexibility of context extension, which flexibly supports ad-hoc extension of diverse context lengths. 2) Strong sample efficiency of training, which enables the embedding model to be learned in a cost-effective way. 3) Superior compatibility with the existing LLMs, where the extensible embedding can be seamlessly introduced as a plug-in component. Comprehensive evaluations on long-context language modeling and understanding tasks verify extensible embedding as an effective, efficient, flexible, and compatible method to extend the LLM's context.

Figures (5)

• Figure 1: Framework. The input data is partitioned into chunks. Each sub-sequence is transformed and down-scaled as extensible embeddings. The new tokens are predicted based on the extensible embeddings from preceding chunks and the token embeddings in the same chunk. The extensible embedder is learned with a fixed downstream LLM.
• Figure 2: Two-Stream AR. In the first pass, the normal embeddings are transformed into extensible embeddings (with a scaling factor $k=3$). In the second pass (window size 10, chunk size 3), the auto-regression is accomplished in two sliding steps: the $x_{1-3}$ and $x_{4-6}$ predicted in the first step, $x_{7-9}$ and $x_{10-12}$ predicted in the second step.
• Figure 3: The extensible embedding trained on LLaMA-2-7B can be directly utilized by LongAlpaca-16K and LongChat-32K, leading to the scaling of their context lengths by $\times16$ and $\times32$ (PPL on PG19). Remarkably, the context of LongChat can be extended to 1 million.
• Figure 4: Comparison between extensible embedding and other context extension methods, including 1) Position Interpolation chen2023extending, 2) NTK-Aware Scaled RoPE ntkaware2023, 3) LongLLaMA tworkowski2023focused. Extensible Tokenization presents a superior long-context language modeling capability, along with better efficency in terms of memory and time. Perplexity is measured on PG19 raecompressive2019 following the method in chevalier2023adapting
• Figure 5: Compatibility of extensible embedding. The metric is perplexity on PG19. The darker the color, the higher the perplexity. And the blank area denotes instance where perplexity $> 10^2$.