Skip to content
Preprint

In-Context Learning as Implicit Policy Gradient

Jul 2026 · 0 citations · 44 references
Computer Science

TL;DR

It is shown that score-conditioned In-Context Learning (ICL) admits a structural correspondence to policy gradient optimization, and an exact upper bound on the distribution shift induced by a bounded attention update is derived, yielding a trust-region-like analogy to KL-constrained policy optimization.

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.

View source

Similar papers

Preprint Jul 2026

ISO: An RLVR-Native Optimization Stack

Reinforcement learning with verifiable rewards (RLVR) is rapidly advancing the reasoning capabilities of language models, yet the optimization layer that converts reward feedback into weight-space updates remains poorly understood. Building on our prior analysis (Zhu et al., 2025), we study this missing layer through the singular structure of model weights and identify spectral inheritance: RLVR can reuse the base model's weight spectra while acquiring new behavior through changes in the associated input and output singular frames. We operationalize spectral inheritance as Isospectral Optimization (ISO), an RLVR-native, fixed-spectrum optimization framework with complementary offline and online instantiations. Offline, ISO-Merger combines the frame changes of shared-base specialists into a single fixed-spectrum model, requiring no post-merge data, rollouts, gradient updates, or on-policy distillation (OPD). It recovers complementary specialist capabilities and achieves the strongest aggregate performance among the compared data-free merging methods. Online, ISO-Optimizer applies a chosen base optimizer, including AdamW and Muon, to the frame variables while keeping the base spectra fixed. Across reasoning and coding tasks ranging from 1.5B to 8B parameters, ISO-Optimizer improves accuracy in the reported runs and reaches matched scores with substantially fewer training steps. On Qwen3-8B-Base, AdamW reaches an aggregate accuracy of 0.495 after 270 training steps. ISO-AdamW reaches the same accuracy after only 100 training steps and improves further to 0.509 after 210 training steps. Together, ISO offers a concrete answer to RLVR's missing optimization layer: rather than inheriting pre-training optimization wholesale, design post-training around the structure of reward-driven adaptation: inherit the spectrum, optimize the frames.

Hanqing Zhu, Wenyan Cong, Zhizhou Sha et al. · 0 citations
#small language model Preprint Aug 2026

Boosting LLM Exploration via Weak-Model Guidance in RLVR

This work empirically study the potential of outer prefixes, revealing the mechanism of the impact of distributional discrepancy to the exploration dynamics in RLVR training and efficiently mitigates entropy collapse without requiring additional SFT, intricate reward designs, or complex prompting.

Xin Shen, Huishuai Zhang, Peng Li et al. · 0 citations
Preprint Jul 2026

Weak-to-Strong Generalization via Direct On-Policy Distillation

Reinforcement learning with verifiable rewards (RLVR) is a powerful recipe for improving language-model reasoning, but it is expensive to repeat on every new strong model because the target model must generate many rollouts during training. As models scale, post-training itself becomes a bottleneck. We study a weak-to-strong alternative: run RL on a smaller model where rollouts are cheaper, then reuse what that RL run learned to improve a stronger target model. Directly distilling the post-RL weak teacher is not enough, because the teacher's final policy mixes useful RL gains with the limitations of the smaller model. We propose Direct On-Policy Distillation (Direct-OPD), which transfers the teacher's RL-induced policy shift instead. Direct-OPD compares the post-RL teacher with its own pre-RL reference and treats their log-ratio as a dense implicit reward for the student. In plain terms, the checkpoint pair tells us which actions RL made the weak model more or less likely to take, and Direct-OPD applies that signal on the stronger student's own on-policy states. This directly reuses the weak model's RL supervision signal without running sparse-reward RL on the target model. Empirically, Direct-OPD consistently leverages weaker teachers to improve stronger target models; notably, it boosts Qwen3-1.7B from 48.3% to 58.3% on AIME 2024 in just 4 hours on 8 A100 GPUs. It outperforms step-matched direct RL and enables the sequential composition of multiple policy shifts. Our results show that RL outcomes can be reused across model scales as implicit reward signals, not merely as final models to imitate.

Shiyuan Feng, Huan Gao, Haohan Chi et al. · 6 citations
Preprint Jul 2026

TAPO: Transition-Aware Policy Optimization for LLM Agents

This work proposes TAPO: Transition-Aware Policy Optimization for LLM Agents, a unified training framework that alternates between policy optimization and transition supervision, and demonstrates that TAPO consistently improves task performance over pure policy optimization baselines.

Cong Li, Peixi Peng, Yisen Zhao et al. · 0 citations
Book Open access Aug 2026

InfRL: Inference-time Reinforcement Learning for Research Idea Optimization

Large language models (LLMs) possess extensive latent knowledge yet remain largely static at inference. Once prompted, their generation policy typically cannot evolve, and post-hoc ''self-reflection'' methods provide no explicit principled learning signals. To address this limitation, we formally model iterative research idea optimization as a finite-horizon Markov Decision Process and propose InfRL (Inference-time Reinforcement Learning), a framework designed for effective policy improvement at inference without updating model weights. InfRL coordinates three specialized LLM agents: (i) a State Transition Agent that proposes candidate ideas; (ii) a Policy Update Agent that learns feedback strategies based on idea trajectories; and (iii) a Reward Agent that assigns normalized, comparative rewards, enabling nuanced reinforcement signals. This empowers the Policy Update Agent to dynamically refine feedback strategies, progressively leveraging latent knowledge encoded within the LLM. We evaluate InfRL on five balanced datasets covering 500 recent research papers from health, genetics, environment, neuroscience, and engineering domains. Compared to a single-pass GPT-4o baseline and a strong self-reflection baseline, InfRL improves the novelty of generated ideas by 3.93%--32.7% and the feasibility of generated ideas by 70.1%--448.6% with GPT-4o, exhibiting consistent improvements across ten inference-time iterations. Ablation studies underscore the critical roles of our reward formulation and modular agent design, while human evaluations confirm alignment between reward trends and perceived idea quality. InfRL thus offers a practical and domain-agnostic approach to harness reinforcement learning during inference, bridging the gap between static prompting and computationally intensive parameter-level fine-tuning. The code and the dataset we use are provided at: https://github.com/amir-hassan25/InfRL

Sikun Guo, Amir Hassan Shariatmadari, Jiuqi Wang et al. · 0 citations