[2603.01335] Provable and Practical In-Context Policy Optimization for Self-Improvement

Computer Science > Machine Learning arXiv:2603.01335 (cs) [Submitted on 2 Mar 2026] Title:Provable and Practical In-Context Policy Optimization for Self-Improvement Authors:Tianrun Yu, Yuxiao Yang, Zhaoyang Wang, Kaixiang Zhao, Porter Jenkins, Xuchao Zhang, Chetan Bansal, Huaxiu Yao, Weitong Zhang View a PDF of the paper titled Provable and Practical In-Context Policy Optimization for Self-Improvement, by Tianrun Yu and Yuxiao Yang and Zhaoyang Wang and Kaixiang Zhao and Porter Jenkins and Xuchao Zhang and Chetan Bansal and Huaxiu Yao and Weitong Zhang View PDF HTML (experimental) Abstract:We study test-time scaling, where a model improves its answer through multi-round self-reflection at inference. We introduce In-Context Policy Optimization (ICPO), in which an agent optimizes its response in context using self-assessed or externally observed rewards without modifying its parameters. To explain this ICPO process, we theoretically show that with sufficient pretraining under a novel Fisher-weighted logit-matching objective, a single-layer linear self-attention model can provably imitate policy-optimization algorithm for linear bandits. Building on this theory, we propose Minimum-Entropy ICPO (ME-ICPO), a practical algorithm that iteratively uses its response and self-assessed reward to refine its response in-context at inference time. By selecting the responses and their rewards with minimum entropy, ME-ICPO ensures the robustness of the self-assessed rewards via majority v...