[논문 리뷰] Provable and Practical In-Context Policy Optimization for Self-Improvement

Tianrun Yu, Yuxiao Yang외 7명|arXiv (Cornell University)|2026. 03. 02.

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.

연구 동기 및 목표

• LLMs에서 자기개선을 이해하기 위해 다회 차 자기 성찰을 맥락 내 정책 최적화(ICPO)로 공식화한다.
• 충분히 사전학습된 단일 층 선형 자체 주의 모델이 선형 밴딧에서의 정책 최적화를 모방할 수 있음을 증명한다.
• 자기 평가 보상을 사용해 맥락 내 응답을 실시간으로 개선하는 실용적인 테스트 시점 알고리즘 ME-ICPO를 소개한다.
• 표준 수학 추론 벤치마크에서 강건성과 경쟁력 있는 성능을 입증한다.

제안 방법

• 과거 행동과 관찰된 보상을 사용해 맥락 내 정책을 업데이트하는 ICPO 프레임워크를 정의한다.
• 피셔 가중 로그-매칭 목적하에서 정책 최적화 업데이트를 모방하는 닫힌 형식의 선형 자기 주의(LSA) 표현을 도출한다.
• LSA 출력을 정책 최적화 로짓과 연결하는 투영된 Fisher 손실 및 KL 대리 관계를 통한 지도학습 사전학습 손실을 제안한다.
• 이론적 보장을 제시한다: 전체 모집단 등가성(Theorem 4.2), 유한 샘플 일관성(Theorem 4.3), 보상 가 perturbations에 대한 안정성(Theorem 4.8).
• 다수결 보상, CoTSummaries, 그리고 최소 엔트로피 선택을 사용해 맥락 내 이력을 업데이트하는 실용적이고 기울기 없는 테스트시점 알고리즘 ME-ICPO를 제시한다.
• ME-ICPO에 대한 알고리즘적 청사진을 제공하고 복잡도 고려사항과 강건성 측면을 논의한다.

실험 결과

연구 질문

• RQ1 ICPO가 파라미터를 업데이트하지 않고도 LLM이 어떻게 자기 성찰을 수행하고 응답을 개선하는지 설명할 수 있는가?
• RQ2 사전 학습 하에서 간단한 단일 층 선형 자기 주의 네트워크가 맥락 내 데이터를 사용하여 정책 최적화 과정을 모방할 수 있는가?
• RQ3 피셔 가중 투영 손실이 ICPO에서 KL 기반 정책 목표에 대한 신뢰할 수 있는 대리손실을 제공하는가?
• RQ4 ME-ICPO가 잡음이 있는 자기 평가 보상에 강건하고 수학적 추론 과제에서 추론 시점에 확장 가능하는가?
• RQ5 기준선과 비교하여 표준 수학 추론 벤치마크에서 ME-ICPO가 달성하는 실증적 이득은 무엇인가?

주요 결과

• 충분한 사전 학습이 주어지면 한 층 LSA가 선형 밴딧에 대한 정책 최적화를 이론적으로 모방할 수 있다.
• 피셔 가중 손실은 혼합 정책 KL의 촘촘한 대리로 작용하며 맥락 내 정책 학습을 뒷받침한다.
• 모집단 등가성이 성립한다: 사전 학습된 LSA는 모든 이력에 대해 PO 정책을 모방할 수 있다(Theorem 4.2).
• 유한 샘플 보장은 실증 학습 데이터가 PO 정책을 높은 확률로 모방하기에 충분함을 보장한다(Theorem 4.3).
• ME-ICPO는 추론 시점의 효율성 이점과 함께 수학적 추론 과제에서 경쟁력 있는 성능을 달성한다.
• ME-ICPO는 훈련이 진행될수록 감소하는 효과를 통해 원샷 보상 교란에 대한 강건성을 보인다(Theorem 4.8).