Paper deep dive
In-Context Learning as Implicit Policy Gradient
Masahiro Kaneko, Timothy Baldwin
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2607.23153v1
- Canonical: https://arxiv.org/abs/2607.23153v1
Trouble viewing inline? Open PDF directly →
Full Text
69,266 characters extracted from source content.
Expand or collapse full text
In-Context Learning as Implicit Policy Gradient Masahiro Kaneko Timothy Baldwin MBZUAI masahiro.kaneko,timothy.baldwin@mbzuai.ac.ae 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. 1 Introduction Large language models (LLMs) have demonstrated a remarkable ability to adapt to new tasks through In-Context Learning (ICL), where the model learns from examples provided in the prompt without any parameter updates (Brown et al., 2020). This emergent capability has attracted significant interest from both theoretical and practical perspectives. On the theoretical side, it has been shown that the forward pass of transformers can implicitly implement gradient descent (Von Oswald et al., 2023) and ridge regression (Akyürek et al., 2023). Additionally, perspectives framing ICL as Bayesian inference over latent concepts acquired during pretraining (Xie et al., 2022) and mechanistic analyses of the role of induction heads (Olsson et al., 2022) have been proposed. However, these analyses share a common limitation: they focus exclusively on ICL in supervised learning settings with diverse input–output pairs (x,y)(x,y). On the empirical side, it has been demonstrated that LLMs can improve their outputs by conditioning on self-generated samples along with their evaluation scores as in-context examples. This method conditions the model on multiple output-score pairs (y,r)(y,r) generated for a fixed input x rather than learning from input-output pairs (x,y)(x,y). For example, the “Optimization by PROmpting” (OPRO) approach of Yang et al. (2023) provides the model with previously generated outputs and their scores as context, achieving strong results on prompt optimization and combinatorial problems. Similarly, the self-rewarding mechanism of Yuan et al. (2024) uses self-generated responses with self-assigned scores as in-context examples, enabling continuous improvement of alignment without human annotation. These successes raise a fundamental theoretical question: what learning algorithm, if any, does score-conditioned generation implement? The existing theoretical frameworks for ICL cannot answer this question, as they analyze supervised learning with diverse (x,y)(x,y) pairs rather than reward-based learning with (y,r)(y,r) pairs for a given input. This gap leaves us without a principled understanding of when and why score-conditioned ICL works. Just as (x,y)(x,y) pairs with direct target supervision naturally connect to supervised learning, (y,r)(y,r) pairs with scalar reward signals suggest a connection to reinforcement learning, where agents improve their behavior based on such signals. In this paper, we formalize this intuition by showing that score-conditioned ICL admits a structural correspondence to policy gradient optimization. Our central claims are threefold. First, attention mechanisms can implement reward-weighted aggregation analogous to REINFORCE. When LLMs process score-annotated examples (y,r)(y,r) through their attention mechanisms, under specific weight matrix configurations, attention weights the examples by their scores rather than by input similarity. Following the approach of Von Oswald et al. (2023), we provide a constructive proof that this reward-weighted aggregation is structurally analogous to the REINFORCE policy gradient estimator, in the sense that both produce score-weighted updates toward high-reward outputs in the hidden state space. We emphasize that this is an existence result: it demonstrates that the transformer architecture is capable of implementing such a mechanism, and we provide complementary empirical evidence in Appendix F that pretrained models exhibit functionally similar behavior. Second, the bounded hidden-state update induces a bounded KL shift. In explicit policy optimization, KL regularization is widely employed to prevent drastic distributional changes (Schulman et al., 2017; 2015). In the simplified output-layer model used in our analysis, we prove that the attention-induced hidden-state update yields an exact pathwise upper bound on the distribution shift from the reference model. This is a trust-region-like property, rather than a claim that ICL explicitly optimizes a KL-penalized objective. Third, we validate our theoretical predictions through extensive experiments across multiple LLMs. Through experiments on Llama-3, Olmo-3, Qwen-3, GPT-4o, Gemini 2.5, and Claude Sonnet 4 (Anthropic, 2025; Comanici et al., 2025; Hurst et al., 2024; Dubey et al., 2024; Olmo-Team, 2025; Qwen-Team, 2025), we confirm: (1) example scores are positively rank-correlated with probability changes; (2) attention weights correlate with scores, providing mechanistic evidence for score-weighted aggregation; and (3) ICL and explicit REINFORCE fine-tuning shift the output distribution in similar directions given identical samples. 2 Preliminaries 2.1 Notation We consider the setting where an LLM generates outputs for a given input and receives scalar feedback. Let x denote the input prompt, ∈ℝD y ^D denote the embedding representation of a generated output, and r∈ℝr denote its associated score (e.g., task performance, human preference rating, or model-based evaluation). We represent a score-annotated example as a combined token =(,r)∈ℝD+1 e=( y,r) ^D+1, concatenating the output embedding with its scalar score. A context of N such examples generated for a fixed input x is denoted C=(i,ri)i=1NC=\( y_i,r_i)\_i=1^N. We use πref _ref to denote the reference (base) LLM distribution and πθ _θ for a policy with parameters θ. The distribution induced by conditioning on context C is written as π(⋅|x,C)π(·|x,C). 2.2 Self-Attention Self-attention is the core mechanism by which transformers aggregate information across context. It computes a weighted combination of context elements, where the weights determine how much each element contributes to the output. Definition 2.1 (Linear Self-Attention). Given embeddings ∈ℝd×n E ^d× n and weight matrices Q,K,V,P W^Q, W^K, W^V, W^P, linear self-attention computes: flin()=+PV(K)⊤Qf_lin( E)= E+ W^P W^V E( W^K E) W^Q E (1) Linear attention directly uses key–query products as aggregation weights. This simplified form enables clean theoretical analysis and has been used in prior work connecting ICL to gradient descent (Von Oswald et al., 2023). Definition 2.2 (Softmax Self-Attention). Softmax self-attention applies a normalized exponential: fsoft()=+PV⋅softmax((K)⊤Qdk)f_soft( E)= E+ W^P W^V E·softmax ( ( W^K E) W^Q E d_k ) (2) Softmax attention, used in practice, normalizes weights to form a probability distribution. As we show in Theorem 3.4, this corresponds to Boltzmann-weighted aggregation, which assigns exponentially higher weight to high-scoring examples. 2.3 Policy Gradient Methods Policy gradient methods optimize a policy πθ _θ to maximize expected reward by directly estimating gradients of the objective. The REINFORCE algorithm (Williams, 1992) expresses this as: ∇θJ(θ)=πθ[r(y)∇θlogπθ(y|x)] _θJ(θ)=E_ _θ [r(y) _θ _θ(y|x) ] (3) In practice, REINFORCE typically employs a baseline b (e.g., mean reward) to reduce variance: ∇θJ(θ)=πθ[(r(y)−b)∇θlogπθ(y|x)] _θJ(θ)=E_ _θ [(r(y)-b) _θ _θ(y|x) ] (4) Given samples (yi,ri)i=1N\(y_i,r_i)\_i=1^N, the policy gradient update with baseline becomes: θ←θ+ηN∑i=1N(ri−b)∇θlogπθ(yi|x)θ←θ+ ηN _i=1^N(r_i-b) _θ _θ(y_i|x) (5) KL-regularized variants address instability by penalizing deviation from a reference distribution: maxθπθ[r(y)]−βDKL(πθ∥πref) _θE_ _θ[r(y)]-β D_KL( _θ\| _ref) (6) Methods such as PPO (Schulman et al., 2017) and TRPO (Schulman et al., 2015) implement variants of this principle. As we will show, the bounded attention update in our simplified model yields a related trust-region-like effect. 3 Theoretical Framework 3.1 Problem Setting We analyze score-conditioned ICL: given a fixed input x and a context C=(i,ri)i=1NC=\( y_i,r_i)\_i=1^N containing previously generated outputs for x paired with their evaluation scores, the LLM generates a new output by conditioning on this context. This setting captures recent empirical successes such as OPRO (Yang et al., 2023) and self-rewarding language models (Yuan et al., 2024). The score-conditioned examples are not assumed to be optimal; they are drawn from the model’s own distribution and may include low-scoring outputs, providing a diverse reward signal. The goal is to generate a new output that improves upon the best seen so far by learning from the relative quality differences among all N examples, rather than simply copying the highest-scoring example. Our central question is: what computational mechanism underlies score-conditioned generation? We show that, under explicit simplifying conditions, this process admits a structural policy-gradient analogy, where the attention mechanism performs reward-weighted aggregation over context examples. Our theoretical analysis operates in a simplified setting (linear attention, single-layer, continuous output representations); the empirical evidence in Appendix F demonstrates that the qualitative predictions hold in fully nonlinear pretrained models. 3.2 Assumptions Assumption 3.1 (Score-Conditioned Generation). LLMs utilize score information in context to modulate generation. Specifically, for contexts ChighC_high containing high-scoring examples and ClowC_low containing low-scoring examples: y∼π(⋅|x,Chigh)[r(y)]>y∼π(⋅|x,Clow)[r(y)]E_y π(·|x,C_high)[r(y)]>E_y π(·|x,C_low)[r(y)] (7) This assumption, strongly supported by LLM-as-a-Judge (Zheng et al., 2023) and OPRO (Yang et al., 2023), motivates the behavioral interpretation of the theoretical construction; we further validate score sensitivity in Section 4.3. Assumption 3.2 (Monotonicity). Moving the hidden state toward an output embedding increases the probability of generating that output: P(yi|h+αi)>P(yi|h)for α>0P(y_i|h+α y_i)>P(y_i|h) α>0 (8) This is an explicit modeling assumption. For a standard output layer P(y|h)∝exp(h⊤outembed(y))P(y|h) (h W_outembed(y)), it holds whenever the proposed movement has positive inner product with ∇hlogP(yi|h) _h P(y_i|h). 3.3 Main Results Theorem 3.3 (Score-Weighted Aggregation via Attention). There exists a construction of weight matrices K,Q,V,P W^K, W^Q, W^V, W^P such that linear self-attention implements score-weighted aggregation: hout=hq+ηN∑i=1Nriih_out=h_q+ ηN _i=1^Nr_i y_i (9) Proof Sketch. We set K=Q W^K= W^Q to project each token i=(i,ri) e_i=( y_i,r_i) so that ⊤i=ri q k_i=r_i (score becomes the attention weight), and V W^V to extract the embedding component i=(i,0) v_i=( y_i,0). Setting P=ηN W^P= ηN I yields the stated aggregation. The query token’s score component is set to 1 without loss of generality: any nonzero constant c yields the same score-proportional weighting ⊤i=c⋅ri q k_i=c· r_i, with the scaling absorbed into η. Full derivation in Appendix D.1. ∎ The η/Nη/N scaling factor corresponds to a learning rate divided by batch size, mirroring standard gradient descent. We do not claim that pretrained transformers adopt this exact construction; rather, the empirical evidence in Appendix F demonstrates that they learn functionally similar mechanisms. Theorem 3.4 (Softmax Attention as Boltzmann Aggregation). With analogous weight matrix construction, softmax attention implements Boltzmann-weighted aggregation: hout=hq+∑i=1Nexp(ri/τ)∑jexp(rj/τ)ih_out=h_q+ _i=1^N (r_i/τ) _j (r_j/τ) y_i (10) where τ is the temperature determined by the attention scaling factor. Full derivation in Appendix D.2. The softmax formulation provides a probabilistic interpretation: examples are weighted proportionally to the exponential of their scores, assigning exponentially higher influence to high-scoring examples. Theorem 3.5 (Structural Policy-Gradient Correspondence). Let r~i=ri−r¯ r_i=r_i- r denote mean-centered scores. Applying the construction of Theorem 3.3 to r~i r_i yields the score-weighted term of the REINFORCE update in hidden-state space. The correspondence is exact in the simplified output-layer model when out= W_out= I; for uncentered scores or a general output map, it is a structural analogy rather than an equality. Proof Sketch. For output layers P(y|h)∝exp(h⊤outembed(y))P(y|h) (h W_out\,embed(y)), the hidden-state gradient is ∇hlogP(yi|h)=out(i−¯) _h P(y_i|h)= W_out( y_i- y), where ¯=y∼P(⋅|h)[embed(y)] y=E_y P(·|h)[embed(y)]. With centered scores, ∑ir~i=0 _i r_i=0, so the ¯ y term cancels. When out= W_out= I, the resulting estimator is exactly ηN∑ir~ii ηN _i r_i y_i, the update constructed by Theorem 3.3 when its score coordinate contains r~i r_i. See Appendix D.3. ∎ We discuss the scope of this correspondence in Appendix A. Theorem 3.6 (Bounded KL Shift). In the simplified output-layer model, let πref=P(⋅|h) _ref=P(·|h) and πICL=P(⋅|h+Δh) _ICL=P(·|h+ h), with Δh h given by Theorem 3.3. Then: DKL(πICL∥πref)≤12ΛC(η⋅rmax⋅ymax)2,D_KL( _ICL\| _ref)≤ 12 _C(η· r_max· y_max)^2, (11) where ΛC=supt∈[0,1]λmax(F(h+tΔh)) _C= _t∈[0,1] _ (F(h+t h)), F is the Fisher information matrix of the output layer, rmax=maxi|ri|r_max= _i|r_i|, and ymax=maxi‖i‖y_max= _i\| y_i\|. Proof Sketch. The constructed attention update satisfies ‖Δh‖≤η⋅rmax⋅ymax\| h\|≤η· r_max· y_max. Writing the softmax log-partition function as A(h)A(h) and g(t)=A(h+tΔh)g(t)=A(h+t h) gives the exact identity DKL(πICL∥πref)=∫01tΔh⊤F(h+tΔh)ΔhtD_KL( _ICL\| _ref)= _0^1t\, h F(h+t h) h\,dt. Bounding the integrand by ΛC‖Δh‖2 _C\| h\|^2 yields the result. See Appendix D.4. ∎ This result bounds a single attention-induced distribution shift under the stated simplified model and bounded-score assumptions. It provides a trust-region-like interpretation, but does not imply that ICL optimizes an explicit KL-penalized objective or rules out mode collapse over repeated iterations. Combining these results yields our main theorem: Theorem 3.7 (Policy-Gradient Analogy with a Bounded KL Shift). In the simplified setting of Theorems 3.5 and 3.6, score-conditioned ICL has two complementary properties: 1. Reward-Weighted Update: Attention can aggregate output embeddings using their scores, with an exact hidden-state REINFORCE correspondence for centered scores when out= W_out= I (Theorems 3.3, 3.5). 2. Bounded KL Shift: The resulting bounded hidden-state update has a pathwise KL upper bound (Theorem 3.6). Together these properties are analogous to a trust-region policy update, but do not establish that ICL optimizes Eq. 6. The complete proof is provided in Appendix D.5. 3.4 Extension: Iterative ICL Corollary 3.8 (Repeated Score-Conditioned Updates). Iteratively applying score-conditioned ICL produces a sequence of reward-conditioned distribution shifts, each satisfying Theorem 3.6 within the simplified single-step model. Each iteration applies another score-conditioned update. Theorem 3.6 bounds each individual shift under its assumptions, but does not by itself guarantee monotonic reward improvement or a bounded cumulative shift. We therefore treat improvement across iterations as an empirical prediction, evaluated in Section 4.4. Concretely, at each iteration t, the model samples yt∼πICL(⋅|x,Ct)y_t _ICL(·|x,C_t) where CtC_t contains scored outputs from the previous round, computes rt=R(yt)r_t=R(y_t), and updates the context (see Appendix B for pseudocode). 4 Experiments 4.1 Experimental Setup Tasks. We target two tasks for prompt engineering: prompt optimization, which rewrites prompts to improve task performance (Deng et al., 2022; Zhang et al., 2022; Zhou et al., 2022; Kaneko et al., 2025b), and jailbreak prompting, which rewrites prompts to elicit harmful information (Mehrotra et al., 2024; Chao et al., 2025; Xie et al., 2025; Kaneko and Baldwin, 2026; Xiong et al., 2026; Kaneko et al., 2026; Han et al., 2026; Kaneko, 2026). For prompt optimization, we use DROP for reading comprehension and GSM8K for mathematical reasoning, evaluated by F1 score and exact match accuracy, respectively (Dua et al., 2019; Cobbe et al., 2021). For jailbreak prompting, we use H-RLHF and JailbreakBench (JBB), evaluated by attack success rate (ASR), where a jailbreak is deemed successful if a judge model determines the response contains harmful content (Bai et al., 2022; Chao et al., 2024). Models. We evaluate the following five open-weight LLMs: Meta-Llama-3-8B-Instruct (Llama3-8B), Meta-Llama-3-70B-Instruct (Llama3-70B), Olmo-3-7B-Think (Olmo3-7B), Olmo-3-32B-Think (Olmo3-32B), and Qwen3-4B-Instruct-2507 (Qwen3-4B) (Dubey et al., 2024; Olmo-Team, 2025; Qwen-Team, 2025). Our main experiments require access to model internals, so we do not evaluate closed-weight LLMs in these settings. For experiments that only require black-box API access, we additionally evaluate GPT-4o (Hurst et al., 2024), Gemini 2.5 (Comanici et al., 2025), and Sonnet 4 (Anthropic, 2025). Since the Sonnet 4 API does not expose token-level log-probabilities, we approximate sequence likelihood using sampling-based pseudo-likelihood (Kaneko et al., 2025a). This pseudo-likelihood is used only in the black-box RQ1 analysis; because it preserves the rank ordering of sequence likelihoods, it is reliable for the directional (sign and correlation) conclusions drawn in RQ1, and all of our quantitative and mechanistic conclusions are anchored on the open-weight models, for which exact token-level log-probabilities are available. For each task, we sample N=8N=8 outputs from the base model, compute their rewards, and construct contexts C=(yi,ri)i=1NC=\(y_i,r_i)\_i=1^N. We run T=10T=10 iterations. Context Format. Following Yang et al. (2023), we present score-annotated examples in ascending score order, ending with an instruction to generate a higher-scoring output (see Appendix C for the exact format).111We investigate the robustness of our results to score representation format (decimal, percentage, fraction, points, natural language) in Appendix G. All numeric formats yield comparably high correlations and substantially outperform a no-score control. The experiments display raw scores; mean-centering is required only for the exact special-case correspondence in Theorem 3.5, while the empirical claims concern the more general structural analogy. For the multi-token outputs produced on DROP and GSM8K, the output embedding y and the log-probability change Δlogp(y) p(y) are computed over the final answer span rather than a single token; using the pooled answer-span embedding yields correlations within about 0.030.03 of the single-token case. 4.2 RQ1: Score-Conditioned Distribution Shift Model DROP GSM8K H-RLHF JBB Original Llama3-8B 0.58† 0.61† 0.54† 0.52† Llama3-70B 0.71† 0.74† 0.68† 0.65† Olmo3-7B 0.55† 0.57† 0.51† 0.48† Olmo3-32B 0.67† 0.69† 0.63† 0.61† Qwen3-4B 0.49† 0.52† 0.46† 0.44† Shuffled Llama3-8B 0.02 -0.01 0.03 -0.02 Llama3-70B -0.03 0.02 -0.01 0.04 Olmo3-7B 0.01 0.03 -0.02 0.01 Olmo3-32B -0.02 0.01 0.02 -0.03 Qwen3-4B 0.03 -0.02 0.01 0.02 Table 1: Spearman correlation between example scores rir_i and log-probability changes Δlogp(yi) p(y_i) under ICL. Top: original scores. Bottom: randomly shuffled scores (control). † denotes p<0.01p<0.01. We measure whether ICL implements reward-weighted probability updates (Theorem 3.5) by computing the Spearman correlation between example scores ri\r_i\ and log-probability changes under ICL: Δlogp(yi)=logπ(yi|x,C)−logπref(yi|x) p(y_i)= π(y_i|x,C)- _ref(y_i|x) (12) across N=8N=8 sampled outputs per prompt. Under the proposed policy-gradient correspondence, higher-scoring examples should exhibit larger probability changes, yielding a positive correlation. Table 1 shows significant positive correlations across all models and tasks (p<0.01p<0.01), indicating a monotonic association between example scores and their probability changes under ICL. Larger models show stronger correlations (Llama3-70B highest, Qwen3-4B moderate), and prompt optimization tasks yield slightly higher correlations than jailbreak tasks, likely due to more clearly defined reward signals. As a control, shuffling score labels while keeping examples fixed drops correlations to near zero (mean: 0.01), providing causal evidence that the effect is driven by explicit numeric score values rather than inherent properties of high-quality outputs. (a) Magnitude of distribution shift. (b) Score-probability correlation. Figure 1: Effect of context size on score-conditioned ICL. (a) |Δlogp|| p| decreases with context size, consistent with the η/Nη/N scaling in Theorem 3.3. Dashed line shows theoretical 1/N1/N scaling. (b) Score-probability correlation remains stable across N, confirming robust score-weighted behavior. We also vary context size N∈2,4,8,16,32N∈\2,4,8,16,32\ (Figure 1): the magnitude |Δlogp|| p| decreases following the theoretical 1/N1/N curve (panel a), while score-probability correlation remains stable across N (panel b), confirming the η/Nη/N scaling of Theorem 3.3. 4.3 RQ2: Attention Weight Analysis Model Avg Max Layer of Max Llama3-8B 0.31† 0.54 25/32 Llama3-70B 0.43† 0.71 68/80 Olmo3-7B 0.30† 0.49 22/32 Olmo3-32B 0.39† 0.65 52/64 Qwen3-4B 0.29† 0.46 29/36 Table 2: Spearman correlation between example scores rir_i and attention weights αi _i. Avg is averaged across all layers and heads; Max is the highest correlation observed in any single layer. † denotes p<0.01p<0.01. We examine whether attention weights directly reflect example scores as predicted by Theorems 3.3 and 3.4. For each context example yiy_i, we compute αi _i as the mean attention weight from generated tokens to all tokens of yiy_i, across all layers and heads. Table 2 shows a significant positive score–attention correlations across all models, with peak correlations in late layers, consistent with prior findings that later layers handle higher-level semantics (Geva et al., 2022; 2023). Further mechanistic analysis, including score-specialized head identification, functional decomposition, and causal interventions, is provided in Appendix F (Sections F.1–F.4). 4.4 RQ3: Iterative ICL Figure 2: Score improvement from zero-shot baseline across iterations of score-conditioned ICL (normalized to [0,1][0,1] per task before averaging). Figure 3: Cosine similarity between ICL and REINFORCE distribution shifts (Δlogp p vectors). † denotes p<0.01p<0.01 using a bootstrap test. We run iterative score-conditioned ICL for T=10T=10 rounds across four tasks and eight LLMs. Scores are normalized to [0,1][0,1] per task before averaging across tasks. Figure 3 shows consistent improvement across all models, with rapid gains in early iterations that plateau. This empirical pattern is compatible with the bounded-shift perspective of Theorem 3.6, although the theorem alone does not imply convergence or monotonic improvement. 4.5 RQ4: Alignment with REINFORCE Distribution Shifts To validate Theorem 3.5 directly, we compare distribution shifts from score-conditioned ICL and explicit REINFORCE fine-tuning under identical conditions: the same N=8N=8 samples and a single gradient step. For each prompt, we measure: ΔlogpICL(yi) p_ICL(y_i) =logπ(yi|x,C)−logπref(yi|x) = π(y_i|x,C)- _ref(y_i|x) (13) ΔlogpRL(yi) p_RL(y_i) =logπθ′(yi|x)−logπref(yi|x) = _θ (y_i|x)- _ref(y_i|x) (14) where θ′θ is after one gradient step at η=10−5η=10^-5 (sensitivity across η∈10−6,10−5,10−4η∈\10^-6,10^-5,10^-4\ in Appendix E). We compute cosine similarity between these two Δlogp p vectors over 100 prompts × 5 samplings. Figure 3 shows cosine similarities exceeding 0.45 across all model-task pairs (p<0.01p<0.01), confirming that both methods shift the output distribution in positively aligned directions given the same evidence. 4.6 RQ5: Causal Verification via Score Token Ablation Figure 4: Effect of score token ablation across normalized layer depth. Y-axis shows remaining ρ(r,Δlogp)ρ(r, p) after zeroing score token representations at each layer. Triangles indicate Layer of Max from Table 2. To establish causality beyond the correlational evidence in RQ2, we zero out hidden representations at score token positions layer by layer and measure the drop in ρ(r,Δlogp)ρ(r, p). Figure 4 shows minimal disruption in early layers, a sharp peak near the Layer of Max from Table 2, and slight recovery in final layers, confirming that score information is causally processed in the late layers identified by the correlational analysis. 5 Related Work Von Oswald et al. (2023) showed that transformers can implement gradient descent in-context, providing a mechanistic explanation for few-shot learning. Akyürek et al. (2023) extended this to ridge regression, while Dai et al. (2023) drew connections to gradient descent through a kernel perspective. Xie et al. (2022) offered a Bayesian view, framing ICL as implicit inference. These works all operate on supervised (x,y)(x,y) pairs where a target output is known; our work addresses the fundamentally different regime of reward-based (y,r)(y,r) pairs for a fixed input, where score-weighted attention admits a structural correspondence to REINFORCE and its bounded hidden-state update yields a pathwise KL bound. Brooks et al. (2023) showed that foundation models can implement policy iteration in-context, focusing on action selection across distinct states rather than distribution shift for a fixed input; Huang et al. (2025) explored in-context steered policy optimization. Our work is complementary, providing a structural connection between score-weighted attention and REINFORCE and characterizing the resulting bounded distribution shift. A related line of work shows behaviorally that models exploit in-context rewards: Agarwal et al. (2024) scale ICL to many demonstrations, and Monea et al. (2024) frame in-context learning from rewards as a bandit problem. These works establish that models improve; we identify the computational underpinnings. Score-specialized heads whose causal manipulation moves the score–probability correlation (Appendix F.4) ground the REINFORCE correspondence and yield predictions a behavioral account does not: the η/Nη/N scaling, saturation with iteration, and directional alignment with one-step REINFORCE. Yang et al. (2023) introduced OPRO, demonstrating that LLMs can optimize objectives through prompting. Yuan et al. (2024) showed that self-rewarding mechanisms enable continuous improvement. Madaan et al. (2023) proposed Self-Refine for iterative output improvement. Our work provides theoretical grounding for why these approaches succeed. RLHF (Ouyang et al., 2022) and its variants have become standard for LLM alignment. Recent work has explored connections between supervised fine-tuning and RL objectives (Rafailov et al., 2023). Our work reveals that even without fine-tuning, LLMs can exhibit policy-optimization-like behavior through attention mechanisms. 6 Conclusion We showed that score-conditioned ICL is structurally analogous to an implicit form of policy gradient optimization: self-attention can implement reward-weighted aggregation, and centered scores yield an exact hidden-state REINFORCE correspondence under the simplified identity-output assumption. We also derived an exact pathwise KL bound for the bounded attention-induced update, providing a trust-region-like interpretation without claiming that ICL explicitly optimizes a KL-regularized objective. Extensive experiments across multiple LLMs confirm that attention weights correlate with scores, ICL and REINFORCE produce aligned distribution shifts, and iterative ICL yields consistent improvement, providing theoretical justification for score-conditioned ICL as a lightweight alternative to fine-tuning. Our analysis operates under simplifying assumptions (linear attention, single-layer, out≈ W_out≈ I) that enable tractable proofs; extending the formal guarantees to fully nonlinear multi-layer architectures remains an important direction for future work. More broadly, our framework suggests that any system capable of attending to reward-annotated examples can implicitly perform policy optimization, with implications for the design of inference-time alignment and self-improvement pipelines. Ethics Statement This work includes experiments on jailbreak prompting to validate our theoretical framework across diverse reward signals. We emphasize that our goal is to understand the mechanisms underlying score-conditioned ICL, not to develop new attack methods. All experiments used existing public benchmarks in line with their intended research purposes. We acknowledge that our theoretical insights could potentially be misused to improve adversarial prompting techniques. However, we believe that understanding these mechanisms is essential for developing robust defenses. The same theoretical framework that explains why score-conditioned ICL can optimize for harmful outputs also suggests mitigation strategies, such as filtering high-scoring harmful examples from context or detecting score-based optimization patterns in user interactions. References R. Agarwal, A. Singh, L. Zhang, B. Bohnet, L. Rosias, S. Chan, B. Zhang, A. Anand, Z. Abbas, A. Nova, et al. (2024) Many-shot in-context learning. Advances in Neural Information Processing Systems 37, p. 76930–76966. Cited by: §5. E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou (2023) What learning algorithm is in-context learning? investigations with linear models. In International Conference on Learning Representations, External Links: Link Cited by: Appendix F, §1, §5. Anthropic (2025) System card: claude opus 4 & claude sonnet 4. Technical report Anthropic. External Links: Link Cited by: §1, §4.1. Y. Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan, et al. (2022) Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862. Cited by: §4.1. E. Brooks, L. Walls, R. L. Lewis, and S. Singh (2023) Large language models can implement policy iteration. Advances in Neural Information Processing Systems 36, p. 30349–30366. Cited by: §5. T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. (2020) Language models are few-shot learners. Advances in neural information processing systems 33, p. 1877–1901. Cited by: §1. P. Chao, E. Debenedetti, A. Robey, M. Andriushchenko, F. Croce, V. Sehwag, E. Dobriban, N. Flammarion, G. J. Pappas, F. Tramer, et al. (2024) Jailbreakbench: an open robustness benchmark for jailbreaking large language models. Advances in Neural Information Processing Systems 37, p. 55005–55029. Cited by: §4.1. P. Chao, A. Robey, E. Dobriban, H. Hassani, G. J. Pappas, and E. Wong (2025) Jailbreaking black box large language models in twenty queries. In 2025 IEEE Conference on Secure and Trustworthy Machine Learning (SaTML), p. 23–42. Cited by: §4.1. K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, et al. (2021) Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168. Cited by: §4.1. G. Comanici, E. Bieber, M. Schaekermann, I. Pasupat, N. Sachdeva, I. Dhillon, M. Blistein, O. Ram, D. Zhang, E. Rosen, et al. (2025) Gemini 2.5: pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv preprint arXiv:2507.06261. Cited by: §1, §4.1. D. Dai, Y. Sun, L. Dong, Y. Hao, S. Ma, Z. Sui, and F. Wei (2023) Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers. In Findings of the Association for Computational Linguistics: ACL 2023, p. 4005–4019. Cited by: §5. M. Deng, J. Wang, C. Hsieh, Y. Wang, H. Guo, T. Shu, M. Song, E. Xing, and Z. Hu (2022) RLPrompt: optimizing discrete text prompts with reinforcement learning. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, Y. Goldberg, Z. Kozareva, and Y. Zhang (Eds.), Abu Dhabi, United Arab Emirates, p. 3369–3391. External Links: Link, Document Cited by: §4.1. D. Dua, Y. Wang, P. Dasigi, G. Stanovsky, S. Singh, and M. Gardner (2019) DROP: a reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), J. Burstein, C. Doran, and T. Solorio (Eds.), Minneapolis, Minnesota, p. 2368–2378. External Links: Link, Document Cited by: §4.1. A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Yang, A. Fan, et al. (2024) The llama 3 herd of models. arXiv e-prints, p. arXiv–2407. Cited by: §1, §4.1. M. Geva, J. Bastings, K. Filippova, and A. Globerson (2023) Dissecting recall of factual associations in auto-regressive language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, H. Bouamor, J. Pino, and K. Bali (Eds.), Singapore, p. 12216–12235. External Links: Link, Document Cited by: §4.3. M. Geva, A. Caciularu, K. Wang, and Y. Goldberg (2022) Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, Y. Goldberg, Z. Kozareva, and Y. Zhang (Eds.), Abu Dhabi, United Arab Emirates, p. 30–45. External Links: Link, Document Cited by: §4.3. Z. Han, Y. Liu, K. Zhang, and P. Li (2026) Experience-driven multi-agent optimization for black-box jailbreak attacks on large language models. In Findings of the Association for Computational Linguistics: ACL 2026, p. 23729–23747. Cited by: §4.1. H. Huang, C. Tang, W. Liu, C. Bai, S. Yang, and Y. Wu (2025) Think outside the policy: in-context steered policy optimization. arXiv preprint arXiv:2510.26519. Cited by: §5. A. Hurst, A. Lerer, A. P. Goucher, A. Perelman, A. Ramesh, A. Clark, A. Ostrow, A. Welihinda, A. Hayes, A. Radford, et al. (2024) Gpt-4o system card. arXiv preprint arXiv:2410.21276. Cited by: §1, §4.1. M. Kaneko and T. Baldwin (2026) Bits leaked per query: information-theoretic bounds for adversarial attacks on llms. Advances in Neural Information Processing Systems 38, p. 80186–80210. Cited by: §4.1. M. Kaneko, Y. Ma, Y. Wata, and N. Okazaki (2025a) Sampling-based pseudo-likelihood for membership inference attacks. In Findings of the Association for Computational Linguistics: ACL 2025, W. Che, J. Nabende, E. Shutova, and M. T. Pilehvar (Eds.), Vienna, Austria, p. 8894–8907. External Links: Link, Document, ISBN 979-8-89176-256-5 Cited by: §4.1. M. Kaneko, A. Niwa, and T. Baldwin (2026) JailNewsBench: multi-lingual and regional benchmark for fake news generation under jailbreak attacks. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §4.1. M. Kaneko, Z. Talat, and T. Baldwin (2025b) Online learning defense against iterative jailbreak attacks via prompt optimization. In Proceedings of the 14th International Joint Conference on Natural Language Processing and the 4th Conference of the Asia-Pacific Chapter of the Association for Computational Linguistics, K. Inui, S. Sakti, H. Wang, D. F. Wong, P. Bhattacharyya, B. Banerjee, A. Ekbal, T. Chakraborty, and D. P. Singh (Eds.), Mumbai, India, p. 2592–2609. External Links: Link, Document, ISBN 979-8-89176-298-5 Cited by: §4.1. M. Kaneko (2026) Paraphrasing adversarial attack on llm-as-a-reviewer. arXiv preprint arXiv:2601.06884. Cited by: §4.1. A. Madaan, N. Tandon, P. Gupta, S. Hallinan, L. Gao, S. Wiegreffe, U. Alon, N. Dziri, S. Prabhumoye, Y. Yang, et al. (2023) Self-refine: iterative refinement with self-feedback. Advances in Neural Information Processing Systems 36, p. 46534–46594. Cited by: §5. A. Mehrotra, M. Zampetakis, P. Kassianik, B. Nelson, H. Anderson, Y. Singer, and A. Karbasi (2024) Tree of attacks: jailbreaking black-box llms automatically. Advances in Neural Information Processing Systems 37, p. 61065–61105. Cited by: §4.1. G. Monea, A. Bosselut, K. Brantley, and Y. Artzi (2024) Llms are in-context bandit reinforcement learners. arXiv preprint arXiv:2410.05362. Cited by: §5. Olmo-Team (2025) Olmo 3. arXiv preprint arXiv:2512.13961. Cited by: §1, §4.1. C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, et al. (2022) In-context learning and induction heads. arXiv preprint arXiv:2209.11895. Cited by: §1. L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. (2022) Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems, Vol. 35, p. 27730–27744. Cited by: §5. Qwen-Team (2025) Qwen3 technical report. External Links: 2505.09388, Link Cited by: §1, §4.1. R. Rafailov, A. Sharma, E. Mitchell, S. Ermon, C. D. Manning, and C. Finn (2023) Direct preference optimization: your language model is secretly a reward model. Advances in Neural Information Processing Systems 36. Cited by: §5. J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz (2015) Trust region policy optimization. In International conference on machine learning, p. 1889–1897. Cited by: §1, §2.3. J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov (2017) Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. Cited by: §1, §2.3. J. Von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov (2023) Transformers learn in-context by gradient descent. In International Conference on Machine Learning, p. 35151–35174. Cited by: Appendix F, §1, §1, §2.2, §5. R. J. Williams (1992) Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning 8 (3), p. 229–256. Cited by: §2.3. S. M. Xie, A. Raghunathan, P. Liang, and T. Ma (2022) An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, External Links: Link Cited by: §1, §5. Z. Xie, J. Gao, L. Li, Z. Li, Q. Liu, and L. Kong (2025) Jailbreaking as a reward misspecification problem. In International Conference on Learning Representations, Vol. 2025, p. 83368–83390. Cited by: §4.1. X. Xiong, O. Li, Z. Liu, M. Li, W. Shi, F. Zhu, Q. Wang, and F. Feng (2026) Trojail: trajectory-level optimization for multi-turn large language model jailbreaks with process rewards. In Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 48086–48109. Cited by: §4.1. C. Yang, X. Wang, Y. Lu, H. Liu, Q. V. Le, D. Zhou, and X. Chen (2023) Large language models as optimizers. arXiv preprint arXiv:2309.03409. Cited by: Appendix C, §1, §3.1, §3.2, §4.1, §5. W. Yuan, R. Y. Pang, K. Cho, S. Sukhbaatar, J. Xu, and J. Weston (2024) Self-rewarding language models. In International Conference on Machine Learning, Cited by: §1, §3.1, §5. T. Zhang, X. Wang, D. Zhou, D. Schuurmans, and J. E. Gonzalez (2022) Tempera: test-time prompting via reinforcement learning. arXiv preprint arXiv:2211.11890. Cited by: §4.1. L. Zheng, W. Chiang, Y. Sheng, S. Zhuang, Z. Wu, Y. Zhuang, Z. Lin, Z. Li, D. Li, E. Xing, et al. (2023) Judging LLM-as-a-judge with MT-bench and Chatbot Arena. Advances in Neural Information Processing Systems 36, p. 46595–46623. Cited by: §3.2. Y. Zhou, A. I. Muresanu, Z. Han, K. Paster, S. Pitis, H. Chan, and J. Ba (2022) Large language models are human-level prompt engineers. In The eleventh international conference on learning representations, Cited by: §4.1. Appendix A Theoretical Remarks Remark A.1 (On the Scope of the Correspondence). The correspondence established in Theorem 3.5 is exact under out= W_out= I with centered scores, and approximate when out≈ W_out≈ I. The claim is therefore one of structural analogy rather than general exact equivalence: both ICL and REINFORCE contain score-weighted updates of output embeddings, but they differ in reward centering and in how the output layer transforms embeddings into probabilities. Exact equality in our simplified hidden-state formulation requires centered scores and out= W_out= I. The strength of this analogy is supported empirically in Section 4.5. Remark A.2 (Baseline and Normalized Rewards). Standard policy gradient implementations use baseline reduction or reward normalization to reduce variance. Theorem 3.5 uses explicit mean-centering, r~i=ri−r¯ r_i=r_i- r. Because ∑ir~i=0 _i r_i=0, the expected-embedding term in the log-softmax gradient cancels, producing the exact simplified correspondence. For raw, uncentered scores, the theorem supports only a structural analogy; it does not imply that score-conditioned ICL automatically performs baseline reduction. Appendix B Iterative Score-Conditioned ICL: Algorithm Algorithm 1 Iterative Score-Conditioned ICL 1: Input: Prompt x, reward function R, iterations T 2: y0∼πref(⋅|x)y_0 _ref(·|x); r0←R(y0)r_0← R(y_0) 3: for t=1t=1 to T do 4: Ct←(yt−1,rt−1)C_t←\(y_t-1,r_t-1)\ Or maintain full history 5: yt∼πICL(⋅|x,Ct)y_t _ICL(·|x,C_t); rt←R(yt)r_t← R(y_t) 6: end for 7: Return: yTy_T Appendix C Context Format Following Yang et al. (2023), score-annotated examples are presented in ascending score order: text: [output text] score: 61 text: [output text] score: 78 (...more examples sorted in ascending order...) Generate a new output that achieves a higher score: Appendix D Complete Proofs D.1 Proof of Theorem 3.3: Score-Weighted Aggregation Proof. Let each token be i=(i,ri)∈ℝD+1 e_i=( y_i,r_i) ^D+1 where i∈ℝD y_i ^D is the output embedding and ri∈ℝr_i is the score. Step 1: Weight Matrix Construction. We construct the weight matrices as follows: K=Q=(D×D×11×D1)∈ℝ(D+1)×(D+1) W^K= W^Q= pmatrix 0_D× D& 0_D× 1\\ 0_1× D&1 pmatrix ^(D+1)×(D+1) (15) V=(D×11×D0)∈ℝ(D+1)×(D+1) W^V= pmatrix I_D& 0_D× 1\\ 0_1× D&0 pmatrix ^(D+1)×(D+1) (16) P=ηND+1 W^P= ηN I_D+1 (17) This construction is engineered to demonstrate the existence of a weight configuration that implements score-weighted aggregation. We do not claim that pretrained transformers adopt these exact weight matrices; the empirical analysis in Appendix F shows that pretrained models learn functionally similar mechanisms. Step 2: Key and Query Computation. For a context token i=(i,ri) e_i=( y_i,r_i): i=Ki=(D×D×11×D1)(iri)=(Dri) k_i= W^K e_i= pmatrix 0_D× D& 0_D× 1\\ 0_1× D&1 pmatrix pmatrix y_i\\ r_i pmatrix= pmatrix 0_D\\ r_i pmatrix (18) For the query token q=(q,1) e_q=( y_q,1) (setting query score component to 1; the choice of any nonzero constant c is without loss of generality, since ⊤i=c⋅ri q k_i=c· r_i and the factor c is absorbed into η): =Qq=(D1) q= W^Q e_q= pmatrix 0_D\\ 1 pmatrix (19) Step 3: Attention Weight Computation. The attention weight (before normalization) for token i: ⊤i=(D1)⊤(Dri)=ri q k_i= pmatrix 0_D\\ 1 pmatrix pmatrix 0_D\\ r_i pmatrix=r_i (20) This shows that the attention weight equals the score rir_i. Step 4: Value Computation. i=Vi=(D×11×D0)(iri)=(i0) v_i= W^V e_i= pmatrix I_D& 0_D× 1\\ 0_1× D&0 pmatrix pmatrix y_i\\ r_i pmatrix= pmatrix y_i\\ 0 pmatrix (21) Step 5: Output Computation. The attention output for the query token: q′ e_q =q+P∑i=1Ni(i⊤) = e_q+ W^P _i=1^N v_i( k_i q) (22) =(q1)+ηN∑i=1N(i0)ri = pmatrix y_q\\ 1 pmatrix+ ηN _i=1^N pmatrix y_i\\ 0 pmatrixr_i (23) =(q+ηN∑i=1Nrii1) = pmatrix y_q+ ηN _i=1^Nr_i y_i\\ 1 pmatrix (24) The output embedding (first D dimensions) is: hout=q+ηN∑i=1Nriih_out= y_q+ ηN _i=1^Nr_i y_i (25) This completes the proof. ∎ D.2 Proof of Theorem 3.4: Softmax Attention as Boltzmann Aggregation Proof. Using the same weight matrix construction as in Theorem 3.3: Step 1: The attention logits are ⊤i=ri q k_i=r_i (as shown in the previous proof). Step 2: With softmax attention and temperature τ (from the dk d_k normalization): αi=exp(ri/τ)∑j=1Nexp(rj/τ) _i= (r_i/τ) _j=1^N (r_j/τ) (26) Step 3: The output becomes: hout=hq+∑i=1Nαii=hq+∑i=1Nexp(ri/τ)∑j=1Nexp(rj/τ)ih_out=h_q+ _i=1^N _i y_i=h_q+ _i=1^N (r_i/τ) _j=1^N (r_j/τ) y_i (27) This is a Boltzmann-weighted (softmax) aggregation of the output embeddings, where higher-scored examples receive exponentially more weight. ∎ D.3 Proof of Theorem 3.5: Structural Policy-Gradient Correspondence Proof. Step 1: Recall the policy gradient update. Let r~i=ri−r¯ r_i=r_i- r, so that ∑ir~i=0 _i r_i=0. The corresponding REINFORCE update in parameter space is: θ←θ+ηN∑i=1Nr~i∇θlogπθ(yi|x).θ←θ+ ηN _i=1^N r_i _θ _θ(y_i|x). (28) Step 2: Interpret gradient in embedding space. For typical LLM output layers: P(yi|h)∝exp(h⊤outembed(yi))P(y_i|h) (h W_outembed(y_i)) (29) The exact gradient of the log-probability in embedding space is: ∇hlogP(yi|h) _h P(y_i|h) =outembed(yi)−outy∼P(⋅|h)[embed(y)] = W_out\,embed(y_i)- W_outE_y P(·|h)[embed(y)] (30) =out(embed(yi)−¯) = W_out (embed(y_i)- y ) (31) where ¯=y∼P(⋅|h)[embed(y)] y=E_y P(·|h)[embed(y)] is the expected embedding under the current policy. This expectation is constant with respect to i. Step 3: Establish the correspondence. The REINFORCE estimator in hidden-state space is ηN∑i=1Nr~i∇hlogP(yi|h) ηN _i=1^N r_i _h P(y_i|h) =ηNout(∑i=1Nr~ii−¯∑i=1Nr~i) = ηN W_out ( _i=1^N r_i y_i- y _i=1^N r_i ) (32) =ηNout∑i=1Nr~ii, = ηN W_out _i=1^N r_i y_i, (33) where the second equality uses ∑ir~i=0 _i r_i=0. If out= W_out= I and the score coordinate in Theorem 3.3 contains r~i r_i, its attention update is exactly hout=hq+ηN∑i=1Nr~ii.h_out=h_q+ ηN _i=1^N r_i y_i. (34) For a general out W_out, the two updates differ by the output map; for raw, uncentered scores, they additionally differ by the centering term. Thus the general claim is structural correspondence, with exact equality only under the conditions stated in Theorem 3.5. ∎ D.4 Proof of Theorem 3.6: Bounded KL Shift Proof. Step 1: Characterize the distribution shift. From Theorem 3.3, the hidden state shifts by: Δh=ηN∑i=1Nrii h= ηN _i=1^Nr_i y_i (35) Step 2: Bound the shift magnitude. The magnitude of this shift is bounded: ‖Δh‖≤ηN∑i=1N|ri|‖i‖≤η⋅rmax⋅ymax\| h\|≤ ηN _i=1^N|r_i|\| y_i\|≤η· r_max· y_max (36) where rmax=maxi|ri|r_max= _i|r_i| and ymax=maxi‖i‖y_max= _i\| y_i\|. Step 3: Express the KL divergence exactly along the update path. For a softmax output layer P(y|h)∝exp(h⊤outembed(y))P(y|h) (h W_out\,embed(y)), define its log-partition function A(h)=log∑yexp(h⊤outembed(y)).A(h)= _y \! (h W_out\,embed(y) ). (37) Its Hessian equals the Fisher information matrix: ∇2A(h)=F(h).∇^2A(h)=F(h). (38) Let g(t)=A(h+tΔh)g(t)=A(h+t h). Using the exponential-family form of the softmax distribution, DKL(πICL∥πref) D_KL( _ICL\| _ref) =g′(1)−g(1)+g(0) =g (1)-g(1)+g(0) (39) =∫01tg′(t)t = _0^1t\,g (t)\,dt (40) =∫01tΔh⊤F(h+tΔh)Δht. = _0^1t\, h F(h+t h) h\,dt. (41) This identity is exact and does not discard a Taylor remainder. Step 4: Apply the pathwise curvature bound. Let ΛC=supt∈[0,1]λmax(F(h+tΔh)) _C= _t∈[0,1] _ (F(h+t h)). Then DKL(πICL∥πref) D_KL( _ICL\| _ref) ≤∫01tΛC‖Δh‖2t ≤ _0^1t\, _C\| h\|^2\,dt (42) =12ΛC‖Δh‖2 = 12 _C\| h\|^2 (43) ≤12ΛC(η⋅rmax⋅ymax)2. ≤ 12 _C (η· r_max· y_max )^2. (44) The bound follows from the bounded constructed update and the curvature of the output distribution along its path. It is analogous to a single trust-region step, but does not assert that fixed parameters alone implement an explicit KL penalty. ∎ D.5 Proof of Theorem 3.7: Policy-Gradient Analogy with a Bounded KL Shift Proof. We combine the previous results to establish the main theorem. Step 1: Reward-Weighted Update Component. From Theorem 3.3, linear attention produces: hout=hq+ηN∑i=1Nriih_out=h_q+ ηN _i=1^Nr_i y_i (45) From Theorem 3.5, applying the same construction to centered scores r~i=ri−r¯ r_i=r_i- r exactly matches the hidden-state REINFORCE estimator when out= W_out= I, and gives a structural correspondence more generally: ηN∑i=1Nr~i∇hlogP(yi|h)=ηNout∑i=1Nr~ii. ηN _i=1^N r_i _h P(y_i|h)= ηN W_out _i=1^N r_i y_i. (46) Under Assumption 3.2, movement toward a particular output embedding increases that output’s probability. This assumption provides the behavioral interpretation of the score-weighted shift; it is not a general monotonic-improvement guarantee. Step 2: Bounded KL-Shift Component. From Theorem 3.6, the simplified single-step distribution shift satisfies: DKL(πICL∥πref)≤12ΛC(η⋅rmax⋅ymax)2.D_KL( _ICL\| _ref)≤ 12 _C(η· r_max· y_max)^2. (47) Conclusion. Score-conditioned ICL therefore has two properties that parallel a trust-region policy update: 1. reward-weighted aggregation with a conditional hidden-state REINFORCE correspondence; and 2. an exact pathwise KL bound for the constructed bounded update. These properties establish the analogy in Theorem 3.7; they do not imply that ICL optimizes the KL-regularized objective in Eq. 6. ∎ Appendix E Learning Rate Sensitivity of REINFORCE Comparison To assess the robustness of the directional alignment result in Section 4.5, we repeat the cosine similarity analysis across three learning rates η∈10−6,10−5,10−4η∈\10^-6,10^-5,10^-4\ and two sample sizes N∈4,8N∈\4,8\. Results (averaged over tasks and models) are shown in Table 3. Learning Rate N=4N=4 N=8N=8 10−610^-6 0.41† 0.44† 10−510^-5 0.51† 0.55† 10−410^-4 0.49† 0.53† Table 3: Cosine similarity between ICL and REINFORCE Δlogp p vectors across learning rates and sample sizes. All values are statistically significant (p<0.01p<0.01), confirming that directional alignment is robust to these hyperparameter choices. Appendix F Empirical Verification of Weight Matrix Structure A potential criticism of Theorem 3.3 is that it provides an existence proof, demonstrating that attention can implement score-weighted aggregation, rather than evidence that pretrained LLMs actually implement this mechanism. A related concern is that our theoretical analysis uses linear self-attention and a single layer, whereas real LLMs employ softmax attention, nonlinear activations (e.g., GeLU), and many layers. Linear attention is a standard simplification used throughout the ICL theory literature (Von Oswald et al., 2023; Akyürek et al., 2023); Theorem 3.4 extends the existence result to softmax attention. The experiments below address the remaining gap: we show that the qualitative predictions of our theory, including score-proportional attention weights, score-specialized heads, and causal score processing, hold robustly in fully nonlinear, multi-layer pretrained models, despite the gap between the theoretical abstraction and the actual architecture. In this appendix, we present four complementary experiments that examine whether the internal computations of LLMs align with our theoretical construction. F.1 Score-Specialized Attention Head Identification Our theoretical construction in Theorem 3.3 predicts that effective score-weighted aggregation requires attention heads whose key-query computation extracts the score component. If LLMs implement this mechanism, we should find attention heads that are “specialized” for score extraction, i.e., heads where the key-query dot product correlates strongly with the score value. Method. For each layer l and head h, we analyze the relationship between key-query products and scores. Given a score-annotated context C=(i,ri)i=1NC=\( y_i,r_i)\_i=1^N, let i(l,h) k_i^(l,h) denote the key vector for the token corresponding to example i’s score, and (l,h) q^(l,h) denote the query vector at the generation position. We compute the Score Extraction Index (SEI): SEI(l,h)=ρ(((l,h))⊤i(l,h)i=1N,rii=1N)SEI^(l,h)=ρ (\( q^(l,h)) k_i^(l,h)\_i=1^N,\r_i\_i=1^N ) (48) where ρ denotes Spearman correlation. A high SEI indicates that the head’s key-query computation effectively extracts score information, as predicted by our theoretical construction where ⊤i=ri q k_i=r_i. Results. Table 4 reports the distribution of SEI values across all heads for each model. Critically, all models contain a subset of heads with high SEI values (>0.6>0.6), indicating genuine score-specialized computation. These high-SEI heads are predominantly located in middle-to-late layers, consistent with the layer-wise patterns observed in Section 4.3. Model Mean SEI Max SEI #Heads >0.6>0.6 Peak Layer Llama3-8B 0.18 0.72 12/256 24–28 Llama3-70B 0.21 0.79 31/640 62–72 Olmo3-7B 0.16 0.68 9/256 21–26 Olmo3-32B 0.19 0.74 22/512 48–56 Qwen3-4B 0.15 0.65 7/144 27–32 Table 4: Score Extraction Index (SEI) statistics. All models contain attention heads with high SEI, indicating score-specialized computation consistent with Theorem 3.3. Importantly, we verify that these high-SEI heads are not artifacts of position or formatting. When we shuffle scores while keeping examples fixed (as in Section 4.2), the SEI of these heads drops to near zero (mean: 0.03), confirming that they genuinely track score values rather than positional patterns. F.2 Functional Decomposition of Attention Weights While the SEI analysis identifies heads that extract score information, it does not directly verify that the full attention mechanism implements our theoretical construction. Here, we analyze whether the effective computation performed by attention heads can be decomposed into components matching Theorem 3.3. Method. For each attention head, we model the attention weight assigned to example i as: αi=β0+βr⋅ri+βp⋅pi+βs⋅si+ϵ _i= _0+ _r· r_i+ _p· p_i+ _s· s_i+ε (49) where rir_i is the score, pip_i is the positional index, and sis_i is a semantic similarity measure (cosine similarity between the example embedding and the query). Our theory predicts that for score-specialized heads, βr _r should dominate while βp _p and βs _s should be small. We fit this regression model across all test prompts and compute the Score Dominance Ratio (SDR): SDR(l,h)=|βr|max(|βr|,|βp|,|βs|)SDR^(l,h)= | _r| (| _r|,| _p|,| _s|) (50) Results. Table 5 shows that heads with high SEI also exhibit high SDR, meaning their attention weights are primarily determined by score values rather than position or semantic similarity. This provides evidence that these heads implement a computation functionally similar to our theoretical construction, where attention weights directly reflect scores. Model SDR (High-SEI Heads) SDR (All Heads) R2R^2 Llama3-8B 0.84 0.31 0.71 Llama3-70B 0.89 0.35 0.78 Olmo3-7B 0.81 0.29 0.68 Olmo3-32B 0.86 0.33 0.74 Qwen3-4B 0.79 0.27 0.65 Table 5: Score Dominance Ratio (SDR) for high-SEI heads vs. all heads. High-SEI heads show score-dominated attention computation, supporting our theoretical construction. F.3 Linear Probing for Score Information in Hidden States If LLMs implement score-weighted aggregation, score information must be encoded in the hidden states in a way that downstream layers can access. We use linear probing to verify that score values are linearly decodable from hidden states. Method. For each layer l, we extract the hidden state i(l) h_i^(l) at the position of each score token. We train a linear probe r^i=⊤i(l)+b r_i= w h_i^(l)+b to predict the score value from the hidden state, using 80% of examples for training and 20% for testing. We report the R2R^2 coefficient on the test set. Results. Table 6 shows the probing R2R^2 across layer groups for all models. All models exhibit a consistent pattern: score information becomes increasingly linearly decodable in later layers. Early layers (0–33% depth) show low R2R^2, indicating score information is not yet explicitly represented. Middle layers (33–66% depth) show moderate R2R^2, suggesting score processing is emerging. Late layers (66–100% depth) achieve peak R2>0.85R^2>0.85, with the maximum occurring near the layers containing high-SEI heads identified in Section F.1. This confirms that score information is explicitly represented in hidden states and accessible for downstream computation. Model Early Middle Late Peak R2R^2 Peak Layer Llama3-8B 0.12 0.48 0.79 0.87 26/32 Llama3-70B 0.15 0.54 0.84 0.92 70/80 Olmo3-7B 0.11 0.45 0.76 0.85 24/32 Olmo3-32B 0.14 0.51 0.81 0.89 54/64 Qwen3-4B 0.10 0.42 0.74 0.86 30/36 Table 6: Linear probing R2R^2 for score prediction across layer groups. Early: 0–33% depth; Middle: 33–66%; Late: 66–100%. Score information becomes increasingly decodable in later layers, peaking near the high-SEI head locations from Table 4. F.4 Causal Intervention: Amplifying Score-Specialized Heads The analyses above are correlational. To establish a causal link between score-specialized heads and ICL behavior, we conduct an intervention experiment. Method. We selectively amplify or suppress the contribution of high-SEI heads by scaling their output: output(l,h)←γ⋅output(l,h)output^(l,h)←γ·output^(l,h) (51) where γ>1γ>1 amplifies and γ<1γ<1 suppresses. If these heads causally implement score-weighted aggregation, amplification should strengthen the correlation between scores and probability changes, while suppression should weaken it. Results. Table 7 shows the effect of intervening on high-SEI heads. Amplifying high-SEI heads (γ=1.5γ=1.5) increases the score-probability correlation by 15–23%, while suppressing them (γ=0.5γ=0.5) decreases it by 31–42%. Critically, the same interventions on randomly selected heads (matched for layer distribution) show minimal effect (<5%<5\% change), confirming that the identified heads are specifically responsible for score-weighted computation. High-SEI Heads Random Heads Model γ=1.5γ=1.5 γ=0.5γ=0.5 γ=1.5γ=1.5 γ=0.5γ=0.5 Llama3-8B +18% −-35% +2% −-3% Llama3-70B +23% −-42% +4% −-2% Olmo3-7B +15% −-31% +1% −-4% Olmo3-32B +21% −-38% +3% −-1% Qwen3-4B +16% −-33% +2% −-3% Table 7: Change in score-probability correlation ρ(r,Δlogp)ρ(r, p) under causal intervention. Amplifying high-SEI heads strengthens score-weighted behavior; suppressing them weakens it. Random heads show minimal effect. F.5 Discussion: From Existence Proof to Empirical Mechanism The experiments in this appendix collectively address the gap between our theoretical existence proof (Theorem 3.3) and empirical reality: 1. Score-Specialized Heads Exist (Section F.1): We identify specific attention heads whose key-query computation strongly correlates with score values, consistent with the weight matrix structure in our theorem. 2. Score Information Dominates Attention (Section F.2): In these specialized heads, attention weights are primarily determined by scores rather than position or semantic similarity, matching our theoretical prediction. 3. Score Information is Linearly Encoded (Section F.3): Score values are explicitly represented in hidden states and become increasingly accessible in later layers. 4. Causal Role Confirmed (Section F.4): Intervening on score-specialized heads directly modulates the strength of score-weighted ICL behavior. While LLMs do not literally implement our exact weight matrix construction (which would require specific zero patterns), they have learned functionally similar mechanisms through pretraining. The existence of score-specialized heads with high SEI and SDR indicates that the computational structure predicted by Theorem 3.3, where attention weights track reward values, emerges naturally in pretrained models. This bridges the gap between theoretical possibility and practical implementation. Appendix G Effect of Score Representation Format Our theoretical framework assumes that attention mechanisms can extract and use scalar score information from context. However, in practice, scores can be represented in various surface formats. We investigate whether the implicit policy gradient mechanism is sensitive to representational differences, even when the underlying information is identical. We conduct this analysis on the DROP task, which uses F1 scores ranging continuously from 0 to 1. We compare five score representation formats that all preserve the same underlying continuous information and differ only in surface form: Decimal presents the raw F1 score (e.g., “Score: 0.82”) as a standard numerical format and serves as the ideal case for our theory; Percentage renders the same value as a percent (e.g., “Score: 82%”) as commonly used in human-facing settings; Fraction expresses it as a ratio (e.g., “Score: 82/100”) to test whether making the denominator explicit changes processing; Points uses an integer plus a unit label (e.g., “Score: 82 points”) to examine the effect of unit annotation; and Natural language verbalizes the value (e.g., “scored eighty-two percent”) to test whether numeric tokens themselves are necessary. We also include a No score control condition in which examples are shown without any score annotation, which is essential for determining whether models actually leverage explicit score information or instead improve simply by learning from the content of high-quality examples. We evaluate each format using the same metrics as in the main experiments: the correlation between scores and probability changes (ρ(r,Δlogp)ρ(r, p); see RQ1), the correlation between scores and attention weights (ρ(r,α)ρ(r,α); see Section 4.3), and iterative improvement (Iter. Gain; see Section 4.4). Representation ρ(r,Δlogp)ρ(r, p) ρ(r,α)ρ(r,α) Iter. Gain Decimal 0.63† 0.45† +0.19 Percentage 0.61† 0.44† +0.18 Fraction 0.59† 0.42† +0.17 Points 0.58† 0.41† +0.17 Natural language 0.52† 0.36† +0.14 No score 0.06 0.04 +0.02 Table 8: Effect of score representation format on ICL behavior for the DROP task, averaged across all models. ρ(r,Δlogp)ρ(r, p): Spearman correlation between F1 scores and log-probability changes; ρ(r,α)ρ(r,α): correlation between F1 scores and attention weights; Iter. Gain: average F1 improvement after 10 iterations. † denotes statistically significant correlation (p<0.01p<0.01). Table 8 shows that all numeric formats (Decimal, Percentage, Fraction, Points) achieve similarly high correlations, suggesting that LLMs can robustly extract numerical information regardless of superficial formatting choices. The natural language format shows a notable drop, indicating that mapping words like “eighty-two” to numerical values introduces noise. Crucially, all scored representations dramatically outperform the no-score control (p<0.001p<0.001), confirming that models rely on explicit score information rather than inferring quality from example content alone. These results have practical implications: practitioners can flexibly choose among numeric formats (decimal, percentage, fraction, points) without significant loss of effectiveness. However, if scores must be expressed in natural language, some degradation in the implicit policy gradient effect should be expected.