Provable and Practical In-Context Policy Optimization for Self-Improvement
This addresses the challenge of enabling models to self-improve at inference time without parameter updates, offering a principled approach for mathematical reasoning, though it appears incremental as it builds on existing self-reflection concepts.
The paper tackles the problem of test-time scaling for models to improve answers through self-reflection at inference, introducing In-Context Policy Optimization (ICPO) and a practical algorithm, ME-ICPO, which achieves competitive, top-tier performance on standard mathematical reasoning tasks while keeping inference costs affordable.
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 voting. Across standard mathematical reasoning tasks, ME-ICPO attains competitive, top-tier performance while keeping inference costs affordable compared with other inference-time algorithms. Overall, ICPO provides a principled understanding of self-reflection in LLMs and yields practical benefits for test-time scaling for mathematical reasoning.