REVIEW 4 major objections 6 minor 47 references
In-Context Learning as Implicit Policy Gradient
T0 review · 4 major / 6 minor · reviewed 2026-08-01 · deepseek-v4-flash
Pith's one-line read Score-conditioned in-context learning is structurally equivalent to a policy-gradient update: attention weights examples by their scores, and under simplifying conditions this matches REINFORCE exactly with a bounded KL shift.
desk verdict Clean existence proof for score-conditioned attention as implicit REINFORCE, with honest hedging—but the empirical bridge is weaker than the title suggests, and one headline result claims support from the wrong theorem. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing machinery is a specific weight-matrix construction for self-attention: W_K = W_Q project each context token (y_i, r_i) onto its score coordinate, so the query–key product equals the score r_i; W_V extracts the embedding component; and W_P = (η/N) I scales the output. This makes the attention update h_out = h_q + (η/N) Σ r_i y_i, a score-weighted aggregation of example embeddings that mirrors a REINFORCE step. The KL bound rests on the exponential-family identity that the Hessian of the log-partition of the softmax output layer is the Fisher information matrix (a curvature measure of the output distribution), giving the exact path-integral expression D_KL = ∫₀¹ t ∆hᵀ F(h+t∆h)
What would settle it
Compute, for a pretrained model given N scored examples, the actual hidden-state shift at the final layer and its cosine similarity with the score-weighted sum of example embeddings (with centered scores). If the cosine similarity is near zero or negative for a model that nonetheless shows score–probability correlation, the structural mechanism is refuted as a general explanation. Alternatively, test the exact equality in Theorem 3.5: measure whether the ICL-induced log-probability change equals the effect of the hidden-state REINFORCE update when the output layer is set to identity; failure w
Extended reading notes
Core claim
The central claim is that score-conditioned in-context learning admits a structural correspondence to policy-gradient optimization. Concretely, self-attention can be configured — via a construction where the key–query product reads off each example's score and the value projection extracts its embedding — to produce the hidden-state update h_out = h_q + (η/N) Σ_i r_i y_i. This is a score-weighted aggregation of example embeddings, which is the hidden-state form of the REINFORCE estimator (a policy-gradient update using a score-weighted sum of gradient directions) when the scores are mean-centered and the output embedding map is the identity. The softmax version becomes a Boltzmann-weighted a
Load-bearing premise
The load-bearing premise is that the simplified single-layer linear-attention construction with identity output map is a faithful model of how pretrained multi-layer transformers move hidden states in response to score-conditioned context; if real attention does not move hidden states along score-weighted output embeddings, or when uncentered scores and a non-identity output layer are used (as in the experiments), the policy-gradient correspondence is merely an analogy rather
Editorial extensions
If this is right
- If the correspondence is correct, a score-conditioned prompt alone — with no weight updates — shifts the model's output distribution along the score-weighted direction, so in-context examples act as data for a hidden-state gradient step.
- In the simplified model, the empirically used configuration (raw scores, general output layer) yields a directional analogy; the exact REINFORCE equality requires centered scores and an identity output map, which is testable.
- Because the update is bounded, a single score-conditioned step cannot push the distribution arbitrarily far: the paper proves D_KL ≤ ½ Λ_C (η r_max y_max)^2, the analog of a trust region.
- Iterative score-conditioned ICL behaves like repeated reward-conditioned shifts; each step is individually bounded, though monotonic improvement is an empirical finding rather than a theorem.
- Pretrained models appear to contain score-specialized attention heads whose causal manipulation strengthens or weakens score-driven behavior, which pins the mechanism to identifiable internal components.
Reading between the lines
- If the simplified model captures real dynamics, one testable consequence is that models whose output layer is closer to an identity map should show stronger alignment with REINFORCE; models whose output layer strongly reshapes embeddings should show weaker alignment. This is an inference, not a claim in the paper.
- A safety-relevant extension of the KL bound is that a single adversarial construction of high-scoring context may not dramatically shift the output distribution on its own; but the paper's own caveat that repeated iterations are not covered means cumulative drift across many turns remains an open question.
- The structural correspondence suggests that training data containing scored samples could implicitly teach attention heads to become score extractors, implying the mechanism might be present in models not explicitly trained with reward-conditioned prompts; the paper does not test this.
- A practical testable extension: if the mechanism is correct, then adding a mean-centered baseline to the scores in the prompt should strengthen the alignment with REINFORCE beyond using raw scores, since the paper shows the exact equality requires centering.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper argues that score-conditioned in-context learning (ICL) is structurally analogous to policy gradient optimization. It proves (Thm 3.3) that with a specific weight construction, linear self-attention computes h_out = h_q + (η/N) Σ r_i y_i; (Thm 3.4) that softmax attention gives Boltzmann weighting; (Thm 3.5) that with centered scores and W_out = I this update equals the hidden-state REINFORCE estimator; and (Thm 3.6) that the induced softmax distribution shift obeys a pathwise KL upper bound. Experiments on Llama-3, Olmo-3, Qwen-3, GPT-4o, Gemini 2.5, and Sonnet 4 measure score–probability correlations, score–attention correlations, iterative gains, causal head interventions, and cosine similarity with one-step REINFORCE. The authors acknowledge that the formal results hold in a simplified single-layer linear model, and treat the transfer to pretrained models as an empirical question.
Significance. The constructive existence proof and the exact KL identity (Appendix D.4) are clean and useful: they show that attention can implement reward-weighted hidden-state updates and that the resulting update has a trust-region-like KL bound. The paper also includes controls that go beyond simple correlation: shuffled-score controls, score-specialized head identification with SEI and SDR, linear probing, and causal head amplification/suppression. If the empirical bridge were established, this would be an important step toward explaining OPRO/self-rewarding behavior. However, the bridge is currently incomplete: the central empirical evidence does not verify the hidden-state displacement direction and is compatible with simpler score-sensitive attention mechanisms.
major comments (4)
- [§4.2, Fig. 1a] The claimed confirmation of 1/N scaling is directed at Theorem 3.3, but the experiments use softmax attention; Theorem 3.4 predicts Boltzmann weighting without an η/N factor. A decrease in |Δlogp| roughly following 1/N is therefore not evidence for Theorem 3.3 as stated. To support the claimed scaling, the paper would need to measure the effective aggregation (e.g., in high-SEI heads) or evaluate a linear-attention variant; otherwise Fig. 1a should be reported as an empirical observation rather than a test of Theorem 3.3.
- [§4.5, RQ4] The cosine-similarity test lacks alternative baselines. Any context-induced shift that is monotonically increasing in score will have positive cosine with the one-step REINFORCE direction; values above 0.45 are not diagnostic. Moreover, Theorem 3.5's exact equality requires centered scores and W_out = I, while the experiments use raw scores in multi-layer softmax transformers with learned output maps. The statement that this experiment 'directly validates Theorem 3.5' is therefore overstated. The experiment should compare against copy-best, rank-based attention, and semantic-matching baselines, and state explicitly which part of the theorem is being tested.
- [§4.2–4.3 and Appendix F] The correlations and causal interventions show that LLMs upweight high-scoring examples, but they do not establish that the hidden state moves along Σ r_i y_i. Appendix F's SDR analysis conditions on attention weights, not on the hidden-state update; a head could weight by score while the output map and residual stream produce a different update. To test the structural correspondence, the paper should measure the direction of the actual hidden-state displacement Δh and compare it with the predicted direction, rather than only correlating attention weights with scores.
- [§3.3 (Theorem 3.5), §4.1, §4.5] The paper's own caveats are in tension with some experimental claims. Theorem 3.5 is exact only under centered scores and W_out = I, and Remark A.1 limits the general claim to a structural analogy; however, Section 4.1 states that raw scores are used while Section 4.5 says Theorem 3.5 is validated directly. This mismatch should be resolved by either adapting the REINFORCE comparison to centered scores with a controlled output layer, or explicitly labeling the result as testing only a weaker monotone-correspondence claim.
minor comments (6)
- [§4.2] The text calls the shuffled-score condition 'causal evidence'; shuffled labels rule out some confounds but are not a causal intervention. Suggest rephrasing to 'control condition'.
- [Appendix B, Algorithm 1] Algorithm 1 sets C_t to the single latest example, while §3.4 and §4.1 describe the context as containing scored outputs from the previous round. Please clarify whether the context accumulates history or uses only the most recent sample.
- [§4.1] The claim that sampling-based pseudo-likelihood 'preserves the rank ordering' of sequence likelihoods is asserted without justification. Since it is used only for black-box RQ1, either provide a reference or proof, or drop the claim.
- [Figure 3] Figure 3 lacks axis labels, error bars, and a description of how the cosine similarity is aggregated across prompts. Please add per-task statistics and confidence intervals.
- [Appendix F.1 vs §4.3] The SEI definition in Eq. (48) uses key-query products at score-token positions, while Section 4.3 defines α_i as attention weights averaged over all tokens of y_i. Reconcile these definitions so that the layer-wise comparison between Table 2 and Table 4 is interpretable.
- [Theorem 3.6 / §4.2] Theorem 3.6 bounds the shift using ∥y_i∥ in embedding space, which is not directly measured by |Δlogp|. The text should note this gap when interpreting Figure 1a.
Circularity Check
No significant circularity; the theoretical results are explicit constructions and algebraic identities, and the only self-citation (pseudo-likelihood) is non-load-bearing.
full rationale
The paper's core derivation chain is self-contained rather than circular. Theorem 3.3 is an explicit existence construction: the paper defines W_K, W_Q, W_V, W_P so that the attention weights equal scores, and then proves the resulting hidden-state update is h_out = h_q + (eta/N) sum r_i y_i. This is a legitimate constructive proof, not a conclusion that was assumed as an input. Theorem 3.5 is an algebraic identity: with centered scores and W_out = I, the hidden-state REINFORCE estimator reduces to exactly the same score-weighted sum. The equality is derived, not fitted, and the paper explicitly limits exactness to these stated simplifying conditions (Remark A.1, Theorem 3.5, Appendix D.3). Theorem 3.6 is a direct bound from the norm of the constructed update and the pathwise Fisher information; no fitted quantity is renamed as a prediction. The empirical sections test independent predictions: score-probability rank correlation with a shuffled-score control, attention-weight correlation, causal head intervention, and comparison with explicit REINFORCE. These tests do not use the constructed weight matrices as fitted parameters, and the causal intervention (Appendix F.4) is an independent manipulation rather than a restatement of the construction. The only self-citation is Kaneko et al. (2025a) for sampling-based pseudo-likelihood; the paper explicitly limits this to the black-box Sonnet 4 RQ1 analysis and states that all quantitative and mechanistic conclusions are anchored on open-weight models with exact log-probabilities, so it is not load-bearing. The skeptic's concerns about softmax vs. linear attention, W_out != I, and the non-diagnostic nature of positive cosine similarity are about evidential strength and generalization, not about circularity under the rubric. The paper also explicitly discloses the relevant limitations (linear attention, single-layer, W_out ~ I, centered-score special case), which further supports that no circular step is hidden.
Assumptions & free parameters
free parameters (3)
- η (attention construction learning rate) =
Not fitted in theory; set to 10^-5 for the REINFORCE comparison in RQ4
- Score-coordinate scale c (query constant) =
1 (gauge freedom, absorbed into η)
- SEI threshold and max-layer selection =
0.6 for SEI; 'Layer of Max' chosen as max over layers
assumptions (4)
- domain assumption Assumption 3.1: LLMs utilize score information in context to modulate generation (E[r(y)|C_high] > E[r(y)|C_low])
- domain assumption Assumption 3.2: Moving the hidden state toward an output embedding increases that output's probability (P(y_i|h+αy_i) > P(y_i|h))
- ad hoc to paper Simplified hidden-state model: single-layer linear self-attention, continuous output embeddings, W_out≈I
- standard math Softmax output layer as exponential family: P(y|h) ∝ exp(h^T W_out embed(y))
Cite this review
Pith. "Pith review of In-Context Learning as Implicit Policy Gradient." pith.science (2026). https://pith.science/paper/XKNODIBJ
@misc{pith2026260723153,
author = {Pith},
title = {Pith review of: In-Context Learning as Implicit Policy Gradient},
year = {2026},
howpublished = {\url{https://pith.science/paper/XKNODIBJ}},
note = {Machine review of arXiv:2607.23153}
}
read the original 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.
Figures
Reference graph
Works this paper leans on
-
[1]
Advances in neural information processing systems , volume=
Language models are few-shot learners , author=. Advances in neural information processing systems , volume=
-
[2]
International Conference on Machine Learning , pages=
Transformers learn in-context by gradient descent , author=. International Conference on Machine Learning , pages=. 2023 , organization=
2023
-
[3]
Advances in Neural Information Processing Systems , volume=
Direct preference optimization: Your language model is secretly a reward model , author=. Advances in Neural Information Processing Systems , volume=
-
[4]
arXiv preprint arXiv:2309.03409 , year=
Large language models as optimizers , author=. arXiv preprint arXiv:2309.03409 , year=
-
[5]
Findings of the Association for Computational Linguistics: ACL 2023 , pages=
Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers , author=. Findings of the Association for Computational Linguistics: ACL 2023 , pages=
2023
-
[6]
Advances in Neural Information Processing Systems , volume=
Self-refine: Iterative refinement with self-feedback , author=. Advances in Neural Information Processing Systems , volume=
-
[7]
International Conference on Learning Representations , year=
An Explanation of In-context Learning as Implicit Bayesian Inference , author=. International Conference on Learning Representations , year=
-
[8]
International Conference on Learning Representations , year=
What learning algorithm is in-context learning? Investigations with linear models , author=. International Conference on Learning Representations , year=
Show all 47 references
-
[9]
Advances in Neural Information Processing Systems , volume=
What Can Transformers Learn In-Context? A Case Study of Simple Function Classes , author=. Advances in Neural Information Processing Systems , volume=
-
[10]
arXiv preprint arXiv:2209.11895 , year=
In-context learning and induction heads , author=. arXiv preprint arXiv:2209.11895 , year=
-
[11]
Advances in Neural Information Processing Systems , volume=
Training language models to follow instructions with human feedback , author=. Advances in Neural Information Processing Systems , volume=
-
[12]
International Conference on Machine Learning , year=
Self-Rewarding Language Models , author=. International Conference on Machine Learning , year=
-
[13]
International Conference on Machine Learning , pages=
Model-agnostic meta-learning for fast adaptation of deep networks , author=. International Conference on Machine Learning , pages=. 2017 , organization=
2017
-
[14]
Advances in Neural Information Processing Systems , volume=
Attention is all you need , author=. Advances in Neural Information Processing Systems , volume=
-
[15]
International conference on machine learning , pages=
Trust region policy optimization , author=. International conference on machine learning , pages=. 2015 , organization=
2015
-
[16]
arXiv preprint arXiv:1707.06347 , year=
Proximal policy optimization algorithms , author=. arXiv preprint arXiv:1707.06347 , year=
-
[17]
2025 , month=
System Card: Claude Opus 4 & Claude Sonnet 4 , author=. 2025 , month=
2025
-
[18]
arXiv preprint arXiv:2507.06261 , year=
Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities , author=. arXiv preprint arXiv:2507.06261 , year=
-
[19]
arXiv preprint arXiv:2410.21276 , year=
Gpt-4o system card , author=. arXiv preprint arXiv:2410.21276 , year=
-
[20]
Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space
Geva, Mor and Caciularu, Avi and Wang, Kevin and Goldberg, Yoav. Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space. Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing. 2022. doi:10.18653/v1/2022....
2022 doi
-
[21]
Dissecting Recall of Factual Associations in Auto-Regressive Language Models
Geva, Mor and Bastings, Jasmijn and Filippova, Katja and Globerson, Amir. Dissecting Recall of Factual Associations in Auto-Regressive Language Models. Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. 2023. doi:10.18653/v1/2023.emnlp-main.751
2023 doi
-
[22]
Machine learning , volume=
Simple statistical gradient-following algorithms for connectionist reinforcement learning , author=. Machine learning , volume=. 1992 , publisher=
1992
-
[23]
arXiv e-prints , pages=
The llama 3 herd of models , author=. arXiv e-prints , pages=
-
[24]
arXiv preprint arXiv:2512.13961 , year=
Olmo 3 , author=. arXiv preprint arXiv:2512.13961 , year=
-
[25]
2025 , eprint=
Qwen3 Technical Report , author=. 2025 , eprint=
2025
-
[26]
Advances in Neural Information Processing Systems , volume=
Jailbreakbench: An open robustness benchmark for jailbreaking large language models , author=. Advances in Neural Information Processing Systems , volume=
-
[27]
arXiv preprint arXiv:2204.05862 , year=
Training a helpful and harmless assistant with reinforcement learning from human feedback , author=. arXiv preprint arXiv:2204.05862 , year=
-
[28]
arXiv preprint arXiv:2110.14168 , year=
Training verifiers to solve math word problems , author=. arXiv preprint arXiv:2110.14168 , year=
-
[29]
DROP : A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs
Dua, Dheeru and Wang, Yizhong and Dasigi, Pradeep and Stanovsky, Gabriel and Singh, Sameer and Gardner, Matt. DROP : A Reading Comprehension Benchmark Requiring Discrete Reasoning Over Paragraphs. Proceedings of the 2019 Conference of the North A merican Chapter of the Associa...
2019 doi
-
[30]
Zheng, Lianmin and Chiang, Wei-Lin and Sheng, Ying and Zhuang, Siyuan and Wu, Zhanghao and Zhuang, Yonghao and Lin, Zi and Li, Zhuohan and Li, Dacheng and Xing, Eric and others , journal=. Judging
-
[31]
arXiv preprint arXiv:2510.26519 , year=
Think Outside the Policy: In-Context Steered Policy Optimization , author=. arXiv preprint arXiv:2510.26519 , year=
-
[32]
Sampling-based Pseudo-Likelihood for Membership Inference Attacks
Kaneko, Masahiro and Ma, Youmi and Wata, Yuki and Okazaki, Naoaki. Sampling-based Pseudo-Likelihood for Membership Inference Attacks. Findings of the Association for Computational Linguistics: ACL 2025. 2025. doi:10.18653/v1/2025.findings-acl.465
2025 doi
-
[33]
Advances in Neural Information Processing Systems , volume=
Large language models can implement policy iteration , author=. Advances in Neural Information Processing Systems , volume=
-
[34]
Advances in Neural Information Processing Systems , volume=
Many-shot in-context learning , author=. Advances in Neural Information Processing Systems , volume=
-
[35]
arXiv preprint arXiv:2410.05362 , year=
Llms are in-context bandit reinforcement learners , author=. arXiv preprint arXiv:2410.05362 , year=
-
[36]
Advances in Neural Information Processing Systems , volume=
Bits leaked per query: Information-theoretic bounds for adversarial attacks on llms , author=. Advances in Neural Information Processing Systems , volume=
-
[37]
The Fourteenth International Conference on Learning Representations , year=
JailNewsBench: Multi-Lingual and Regional Benchmark for Fake News Generation under Jailbreak Attacks , author=. The Fourteenth International Conference on Learning Representations , year=
-
[38]
Online Learning Defense against Iterative Jailbreak Attacks via Prompt Optimization
Kaneko, Masahiro and Talat, Zeerak and Baldwin, Timothy. Online Learning Defense against Iterative Jailbreak Attacks via Prompt Optimization. Proceedings of the 14th International Joint Conference on Natural Language Processing and the 4th Conference of the Asia-Pacific Chapte...
2025 doi
-
[39]
arXiv preprint arXiv:2601.06884 , year=
Paraphrasing Adversarial Attack on LLM-as-a-Reviewer , author=. arXiv preprint arXiv:2601.06884 , year=
-
[40]
RLP rompt: Optimizing Discrete Text Prompts with Reinforcement Learning
Deng, Mingkai and Wang, Jianyu and Hsieh, Cheng-Ping and Wang, Yihan and Guo, Han and Shu, Tianmin and Song, Meng and Xing, Eric and Hu, Zhiting. RLP rompt: Optimizing Discrete Text Prompts with Reinforcement Learning. Proceedings of the 2022 Conference on Empirical Methods in...
2022 doi
-
[41]
arXiv preprint arXiv:2211.11890 , year=
Tempera: Test-time prompting via reinforcement learning , author=. arXiv preprint arXiv:2211.11890 , year=
-
[42]
The eleventh international conference on learning representations , year=
Large language models are human-level prompt engineers , author=. The eleventh international conference on learning representations , year=
-
[43]
2025 IEEE Conference on Secure and Trustworthy Machine Learning (SaTML) , pages=
Jailbreaking black box large language models in twenty queries , author=. 2025 IEEE Conference on Secure and Trustworthy Machine Learning (SaTML) , pages=. 2025 , organization=
2025
-
[44]
Advances in Neural Information Processing Systems , volume=
Tree of attacks: Jailbreaking black-box llms automatically , author=. Advances in Neural Information Processing Systems , volume=
-
[45]
Findings of the Association for Computational Linguistics: ACL 2026 , pages=
Experience-Driven Multi-Agent Optimization for Black-Box Jailbreak Attacks on Large Language Models , author=. Findings of the Association for Computational Linguistics: ACL 2026 , pages=
2026
-
[46]
Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages=
Trojail: Trajectory-level optimization for multi-turn large language model jailbreaks with process rewards , author=. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages=
-
[47]
International Conference on Learning Representations , volume=
Jailbreaking as a reward misspecification problem , author=. International Conference on Learning Representations , volume=
Reviewed August 1, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.