A Unified Framework for In-Context Learning with Causal and Masked Language Models
- URL: http://arxiv.org/abs/2607.04081v1
- Date: Sun, 05 Jul 2026 02:23:18 GMT
- Title: A Unified Framework for In-Context Learning with Causal and Masked Language Models
- Abstract summary: In-context learning (ICL) has emerged as a central capability of pretrained language models.<n>Masked language models recover masked tokens from bidirectional context.<n>ICL behavior is not specific to causal language models.
- Score: 9.893068784551879
- License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
- 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.
Related papers
- The Maskability Index: Predicting Task-Objective Alignment in Pretrained Language Models [0.0]
Maskability Index (MI) is a metric that estimates whether a knowledge relation is better suited to masked-style prompting or prefix-style prompting in few-shot generation.<n>MI can help select appropriate prompting templates and adaptation strategies for extracting relational knowledge from pretrained language models.
arXiv Detail & Related papers (2026-07-22T15:19:23Z) - On Predicting the Post-training Potential of Pre-trained LLMs [60.07459271263409]
We introduce a new task of predicting post-training potential - forecasting a base model's performance before post-training.<n>We propose RuDE, a unified framework that bypasses the generation gap of base models by leveraging response discrimination.<n>Experiments demonstrate a correlation greater than 90% with post-training performance.
arXiv Detail & Related papers (2026-05-12T11:33:49Z) - Is my model perplexed for the right reason? Contrasting LLMs' Benchmark Behavior with Token-Level Perplexity [6.285055089721361]
We introduce a principled interpretability framework based on token-level perplexity to test whether models rely on linguistically relevant cues.<n>Our method enables precise, hypothesis-driven analysis without relying on unstable feature-attribution techniques.
arXiv Detail & Related papers (2026-03-31T08:00:55Z) - Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models [102.20309135516186]
Cross-entropy (CE) training provides dense and scalable supervision for language models.<n>We introduce a feature-matching objective for language-model fine-tuning.<n>To optimize this objective efficiently, we propose energy-based fine-tuning.
arXiv Detail & Related papers (2026-03-12T17:57:50Z) - KBQA-R1: Reinforcing Large Language Models for Knowledge Base Question Answering [64.62317305868264]
We present textbfKBQA-R1, a framework that shifts the paradigm from text imitation to interaction optimization via Reinforcement Learning.<n>Treating KBQA as a multi-turn decision process, our model learns to navigate the knowledge base using a list of actions.<n>Experiments on WebQSP, GrailQA, and GraphQuestions demonstrate that KBQA-R1 achieves state-of-the-art performance.
arXiv Detail & Related papers (2025-12-10T17:45:42Z) - Context-level Language Modeling by Learning Predictive Context Embeddings [79.00607069677393]
We introduce textbfContextLM, a framework that augments standard pretraining with an inherent textbfnext-context prediction objective.<n>This mechanism trains the model to learn predictive representations of multi-token contexts, leveraging error signals derived from future token chunks.<n>Experiments on the GPT2 and Pythia model families, scaled up to $1.5$B parameters, show that ContextLM delivers consistent improvements in both perplexity and downstream task performance.
arXiv Detail & Related papers (2025-10-23T07:09:45Z) - In-Context Learning Is Provably Bayesian Inference: A Generalization Theory for Meta-Learning [51.56484100374058]
We introduce a principled risk decomposition that separates the total ICL risk into two components: Bayes Gap and Posterior Variance.<n>For a uniform-attention Transformer, we derive a non-asymptotic upper bound on this gap, which explicitly clarifies the dependence on the number of pretraining prompts.<n>The Posterior Variance is a model-independent risk representing the intrinsic task uncertainty.
arXiv Detail & Related papers (2025-10-13T03:42:31Z) - How Does the Pretraining Distribution Shape In-Context Learning? Task Selection, Generalization, and Robustness [6.723482324209954]
We show how statistical properties of the pretraining distribution shape ICL on numerical tasks.<n>We develop a theoretical framework that unifies task selection and generalization, extending and sharpening earlier results.<n>We then empirically study how ICL performance varies with the pretraining distribution on challenging tasks.
arXiv Detail & Related papers (2025-10-01T17:52:29Z) - Mitigating Spurious Correlations in LLMs via Causality-Aware Post-Training [57.03005244917803]
Large language models (LLMs) often fail on out-of-distribution (OOD) samples due to spurious correlations acquired during pre-training.<n>Here, we aim to mitigate such spurious correlations through causality-aware post-training (CAPT)<n> Experiments on the formal causal inference benchmark CLadder and the logical reasoning dataset PrOntoQA show that 3B-scale language models fine-tuned with CAPT can outperform both traditional SFT and larger LLMs on in-distribution (ID) and OOD tasks.
arXiv Detail & Related papers (2025-06-11T06:30:28Z) - CSE-SFP: Enabling Unsupervised Sentence Representation Learning via a Single Forward Pass [3.0566617373924325]
Recent advances in pre-trained language models (PLMs) have driven remarkable progress in this field.<n>We propose CSE-SFP, an innovative method that exploits the structural characteristics of generative models.<n>We show that CSE-SFP not only produces higher-quality embeddings but also significantly reduces both training time and memory consumption.
arXiv Detail & Related papers (2025-05-01T08:27:14Z) - In-context Learning Generalizes, But Not Always Robustly: The Case of Syntax [36.98247762224868]
In-context learning (ICL) is now a common method for teaching large language models (LLMs) new tasks.
Do models infer the underlying structure of the task defined by the context, or do they rely on superficial generalizations that only generalize to identically distributed examples?
In experiments with models from the GPT, PaLM, and Llama 2 families, we find large variance across LMs.
The variance is explained more by the composition of the pre-training corpus and supervision methods than by model size.
arXiv Detail & Related papers (2023-11-13T23:52:43Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.