A Unified Framework for In-Context Learning with Causal and Masked Language Models

In-context learning (ICL) has emerged as a central capability of pretrained language models, yet its theoretical analysis has focused primarily on causal language models trained by left-to-right autoregressive prediction, such as GPT-style models.

Masked language models instead recover masked tokens from bidirectional context, and their role in ICL remains less understood.

We develop a statistical learning framework that represents the context examples by their empirical measure and models prediction as a function of the context and the query.

This formulation places autoregressive and masked pretraining objectives within a common excess-risk analysis.

Under Wasserstein-type regularity conditions, we relate pretraining with T tasks and N samples per task to k-shot excess risk at inference, obtaining same-order upper bounds for masked and autoregressive objectives.

We also study task-distribution shift, where pretraining tasks are sampled from P and inference tasks from Q; the resulting bound contains an additional term controlled by the lifted Wasserstein distance between P and Q.

The bounds further imply an order-optimal allocation under a fixed pretraining data budget and refined rates under intrinsic low-dimensional structure.

Experiments on controlled function-learning tasks show that the Masked Pair Encoder (MPE) can achieve performance comparable to GPT-2-style causal Transformers, suggesting that ICL behavior is not specific to causal language models.