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

arXiv: https://arxiv.org/abs/2607.04081 · HTML: https://mlsanity.com/paper/2607.04081/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 41 of its 61 references that are in the map.
- First posted on arXiv 5 Jul 2026.

## Builds on (most cited first, 30 of 41)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://mlsanity.com/paper/2006.03654/) (2020)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Reconciling modern machine learning practice and the bias-variance trade-off](https://mlsanity.com/paper/1812.11118/) (2018)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [The Intrinsic Dimension of Images and Its Impact on Learning](https://mlsanity.com/paper/2104.08894/) (2021)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (2023)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [The Benefit of Multitask Representation Learning](https://mlsanity.com/paper/1505.06279/) (2015)
- [Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression](https://mlsanity.com/paper/2306.15063/) (2023)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?](https://mlsanity.com/paper/2310.08391/) (2023)
- [In-Context Learning through the Bayesian Prism](https://mlsanity.com/paper/2306.04891/) (2023)
- [Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression](https://mlsanity.com/paper/2310.17086/) (2023)
- [Learning Theory for Distribution Regression](https://mlsanity.com/paper/1411.2066/) (2014)
- [Transformers are Minimax Optimal Nonparametric In-Context Learners](https://mlsanity.com/paper/2408.12186/) (2024)
- [Can In-context Learning Really Generalize to Out-of-distribution Tasks?](https://mlsanity.com/paper/2410.09695/) (2024)
- [In-Context Learning Is Provably Bayesian Inference: A Generalization Theory for Meta-Learning](https://mlsanity.com/paper/2510.10981/) (2025)
- [Distribution-Free Distribution Regression](https://mlsanity.com/paper/1302.0082/) (2013)
- [Intrinsic Dimension Estimation Using Wasserstein Distances](https://mlsanity.com/paper/2106.04018/) (2021)
- [Data Management For Training Large Language Models: A Survey](https://mlsanity.com/paper/2312.01700/) (2023)
- [When can in-context learning generalize out of task distribution?](https://mlsanity.com/paper/2506.05574/) (2025)
- [A survey and taxonomy of loss functions in machine learning](https://mlsanity.com/paper/2301.05579/) (2023)
- [BERTs are Generative In-Context Learners](https://mlsanity.com/paper/2406.04823/) (2024)

## Abstract (arXiv)

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.