Context-level Language Modeling by Learning Predictive Context Embeddings

TL;DR

ContextLM addresses the limits of token-level next-token prediction by introducing a next-context prediction objective that learns predictive context embeddings within a hierarchical framework. A Context Predictor, fed by token representations, forecasts future context chunks which are then fused into token-level decoding, enabling joint modeling of fine-grained lexical patterns and higher-level semantics without changing the autoregressive evaluation regime. Empirical results on GPT2 and Pythia up to 1.5B parameters show consistent perplexity reductions and stronger downstream and instruction-following performance, with modest computational overhead. The approach yields improved long-range coherence and more efficient use of compute, suggesting a scalable path toward richer contextual modeling in large language models.

Abstract

Next-token prediction (NTP) is the cornerstone of modern large language models (LLMs) pretraining, driving their unprecedented capabilities in text generation, reasoning, and instruction following. However, the token-level prediction limits the model's capacity to capture higher-level semantic structures and long-range contextual relationships. To overcome this limitation, we introduce \textbf{ContextLM}, a framework that augments standard pretraining with an inherent \textbf{next-context prediction} objective. This mechanism trains the model to learn predictive representations of multi-token contexts, leveraging error signals derived from future token chunks. Crucially, ContextLM achieves this enhancement while remaining fully compatible with the standard autoregressive, token-by-token evaluation paradigm (e.g., perplexity). Extensive experiments on the GPT2 and Pythia model families, scaled up to $1.5$B parameters, show that ContextLM delivers consistent improvements in both perplexity and downstream task performance. Our analysis indicates that next-context prediction provides a scalable and efficient pathway to stronger language modeling, yielding better long-range coherence and more effective attention allocation with minimal computational overhead.

Figures (6)

• Figure 1: Scaling performance comparison across three dimensions: parameters, training tokens, and training FLOPs for GPT2 and ContextLM-GPT2.
• Figure 2: Overview of the ContextLM architecture. The model extends a standard NTP-based backbone by introducing a context predictor $\mathcal{P}$ which generates context embeddings that are fed into the token decoder $\mathcal{D}$ to guide token generation. The pseudocode details the inference procedure, showing the integration of context prediction into the token generation loop.
• Figure 3: Instruction-following evaluation on MT-Bench across multiple subtasks.
• Figure 4: Validation loss on OpenWebText for baseline GPT2 and ContextLM-GPT2 across chunk sizes, context predictor depths, and token encoder/decoder depths.
• Figure 5: Analysis of long-range modeling ability. The $\Delta Loss$ between ContextLM-GPT2 and GPT2 grows as block size increases.