Context-level Language Modeling by Learning Predictive Context Embeddings
This addresses the problem of capturing long-range contextual relationships in language models for AI researchers, representing an incremental enhancement to existing pretraining methods.
The paper tackles the limitation of token-level prediction in large language models by introducing ContextLM, a framework that augments pretraining with next-context prediction to capture higher-level semantic structures, resulting in consistent improvements in perplexity and downstream task performance for models up to 1.5B parameters.
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.