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

Abstract: In-context learning (ICL) has emerged as a central capability of pretrained LLMs, yet its theoretical analysis has focused primarily on causal LLMs trained by left-to-right autoregressive prediction, such as GPT-style models. Masked LLMs 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 LLMs.

• The paper demonstrates that both masked and autoregressive models achieve similar excess-risk convergence in in-context learning through a unified statistical formulation.
• It introduces a measure-theoretic framework that quantifies task diversity, within-task sampling, and out-of-distribution shifts using Wasserstein distances.
• Empirical results on synthetic function learning show that masked-pair encoders can match GPT-style models, informing optimal strategies for pretraining data allocation.

A Unified Statistical Framework for In-Context Learning with Causal and Masked LLMs

Introduction and Motivation

In-context learning (ICL) is a core capability of large pretrained LLMs, enabling adaptation to new tasks from input–output demonstrations at inference, without parameter updates. Theoretical understanding of ICL has predominantly focused on causal LLMs (CLMs), such as GPT-style left-to-right Transformers, with little attention to masked LLMs (MLMs), such as BERT, despite their empirical ICL behavior and centrality in pretraining paradigms. This paper ["A Unified Framework for In-Context Learning with Causal and Masked LLMs" (2607.04081)] bridges this gap by providing a unified statistical learning framework for ICL that accommodates both training objectives on equal theoretical footing. A central aim is to determine how the choice of masked versus autoregressive objective influences excess risk, scaling behavior, and transferability in ICL.

Unified Empirical-Measure Formulation

The core abstraction represents the in-context prompt—comprising $k$ input–output examples and a query—via an empirical measure $\widehat \rho^{\,k}$ on the data space and models the predictor as a measurable function $f(\rho, \mathbf{x})$. This reframes ICL as distribution-to-function regression: the context provides empirical task information (via $\widehat \rho^{\,k}$), while the query specifies where to make predictions. The framework is agnostic to the specifics of the learning architecture and is readily instantiated by both CLMs and MLMs under suitable pretraining objectives.

Both pretraining paradigms are covered:

• Autoregressive (causal) objective: for each target, the context is a prefix empirical measure, excluding the target and all subsequent examples.
• Masked objective: for each target, the context is a leave-one-out empirical measure, using all examples except the target.

This representation aligns with recent measure-theoretic formulations of Transformer architectures and enables analysis of prompt scaling, task diversity, and out-of-distribution (OOD) generalization without limiting assumptions on model structure.

Statistical Learning Theory for ICL

Excess-Risk Decomposition

Under standard Wasserstein regularity (Lipschitz) conditions for both the loss and the predictor, the framework yields an excess-risk decomposition for pretraining with $T$ tasks and $N$ examples per task:

• Terms for estimating the meta-distribution of tasks ($T$-level)
• Terms for within-task estimation ($N$-level, i.e., empirical approximation of task distributions)
• Terms for context reduction specific to masked or causal objectives

Explicitly, for both masked and autoregressive pretraining under joint Lipschitz regularity (as satisfied by Transformer architectures without positional encodings), the $k$-shot excess risk at inference admits the upper bound:

$$\mathbb{E}[R^\ell_{k, \mathbb{P}}(f_{T,N})] - R^\ell_{k, \mathbb{P}}(f^*_k) \lesssim T^{-1/s} + N^{-1/(d_x + d_y)} + k^{-1/(d_x + d_y)} + N^{-1}$$

where $\widehat \rho^{\,k}$0 is the intrinsic meta-distribution complexity, and $\widehat \rho^{\,k}$1 are input/output dimensions. Notably, the same convergence order is obtained for both pretraining objectives, implying that the statistical benefits of in-context learning can be realized by both MLMs and CLMs.

Task-Distribution Shift and Transferability

The analysis generalizes to the OOD setting, where pretraining tasks are drawn from meta-distribution $\widehat \rho^{\,k}$2 and test tasks from $\widehat \rho^{\,k}$3, possibly distinct. The performance degradation under shift is linearly controlled by the lifted Wasserstein distance $\widehat \rho^{\,k}$4:

$\widehat \rho^{\,k}$5

This provides a rigorous measure for transferability and OOD generalization.

Optimal Allocation Under Budget Constraints

Given a fixed total data budget $\widehat \rho^{\,k}$6, the framework yields a principled allocation strategy between task diversity and within-task sampling:

$\widehat \rho^{\,k}$7

Order-optimal allocation minimizes excess risk, balancing gains from increased diversity and from improved within-task estimation.

Intrinsic Low-Dimensional Structure and Clusterability

If the within-task distribution supports are concentrated on a $\widehat \rho^{\,k}$8-dimensional submanifold or exhibit clusterable structure, the convergence rates in $\widehat \rho^{\,k}$9 and $f(\rho, \mathbf{x})$0 can be replaced by corresponding intrinsic-dimensional rates (or cluster rates), reducing the curse of dimensionality for appropriately structured data.

Empirical Verification: Synthetic In-Context Function Learning

The theoretical findings are substantiated by controlled synthetic experiments comparing the Masked Pair Encoder (an MLM-inspired encoder) to a GPT-2-style CLM across a variety of function-learning scenarios:

• Linear regression: Masked Pair Encoder matches the causal Transformer and the OLS estimator, saturating at the minimax optimal threshold for in-context sample size.
• Noisy linear regression: Both architectures exhibit least-squares behavior, robust to label noise and capturing double descent phenomena.
• Decision trees: Both methods outperform task-agnostic baselines in nonlinear settings.
• Two-layer ReLU networks: Comparable performance is observed on highly nonlinear function classes, tracking the optimal fit by a two-layer neural network fit per prompt.

Empirically, ICL behavior is not exclusive to causal models: MLMs equipped with appropriate masked-pair objectives can match the in-context learning performance of CLMs across diverse function classes.

Figure 1: Synthetic in-context function learning—Panel (a) shows linear regression, (b) noisy linear regression, (c) decision trees, (d) two-layer ReLU networks. The Masked Pair Encoder matches GPT-2-style causal baselines and is competitive with task-specific reference methods.

Implications and Future Directions

The unified statistical learning theory for both CLMs and MLMs in ICL delineates the precise scaling behaviors and transfer properties as a function of task diversity, context length, underlying data geometry, and OOD generalization. This has tangible implications:

• Pretraining strategy: Both MLM and CLM objectives can support scalable, transferable ICL under sufficient data and regularity, guiding architecture and objective selection in practical LLM pretraining.
• Data curation: The allocation results inform pretraining data collection, balancing diversity vs. per-task density.
• Intrinsic dimension: For data with strong geometric or cluster structure, the curse of dimensionality in ICL can be substantially mitigated by leveraging this structure in modeling and pretraining.

Open questions include the extension of these results to real language or code datasets, aggregate modeling of architectures with positional (or rotary) encoding intricacies, and characterization of the limits of OOD generalization in more complex task families. The deployment of these theoretical guarantees in neural network discovery of latent low-dimensional structure also merits investigation.

Conclusion

This paper provides a comprehensive statistical foundation for in-context learning across both causal and masked pretraining paradigms. The unified excess-risk analysis, supported by empirical evidence on synthetic function learning, establishes that MLMs can match CLMs in ICL capability under analogous objectives and architecture regularity. Transferability, scaling behavior, and data allocation strategies are formalized, providing actionable guidance for continued advancement in LLM pretraining and deployment.