In-Context Learning as Implicit Policy Gradient

2026-07-25

AI Summary

This work elucidates the theoretical mechanism by which large language models optimize outputs through score-conditioned in-context learning. By establishing a structural correspondence between score-conditioned in-context learning and policy gradient methods such as REINFORCE, the study provides the first theoretical proof that, under specific weight configurations, the self-attention mechanism can implement reward-weighted aggregation. The authors introduce a directional alignment in the hidden state space, derive an upper bound on the distributional shift induced by attention updates, and validate their findings through self-attention analysis, constructive proofs, and simplified state-space models. Experimental results demonstrate that the model effectively leverages score information to adjust its output distribution, with attention weights exhibiting strong correlation with example scores.

Abstract

Recent work has shown that large language models (LLMs) can iteratively improve their outputs by incorporating generated samples and their corresponding evaluation scores as in-context examples. Despite these empirical findings, the theoretical foundations underlying this phenomenon remain poorly understood. In this paper, we show that score-conditioned In-Context Learning (ICL) admits a structural correspondence to policy gradient optimization. We first provide a constructive proof that self-attention mechanisms can implement reward-weighted aggregation analogous to the REINFORCE algorithm under specific weight matrix configurations, and discuss the relationship between this construction and the behavior of pretrained transformers. The correspondence is directional in hidden-state space and holds exactly only under the stated simplifying conditions; we quantify its strength empirically. Within our simplified hidden-state model, we furthermore derive an exact upper bound on the distribution shift induced by a bounded attention update, yielding a trust-region-like analogy to KL-constrained policy optimization. We validate our theory through extensive experiments across multiple LLMs, demonstrating that LLMs effectively utilize score information to shift output distributions toward high-scoring exemplars, and that attention weights exhibit a strong correlation with example scores.

Problem

• In-Context Learning
• Policy Gradient
• Large Language Models
• Theoretical Foundation
• Score-Conditioned Learning

Innovation

• In-Context Learning
• Policy Gradient
• Self-Attention
• REINFORCE
• Distribution Shift