Preprint

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

Chenrui Liu Chuanlong Xie Falong Tan Yicheng Zeng Li-Xing Zhu

Jul 2026 · 0 citations · 61 references

Computer Science Mathematics

TL;DR

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

Abstract

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.

arXiv Semantic Scholar

View source

Similar papers

Jul 2026

Induction in Both Directions: A Mechanistic Analysis of In-Context Learning in Masked Diffusion Language Models

While the internal mechanisms of autoregressive (AR) transformers have been studied extensively, much less is known about diffusion language models (DLMs), an emerging alternative that generates text by iterative denoising. In this work, we study how DLMs implement induction, a mechanism behind in-context learning in which the model finds a repeated context and copies the token that followed it. Our analysis compares attention-only AR models and absorbing-mask DLMs with matched architectures. We find that DLMs learn a bidirectional induction circuit, where previous-token and next-token heads write local context into the residual stream and later induction heads use it to find and copy the answer from the matching source position. The circuit is direction-symmetric, working whether the source appears in the past or in the future. When only left context is visible, matching what an AR model sees, the DLM does not outperform its AR counterpart in induction capabilities. However, we observe it has stronger induction when both sides of the masked token are visible, pointing to bidirectional context access rather than a stronger one-sided mechanism. Beyond induction, we provide causal evidence that DLMs compute the global fraction of masked tokens and use it as an implicit timestep, even though they are given no explicit timestep embedding.

Andy Catruna, Emilian Radoi · arXiv.org · 0 citations

Jul 2026

Knowledgeless Language Models: Suppressing Parametric Recall for Evidence-Grounded Language Modeling

The results demonstrate that suppressing entity-linked supervision during pretraining induces a shift in epistemic behavior: KLLMs rely less on parametric knowledge and more on external evidence, leading to improved reliability under realistic conditions, suggesting that pretraining-time control over knowledge acquisition can complement retrieval-augmented and tool-based systems by providing a more evidence-sensitive base model.

Roi Cohen, Yvan Carr'e, Nick Lechtenborger et al. · arXiv.org · 0 citations

#artificial intelligence Review Aug 2026

The Illusion of Replacement: Rethinking Specialized Machine Learning Models in the Foundation Model Era

Although performance of language-based models is improved by scaling, whether the gap to a structure-aware architecture can eventually be eliminated remains untested.

Kiyan Rezaee · 0 citations

Preprint Aug 2026

Attention-Path Fragility as an Uncertainty Signal in Large Language Models

It is proposed that a model's uncertainty about a token is reflected not only in the breadth of its output distribution but also in whether a confident prediction is \emph{fragile} under perturbation of its attention pathways, a training-free estimator that masks attention heads and measures the BALD mutual information among the resulting subnetworks.

Minsoo Kim, Sungyoung Ji, Kisung Moon et al. · 0 citations

Jul 2026

In-Context Learning as Implicit Policy Gradient

It is shown that score-conditioned In-Context Learning (ICL) admits a structural correspondence to policy gradient optimization, and an exact upper bound on the distribution shift induced by a bounded attention update is derived, yielding a trust-region-like analogy to KL-constrained policy optimization.

Masahiro Kaneko, Timothy Baldwin · arXiv.org · 0 citations

Jul 2026

A Better Start for Language Models: Domain-Conditional Position Offsets

A domain-conditional position offset is shown to improve retrieval reranking and domain classification when decisions depend on early in-domain tokens, and to be a lightweight, hot switchable tool for short in-domain scoring and calibration.

Ye Qiao · arXiv.org · 0 citations