Provable and Practical In-Context Policy Optimization for Self-Improvement

TL;DR

This paper introduces ICPO, a method enabling language models to improve their responses through self-reflection at inference time, backed by theory and practical algorithms that enhance reasoning performance.

Contribution

The paper presents a novel theoretical framework for self-reflection in language models and proposes ME-ICPO, a practical algorithm for in-context policy optimization during inference.

Findings

ME-ICPO achieves top-tier performance on mathematical reasoning tasks.

Theoretical proof shows single-layer attention can imitate policy optimization algorithms.

ME-ICPO maintains inference efficiency while improving reasoning accuracy.

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 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.