In-Context Learning Strategies Emerge Rationally

2025-06-21

AI Summary

This paper investigates the origins of strategy diversity in large language models’ in-context learning (ICL), particularly why models adopt distinct ICL strategies across mixed-task training regimes. Method: We propose a unified explanatory framework based on a hierarchical Bayesian model: pretraining updates the posterior distribution over strategies, while inference performs a weighted average over strategy-specific predictions. Crucially, we introduce a rational analysis perspective—formalizing strategy selection as an optimization trade-off between prediction loss and strategy complexity. Contribution/Results: We theoretically derive that increasing task diversity induces a superlinear growth in the transition time from memory-based to generalization-based strategies. The framework requires no access to model weights; instead, it predicts Transformer token-level behavior and diverse ICL phenomena with high accuracy using only task structure and training dynamics. It provides the first quantifiable, predictive theory of ICL strategy evolution.

Technology Category

• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Recent work analyzing in-context learning (ICL) has identified a broad set of strategies that describe model behavior in different experimental conditions. We aim to unify these findings by asking why a model learns these disparate strategies in the first place. Specifically, we start with the observation that when trained to learn a mixture of tasks, as is popular in the literature, the strategies learned by a model for performing ICL can be captured by a family of Bayesian predictors: a memorizing predictor, which assumes a discrete prior on the set of seen tasks, and a generalizing predictor, wherein the prior matches the underlying task distribution. Adopting the lens of rational analysis from cognitive science, where a learner's behavior is explained as an optimal adaptation to data given computational constraints, we develop a hierarchical Bayesian framework that almost perfectly predicts Transformer next token predictions throughout training without assuming access to its weights. Under this framework, pretraining is viewed as a process of updating the posterior probability of different strategies, and its inference-time behavior as a posterior-weighted average over these strategies' predictions. Our framework draws on common assumptions about neural network learning dynamics, which make explicit a tradeoff between loss and complexity among candidate strategies: beyond how well it explains the data, a model's preference towards implementing a strategy is dictated by its complexity. This helps explain well-known ICL phenomena, while offering novel predictions: e.g., we show a superlinear trend in the timescale for transition to memorization as task diversity is increased. Overall, our work advances an explanatory and predictive account of ICL grounded in tradeoffs between strategy loss and complexity.

Problem

• Understand why models learn diverse in-context strategies
• Predict Transformer behavior via Bayesian framework without weights
• Explain ICL phenomena through strategy loss-complexity tradeoffs

Innovation

• Hierarchical Bayesian framework predicts Transformer behavior
• Pretraining updates posterior probability of strategies
• Tradeoff between strategy loss and complexity