[ARCH-PROP] ContextLM

[dbylynn]

Architecture Name

ContextLM

Parent issue

#1

Motivations

• Scaling laws offer predictable progress, but demand exponentially growing data, parameters, and FLOPs.
• From bidirectional to autoregressive modeling, and from byte-level to token-level prediction, progress has often come from more challenging pretraining tasks.
• Following this direction, and inspired by JEPA, ContextLM predicts future context in latent representation space.

Proposed Architecture

ContextLM augments standard autoregressive language modeling with an intrinsic next-context prediction objective. Instead of only learning token-level hidden states for immediate next-token prediction, ContextLM learns predictive context embeddings that span multiple tokens and guide token-level generation.

The architecture decomposes a standard Transformer into three components:

• Token Encoder: encodes input tokens into token-level hidden states.
• Context Predictor: aggregates token hidden states into chunk-level context embeddings and autoregressively predicts future context embeddings.
• Token Decoder: combines token-level states with predicted context embeddings and performs standard next-token prediction.

Relative to the parent Transformer architecture, the key innovation is the insertion of a context-level autoregressive pathway between token encoding and token decoding. This pathway models the temporal evolution of latent context states, while preserving the same token-by-token evaluation paradigm used by standard language models.

Relative to MTP, ContextLM does not simply add more future token targets. It uses aggregated error signals from future tokens to optimize predictive context embeddings. This provides context-level gradient signals to train higher-level representations, allowing the model to learn abstractions beyond local token dependencies while still optimizing the standard token-level cross-entropy objective.

During training, ContextLM keeps the standard token-level CE loss. Each predicted context embedding receives gradients from tokens in its corresponding future chunk. This creates an intrinsic next-context prediction signal in latent space, where the context predictor is trained to produce representations that are useful for subsequent token generation.

At inference time, ContextLM remains compatible with standard autoregressive decoding. The model still generates tokens one by one, and can be evaluated directly with perplexity and downstream task metrics without changing the evaluation protocol.

Preliminary Results (if any)

Existing GPT-2 experiments show that ContextLM improves language modeling efficiency across multiple scaling dimensions.

On GPT-2 backbones trained on OpenWebText with a 10B-token training budget, ContextLM demonstrates consistent scaling effects up to the 1.5B-parameter scale. Compared with standard next-token prediction, ContextLM shifts the Pareto frontier of scaling laws across parameters, training tokens, and training FLOPs.

Under the same training data budget, ContextLM achieves the same perplexity with 39% fewer parameters. Under the same parameter budget, it achieves the same perplexity with 23% less pretraining data. Under the same parameter and training data budgets, it achieves the same performance with 20% less compute.

At the GPT2-XL scale, ContextLM outperforms both next-token prediction and multi-token prediction baselines. Compared with GPT2-XL-NTP and GPT2-XL-MTP, ContextLM-XL achieves better perplexity and stronger average downstream task accuracy on OpenWebText.

Experiments Plan

We will further validate ContextLM on larger and more diverse backbone families, such as Pythia, to test whether the GPT-2 scaling results transfer beyond a single architecture and pretraining setup.

Research questions

1. Can ContextLM maintain its scaling advantage when applied to larger backbone families?
2. Does ContextLM improve perplexity and downstream task performance under the same pretraining data and comparable model scales?
3. Can ContextLM improve instruction-following ability after fine-tuning?
4. Do the gains remain stable across model sizes up to the 1.5B scale?
5. How sensitive are the results to chunk size, encoder/decoder depth allocation, and context prediction design?

Comparisons

• Standard backbone baselines trained with next-token prediction.
• ContextLM variants built on the same backbones, trained with the same data and comparable model scales.
• Multi-Token Prediction baselines at comparable scale where applicable.
• ContextLM variants with different chunk sizes.
• ContextLM variants with different Token Encoder, Context Predictor, and Token Decoder allocations.

Evaluation

• Train and evaluate ContextLM models on larger backbones, up to the 1.5B-parameter scale and 300B training tokens.
• Measure validation perplexity under the same pretraining setup as the corresponding backbone baselines.
• Evaluate downstream task performance on standard language understanding benchmarks.
• Fine-tune ContextLM and backbone baselines under the same instruction-tuning setup.
• Evaluate instruction-following ability using MT-Bench-style multi-subtask evaluation.
• Analyze whether the context-level prediction gains observed on GPT-2 persist across larger-scale backbone settings.

Success criteria

• Achieve lower perplexity than backbone baselines under the same model size and training data.
• Improve downstream task accuracy across model sizes.
• Improve instruction-following performance after fine-tuning.
• Demonstrate stable gains across a larger scaling ladder up to 1.5B parameters and 300B training tokens.
• Confirm that context-level prediction remains effective beyond the GPT-2/OpenWebText setting.

[JT-Ushio]

The motivation, architecture, and validation plan are clear. ContextLM introduces a context-level autoregressive pathway while preserving standard token-by-token decoding, and the proposed comparisons against NTP and MTP directly address its claimed advantage over stronger pretraining objectives.

The proposal can proceed to validation and has been marked as in-progress.