[论文解读] Provable and Practical In-Context Policy Optimization for Self-Improvement

Tianrun Yu, Yuxiao Yang+7 more|arXiv (Cornell University)|Mar 2, 2026

Reinforcement Learning in Robotics被引用 0

一句话总结

该论文理论上证明单层线性自注意力变换器在有上下文数据时可以模仿策略优化，并提出 ME-ICPO，一种在测试时进行最小熵选择的自我提升数学推理的实用自我改进算法。

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.

研究动机与目标

• 将多轮自我反思形式化为上下文中的策略优化（ICPO），以理解LLM的自我提升。
• 证明充分预训练的单层线性自关注模型可以在线性赌博问题上模仿策略优化。
• 引入 ME-ICPO，一种在测试时通过自我评估的奖励在上下文中细化回答的实用算法。
• 在标准数学推理基准上展示鲁棒性和具有竞争力的性能。

提出的方法

• 定义一个 ICPO 框架，其中模型使用过去的行动和观测到的奖励来更新其上下文策略。
• 推导一个闭式、线性自关注（LSA）表示，能够在 Fisher 加权对数匹配目标下模拟策略优化更新。
• 提出一个有监督的预训练损失，通过投影的 Fisher 损失和 KL 替代关系将 LSA 输出与策略优化对数（logits）关联起来。
• 建立理论保证：总体等价性（定理 4.2）、有限样本一致性（定理 4.3）以及对奖励扰动的稳定性（定理 4.8）。
• 给出 ME-ICPO 的实用、无梯度测试时算法，利用多数票奖励、CoT 汇总和最小熵选择来更新上下文历史。
• 提供 ME-ICPO 的算法蓝本，讨论复杂性考量与鲁棒性方面。

实验结果

研究问题

• RQ1 ICPO 是否能解释在不更新参数的情况下，LLM 如何进行自我反思并改进回答？
• RQ2 在预训练条件下，简单的单层线性自关注网络是否能利用上下文数据模仿策略优化过程？
• RQ3 在 ICPO 中，Fisher 加权投影损失是否能为 KL 基的策略目标提供可靠的代理？
• RQ4 ME-ICPO 在面对嘈杂的自我评估奖励时是否鲁棒，且在推理时对于数学推理任务是否具备可扩展性？
• RQ5 相较于基线，ME-ICPO 在标准数学推理基准上的经验提升有哪些？

主要发现

• 在充分预训练的前提下，单层 LSA 可以证明性地模仿线性带宽问题的策略优化。
• Fisher 加权损失为混合策略 KL 提供紧致代理，支持上下文中的策略学习。
• 总体等价性成立：预训练的 LSA 可以模仿对所有历史的 PO 策略（定理 4.2）。
• 有限样本保证确保经验训练数据在高概率下足以模仿 PO 策略（定理 4.3）。
• ME-ICPO 在数学推理任务上表现具有竞争力，并在推理时具备效率优势。
• ME-ICPO 对一次性奖励扰动具有鲁棒性，随着训练进展呈现减弱效应（定理 4.8）。