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
Venue: ICLR 2026 Poster (2026)
Keywords: in-context learning
self-reflection
policy optimization
FTRL
bandits
large language models
reasoning
Primary area: foundation or frontier models, including LLMs
TL;DR: We provide a provable and practical in-context policy optimization for test-time scaling

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.

Published as a conference paper at ICLR 2026
PROVABLE AND PRACTICAL IN-CONTEXT POLICY OP- TIMIZATION FOR SELF-IMPROVEMENT
Tianrun Yu Yuxiao Yang", Zhaoyang Wang Kaixiang Zhao', Porter Jenkins', Xuchao Zhang', Chetan Bansal Huaxiu Yao", Weitong Zhang² 'Brigham Young University, University of North Carolina at Chapel Hill, "Microsott "Equal contribution. tianruny&byu.cem {yxyang, weitongz] (esttongg@@un...... @unc. edu
ABSTRACT
We study test-time scaling, where model improves its answer through multi- 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 novel Fisher-weighted logit- 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 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-ting algorithms. Overall, ICPO provides a principled understanding of self-reflection in LLMs and yields practical benefits for test-time scaling for mathematical reasoning.
INTRODUCTION
Recent years have witnessed a growing capacity for large language models (LLMs) with rising abilities in mathematical reasoning (Yang et al., 2024a; Wei al., 2022), problem solving (Rein et al., 2024; Zhou et al., 2024) and tool use (Yao et al., 2023b). Among these new abilities, the emergence of test-time scaling has been playing an important role, where the L.MMs progressively improve their response through multi-round self-reflection without parameter updates. This test-time scaling has demonstrated a strong ability to enable LLMS to perform post-training search (Yao et al., 2023a; Besta et al., 2024), self-reflection and self-rewarding (Madaan et al., 2023; Shinn et al., 2023; Lightman et al., 2023) and Chain-of- Thoughts (CoT, Wei et al. 2022). The key part of this process hinges on the model's ability digest the in-context information improve its response. Such in-context information can be the previous response with users' finetuning instructions, or the CoT process with self-assessed rewards. However, despite repeated empirical validation, the mechanism underlying such in-context self-improvement remains under-explored in the literature. Existing works (Park al., 2024) usually assume this ability for conducting the posterior sampling or policy optimization intrinsically within LLMs without answering why this ability emerges during the pretraining process.
ofhiinoo-Thhuught
On the other hand, recent works have attempted understand the in-context learning for supervised learning (c.g., linear regression Zhang et al. 2024b; Garg et al. 2022) and reinforcement learning (e.g.. TD learning Wang et 2024) that shows that some carefully designed transformers can learn these algorithms with sufficient pretraining. Yet, most of these works consider empowering the L.MMs to predict the output based on the input, while it is vacant in literature understanding how transformers learn to optimize its behavior by optimizing its policy towards maximizing the response y. In addition, there is a huge gap between the current theoretical understanding of the in-context learning and the empirical implementation of the in-context test-time scaling. Witnessing these lacks of the- oretical understanding of the in-context policy optimization and the missing of how to leverage these in-context information iteratively in the test-time scaling for reasoning tasks, we would like to ask:

---

Published as a conference paper at ICLR 2026
PUO
Cheetton many direars
Iteration t Begins
Fiuk
F22
Rewarding
x
X
Ranking Selecting
P
Sampling
X
Agent
X1
Ex
FER
x
Ti
Append
The In-Context Policy Optimization (ICPO) framework. At each round t, the agent leverages its history of past attempts with bandit feedback ((x,,r))....,, (x1, to improve its response X in order to maximize the received reward Tt-
Figure
Can understand the self-reflection process of LLM from the in-context learning that inspires a test-time scaling for reasoning?
In this paper, we answer this question affirmatively by providing the In-Context Policy Optimization (ICPO) framework which considers how LLMs leverage the in-context actions and response to improve their response x instead of predicting some certain outcomes. As illustrated in Figure 1, the ICPO process considers the transformer (LLMs) generating its response x; and receives the reward given by user or self-assessment Yt and then improve its response by generating X Theoretically, show that with sufficient pretraining, one-layer transformer is sufficient to execute a policy- optimization framework that gradually improves its response x using the observed rewards r. When applying the ICPO framework into the practical mathematical reasoning, ICPO framework is robust enough to take the self-accessed reward into its policy optimization process and to gradually improve its response. Together with solid theoretical results on the ICPO process and a carefully designed practical algorithm ME-ICPO, we provide a provable and practical in- in-context learning framework for test-time scaling for mathematical reasoning. Together, our contributions are:
empirically show that the
We formulate the multi-roumt round self-reflection mechanism as in-context policy optimization (ICPO) framework where the agent generates and improves its response with the received feedback. The ICPO framework extends the current in-context learning framework from supervised learning to policy optimization with bandit feedback. ICPO builds theoretical foundation understand the self-reflection and self-improvement for LLM reasoning.
We prove that, under an explicit design of the Linear Self-Attention (LSA) transformer, when the LSA is sufficiently pre-trained trajectories generated by special policy-optimization framework, it provably mímics the underlying policy optimization even under previously unseen reward functions. To the best of our knowledge, this is the first directly derived mechanistic account in-context policy optimization that provides detailed structural characterization.
Empirically, we propose ME-ICPO, practical algorithm grounded in our theory that yields sub- stantial improvements over base models on mathematical reasoning tasks. ME-ICPO demonstrates that the ICPO framework can leverage self-an assessed feedback, using entropy-regularizee response selection to ensure robust policy updates.
Together, work shows the first in-context optimization mechanism to help understand how LLMs can improve their response with self-reflection, with strong empirical performance in various mathematical reasoning tasks.
Notation. In this paper, we use plain letters such as x to denote scalars, lowercase bold letters such as denote vectors, and uppercase bold letters such to denote matrices. Functions denoted by bold symbols such as f. Sets and classes are denoted by the calligraphic font such as F. For vector x, ||x|| denotes its fo-nr.. For matrix A. ||A llop denotes its operator (spectral) norm, i.c., ||A||e llop = sup]|>]]== |Ax|2 max (A). F(p) == Diag(p) pp For a vector E g Diag(x) I K denotes the diagonal matrix with [Diag(x)] = I; and off off-diagonals zero. For a
positive integer N, use [N] to denote (1,2, ...,....A N..
2

---

Published as a conference paper at ICLR 2026
2 RELATED WORK
We introduce the works on test-time scaling, self-reflection and in-context learning in this section. Test-Time Scaling, Test-time scaling (or inference-time scaling) refers to the phenomenon where allocating more resources during test-time can improve the LLM's ability in reasoning and have been widely adopted in practice (Jaech et al., 2024; Guo et al., 2025). The earliest test-time scaling can be dated back to the few-shot Chain-of- Chain-of-Thought (Wei ct al., 2022) where the provided few-shot demonstrations can improve the LLM reasoning ability. Other test-time scaling works focus on the post-training search algorithms, including the Monte-Carle Tree Search (Zhang et al., 2024a), Best- of-N (Wang al., 2022; Huang al., 2025), Tree of Thoughts (Yao et al., 2023a). Following up with these works, TTRL (Zuo et al., 2025) provides a gradient- -baed update based on the self-assessment during the test-time and improves the LLM's response by updating its parameters.
Self-reflection and self-assessment. At the core of the test-time scaling lies the self-reflection and self-assessment where the LLM evaluates its own response via the self-rewarding (Madaan et al., 2023; Shinn et al., 2023). In particular, LLM-as-a-Judge and Majority-Judgnan inexpensive but noisy preference signals, and self consisteccy can be converted into preferences or rankings (Prasad et al., 2024). However, self-evaluation suffers from prompt/position sensitivity and stylistic bias, calling for calibration (symmetric prompting, position shuffling, executabilityfoormat checks) (Wang et al., 2025). In parallel, process supervision (PRMs (Wang et al., 2023a; Chen et al., 2024b), step-wise verifiers (Lightman et al., 2023)) shifts supervision from outcomes to intermediate steps, reliably filtering errors across rounds (Lightman et al., 2023). Recent analysis of the generation- verification (GV) gap shows iterative improvement succeeds when verification is substantially easier than generation, motivating robust filtering and feedback (Song et al., 2024). Beyond these empirical works, recent theoretical works focus on the posterior sampling of LLM (Bai et al., 2023; Von Oswald al., 2023) by directly assuming the LLM's ability for estimating the posterior distribution.
(MJ) provide
In-Context Learning and In-Context Reinforcement Learning. Besides the empirical advances, a line of theory clarifies regression as core sandbox for ICL. For ridge linear regression, trained linear self-attention can implement preconditioned gradient descent in context, with model depth matching the number of implicit update steps and geometric convergence under standard assump- tions (Von Oswald et al., 2023). From a training- -oeniinggddnmmisss perspective, prior work shows that, after training, a single-head linear attention layer effectively performs one step of gradient descent on the contextual least-squares objective (Zhang et al., 2024b). In sparse settings, multi-head constructions can recover sparse signals and carry out sparse linear regression in context (Chen et al., 2024a); recent work further identifies layered mechanism in which first-layer heads preprocess the context and later layers carry out simple iterative optimization, together yielding excess risk guarantees that improve over naive gradient descent and ridge baselines (Chen et al., 2024a). Besides these works in understanding the regressions, more recent work has pushed forward the understanding of in-context learning meta-reiinforcement learning process. In particular, (Lin et al., 2023) proved that multi-layer transformer structure can imitate bandit/RL-style updates by pretraining on trajectories, and (Wang et al., 2024) shows that the linear regression for the in-context learning can be extended to the TD learning used in RL. Despite these, rare recent literature has covered the policy optimization which directly optimized the output x given the historical information.
3 PRELIMINARIES
We consider a multi-arm bandit abstraction which is a standard theoretical framework for sequential decision making. We consider K-armed bandit and at each round t [T], the agent selects an action A, E [K] and plays the corresponding action written as the one-hot vector x, = CA, E R The agent then receives scalar reward generated from a linear model with an unknown task vector w E R by r. = (w,x.) +,,, where e, a zero-mean uu-subb-aussia random variable. The agent overall goal is to optimize its policy x, by maximizing the expected return (w.x.).
Policy Optimization Framework. We consider the pretrained dataset is generated from the policy optimization process in meta reinforcement learning. In particular, we start with the mirror descent that is similar to the Follow-the- -eegular--eell Leader (FTRL, Shalev-Shwartz 2007; McMahan 201la) framework given by
P++1 arg PEAR Eí / Is
Poias x,.p) R(p),
3

---

Published as a conference paper at ICLR 2026
in the ICPO framework, we consider a practical solution in optimizing the log- log-likelihood of the policy defined by x log p. In the following of this paper, we consider the policy optimization written by
1 1 arg m Σ. (rs St) > E (x,;
S Hs.
(3.1)
2
where we implement the entropy regularizer R(p) Ts Ts/Pa,A, used FTRL with a Lagrange multiplier X Ei= (X,, s:) to penalize the frequently visited arms. The closed form solution for equation 3.1 yields a linear structure written by
Hs and replace the unbiased estimator
St+1 - m (Uge + Vn,), where U = - = - g: -
ToX,, -...==-.-- Xs (3.2)
where we set 7/2 = c/t and the optimized policy is then given by softmax policy mixed with y-greedy random exploration
PO P++1 softmax(s:+1), Pt+1 (1 11--7/ (1-7)p + + 7 7 10,
(3.3)
Supervised Pretraining Data Generation. We generate a pretraining dataset by running the expert policy optimization algorithm. We sample B independent trajectories. For each trajectory E [B], task vector w, (N, r x)) is sampled from the prior. The teacher is then executed for N steps to generate a complete history of interactions H.,N = (,,1), corresponding sequence of expert logit vectors {s P N
(x,,N, ,,,,) and the [t=l is updated in equation 3.2. From these
trajectories, we construct a supervised training dataset, D consisting of pairs of history prefixes and the [E[N-1] teacher's next-step logits. The final dataset is the set of all such pairs = {(Hr.,, S s )))
where Hr, is the history prefix of trajectory T up to and including round t. The following assumption is made on the data coverage on the supervised pretraining data D, which is a standard diversity assumption in linear bandits (Papini et al., 2021; Hao et al., 2020; Wu et al., 2020).
Assumption 3.1 (Data Coverage and Signal Dominance). We assume that in the pretrained dataset, the coverage of the task 7 and the FTRL exploration parameter y wide enough cover the reward noise. In particular, define the learning rate 1/:- = c/t, we assume the coverage rate is strictly positive
CA = Tw - (1 - llop / O.
(LSA) is simplified variant of the standard self-attention mechanism, which has been established as a useful model for the theoretical analysis of transformers and in-context learning (Von Oswald et al., 2023; Zhang et al., 2024b). An LSA layer takes sequence of input embeddings, represented matrix E E RdxN where d is the embedding dimension and N is the sequence length. It produces an output matrix of the same dimension through the following computational form:
Linear Self-Attention (LSA). The Linear
fia (E; 0) = + PVV. E (E W w<0E//).
(3.4)
where 0 = (WKQ, WPV are learnable parameters (matrices) and p is a normalization factor for the attention matrix. This operation allows for interactions between all elements in the input sequence, mediated by the Gram matrix term E W KQ E, to update the initial embeddings.
THEORETICAL FRAMEWORK FOR ICPO
In this section, we provide a theoretical justification for in-context policy optimization based on an inspirational analysis in Linear Self-Attention (LSA) network. Through this minimal LSA model, we theoretically demonstrate that a pretrained LSA can imitate an expert policy optimization algorithm using in-context data. We then present our main theoretical guarantees, which establish that this learning is not only possible in principle but also efficient with a finite amount of data, and robust to perturbations test time.
The ICPO Forward Pass.
We start with introducing the forward pass of ICPO framework. The LSA model parameterized by 0 starts with an empty embedding ((0) = (qxz gr) where qa = 1x, 4r = 0 the placeholder for next-token generation. For each time step E [T]. the LSA model updates its policy according to the logits updates from the next-token generation of LSA described by
a = [Jiss (E(t-1); 00) 1:Ke
P: = (1 - softmaxxx)(() + R
4

---

Published as a conference paper at ICLR 2026
where (liss(E(t11)) ;0)]] 1: stands for the corresponding K dimensions of the newly generated token indexed with 7 is the same exploration factor in the implementation of FTRL. With this policy Pa E R2, the LSA model selects its action A, and receives the reward by
A Pt. X CA. Tt (w, xi) Ce.
Finally the sequence of token is updated by inserting the observed reward ra and embedding x, by
X1
Xr qx
1i
Tt
9r
and update the normalizing factor in LSA as p ensure the attention matrix is upper bounded by 1. The forward pass of ICPO then move to the next round f +- ++11.
Training Objective. The supervised pretraining is conducted the dataset 'D by matching the logits from the model output & = fLSA (E(t-1),0) with the logits from policy optimization s!"" by minimizing the projected weighted loss by
( = E
E
Proj - 57,4 1
(4.1)
where projection Proj defined as II: :=
π 11 which removes the constant bias Ts from the logits, since such shifts do not affect the policy p exp(s). The expected matrix I is inspired by the design of Natural Policy Optimization (Kakade, 2001) defined by
r P=N-EE NIIEEED "GE='Diag(P...) Pr.PP..
The Fisher-weighte loss provides new loss for the supervised pretraining. We show by the following theorem that the common KI. loss between the pretrained data p/41 and the LSA's output Pi+1 is sandwiched by the loss ((O) up to constants.
Theorem 4.1 (mixed-policy KI is controlled by the Fisher-projected quadratic loss). Assume both teacher and student use m-mixture exploration with y (0, 1) as described in equation 3.3, and let N denote the trajectory length of the sample inside the expectation. Then,
(1 7) (
< E
N -E
DKL (pf PI+1) & L()).
K 4y
4
Theorem 4.1 suggests that the widely used KL loss is good surrogate the Fisher weightee loss and explains that even in a linear self-attention layer, using the KI. loss enables the transformers to learn self-reflection and improve it's response.
4.1 THEORETICAL GUARANTEES FOR ICPO
We now present our theoretical results considering the empirical Fisher weighted loss defined by
Σ(0) 2B 1) ETeD (Proj(5,.41-s,00
sPO s..21) f. := BN (Diag(p...) -opa(P..)-Pr.PP... Pr.eP..t.
LE|T
We denote the population optimizer as 0" = arg ming (()) and it's empirical solution 0 arg argmin min (()). Then the first theorem suggests that the population optimizer 0" is exactly imitating the policy optimization algorithm we described in equation 3.3.
Theorem 4.2 (Population Equivalence). Under Assumption 3.1, consider a one-layer LSA with parameter 0" minimizing the population loss L(0) will imitate the policy optimization behavior for all possible history trajectory Hy:, i.e. P (H,; 0*) = (H:).
Theorem 4.2 suggests that the population loss L(0) is precise and informational enough to guarantee that the parameter 0* will drive the LSA exact imitate the policy optimization framework leveraging any in-context data. Then simple concentration analysis suggests that the empirical estimation will also yield similar result with high probability:
Theorem 4.3 (Finite sample result). Under Assumption 3.1, let the one-layer LSA be trained on B i..d. trajectories (H,,,N ) generated by the policy optimization process in equation 3.3. Define M == B - 1). Using all prefixes t E {1....,N N - 1} from each trajectory (allowing within- trajectory dependence), form the empirical Fisher fisher-weighted loss ((0) and let the global optimizer be i arg ming (()). For any E (0, 1), with probability least 1- 5, if M Z 12
5

---

Published as a conference paper at ICLR 2026
then for any fixed test history He we have Pu-1(H(;;) = -Pi41 (H4). In addition, the expected behavioral mismatch is bounded by
Berain
PPt1 - p (Hl )12
I 2(1 2)
where Etrain is taken over the randomness of the B trajectories.
We would like to summarize the theoretical results by the following remark.
Remark 4.4. Theorem 4.3 suggests an ( D(N2K/c3) K/(²) sample complexity for the supervised pretrained data to guarantee the LSA is well imitating the pretrained policy optimization trajectory. Such an constant sample complexity is because of the Assumption 3.1 which is similar with the diversity assumption used in Papini et al. (2021); Hao et al. (2020); Wu et al. (2020) which suggests that 7-greedy exploration suffices for a constant regret in linear bandits.
Our analysis and framework share commons and significant difficulties compared with Lin et al. (2023); Park et al. (2024)
Remark 4.5. Park et al. (2024) analyze the regret of the LLMs assuming the LLM can conduct the posterior sampling without the structural analysis. In contrast, our analysis is built on slightly modified policy optimization framework inspired by inserting an Lagrangian to solve the FTRL mirror descent. With a newly designed supervised learning loss, we show that an linear self attention trans- former can structurally imitate the policy optimization framework. Compared with the unsupervised loss proposed in Park et al. (2024), we show by Theorem 4. that the newly proposed Fisher-weighted loss is a nice surrogate of the practical KL loss, which better justices that the transformers can learn policy optimization with sufficient pretraining.
Remark 4.6. We note that Lin et al. (2023) studies algorithm distillation and proves that ((VT) layers of ReLU-activated transformers can mimic (soft) Linear UCB and Linear Thompson Sampling. In contrast, we target understanding policy optimization instead of its value- counterpart, which is more suitable for analyzing the behavior of L.LMs in improving their responses. In addition, our analysis is built for a one-layer linear self-attention framework so that the network does not need to change the context grows larger, and is more aligned with practical long-context scenarios, whereas Lin et al. (2023) requires the number of network layers grow on the order of vT, where T's is the length of the in-context sequence.
In addition, as frequently diseussed in previous literature (McMahan, 201 lb; Shani et al., 2020), policy optimization methods such as FTRL are known to be robust to adversarial or misspecified rewards, which highlights their applicability to ICPO with self-assessed, noisy, or perturbed rewards. A crucial property for practical self-refinement is that the learned policy is also stable. We analyze the robustness of the learned ICPO loop at test time by examining its response to a single-shot reward shock, which we present in the next theorem. We will start with the definition of the CRN-coupled setup (Glasserman & Yao, 1992).
Definition 4.7 (s-CRN coupled trajectories). We say two trajectories are s-CRN coupled when they share a common random number (CRN) and they are identical up to round 1 in trajectory H. and Hy. We denote F, as the filtration that includes all events happen before observing the reward at round s. At s-th round, the reward is shifted by 6,. We define the drift cause by this reward shift by
Api+1 (0) = p.+1 (H4: 0*) - (P4; 0*) E R
Then we are ready to present the following theorem indicating that the impact from any one-time reward perturbations will be decreasing with well-designed learning rate 7/t =c/l.
Theorem 4.8 (Stability to One-step Reward Perturbations). Under Assumption 3.1, assume the test- time ICPO loop runs with fixed, population-optimal parameters 0* in the s-CRN-coupled trajectories with task w, define
c( - 2
b:= c(1--1)
K ( + Diag(w)))lpp + / 2
U
V
2
that does not grow with s or f, let C₁ = fbb as another absolute constant, for any 1 S 8 S t,
S a(l+0s) (9)²~11 |6-1-
In particular, let the learning rate n, = c/t be sufficiently small such that < 1, the one-time reward shift in the s-CRN-coupled trajectories is decreasing to 0, i.e lim E. [[] A 1//2 F.-1] = 0.

---

Published as a conference paper at ICLR 2026
Algorithm 1 ME-ICPO: Minimum-Entropy In-Context Policy Optimization
Input: Question Q: number of rounds N; candidates per round k;
Input: System prompt SysPrompt: Summarizer (prompt) Summ
1: Ho {SysPrompt, Q}
2: for = 1 to n do
3: Sample response A P- H4-1) with their answer aj for all E [k]
4: Assess 0e MajorityVote (a) (16) } 11=1 let r I(a)) @e) for all j [k]
(t)
(t) 5: Summarize 5 Summ (A)t)) and H, H (al (9)
: ()) for all j E [M]
6: Select the minimum entropy response j arg Ji H(H)) (1)
7: Update in-context list Hr < He-1 U (x r"))
19)
8: end for
Output: Response sampled from P.+1(-1 H.)
5 MINIMUM-ENTROPY IN-CONTEXT POLICY OPTIMIZATION
Based on the theoretical understanding of the ICPO framework, we now present a practical algorithm leveraging the in-context policy optimization and the self-accessed reward to improve its reasoning ability via test-time scaling. To begin with, we adopt the ICPO notation presented in our theoretical analysis. The agent starts with the historical prompt Ho {Q} that only contains the question Q For each time E [7], the model outputs it response x2 ~ ((-1 and then observe the reward rt via self-accessed rewards. Then the agent updates its history H. ( (Q, (x,,
(x,, and move on to the next time step & <- t + 1. However, directly applying the ICPO framework presents two significant challenges. The first challenge is related the length of the contexts preventing the agents from cumulate and conduct effective reasoning process based on a prolonged context through the in-context policy optimization process. The second is the trustworthiness of the eeff-assessed rewards, since the self-evaluation can be noisy and inaccurate. To tackle these two challenges, we present our test-time in-context scaling algorithm: Minimum--Entropy In-Context Policy Optimization (ME-ICPO). The ME-ICPO algorithm works in the following three major procedures.
Response Generation and Self-Assessment. Line 3, the agent samples & responses A using (4) the historical in-context data H4-1 with their final answer d appearing in the final boxed() block (Hendrycks et al., 2021). Then the majority voting as conducted in Wang et al. (2022) is conducted over the answers in Line 4 for assessing the accuracy of the responses.
Mhain-of-Thoughee Summarization. In order to compress the information from the output response and condense the in-context texts, we summarize the responses according their Chain-of- Thouuhhh (CoTs) (+) and ignore the detailed problem-solving process in Line 5. The motivation that the detailed numerical processing is expected to be easier than the global CoTs.
Minimum Entropy Response Selection. Similar to the tree-search-style algorithms Yao et al. (2023a), we select response x0)") and put it and its response into the in-context history H++1 in Line 7. Specifically; unlike the tree-search Yao et al. (2023a), or Best-of-N Huang al. (2025) algorithms that are designed for multi-step reasoning where 2; is the one-step response, we consider the I, to be a CoT description of solving the whole problem. Therefore, instead of selecting the xe with the highest reward, we instead follow "pessimism" in offline reinforcement learning by selecting
the se that leads to the minimum entropy in the future response, i.e., <- arg min JE(K] (H)
as conducted in Line 6. Intuitively, this minimum entropy selection will avoid the agent selecting the corrupted response that may drive the agent to produce random answer. In addition, this procedure will also encourage the agent to select the diversified responses xy- that would help further reduce the entropy.
important to distinguish our test-time approach from methods that train a student model on trajectories from fixed teacher algorithm. Since ME-ICPO performs no parameter updates at test time, we do not claim that the deployed LLM is uniquely equivalent to any specific policy- optimization algorithm. Rather, theoretical analysis shows that the LSA architecture possesses a strong inductive bias for performing such updates. ME-ICPO is designed to provide a usable interface-via reward-aware prompting and principled feedback selection-thal effectively leverages this
7

---

Published as a conference paper at ICLR 2026
inherent computational capability of the model at test time without requiring gradients (Madaan et al., 2023; Shinn et al., 2023). For complexity derivations and prefactor discussions, see Appendix C.
EXPERIMENTS
We present the experiment results to validate our theoretical claims in Section 4 and the performance of ME-ICPO in this section.
6.1 VALIDATION EXPERIMENTS
To verify the Theorem 4.2 and Theorem 4.8 results, we run two controlled checks on single LSA. Figure shows, respectively, the teacher-student policy matching error (Top) and the stability of the mixed policy under one-time reward shock together with the instantiated analytical bound (Bottom).
Teacher-student policy matching. We fix the meta configu- ration K 10, NN- = 30, B = 100, ? 0.2, stepsize T = c/t with c = 1.0, task prior w ~ NN T with Tw = 1.0, and re- ward noise 6 N (0, of) with og 0.5. At test time we draw Bies 64 fresh tasks from the same generative process. At each round we hold the realized history H/ fixed, compute the teacher mixed policy p and the model's 's mixed policy P as defined in Section 4. and then continue the closed loop using the model's policy. Aggregating over tasks, we report the mean and one-standard-deviatio band of E p.p - 11 As shown in Figure 2, the gap rapidly drops to numerical precision and decreases with t, consistent with the population equivalence and finite-sane samile guarantees.
x10-3
Policy time
x10-1
1.0
Eryprniall FANe
Theary
<
Stability under a single reward shock. We use an LSA trained via supervised imitation on PO rollouts with K = 5, N 5, 0.8, 0.5, Tur 0.5, and 0g = 0.1. At test time we keep the data model and y unchanged, extend the horizon to N = 10. and evaluate Best = 256 tasks. For each task we run two s-CRN-coupled trajectories: a baseline and a perturbed path that injects single reward shock ôr, 1.0 at = 2. Following Def. 4.7, we record Ap; and plot the mean and compute b [0.1236, 0.2127] and thus the
5
Figure 2: Validation of ICPO the-
ory. (Top): Policy Matching. (Bot- tom): Reward-Shool Shock Stability.
seeaanddrdddvviaiion band of E|| A 2|2- We non-amplification condition. Figure shows a brief post-shock bump followed by steady decline without amplification over time, with the analytical curve providing a conservative upper bound.
6.2 LLM EXPERIMENTAL SETUP
We evaluate ME-ICPO on standard mathematical QA benchmarks (AIME 2024, AMC, and MATH- 500, split into five difficulty levels following TTRL) (aop, 2024; Mathematical Association of America, 2025; Li et al., 2024; Hendrycks et al., 2021) using representative backbones (Qwen2.5- Math-1.5B and Qwwe2.5-Math-77B (Yang et al., 2025). For reference, also report the perfor- mance of 1emma-3.--B--Intteeee
(Grattafiori et al., 2024) and DeepSSeek-R--Distil--Aammm al., 2025) AIME 2024. Full hyperparameters, baseline specifications, hardware/software environment, and prompt details are provided in Appendices B.2, B.3, and D. Following Guo al. (2025), we generate 1-16 responses per question (T'=0.6, top-p=0.95). We report Mean k [DT EeED Σ cá()), where cı(g) = 1 (a:(4) = a" (4)), and Accurucy, the probability of answering correctly with one attempt.
(Guo
=
6.3 MAIN RESULTS
Our main results are presented in Table 1, reporting the mean and standard deviation (over 5 seeds) of Mean @ 16 (%) and Accuracy (%) across all tasks and models. The results show that ME-ICPO consistently improves performance through its gradient-free, in-context optimization process. These improvements hold for both the larger model (Qwen2.5-Math-77B) and the smaller model (Qwen2.5- Math-1.5B), demonstrating that ME-ICPO is effective across different model scales. Qualitatively curated prompt examples are provided in Appendix D. We also evaluate an oracle-reward variant, with results summarized in Table 2.
We further report the performance of additional models -lama-3.1-8B-Insttee (Grattafiori et al., 2024) and DeepSeek- -1aeek--R1Da
(Guo et al., 2025) on AIME 2024 in Figure 3.
8

---

Published as a conference paper at ICLR 2026
Maam
Mag@16
40
Meang11 aut ME-ICPO Me
DeepSeck- -3 Dettil Lama-U8
Qwen-25-Mathhi
Uama-33 3.1 instr
Figure 3: Performance comparison of backbone models before and after ME-ICPO.
Table 1: Performance comparison on different models with Mean 16 and standard Accuracy (%).
Model Benchmark
AIME 2024
AMC
MATH-LI
MATH--2
MATH--
MATH-LA
MATH-1S
1e (%)
Owen2.5-Math-7B
11.04 1665
41.42 0.99
02.62 1.38
0.97
49.88 0.88
43.00
0.580
30.58
W/ ME-ICPO
30.12 1.81
17.06 = 0.81
62.35 : 1.27
61.11
1123
0.78
19.31
0.97
0.86
+5.64
+9.73
+13.55
+8.99
wwen2.5-Math-L.58
6.46--0.90
30.42 + 0.58
49.27 0.81
48.54 0.67
45.42 = 0.86
40.28 + 0.57
25.23 + 0.45
ME-ICPO A
9.79 1.11
33.56
0.63
61.19 0.77
61.93
52.08
46.41 0.64
29.85 0.70
+3.33
+3.16
+11.22
+6.39
+6.66
+6.16
+4.62
Accuracy (%)
Qwen2.5-Mii -wen2.5-Math-7B
3.227
41.33
1.97
46,98 + 2.71
42.67 + 1.86
43.71
1.75
37.79
1.65
26. 13 1.54
ME-ICPO A
30.05 3.02
47.20
67.32
2.37
54.74
2.04
51.90
1.67
40.84
31.71
+18.92
+5.87
+10.34
+12.07
+-8.19
+33.05
+-5.58
Qwen2.5 5 -.5th-L-5B
6.51 1.94
30.25 1: 1.07
44.68
1.83
39.95
1.32
39.77
0.99
34.97 1: 0.68
20.89 :: 1.35
ME ICPO
9.82 2.15
33.73
0.96
57.06
1.70
47.60
1.000
47.81
1.17
41.55 4 0.72
24.83 4: 1.29
+3.31
43.48
+12.38
47.65
+8.04
+66.88
+33994
We also consider the Maj@k metric (Wang et al., 2023b), which aggregates k responses by majority vote ā(g) = mode(ai())...... (q)..... (g)) (ties broken uniformly) and computes Maj@k EgED 1 ((ā(g) a (q)). As shown in Figure 3, our experiments indicate that ME-ICPO's average performance (Mean @ @16) can surpass the anticipated upper bound given by the base model's majority voting, similar to observations in TTRL. Furthermore, applying majority voting top of ME-ICPO's outputs leads to additional gains.
6.4 ANALYSIS AND ABLATION STUDIES
We further analyze the AIME 2024 benehmark using Qwtt2.5-Math-7B assess the contributions of ME-ICPO's core components and its sensitivity to hyperparameters, detailed in Appendix B.1, along with the computational cost analysis in Appendix C.
Ablation Study. To isolate the contributions of our core eenpoonnnss-entooy-bbbeeee se- lection and explicit reward sig- nals-we conduct an ablation study, with results presented in Table 2. The results clearly demonstrate that the minimum- entropy selection criterion is the most critical component of our algorithm. Removing this greedy selection mechanism (w/o Entropy) causes dramatic performance collapse. The explicit reward signal, made legible by our system prompt, also plays crucial role. While keeping entropy selection active, removing the reward tags (w/o Reward) still results in significant drop.
Table Ablation study of ME-ICPO. (Oracle results use groundtruth labels for reward and are shown for reference only.)
Method
Accuracy (%)
Mean @ 16 (%)
w/o Reward
19.30
19.17
w/o Entropy
5.77
5.83
w/o Entropy Reward
6.21
6.46
ME-ICPO (full)
30.05
30.42
ME-ICPO Oracle
38.19
38.12
7 CONCLUSION
In this work, we studied the test-time scaling and self-reflection phenomenology in LLM reasoning by introducing theoretically grounded In- In-Context Policy Optimization (ICPO) framework. We have shown that, under the ICPO framework, a single-layer linear self-attention transformer can provably imitate policy-optimization algorithm, which provides a theoretical proof-of-concept for how self- reflection can be implemented within LLMss. Based on the ICPO framework, we propose a practical and effective algorithm, Minimum Entropy In-Context Policy Optimization (ME-ICPO), which provides a test time scaling pipeline with self- assessed rewards and in in-context response selection. Extensive empirical studies demonstrate that our improved performance across diverse benchmarks.

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We have made significant efforts to ensure the reproducibility of our results. We detail all datasets, model configurations, hyperparameters, and training procedures Appendix B. To enable faith- ful reproduction, we release our full codebase, experiment scripts, and configuration files at https::/gá thub.
The repository includes exact seeds, environment specifications. Our source code is provided to facilitate faithful reproduction of our experiments in the supplementary materials.
ETHICS STATEMENT
We have carefully reviewed the Code of Ethics and find that work does not raise any significant ethical concerns. Our research does not involve human subjects, sensitive data, or potentially harmful applications. We believe our methodology and contributions align with principles of faimess, transparency, and research integrity.
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their helpful comments. Part of this research is conducted during ZW's internship at Microsoft Research. This research was also supported by WZ's s startup funding provided by the School of Data Science and Society at UNC Chapel Hill.
REFERENCES
Aime problems and solutions, 2024. Accessed: 2025-09-13.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36:57125-5721,, 2023.
Mislav Balunović, Jasper Dekoninck, Ivo Petrov, Nikola Jovanović, and Martin Vechev, Matharena: Evaluating Ilms on uncontaminated math competitions. arXiv preprini arXiv:2505.2325 2025.
Maciej Besta, Nils Blach, Ales Kubicek, Robert Gerstenberger, Michal Podstawski, Lukas Gianinazzi, Joanna Gajda, Tomasz Lehmann, Hubert Niewiadomski, Piotr Nyczyk, et al. Graph of thoughts: Solving elaborate problems with large language models. In Proceedings of the AAAI conference on artificial intelligence, volume 38, pp. 17682-17690, 2024.
Xingwu Chen, Lei Zhao, and Difan Zou. How transformers utilize multi-head attention in-context learning? case study on sparse linear regression. Advances in Neural Information Processing Systems, 17:119573-119611 2024a.
Zhaorun Chen, Zhuokai Zhao, Zhihong Zhu, Ruiqi Zhang, Xiang Li, Bhiksha Raj. and Huaxiu Yao. Autoprm: Automating procedural supervision for multi-step reasoning via controllable question decomposition. arXiv preprint arXivv2402. 1445, 2024b.
Gheorghe Comanici, Eric Bieber, Mike Schaekermann, Ice Pasupat, Noveen Sachdeva, Inderjit Dhillon, Marcel Blistein, Ori Ram, Dan Zhang, Evan Rosen, al. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv preprint arXiv:2507.06260 2025.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study simple function classes. Advances neural information processing systems, 11:30583-30996 2022.
Paul Glasserman and David Yao. Some guidelines and guarantees for common random numbers. Management Science, 38(6):884-908 1992.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXi::2407.2178. 2024.
Daya Guo, Dejian Yang, Haowei Zhang. Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseck-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint 2tti::2501.124444 2025.
10

---

Published as a conference paper at ICLR 2026
Botao Hao, Tor Lattimore, and Csaba Szepesvari. Adaptive exploration in linear contextual bandit. In International Conference on Artificial Intelligence and Statistics, pp.
PMLR, 2020.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song. and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint urXiv:2203.038// 2021.
Audrey Huang, Adam Block, Qinghua Liu, Nan Jiang, Akshay Krishnamurthy, and Dylan J Foster. Is best-of-n the best of them? coverage, scaling, and optimality inference-time alignment. arXiv preprint arXiv:2503.22 0t..:2203....8.. 2878, 2025.
Aaron Jacch, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, et al. Openai ol system card. arXiv preprint arXiv:2412.1672) 2024.
Sham M Kakade. A natural policy gradient. Advances in neural information processing systems, 14, 2001.
Jia Li, Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Huang, Kashif Rasul, Longhui Yu, Albert Q Jiang, Ziju Shen, al. Numinamath: The largest public dataset in ai4maths with 860k pairs of competition math problems and solutions. Hugging Face repository, 13(9):9, 2024.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, 2023.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforce- 2023.
ment learning via supervised pretraining. arXiv preprini
Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, et al. Self-refine: Iterative refmement with self self-feedback. Advances in Neural Information Processing Systems, 36:46534 46594, 2023.
Mathematical Association of America. American mathematics competitions, 2025. Accessed: 2025-09-13,
Brendan McMahan. Follow-the whe.-tth-r-gllrized--eeee leader and mirror descent: Equivalence theorems and 11 regularization. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pp. 525-533. JMLR Workshop and Conference Proceedings, 201 la.
H. Brendan McMahan. Follow-the-reguue tetio--reuuuaazdd-leeeeeee and mirror descent: Equivalence theorems and 11 regularization. In AISTATS, 2011b.
Matteo Papini, Andrea Tirinzoni, Marcello Restelli, Alessandro Lazaric, and Matteo Pirotta. Lever- aging good representations linear contextual bandits. International Conference on Machine Learning, pp. 8371-8380. PMLR, 2021.
Chanwoo Park, Xiangyu Liu, Asuman Ozdaglar, and Kaiqing Zhang. Do Ilm agents have regret? case study in online learning and games. arXiv preprint arXiv:2403.1684 2024.
Archiki Prasad, Weizhe Yuan, Richard Yuanzhe Pang. Jing Xu, Maryam Fazel-Zarandi, Mohit Bansal, Sainbayar Sukhbaatar, Jason Weston, and Jane Yu. Self-consistency preference optimization. arXiv preprint arXiv:2411 1..009, 2024.
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel Bowman. Gpqa: graduate-level google-proe proof q&a benchmark. In First Conference on Language Modeling, 2024.
Shai Shalev-Shwart Online learning: Theory, algorithms, and applications. Hebrew University, 2007.
Lior Shani, Yonathan Efroni, Aviv Rosenberg, and Shie Mannor. Optimistic policy optimization with bandit feedback. In International Conference Machine Learning, pp. 6614-8613. PMLR, 2020.
11

---

Published as a conference paper at ICLR 2026
Noah Shinn, Federico Cassano, Beck Labash, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: Language agents with verbal reinforcement learning, 2023. URL https:/aaxiv. org/abs/2333.11 1366, 1, 2023.
Yuda Song, Hanlin Zhang, Carson Eisenach, Sham Kakade, Dean Foster, and Udaya Ghai. Mind the gap: Examining the self-improvement capabilities of large language models. arXiv preprini arXiv:2412.026744 2024.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp- 35151-35174, PMLR, 2023.
Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, and Shangtong Zhang. Transformers can learn temporal difference methods for in-context reinforcement learning. arXiv preprint arXiv:2405. 13866,, 2024.
Peiyi Wang. Lei Li, Zhihong Shao, RX Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce ilms step-by-step without human annotations. arXiv preprint arXiv:2312.089935 2023a.
Ruiqi Wang, Jiyu Guo, Cuiyun Gao, Guodong Fan, Chun Yong Chong, and Xin Xia. Can Ilms replace human evaluators? an empirical study of llm-as-a-judge in software engineering. Proceedings of the ACM Software Engineering, 1ISSTAA::155511999 2025.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdh- ery. and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. arXiv preprini arXivv2203. 2022.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdh- ery. and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In ICLR, 2023b.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. ofain-of-thouuhh prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824 24837, 2022.
Weiqiang Wu, Jing Yang, and Cong Shen. Stochastic linear contextual bandits with diverse contexts. In International Conference on Artificial Intelligence and Statistics, pp. 2392-2401. PMLR, 2020.
An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jian- hong Tu, Jingren Zhou, Junyang Lin, et al. Qwen2. 5-math technical report: Toward mathematical expert model via self- celf-improvement. arXiv preprint (-Xii:2449.12222 2024a.
An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jian- hong Tu, Jingren Zhou, Junyang Lin, et al. Qwen2. 5-math technical report: Toward mathematical expert model via self-improvement. arXiv preprint arXiv:2409. 12122, 2024b.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09338. 2025.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Advances in neural information processing systems, 36:1 1809-18910000 11822, 2023a.
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao.
React: Synergizing reasoning and acting in language models. In International Conference Learning Representations (ICLR), 2023b.
Di Zhang, Xiaoshui Huang, Dongzhan Zhou, Yuqiang Li, and Wanli Ouyang. Accessing gpt-4 level mathematical olympiad solutions via monte carlo tree self-refine with llama-3 8b. arXiv preprint arXiv:2406.07394, 2024a.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 1554 2024b.
12

---

Published as a conference paper at ICLR 2026
Pei Zhou, Jay Pujara, Xiang Ren, Xinyun Chen, Heng-Tze Cheng, Quoe V Le, Ed Chi, Denny Zhou, Swaroop Mishra, and Huaixiu Steven Zheng. Self-discover: Large language models self-compose reasoning structures. Advances in Neural Information Processing Systems, 37:126032--26055 2024.
Yuxin Zuo, Kaiyan Zhang, Li Sheng, Shang Qu, Ganqu Cui, Xuckai Zhu, Haozhan Li, Yuchen Zhang, Xinwei Long, Ermo Hua, et al. Ttrl: Test-time reinforcement learning. arXiv preprint arXiv:2504.1684, 2025.
13

---

Published as a conference paper at ICLR 2026
APPENDIX
Proofs and Technical Details for Section
14
A.1 Preliminaries and Key Lemmas
14
A.2 Reward-shock Robustness
A.3 KL Divergence vs. wisher-wiighted Duadratic
22
A.4 Closed-Loop Imitation of Policy Optimization
23
A.5
Convergence of the Fisher-trained Two-channel LSA
24
Experimental Details
25
B.1
Hyperparameter Sensitivity....
....aaai........................
25
B.2
Hyperparameter and Implementation Details
25
B.3 Hardware and Environment Configuration
..t.aa..........................
26
B.4
Supplementary Experimental Results
26
Detailed Complexity Analysis
27
C.I Theoretical Time and VRAM Complexity Derivations
27
C.2 Empirical Cost
D Prompt Templates and Qualitative Examples
31
D.I Dataset-Specific System Prompts
31
................................
D.2 Summarization Prompts
................................
32
D.3 Qualitative Case Studies
A PROOFS AND TECHNICAL DETAILS FOR SECTION 4
A.1 PRELIMINARIES AND KEY LEMMAS
In this subsection, establish the minimal set of tools and lemmas required for the subsequent proofs. We adopt all notation and settings from the main text.
Lemma A.1 (Mixture curvature and Fisher bounds on 1 1-). If the teacher uses mixture exploration
p = (1 -1--77 7 1/K, then the Fisher matrix F(p) == Diag(p) - pp is bounded on 14:
2 I 3 F(p) K
- I.
2
Proof. Using linearity Diag(-) and expanding pp T.
F(p) Diag (1 7)P 7 } ((1 - 7)p 1/ ) -pp k)²
= (1 - F + { + 7(1 () ) )
The last rank-one term and F(P) are PSD, hence
F(P) 2 7 F( (
On 1 ( = Diag(E)
TT) acts as (1) the lower bound follows:
F(p) = ? I. For the upper bound, for any x E 1 with ||x||2= =1, 1,
(V2)² 2
(max; I3 min: x2,2
x F = Vara<p/z;] <
=
4
4
2'
where we used Popoviciu's inequality and that among vectors with x 1 0 and ||x|| 1, the range maximized by placing mass +1/V2 /2 on two coordinates. Taking the supremum over such yields F(P) 1 3
[
14

---

Published as a conference paper at ICLR 2026
Lemma A.2 (Softmax is 1/2-Lipsschiii on 1 ) For any R with w
||softmax(u) - softmax(v) S u v/12-
Proof. By the mean value theorem for vector maps, there exists & on the segment (u, u| such that
softmax(u) softmax(b) = J(E) (u - v),v),
where J(] is the Jacobian of softmax. It is well known that J({) F()) with p softmax(E), and J(E) = Since u - E 1 we can restrict to 1 and apply Lemma A.I:
||softmax (u) softmax(v)) (v)ll2 ( v)ll2 s (1 llop ||u vl12 s 1 I|u vll2-
Lemma A.3 (Query-column closed form for one-layer LSA). The next-step logit vector admits the closed form
1 i G.G.
s qx
(A.1)
where PW 11:K; the row-selector matrix, b = w q is the transformed query vector,
q :=
9r 4r
E g++
with qu
R and qr R. and G; E( ( is the history Gram
matrix.
Proof. By the LSA definition in equation 3.4,
841
[fsa (E(); 0)]
1:K,41
(EO) ((
E(9) W E
t w PV ( (E(()) wXQE(0) 1:K.+11
1:K,
[E]
1:KK r+1 + t
(a)
9. W ( (E wwk
1:K
(b)
9r + W (()) (E KQ ee+
1:K
(c)
4. WVV E (E
1:K
q
R.
Justifications: (a) column slicing equals right-multiplying by the standard basis et+1 RR4+1 and [E("]. II.K, ++1 4r: (b) associativity isolates E(0e (c) substitute ((0ert1 = q (q gr)
and regroup (E (E()) (d) apply the concise definitions R [WPV G, (E
and b wkQq.
Lemma A.4 (Two-channel projected logits: exact equality). Assume the architectural normal form: q. = 0, qu E span{1}, and the final column of the action-logit projection is parallel to 1 (i.c., PV || 1). Let the historical statistics up to round be the count vector 12
n
e
1=1
#=1
and the reward vector
g:
T,X
Ts eA.-
Let WPV WKQ E ((K+))((K+) be porti-patttiooned as
( WPP (W T
w PV
WKQ (w
w KK Ko
PVV
12 W
WKQ_
KQJT
Wazz
15

---

Published as a conference paper at ICLR 2026
with wil,, whl R x and w w
KQ KQ RK For the transformed query
W12 W2
wkQ
Φ1
Φ ERR 02 ER,
@2
the projected LSA logits satisfy, for any nonzero Ø 22.
Proj (8++1)
W., n + W WgB:).
(A.2)
1
with effective operators
Wn == Proj WI Diag(oi)).
W = Proj 02 W
Proof. Using Lemma A.3 we have S441 = dr ,RG, b with R [W [W and
I1:K: = [W w
G: EE (E
C (C 9
r
gr
C, = Diag(n), r: r== TT)
g
Projecting onto 1 and expanding block multiplications,
1 qu
Proj 1
Proj
+
t
Proj(R G)
( Proj (WPY w
C2+gga
r r (b2) (d2))
g: [b))
g
Proj( ((C, + 9. 1 gg)02)
w (g. Φi + (r (TTr))))
(e)
Proj(W C. Φ1 + WPF g: 02)
( (Proj(WPP Diag(())) n. + Proj( WPP) g:)
1 W, n. + Wg e
Explanation of steps: (a) apply Proj to the closed form; (b) Projq, 0 since qu span (c) block multiplication with R. G., b; (d) evaluate the product explicitly; (e) both Proj ( 11 q.q Φ1) and Proj(w (PV) vanish because qr, wPy
(f) use (A (i)g(())u y (A Diag(u)) y with ny
v = Φ1 (g) identify W., W, as stated.
Lemma A.5 weighted quadratic in the two-channel operator). Let the normalized historical statistics be n. := ne/t and g: :- g:/t. Define the concatenated operator W := [W, W, e K and the concatenated normalized statistic 2, == (n. g. E R2K Let the population second-moment matrices be
Σ Enn Σ gn
Eng Egg
Σ = E(2,2,
Σn = E(y+11, with Yt+1 = Proj (sP))
With r ([F() the Fisher Fishhr-weighted loss from Eg. equation 4.1 admits the quadratic form
1 (( (r(WEWT) tr( W +
1 tr tr())
2
where Eng := Ely
Proof. By Eq. equation 4.1, averaging uniformly over training pairs (r,t),
(())
E Proj (S 5 1
By Lemma A.4, the projected student logits admit the two-channel form
Proj ($441) (W,, + Wgg.) = W.n. + Wog Wa.
16

---

Published as a conference paper at ICLR 2026
Let Y(+1 Proj (5.42). we obtain
Using |bb] b. b',,, linearity of expectation, and tr(ABC) tr(CAB).
1
[Wz yo+1 1|:]
(())=
E
2
1
(W2: y(+1) T(W2: - y++1)]
E
2
L W rwzi] E[z, Ty411
+ 1
2
1
tr W
HTWEEZZ,
E [ tr Ely++122 2,
WT) (r(TElye
2
1
(T(WEW") tr( W tr Er()
2
Lemma A.6 (Empirical weihht- quadratic form). Let M == B(N 1) be the number of training pairs and index them by (r, () with 7 E [B] and E [N - 1]. Define the empirical Fisher weight and empirical second moments by
f
Diag(p PO) - PO P (T)
Σ =
2 22, M
Σ
M
T,
S y2
Σ T,
Σ yy = M
Σ Fr. 184111
M
where Z (n. B T.
Dr,e/l, Br.t gr,//l, and yr,41 Proj(s Let W [,, W be in Lemma A.5. Then the empirical Fisher-weighted loss from Eq. equation 4.1
Dr,
admits the quadratic form
2 Lr (TWW -
ur WT) +
and in particular the last term is independent of 9.
Proof. By the definition in the main text,
((0) 2M
PO 57,111 67,011 - Σ
7.8
By Lemma A.4, Proj(Sr.a41) = W Zr,1. hence
Σ(0)
1 Σ (Wz,, Vr,441) T (Wz,,, - Vr.4t)).
2M
Expanding the quadratic and using linearity of trace with b Ab = (tr(AbbT)
2(0) tr
1
[fw
W
1
Σ
Z,, WT) tr (r(ff 1
Σ Nr.ttiz, WT)
M
M
+ tr
1
f
1
Σ
M
7,
1
(fwwwT)
tr
1 Er WT) + 2
tr((f Egw)
Lemma A.7 (Population second moment is positive definite on 5). Let 2, := (n:) E R2K with
n, : X. and g: Σ. TxX,, where T. w Tx, + Assume w NO 1)) and is independent of {x,} and { and {e,) is a conditionally zero-mean a, o--sub-Gaussian martingale difference sequence. The agent uses m-mixture exploration. Define the population second-momen
Σ : E(z,2,
17

---

Published as a conference paper at ICLR 2026
Suppose the teacher's parameters (step size co. reward preconditioner U) and problem parameters satisfy
(1-2) 7) co |Ul of Tw K
(A.3)
Then for any 2 2, the restriction Σ ş is strictly positive definite on S 1 @ 1'; in particular,
Amin| (E|s) > 0,
and Eg is invertible on S.
Proof.
We work on = 1- E 1-. Write
Enn Egn
Eng Zgg
Σ = E(Z,2,
Set
Amin Enn 11),
c:= Amin E 1
b : Eng 1lop
Decompose =p, te, with Ps == E[x, F and E
Fa-1l = F(p.). For any u E 1 !
ng
z
t2 U FP..)ul.
By Lemma A.1 and mixture exploration, F ((p))1 - (Y/K)/ I, hence
a
2 2 s=1
(A.4)
K
Kt
Split g: = gt
(sig)
+ gr
(noise)
with ge
(sig)
= { Σ.m (w xx) xx)xx = Diag(ū;) w. For any v E 14,
K
(u ge
2 ) = Th [ Diag [1l) = - Σ
(six)
E Fl of
Let No. = E 1(A, = i}. Since No. 2 Ne.i E[n = B/N > 1/N..1/ = =//,, Mixture exploration yields [R(,,,)) 2 7/ K, so
(A.5)
w Kt
Adding the PSD noise covariance only increases Egg- thus equation A.5 holds. We have
Eng Elngg
(g (ii)) ) J] (g (noise) TJ.
=0 The first term vanishes because w is independent of (nt, (c))) and E [[w] = O. For the second term,
ne (g)
(noise)
1
E[x, Eo x
E|x, es x!]]
uma41
where we used Ee(e [Fs-1= eliminate 8. Fix 8 < u. Consider "world 0" where e, is set to 0, and "world 1" the true world. Then
E[x, G, = E (p
P (0) G x!
since E(p (0) e, x
0. With th, co/u, the projected logits satisfy (one-step normal form)
yu Co
V + ga-1) 3 Δyu = Co
U (c, x,)
when only e, is perturbed. By Lemma A.2 and the (1 - 7) factor from equation 3.3,
[p! (1)
p:0 ll - (1 (1-7) 7) || < (1 7) 1 u
| le.l.
18

---

Published as a conference paper at ICLR 2026
Hence
[ Eo x xllppp
- 1[Il - - 1l2 |ca1
(1 7) co
op
B|e
2 v
(1 - 7) co
op o²
2
Summing u = s + 1..... .,t and = I....,t t 1 gives
1-1
(1
(1-7)
b = Engl 1- llop
U
up o?
22
2
-1
s=1 x=
(1 - 7)
2 t-1
Cu
U
up 0
2
2
C
(A.6)
t
with Cs == [جاا cn | | oZ.
For any unit (u, v) S.
2 a a||u|1² - 2b|lu|| ||||| + c|llal2 > Amin -b -).
Σ
Thus
a
v (a c)² 402
Amin E
(A.7)
2
By equation A.4, equation A.5, and equation A.6,
7 Tw K
C. t
ac
b
5
t.
The small cross-block condition equation A.3 implies vac (for all 2), hence the x matrix
o positive definite and the right-hand side of equation A.7 is strictly positive. Therefore Σ Σ1g>1 0.
Lemma A.8 (Sample second-moment moment concentration for 21). Under Assumption 3.1 and (O IR), ). the normalized statistic ž, := (n, g. T R2 is sub-Gaussian with dimension- free 4/2-norm := 12rl1 depending only on (Twi oe, 7) (and not on K). Let
M
Σ m=1
Z
(M)
= E|Z,Z,11
M
where (z("m)
(mm))1. M
mm=1 are i.i.d. copies of 2, generated by independent rollouts of the same population
process (fixed 7). Let g > 0 be the population lower bound from Lemma A.7, i.e., Amin( [1gg) g) 2 g on S = 1 ① 1 Then there is universal constant C' > O such that, for any 8 E (0, 1). with probability least 1 5,
2K + log(1/8)
2K - log(1/0)
SS llpp
S CCL
(A.8)
M
M
In particular, if
2K + log(1/6)
M 2 Cme
with
Cmc = 1 14,
a2
then Σ
< 2 σ, and consequently
op
Amin
Σ
>
2 g.
(A.9)
8)
19

---

Published as a conference paper at ICLR 2026
Proof. The Sub-Gaussianity of follows from: (i) n. = E X2 is average of one-hot vectors from a "y-mixture policy, hence coordinate- wise sub- uab-Gaussian with th- norm bounded by a constant depending only on 7; (ii) g: = - TxX; with rs w x, + 6x. where w Gaussian independent of (x,) } and {(,), and {c,} } is conditionally aub-Gaussian. Thus gi is sum of sub- Gaussian vectors with 2/2 2 controlled by (Tw, Ge. 7); concatenation preserves uub-Gaussianity with L | (1) in (Tw,
For i.idd aab-Gausiian samples in R. with d 2K, the standard sample covariance operator-norm deviation bound (e.g., matrix Bernstein/ / sub-Gaussian covariance concentration) yields equation A.8:
d + log(1/5) d d+ log(1/)
S - op CT.
M
M
If M 2 m2K + loggl//)))//////// o², then Σ ---lon to S,
Elop op I 0/2. By Weyl's inequality on the restriction
Amin
2 Amin
Σ
E-E
1 o,
9
op
which is equation A.9.
Lemma A.9 (Fisher two-channel quadratic: gradient and Hessian), Consider the population Fisher- weighted quadratic from Eq. equation 4.1 written in the two-channel parameterization
((W) = to (TWWEWT) tr([ Eys W + const,
(A.10)
where W E Kx2K Σ E/Z,z, 1. Eys = Ey and r E (P
).. Then
Vw T(WE- Egz) =: E(W) E Kx2K
(A.11)
V2L(W) = PAΣ (AERKX2K).
(A.12)
Proof. Write = C C2 + const with
(.(W) := 1/2 tr(T W E WT), C2(W) == tr(T Ens W
Using the Frobenius inner product (A, B) = tr(AT B) and d(W E WT) (dW) E w
W E (dW) with Σ = Σ, r =,
dL T dW), dC ( r dW).
so dl- = (T(WE- - Eyæ), dW), which proves equation A.11. For the Hessian, with W(e)
W CA
d de (w =
establishing equation A. Equivalently, d C(A A = tr =
2 0. 1
Lemma A.10 (PL inequality on 1+), Under mixture exploration (Lemma A. so that Lemma A.7
holds, the population Fisher-weighted quadratic
L(W) 1/ (r(r W Σ WT) Ey WT)
const.
satisfies, for any global minimizer W
((W)- L(W') < [W])(
1
/ A (T), (T), X X
Σ s)s)
(A.13)
2μ
where S = - +1 Moreover, by Lemma A.I and Lemma A.7, xt (T) (T)11) 2
7 K
X Σ (s) 2 a > 0,
1 2
7
(A.14)
K
Proof. Let = W- W and define X = 1/2 A 11/2 By Lemma A.9,
CWC(W) = A Σ = 1 x
L(W) (WW*) = 1/2 tr ( (TTTAAT) = 1 I
For any PSD A, B and any matrix X. || A 1A12XBB12 1/2 12 = tr (AAB) >
Amii (A) (
| [ p... Applying this with A r and B Σ g. we obtain
V w ((W) 1/11 which is equivalent to the PL inequality equation A.13. The bound equation A.14 follows from
Z (T) X t (51)) | 2p (C(W)-
Lemma A.I (giving r (r/K) and Lemma A.7 (giving E al).
20

---

Published as a conference paper at ICLR 2026
A.2 REWARD-SHOCK ROBUSTNESS
Proof of Theorem 4.8. Under the s-CR s-CRN coupling (Def. 4.7), the baseline and perturbed runs share the same task w, the same uniforms {Ui} } for sampling, and the same noise sequence { they differ only by feeding a shocked reward T. Ta 6, at round Define
Apt+1 == Pi+1 Pi+1. An, n. nc
Ag. == gr gi, a++1 := El | Fa-1].
and the normalized accumulators n, == nt, Be == gt/t (and their deltas analogously).
Pathwise relations for counts and rewards. With x; es, and (w,x1) + Et.
X.i
Ta X = Di) w +
Σ X..
u=1
u=l
u=l
Let h4 == E === Cu x and h2 fed reward, hence
Eu=
Cu X For the perturbed run, only round changes in the
g: (ū:) w he + 1ft 2 8) x..
Therefore, for 1 8,
Ag: 1
1 Diag(w) w) Ane + Ahe + o Xs.
Ah := h4 he.
(A.15)
t
Two-channel linearization of projected logits. Using the (population-optimal) one-step normal form aligned with equation 3.2 and projecting away the softmax gauge,
Ay-- == Proj (onoss(prrt pertt) Sp41
$/41
(VAn, + UAg. cc An, + UA
(A.16)
Softmax 1/2-Lipschitz on and the 7-mixture equation 3.3 scales the sensitivity by (1 - hence
-7),
1
Ap $ (1 (1-7),
2 |
(A.17)
Combining equation A.15- A.17 and conditioning on Fs-1 yields c(1 - 7)
U 6r x Fo-1
E |IVV+ UDiag(w)) Añ, + U U
a1+1
t
c(1 7)
(IV + DDagg( |lop Blll Anell Fo-1l
2t
+ ||op Bi|l Ahell Fa-1] + U lopp (6,1
(A.18)
CRN coupling controls E An, || and E| Ahel- By inverse- CDF coupling, Vk
P(Au + A. F.-1)) Apull
|A
2
Since | || S V2.
E|||,x, Fa-1] S
V Appll
(A.19)
Summing equation A.19 from u = to f and conditioning on Fa
| F,-1]
(A.20)
For the noise accumulator, using the zero-mean a. assumption (hence E||ca| Fu-1) <
C. 22/T o.) gives
|| [F-i] C.
au.
(A.21)
A Grönwall-type recursion and its solution. Plugging equation A.20 4.20-equation A.21 into equa- tion A.18, define
c(1 (1)
b= c(1
7)
K 2
U Ulop
||V + Diag(w)llop C. (lUlop))
2 V
21

---

Published as a conference paper at ICLR 2026
and obtain, for t 8,
a++1 $ ? 16-1 + Σ au.
a
(A.22)
u
Let St = = Ou with S,-1 = 0. Then
St11 & (1 (1+7)s: St + a |6
(A.23)
Unrolling equation A.23 and using the standard Gamma-ratio bound yields constant C2 >> (depending only on 6) such that
St S a
aC
|6,
Σ
aCb b
(9) 18
u=
Returning to equation A.2
Cb
(
a(1 Cb) ( (-1 b-1
(1)011 16,1-
a++1 S |, St S
|pp
|,,|
8
Therefore, for any st.
a(1
Co)
(
6-1
E [1APi+11l2 F,-1] s
16,1-
(A.24)
In particular, if the learning-rate constant is small enough so that 1, then (2) (1)=-1 a - oc. and the one-shot impact decays to zero:
lim B 1114Ai++ll2
|
F,-1] =0
+0000
This completes the proof.
A.3 KL DIVERGENCE VS. FISHER-WEIGHTED DUADRATIC
Proof of Theorem 4.1, Recall that denotes the trajectory length the dataset construction of
Sec. 4, and expectations below average over (r, () with / E (I....,N - l} and T E D. Let
p(s) == softmax(s) and define the projected logit error
A = (roj s/41 E E
Proj := * 11
Write the (softmax) Fisher matrix as F(p) = Diag(p) - pp and recall from Eq. equation 4.1 that
N-1
N Σ
( = E
AT A
r=
By Lemma A.l (mixture curvature), along the teacher's 7-mixture policy we have the spectral sandwich on 1
7 I = F(!
I.
(A.25)
K
By convexity of the PSD cone, the same bounds transfer to r on 1- Upper bound. Mixing with the uniform distribution is Markov kernel; by data processing for f-divergences,
DK P++1) s Dku(p(sPM) p()++)) Dku(P(s (p)))) (0 A441)).
Let ((s) == log E, e*;; then VΦ softmax and V² F(F(s)). The Bregman integral form of KL yields, for any 1-1
DkI(P(s) Ips A)) = ( (1- ( +2)) Adr.
(A.26)
Using ||F()) llop 1 and (1 r) =
Dka(Prt
PO (P++1) < 1 14
22

---

Published as a conference paper at ICLR 2026
From equation A.25, 4 A+-1 over and taking expectation,
xl hence | I
K A A Averaging
K L()).
E
DRu(Pit) PO (143)
S 47
N
Lower bound. Pinsker's inequality implies DkL (plla) 2 31pp qll 2 - qll2 Let f(s) = (1 P(s) + 1/K K so that p. = f(s and pi+1 = (((+1). ). By the mean-valus integral,
Pi+1 - D1+1 = / V [ + 11+11 A dr = (1 (1-n) / P(p(s + + D dr,
whence
||p/71 Dr+1l12 S (1 - 7) 1 | and thus
| 112 > 4
||
PO PD+1
7)
Using the upper side of equation A.25, A
r
D++1 S 1| 113, we obtain
(()) S E
1
Σ
1
12
S
2 E
P141
PO
p
2
N
(1
7)
N
Finally, combining with Pinsker yields
1
(1 (1- (()).
E
PRI(PER 11-1)
E
|
N
4
This proves the two-sided bound.
A.4 CLOSED-LOOP IMITATION OF POLICY OPTIMIZATION
Proof of Theorem 4.2. By Lemma A.4, the projected student logits are linear in the normalized statistics:
Proj
W Ze,
W (n, Wol Kx2K
Di gr
E 22
For the teacher generated by equation 3.2 with 1/t = c/t, likewise have PO Proj V n, -gg) W, 2t,
Yt+1 Proj
W, = Proj V U] R22
so the labels are realizable by the same two-channel structure.
Lemma A.5 gives the population weighte quadratic form
(( tr(T NETWEWT) E WT) const,
r F (p ]] Σ -[ZZ, ].
Using realizability ye+1 W W...o we have Ey Ely ing gives the cempleted-square form
W WZZ
W Σ. Substitut-
( tr (T(TWW..) W.) W.)") + const.
By Lemma A.1 (mixture curvature), r 2 (r/K) ) 7: by Lemma A.7, Σ 2 2g on S = 1 1
under Assumption 3.1. Hence the quadratic is strongly convex in W on S and has the unique minimizer W W.. Therefore, for any history prefix He,
Proj (H, 0*) = W, 2, = Proj P (H+).
Since softmax is invariant to additive constants (the projected logits remove exactly the span(1) [1] gauge), the induced policies coincide: B.41(He; 0*)
23

---

Published as a conference paper at ICLR 2026
Proof of Theorem 4.3. By Lemma A.4, the projected student logits are linear in the normalized statistics: Proj S = W 2 with W = [W,, We
and & = (n, S. )" Let Yt+1 := Proj (s Lemma A.5 and Lemma A.6 give the population and empirical weighte quadratic forms, and Theorem 4.2 (under Assumption 3.1) guarantees realizability: there exists a unique W, such that yt+1 W.z, for all prefixes. Moreover, Lemma A.1 implies | (T/K)T, and Lemma A.7
implies L|g = ol with o x cx/t on S == @ so the population quadratic in W is strongly convex on S with unique minimizer W. For the empirical problem, Lemma A.8 yields (for i.idd rollouts)
- Elop I CV 2K+Hog((/a)) A
+
2K+Hog(1/d)
Taking M 2 C2 (log( (for large enough absolute C depending C1, L) guarantees E = } gl, while r 1. (7/K)z by mixture curvature. Plugging the realizable labels y(+1
W. By into Lemma A.6 gives
((w) 1/ tr( (W (WW.W)))2.. (W-..") + const,
W.)
so the empirical minimizer satisfies W W. whenever Σ > (which holds with probability
8). Thus ProjS,-1- Proj Po and consequently (P4; 0) p (H4) for the fixed test history He. For the expected mismatch, on the high-probabilityy event the error is O, and on its complement crude bound together with the (1-y) mixture factor yields |[P 1 s 2(1-7)² hence
Eraai |1 1
2(1 2)
This proves the theorem.
A.5 CONVERGENCE OF THE
TWO-CHANNEL LSA
We amalyze the continuous-time gradient flow on the two-channel operator W for the Fisher-weighted quadratic objective in equation A.10.
Theorem A.11 (Exponential convergence of the W-flow). Assume mixture exploration and that Assumption 3.1, Lemma 1, and Lemma A.7 hold. Consider the gradient flow for the population Fisher-weighted quadratic L:
W(t) (W(()) -
(A.27)
initialized at any W(0). Let W be global minimizer of L. Then ((W(t) is strictly decreasing along the flow and
((W()) L(W*) s exp(-2ptt (C(W()) -((")),
(A.28)
with
=
x (T 2-) X
min! E gg))
S $ 1+
and, by Lemma A.1 and Lemma A.7,
mni (T) 2 K
min [1s) 2 a
0,
3
2
7
K
Proof of Theorem A.HI. By Lemma A.9, we have
V Vw r(wE- - Eys).
so the gradient flow W(t) V (WW(((t)) is well-defined. Along the trajectory,
d ((((t) = -V (W(L(W(t)), W()) = [w s 0,
d
hence L( (()( is nonincreasing and strietly decreasing whenever (w(((()) + Next, apply the Folyak-Lojasiewva (PL) inequality on the restricted subspace S == 1- $ 1- (Lemma 10). Let
x (
(T/
A Σ (s
Jr
24

---

Published as a conference paper at ICLR 2026
Then for any global minimizer W*
L(W) - L(W") & 2μ
[wL(W)]
Combining this with the energy decay identity gives, for all 20, 0.
d (L(W(t)) L(W") Vw | - (
dt By Grönwall's inequality,
((W())) L(W*) z exp(-2pt) exp(-2pt(C((W(0) L(W")).
Finally, by mixture exploration (Lemma A.1) and the lower bound on the population second moment on (Lemma A.7),
( K X
0,
which implies μ 2 (7/K)T. a. Strict decrease of the loss holds except at stationary points: by Lemma A.9 together with the restricted positive definiteness, these coincide with the global minimiz- ers on 1
B EXPERIMENTAL DETAILS
This section provides additional details regarding the experimental setup, including specific hyperpa- rameters for method and the baselines, as well as the hardware and software environment used for all experiments. We also present supplementary experimental results that complement the main text.
B.1 HYPERPARAMETER SENSITIVITY.
Figure 4 shows the impact of varying the number of in-context optimization rounds (n) and candidates per round (k) on AIME 2024. As shown in Figure 4a, performance improves substantially as the number of candidates (k) increases from saturates beyond this point. Increasing the number of rounds (n) is broadly beneficial up to a peak near n = 5. complementary grid and the corresponding latency/VR RRMM measures are reported in Appendix C.2 (Tables 8).
64, with Maj @ 16 accuracy more than doubling, then
5-wen225-Math--B Sensitivity
@wen2.5-Math-2 (AIME): Sensitivity
Meang116
Maj@
Accuracy (16)
20
Meang916 Majas 16
24 8
15
32
64
Candidetes (fixed
Rounds Ifieed k 16)
(a) Varying candidates (fixed 5)
(b) Varying rounds (fixed k 16)
Figure 4: Hyperparameter sensitivity of ME-ICPO on AIME 2024 with Owen2.5-Math-77.
B.2 HYPERPARAMETER AND IMPLEMENTATION DETAILS
ME-ICPO (Our Method), For our method, we use the primary settings described the main text (n = 5 rounds of optimization, & 16 candidates per round). Candidate chains-of-thought are sampled using a temperature of 0.6 and a top-p value of 0.95. The summarization step, which is crucial component for managing context length, uses greedy decoding and is prompted to produce a concise summary of approximately 100 tokens, with a hard limit of 500 tokens, focusing on the reasoning strategy. The predictive answer distribution, used for the entropy calculation, is also estimated with a temperature of 0.6. The full text for all system and summarization prompts is provided in Appendix D.
25

---

Published as a conference paper at ICLR 2026
B.3 HARDWARE AND ENVIRONMENT CONFIGURATION
All experiments were conducted on single server node equipped with 8 NVIDIA L40S GPUs, each with 48GB HBM2e memory. Our implementation uses PyTorch 2.2, vLLM O. 0.10.0, and Transformers 4.55.0. The operating system is Red Hat Enterprise Linux 9.6, and the environment is managed via Conda with Python 3.11.7 and CUDA 12.9.
SUPPLEMENTARY EXPER IMENTAL. RESULTS B.4
This section contains supplementary results and comparisons added during the rebuttal phase to address specific computational and scalability questions.
B.4.1 COMPUTATIONAL COST AND BASELINE COMPARISON
To assess efficiency, we compare ME-ICPO against standard prompt-style test-time search base- lines-Tree of Thoughts (ToT) (Yao al., 2023a) and Monte-Caria Tree Refinement (MCTR) (Zhang et al., 2024) and the training-based test-time scaling method TTRI. (7uo et al., 2025). For fair- ness. ME-ICPO uses fixed configuration of N=5 refinement rounds and k=16 samples per round, averaged over 5 seeds. We report average inference time (seconds per question) on AIME 2024. Because ToT/MCTR employ different search depths and branching factors, their runs are not strictly time-matched, for TTRL, we control the experiment to use similar GPU hours as ME-ICPO. As shown in Table and Table 4, ME-ICPO achieves top-tier accuracy and Mean @ 16 at competitive compute budget, reflecting principled gains from the underlying Policy Optimization mechanism. Shallow-search prompt methods (ToT/MCTR) can be faster when search depth is limited but lag notably in accuracy. while under matched wall-clock budgets ME-ICPO attains higher Acc/Mean @ 16 than TTRL, demonstrating more efficient scaling; importantly. ICPO provides a mechanistic account of test-time self-refinement policy optimization rather than heuristic leaderboard tweak.
B.4.2 FRONTIER MODEL SCALABII.IT
To demonstrate the versatility our framework, we evaluate ME-ICPO on recent frontier models, including Qwen3-44B (Yang et al., 2024b) (a long-CoT specialized model), and the Gemini- 2.5 series (Pro and Flash) (Comanici et al., 2025). Tables show that ME-ICPO consistently enhances performance across these diverse architectures, confirming its scalability.
B.4.3 HARDER TASK GENERALIZATION
We further evaluate performance on exceptionally difficult, high-school/collegiaate level math com- petition tasks, specifically the Harvard-Ma Mathematics Tournament (HMMT) (Balunović et al., 2025) and the APEX-sio -hortlisst dataset (Balunović et al., 2025). Results in Table and Table 7 indicate that ME-ICPO provides robustness even on tasks with low baseline solve rates.
Table Computational Cost and Performance Comparison (Mean @ 16, Qwen2.5-Math-78 against Inference and Training Baselines.
Method
AIME-2024
AMC
MATH(Avg)
Time(s/question)
ToT (self eval)
4.38
16.19
12.51
708
ToT (Maj vote)
19.58
29.37
35,63
363
MCTR
4.60
1.20
17.20
1758
TTRL
27.20
45.18
46.83
1253
ME-ICPO (ours)
30.42
47.06
54.71
1152
Table Computational Performance Comparison (Accuracy, Qwen2.5-Math-75
Method
AIME-2024
AMC
MATH(Avg)
ToT (self-eval)
4.40
18.10
10.74
ToT (Maj-vote)
19.30
29.40
35.91
MCTR
23.30
2.40
33.82
TTRL
30.00
43.37
45.11
ME-ICPO (ours)
30.05
47.20
47.30
26

---

Published as a conference paper at ICLR 2026
Table Results on frontier/long-Co CoT models on AIME 2024.
Model
Method
Mean @ 16 (%)
Acc (%)
Qwen3-4B-Insten
Base
20.62
20.59
Swen3-4B-Instmun
ME-ICPO (ours)
57.71
57.67
Gemini-2.5-DP P.o-Poo
Base
58.54
56.60
Gemini-2..-Prr
ME-ICPO (ours)
79.17
80.00
Gemini-2.5-Flaas Base
35.21
35.42
Gemini-2.55-Flan
ME-ICPO (ours)
76.46
76.47
Table 6: Results on Harder Benchmarks (HMMT APEX) using Qwen2.5-Math-
Method
HMMT Mcan 16 (%)
HMMT Acc(%)
APEX
APEX
Mcan 16 (%)
Acc(%)
Base ME-ICPO (ours)
1.04
0.67 1.33
2.55 4.59
2.61
0.42
4.57
Table 7: Results on Harder Benchmarks (HMMT) / APEX) using 2emini-2.5-44
Method
HMMT Mcan @ 16 (%)
HMMT
APEX
APEX
Acc(%)
Mean @ 16 (%)
Acc(%)
Base
14.79 43.12
14.76 43.14
14.68 17.18
18.33
ME-ICPO (ours)
20.00
C DETAILED COMPLEXITY ANALYSIS
C.I THEORETICAL TIME AND VRAM COMPLEXITY DERIVATIONS
We formalize the compute model for decoder-only Transformer and provide exact asymptotic derivations for ME-ICPO (forward-only, summary-aware history) and TTRL (Zuo al., 2025) (backprop-based test-time RL). The statements and proofs below match the main text analyses verbatim; we only add brief connective narration.
Setup and primitive costs. We analyze decoder-only transformer with 7. layers and width d (parameter count [0] = I.d).. For single test instance, the initial prompt length (problem statement template) is sg. Each sampled chain-of though (CoT) has average length f. Per round we sample k candidates; the total number of rounds n. Let 3 ;= E + denote the number of tokens appended per (x, r) pair (the reward stub is O(1), 3 = E). The prompt the beginning of round t therefore has length T;-1 = Su + { - 1). We use K [2, 3] for the backward/forward FLOPs ratio (one backward costs times forward) and for the number of mniicy-opiimization steps per round in TTRL. Primitive costs (suppressing constants) are:
full-sequence forward at length T';
CIwa(T) = G(L(T LU d(
autoregressive decoding of ( tokens from prefix T: (T, ()
one training step ferecie--orring fwd bwd) at length T Cirain (T) = (1 (1++ 1²) Crea (T).
For ME-ICPO, the one-step lookahead score is computed on the just-generated branch by appending constant number of tokens, so cost is an incremental forward
Ciac (T, E) ((L((T + +)
rather than full recomputation at length T + l. C.1.1 ME-ICPO
Theorem C.1 (Time Complexity of ME-ICPO). With one shared prefill length Ti- 1 per round, candidate decodes from that prefix, and incremental on-branch scoring for each candidate, the total time over rounds satisfies
TMP (Ld Ld 12 n + l Ld n² + Led B + k Ld 8n).
27

---

Published as a conference paper at ICLR 2026
and, using x l
T O(dd + O(Ld kln)).
Proof. The total time complexity, TME. is the sum of the costs for prefilling the context (Siwa), decoding the candidates (Sdrr), and scoring each candidate (Secore) over all rounds.
" TME Ctwa + Cdec +
Σ ls Cire
(T) =: Stwd
Saec Sacore.
We analyze each component by first summing over the rounds and then identifying the leading-oree order terms in The prompt length round t is 7,-1 So + B( 1), and use the standard sums min-12- (n²²) and Er= r² (x-1)m(2n-1) (nn)).
E r
The prefill cost, Siwd. is given by:
Stud
Crud (Ti-1)
Ld T 1--1 + Lat Σ To-1
1=1
1=1
Ld nsu s03m(n 1) + 32 (m θ
1)m(2n 11) 1)
Ld 1250 3 n - 1)
6
2
@ (Td + @ n²²)
where (a) substitutes the definition of (T) = @(L(T Td))) and pulls constants out of the sum: (b) substitutes the exact formulas for the sum of linear and quadratic sequences; and (c) identifies the highest-order terms in for each part of the expression.
The decoding cost, S is given by:
Saec
t (Te-1,)
n
kLd
E T nl2 &.I.d(nn
(c)
kLd
(e(nNo
1880
n(n -11)
+ n12
+ k.T.ee)
2
@kLL BEn?) + @kLd en).
where (d) substitutes the definition of (T, O(L((T++ (dd++++++++ kd)); (e) substitutes the
sum for T;-1; and (f) identifies the dominant term in for the Ld component as the quadratic term (kLLd 3ln²).
The incremental scoring cost, Secore. is given by:
Secore
Σ
k C (T)-1,
1=1
@ θ kLd Σ Ti-1 + =nl]
+kLdPn
1=1
- (kId (k.aPn))
where (g) substitutes the definition of Ciue (T, O(L((T + E)d + +d); and (h) identifies the
dominant term (k:Ld 3n²).
Combining the seading-order terms for the three components, we have:
TMME Stwd + Sacc + Sucore
(Ld 32 3 Ldf Bn²) (kLd kLd (n) + @(LddBi Bn² kLd
28

---

Published as a conference paper at ICLR 2026
Grouping the terms by their dependence d and d?, and noting that the scoring cost terms dominated by or are asymptotically equal to the decoding cost terms (since 2 1), the total complexity simplifies to:
TMM = (Ld + θ (ILLdBen)) + + Bn²) + θ (k.Ldd en)
from Spud
from dec
from Stwd
from Siee
Using the approximation 5 (e). we arrive the final expression stated in the theorem. This completes the proof.
Theorem C.2 (VRAM Complexity of ME-ICPO). If candidates are decoded sequentially (no k-way parallelism), the peak memory over n rounds satisfies
MME = ((10)) + (Ld Ld (so + Bm)) - O(())) + ((Ld + En)).
6 < k candidates are decoded concurrently, the second term is multiplied by b.
Proof. The peak VRAM complexity is the sum of the static memory for model parameters and the maximum dynamic memory for the attention KV cache. As ME-ICPO a forward-only method, it requires memory for gradients optimizer states. The model weights occupy @(10) space. For decoder-only transformer with L layers and width d, the KV cache for sequence of length T' requires M2) ((LdT)) memory.
The context length at the beginning of round is T;-1 = + 3(t - 1). During the decoding of a candidate of length the sequence grows monotonically increasing with t, the global peak occurs during the final round (t n), giving maximum sequence length of Tn-1 + l 8q + 8(n 1) ((so + fnn).
maximum of T--1 L. Since this length
Therefore, the total peak memory for sequential decoding is the sum of the static and maximum dynamic components:
MMF 0(10)) + Mkv (0(6o + Bn)) O ( + (Ld((s + Bn)).
Using the approximation 3 (E) yields the equivalent form. If b < k candidates are processed concurrently, each parallel branch maintains its own KV cache, so the activation memory term scales linearly with b.
C.1.2 TTRL
Theorem C.3 (Time Complexity of TTRL). each round, TTRL (i) samples candidates from prefix and (ii) performs policy-optimization steps using feacher-forcing on sequences of length 8o + The total time over n rounds is
TTTRL = 6 (1+ L( 2) (so + +)dd))
and the per-round prefill and sampling costs, @(L(sed + sod?)) and ((kL((ssU (2)d+ lower-order whenever 2 1.
Proof. The total time complexity, TITRI, the sum costs over n rounds. each round, the process performs one shared prefill of the prefix sp. followed by k candidate sampling operations, and finally training steps for each of the & candidates. The cost for a single round is thus Ctwa(so) - l Cdee (s0, E) + gk Cirain (so + f). The total cost over rounds is:
TTTRL = n (Cowa(so) + & (So. E) + gl. Ctruin (So + +))
We substitute the standard complexity formulas for a decoder-only Transformer, using Ctrain (T') + K) (T'), where is the backward/forward FLOPs ratio. This yields three terms:
TITRL = nL (n sud²)) e (nkL((sol + EdP)
Prefill Cost
Sampling Cosl
θ (ngk(1 x) L((s0 e) (5o E)d²)
Training Cusl
29

---

Published as a conference paper at ICLR 2026
To determine the tight asymptotic bound, we identify the dominant term. The training cost term scales with the number of optimization steps and the backpropagation factor (1 + k). Furthermore, its self-attention cost is quadratic in the longer sequence length, So + f. In contrast, the sampling cost's attention component is only linear in the prefix length, (s)(). and the prefill cost is computed on the shorter prefix so and lacks the multiplicative factors of g and it.
Consequently, for any 2 1, the training cost term is of a higher order than both the prefill and sampling costs. Therefore, the lower-order terms are absorbed into the --notation, giving the final tight bound:
TITRL = θ (ngk (1+ ())+))++))) L((0 + f) + (5o +)dd))
Theorem C.4 (Vram Complexity of TTRL). During training sequence length T 80 + a TTRL. step must hold (at least) model weights, gradients, optimizer states (e.g., SGD momentum Adam's s first/second moments), and backward activations. Consequently, for effective batch size batch, the peak memory obeys
MTTRI = +101)+ 0(10) + ((10))0(210)) +
+
-
weights
gradients
optimizer states
backward activations
( + e L (so + +) batch).
Proof. The peak VRAM complexity of a TTRL training step is the sum of two primary components: parameter-resident memory and activation-resident memory. The parameter-resident portion consists of the model weights, their corresponding gradients, and the optimizer states (e.g.. first and second moments for Adam). Since each of these scales linearly with the number of parameters, 0). their combined memory requirement is compactly expressed as Mparam @(1)).
The activation cettrrattioo-rrientt memory is required for the backward pass. For a training sequence of length T + f and decoder-only Transformer with L layers and width d, standard (non- con-cheecpoiitte) backpropagation pass must cache activations (such as hidden states and MLP intermediates) of size (dd) for every token in the sequence. Aggregating this over I. layers and effective batch size of batch, the activation memory scales as Mact = O(L- d T batch).
The total peak memory is the sum of these components, MTTRI- Mparam + Mact Substituting the derived complexities yields the final expression stated in the theorem:
MTTRL @(10)) + ( (50 + e) batch).
C.1.3 COMPARISON WITH TTRL (ZUO ET AL., 2025)
Proposition C.5 (Time Complexity Threshold). ME-ICPO is computationally faster than TTRL. when the number of in-context optimization rounds, n, is below threshold n². This threshold is given by:
1g +K)
S
*
9(1
+
KK)
(s0 (xo+e)
if >> n
e
In practical settings with small number of rounds (e.g., n 10), this condition is typically met, making ME-ICPO the more time-efficient approach.
Proof. The threshold n² is found by equating the leading-order terms of the time complexities derived in Theorem C. and Theorem C.3. We consider two regimes based on the dominant term in the complexity of ME-ICPO.
Case (k < n): The dominant term for ME-ICPO is the prefill cost, TME x Lnd. The TTRL cost FITRL x ngg(11 + x) L(so + dd Equating them and solving for yields:
n² = gk(1 + ks E)²
Ld =1 + k) L(so + e)2d
8
n² su +
V gk((1 + K)).
C
Case 2 (k: n):: The dominant term for ME-ICPO the decoding cost, TMM k En². Equating this with the TTRL cost gives:
LLd ,n ngg(1 x:) T( 0)2d
en= 9 + <50 0)?
30

---

Published as a conference paper at ICLR 2026
(80 + e)2
n = g(1 + k)
In both cases, for n < n,,, the complexity of ME-ICPO is lower.
Proposition C.6 (Memory Complexity Comparison). Let bæI be the number of concurrently decoded candidates in ME-ICPO, and b be the TTRL training batch size. ME-ICPO achieves lower peak VRAM than TTRL if the number of rounds n. is below a threshold. A sufficient condition is:
b(00 + 2) bME 80
param (0
Y2 <
+
bME &
bME Cky Ldd
where Aprrm represents the additional parameter-resideee memory (gradients, optimizer states) required by TTRL, and Okv is a constant. Given that TTRL requires strictly more parameter-side memory and typically uses larger batch size b, this condition holds for all practical values of n.
Proof. We establish the condition by comparing the upper bound on ME-ICPO memory from Theorem C.2 with the lower bound on TTRL memory from Theorem C.4. We seek the condition on for which MME < MTTRL:
ME [0 + Okv bME L d (su + (n) < (cw + cg + cppt) |0| + ckv b d (so + +().
cer
MAR
MTTRI.
Let A param == (ct - cc + copt ) > 0 be the constant factor for the additional parameter-sized tensors (gradients, optimizer states) that TTRL requires. Rearranging the inequality to solve for
Ckv buz Ld < Ckv b L. (su + E) Cev bME L d su + Aparam [
(su + SMESSO
Aparam [6]
n
bMEl
Okv bME Ldd'
Since A Aparam > O, the second term is strictly positive, providing further margin. For typical use-cases like sequential decoding (bME = 1) and b 2 1, the first term large and positive, making the condition true for any practical number of rounds n.
C.2 EMPIRICAL COST
We complement our derivations with controlled latency and memory study on AIME 2024 using Qwen2.5-Math-7B. We vary the number rounds n and the number of candidates per round for ME-ICPO to examine its computational characteristics and to validate the asymptotic trends predicted by Theorem C. We sweep n E (1,3, 5, 7} and E {4, 16, 32). Unless otherwise noted, ME-ICPO uses main protocol: temperature 0.6, top-p 0.95, summary cap of 500 tokens, and entropy lookahead with 1==16 short samples. All runs are performed on 2 (48GB) GPUs using b precision, with candidates decoded sequentially (bMr=1).
metrics. (1) wall time per question (s), averaged over the test set; (2) token usage per question (generated tokens, in thousands); (3) peak vram (GB) from vllm memory statisties.
results. Table 8 reports wall time, token usage, and peak VRAM across the (n, k) grid. Two trends emerge. (i) For fixed k. wall time increases superlinearly with and is well described by prefill- dominated scaling close to ( consistent with Theorem C.I. (ii) For large k, the cost transitions toward Ok n2) as candidate evaluation/selection becomes the bottleneck. Token usage grows with both and & (longer multi-round traces and more candidates), and peak VRAM remains stable in the 81-82.3 GB range across settings, indicating that memory primarily governed by the base model residency and KV-cache size under sequential decoding.
D PROMPT TEMPLATES AND QUALITATIVE EXAMPLES
D.1 DATASET-SPECIHIC SYSTEM PROMPTS
These system prompts define the semantics of the reward tag and are placed once at the beginning of the context.
31

---

Published as a conference paper at ICLR 2026
Table 8: Empirical cost vs. (n,k) on AIME 2024 with Qwen2.5-Math-7B.
k
Wall Time (sqquetioon
Tokens (k/question)
Peak VRAM (GB)
1
197.52
99.59
81.50
I
8
228.89
112.25
81.43
16
286.74
379.82
82.24
32
468.36
852.47
82.27
3
4
406.65
247.76
81.43
3
469.79
272.56
81.43
3
16
715.85
1070.98
82.18
3
32
1163.51
2111.85
82.30
5
609.75
372.18
81.42
5
8
713.25
506,68
81 31.43
5
16
1152.67
1613.49
82.24
5
1726.70 32
3412.94
82.30
7
4
820.60
559.20
81.43
7
1156.62
566.82
82.30
16
1371.78
1753.71
82.27
7
32
2643.74
5123.67
82.29
System Prompt: AIME AMC (Numeric Answer)
You Al mathematiciin. All content you output MUST be in English.
Below compressed solution ideas from previous attempts; cach idea tagged with reward 1 (correct) reward (incorrect). Use the question and these idcas deduce the correct numeric answer. Finish all your reasoning, then on a NEW line output exactly one number (the answer) and nothing else.
Your final output MUST be in the format boxed boxed(<<nmber> <number>>) where <number>> is the final numeric answer only (no expressions, variables, or additional text). content inside boxed (enembern)bbb)) must be decimal number, not a fraction or any other form.
System Prompt: GPQA (Multiple Choice)
You are an Al mathemattcian. All content you output MUST be in English. Below are compressed solution ideas from previous attempts; each idea is tagged with reward I (correct) or reward 0 (incorrect). Use the question and these ideas deduce the correct choice. Finish all your reasoning, then NEW line output exactly letter (the answer) and nothing else. Your final output MUST the format boxed <leeter>>) where <letter> is exactly one of A, C, D.
System Prompt: MATH (Free-form Answer)
are an Al mathematician. All content you output MUST be in English. Below are compressed solution ideas from previous attempts; each idea is tagged with reward I (correct) or reward 0 (incorrect). Use the question and these ideas deduce the correct answer. Finish all your reasoning, then NEW line output exactly answer (the answer) and nothing else. Your final output MUST in format boxed( <answer>).
gives an explicit, task-level meaning to the reward tags (r E {0, 1)), telling the model to learn from high-reward ideas and discount low-reward ones so the model can use feedback without any gradient updates test time.
D.2 SUMMARIZATION PROMPTS
These prompts compress full CoT into short summary that retains the high-level strategy.
Summarization Prompt: AIME AMC
Provide concise summary of the reasoning in the answer below. Do NOT add any introductory phrases extra explanations. Omit all numerical calculations. The summary must be self-contained, no more than 100 tokens.
32

---

Published as a conference paper at ICLR 2026
If there is final numeric result, include it the end in the format boxed <number>> (decimal only, no fractions, variables, or text). If there numeric answer, do not output boxed [Answer start] (... raw model answer.. ... [Answer end] Summary:
Summarization Prompt: GPQA
Provide a concise summary of the reasoning the answer below. Do NOT add any introductory phrases or extra explanations. Omit all numerical calculations. The summary must be self-contained, no more than 100 tokens.
there a final answer choice, include at the end in the format boxed (must be exactly one A, B, C, D, no extra lext). there final answer, not oulput boxed [Answer start] (... raw model answer . [Answer end]
Summary:
Summarization Prompt: MATH
Provide a concise summary of the reasoning in the answer below. Do NOT add any introductory phrases extra explanations. Omit all calculations unless essential understanding. The summary must be self-contained, more than 100 tokens. there a final answer, include at the end the format boxed (<aseed<a [Answer start] (... raw model answer... } [Answer end] Summary:
They replace long CoTIs with short strategy summaries, keeping only decision-relevant logic while dropping arithmetic details, so the history fits in-context and remains useful across rounds.
D.3 QUALITATIVE CASE STUDIES
Total Prompt
You are an Al mathematician. All content you output MUST be in English. Below are compressed solution ideas from previous attempts; each idea is tagged with reward I (correct) or reward 0 (incorrect). Use the question and these ideas deduce the correct numeric answer. Finish all your reasoning, then on NEW line output exactly number (the answer) and nothing else.
Your final output MUST be in the format boxed <number>> ). where <number>> is the final numeric answer only (no expressions, variables, or additional text). The content inside boxed <namber>] must be decimal number, not fraction any other form.
Every morning Aya goes for 9- 9-kilometer-long walk and stops coffce shop afterwards. When she walks at a constant speed of kilometers per hour, the walk takes her 4 hours, including minutes spent in the coffee shop. When she walks s + kilometers per hour, the walk takes her hours and 24 minutes, including t minutes spent in the coffee shop. Suppose Aya walks at + 1/ kilometers per hour. Find the number of minutes the walk takes her, including the minutes spent in the coffee shop.
ideas (reward 0):
[0]- When Aya walks at coffee shop, is approximately 348 minutes.
kilometers per hour, the total time taken, including the t minutes spent in the
Thus, the answer is 348.0
good ideas (reward I):
[0]- To determine how long it would take Aya walk from her house the park and back, including the time spent in the coffee shop, we set up two equations based on the given information. By solving these equations, we found the speed and the time i spent in the coffee shop. Then, we calculated the time required when Aya walks at + 2 km/h, which turned out to be 204 minutes.
Thus, the answer is 204.0
[1]- By analyzing the given conditions and using algebraic manipulation, deduced that the value of z is which when multiplied by 340 gives the final answer of 204.
Thus, the answer is 204.0
[2]- solved the system of equations to find the walking speed s the time spent the coffee shop. Then, we calculated the total time required when Aya walks at speed of s + km/h, including the time spent in the coffee shop. The final answer is 204.0 minutes.
Thus, the answer
204.0
33

---

Published as a conference paper at ICLR 2026
Blue marks the system instructions that enforce output schema and reward semantics. Orange marks the problem statement. Red marks reward-0 (incorrect) ideas retained as counterexamples. Green marks reward-1 (correct) ideas that ICPO prioritizes during selection.
How ICPO appears in this example. ICPO samples multiple CoTs, summarizes them, and assigns rewards based on self-consistency. The green items represent low-entropy agreement on the solution path that eliminates the fixed coffee time and yields the evaluation at as 204.0 minutes. Entropy- based filtering downweights the red outlier (= 348 minutes) that improperly scales total time and ignores the fixed offset. The final evaluator, conditioned on the schema and the curated ideas, outputs the correct numeric answer, 204.0
USE OF LARGE LANGUAGE MODELS
We used LLMs solely as a writing assistant for minor grammar and phrasing corrections during manuscript preparation. LLMs were not involved in research ideation, experiment design, data analysis, or result interpretation.
34


Reviews:
[{"review_id": "W8QLsFBkyu", "reviewer": "Reviewer_gZmq", "summary": "This paper introduces In-Context Policy Optimization (ICPO), a theoretical framework that explains how large language models can self-improve during test-time by iteratively refining their responses without parameter updates.The paper formulates multi-round self-refinement in LLMs as a form of in-context policy optimization, where the model treats its previous responses and associated rewards as contextual experience to adjust future outputs. This extends existing in-context learning theory from supervised prediction to policy optimization with bandit feedback. One the theoretical side, the authors prove that a single-layer linear self-attention transformer, when pretrained using a Fisher-weighted logit-matching objective, can provably imitate a policy optimization algorithm for linear bandits, thereby establishing a mechanistic explanation for the emergence of self-reflection in LLMs. Based on the theory, the paper proposes ME-ICPO, a practical inference-time algorithm that performs iterative response refinement using self-assessed rewards and entropy-based selection to ensure robustness to reward noise. Across standard mathematical reasoning benchmarks, ME-ICPO achieves competitive and often state-of-the-art test-time performance while maintaining affordable inference cost, demonstrating that test-time scaling can be improved without parameter fine-tuning.", "questions": "How can we handle the situation that the model itself cannot score or rank its own responses? How can we handle mis-aligned reward heuristics that the incorrect reasonings are being rewarded or reinforced? How to mitigate such caveats?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors derive provable guarantees showing that a linear self-attention transformer, when trained under a Fisher-weighted objective, can imitate the behavior of a policy optimization algorithm in a linear bandit setting. This is a novel result from the theoretical perspective. \n\nThe paper proposed Minimum-Entropy ICPO (ME-ICPO) algorithm which demonstrates a practical and implementable version of in-context policy optimization. It integrates entropy-regularized response selection and self-assessed rewards, leading to consistent empirical improvements in mathematical reasoning tasks. The experimental results are strong and align with the theoretical insights. \n\nBy modeling the self-reflection and iterative response refinement as In-Context policy optimization problem, the paper offers a clear mechanistic and mathematically grounded explanation for self-improvement phenomena observed in LLMs.", "weaknesses": "The effectiveness of ME-ICPO depends on choices such as number of refinement rounds, sample count per round, and entropy thresholds. Tuning those hyperparameters are non-trivial and might heavily depend on model sizes and datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iExplonOUP", "reviewer": "Reviewer_XA9N", "summary": "This paper introduces ICPO (In-Context Policy Optimization), a framework showing that a one-layer linear self-attention model can imitate a policy-optimization algorithm under a Fisher-weighted training objective. Motivated by insights from the theoretical results, the paper also proposes a practical algorithm, ME-ICPO, which performs multi-round generation, self-assessment with majority voting, chain-of-thought summarization, and minimum-entropy response selection to enable reward-aware prompting and principled feedback selection without further training. The empirical results show improvements on mathematical reasoning benchmarks.", "questions": "1. How does the method perform on recent long-CoT models, for example Qwen3-4B-Instruct? And since this is a training-free method, how does it perform even on frontier models, such as GPT-5 or Gemini-2.5-Pro?\n\n2. How does ICPO extend to harder tasks—for example HMMT, APEX-shortlist tasks—or tasks without final-answer executability, or where the final answer is not discrete?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. It is interesting to formulate ICPO as a bandit-style policy optimization approach. The theoretical grounding for in-context self-refinement is potentially impactful if the claims hold in more realistic settings. \n\n2. The framework and algorithm diagrams are well-organized, and the writing is mostly easy to follow.", "weaknesses": "1. The theoretical framework in Section 4 uses a linear bandit abstraction and a simplified linear self-attention model, whereas ME-ICPO is demonstrated with models like Qwen2.5-Math-7B. It is not clear how these theoretical assumptions connect to the practical model choices.\n\n2. ICPO requires iterative sampling, which implicitly increases inference compute. The paper only compares with the base model; since this is technically a prompting technique, it is unclear how this improvement differs from test-time scaling methods such as Tree-of-Thoughts, ReAct, and Monte-Carlo Tree Refinement, or from lightweight training methods such as GRPO and TTRL.\n\n3. The ME-ICPO also seems limited. Majority voting requires that (1) the model has sufficient capability to solve the task, (2) reasoning verification is cheap and easier than generation, and (3) the majority answer correlates with correctness. These appear to be strong assumptions that many real tasks may not satisfy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mkkCEXnuKY", "reviewer": "Reviewer_4bex", "summary": "- The paper introduces ICPO, framing multi‑round self‑reflection at inference as **in‑context policy optimization** that uses self‑assessed or external rewards without parameter updates. The authors prove that, under **Fisher‑weighted logit‑matching** pretraining, a **single‑layer linear self‑attention (LSA)** model can imitate a policy‑optimization algorithm for linear bandits. Building on this, they propose **Minimum‑Entropy ICPO (ME‑ICPO)**, a practical test‑time algorithm that iteratively samples candidates, assigns self‑assessed rewards, and selects low‑entropy, high‑confidence responses; **majority voting** is used to robustify the reward signal.\n- Experiments use **Qwen2.5‑Math‑7B** and **Qwen2.5‑Math‑1.5B** across **AIME‑2024, AMC, and MATH L1–L5**, etc., demonstrating the power of ME-ICPO.", "questions": "1) Although the proofs target single‑layer LSA, **can ME‑ICPO be safely applied to general (non‑LSA, multi‑layer) Transformers in practice**?\n2) The paper should clearly articulate the scope and the required setup for ICPO vs ICPO with LSA. In Section 4, it seems ICPO is only defined with LSA, is this correct? If so, is ME-ICPO only defined for LSA models?\n3) \"that with sufficient pretraining under a novel Fisher-weighted logit-matching objective, a single-layer linear self-attention model can provably\nimitate policy-optimization algorithm for linear bandits\", does ME-ICPO described in Algorithm 1 **require** such pretraining to work? \nIt is not clear whether ME-ICPO can be used as a test-time only method OR it has to be bundled with the specific pretraining procedures.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1) **Clear mechanistic link:** a theoretically grounded account connecting pretraining under a Fisher‑weighted objective to in‑context policy‑optimization behavior in an LSA. \n2) **Practicality:** ME‑ICPO yields strong math‑reasoning gains with gradient‑free test‑time optimization; **Mean@16 can surpass the base model’s majority‑vote upper bound**, and adding majority vote on ME‑ICPO output brings further gains.", "weaknesses": "- **No variability reported in Table 1.** Table 1 reports only point estimates (Accuracy and Mean@16) with no variability across multiple runs; please add mean±std over, e.g., 5 seeds. \n- **Theory scope.** Guarantees apply to a **single‑layer LSA** and **linear bandits**; practical models may not be LSA, so the theoretical guarantees do not directly cover the standard non-LSA archetictures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7Y8OCkaPT7", "reviewer": "Reviewer_jkwj", "summary": "This paper introduces the In-Context Policy Optimization (ICPO) framework for understanding the phenomenon of test-time scaling in LLMs, i.e., when models improve their responses through multi-round self-reflection without parameter updates. The main idea is to model this process as an agent that optimizes its response (action $x$), based on a history of in-context action-reward pairs $(x_t, r_t)$, effectively performing policy optimization within the context window. The main theoretical contribution proves that a single-layer Linear Self-Attention (LSA) transformer can imitate a specific policy optimization algorithm: a variant of FTRL for linear bandits. This imitation is achieved when the LSA model is pre-trained using a Fisher-weighted logit-matching objective. The authors present this as a foundational proof of how an attention-based architecture can learn to perform in-context optimization. Building on this theoretical result, the authors propose a practical, gradient-free inference-time algorithm called Minimum-Entropy ICPO (ME-ICPO), which they leverage for complex mathematical reasoning tasks. At each round, ME-ICPO generates multiple candidate solutions, assigns self-assessed rewards using majority voting on the final answers, summarizes the reasoning paths (Chain-of-Thought) to manage context length, and selects the next reasoning step to add to the context based on a minimum-entropy criterion. The experiments show that ME-ICPO achieves competitive, and in some cases state-of-the-art, performance on benchmarks like AIME, AMC, and MATH.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. To best of my knowledge, the paper introduces a first framework to formally model in-context self-improvement as a policy optimization problem. The authors provide a new and principled perspective on the mechanisms underlying test-time scaling.\n\n2. The main theorems establish population-level equivalence to an FTRL-like algorithm and provide finite-sample guarantees for learning this algorithm from data. They also analyze the stability of the learned policy to reward perturbations. \n\n3. The proposed ME-ICPO algorithm is an effective method. It demonstrates substantial and consistent performance improvements over specialized base models on multiple mathematical reasoning benchmarks.\n\n4. ME-ICPO is a gradient-free, inference-time-only algorithm. This makes it significantly more computationally efficient (particularly in terms of VRAM) than methods that require test-time backpropagation, such as TTRL. This practicality makes it a more accessible method for improving LLM performance.", "weaknesses": "Please respond to weaknesses, I will consider raising my score from 6 to 8 if all weaknesses are addressed -- this is a good paper!\n\n1. The most significant weakness is the large abstraction gap between the theoretical model and the practical application. The theory is built on a single-layer Linear Self-Attention model solving a linear bandit problem, whereas the experiments are run on deep, multi-layer, non-linear transformers performing complex, structured reasoning. The paper's claim to \"explain\" the mechanism of self-reflection is an overstatement. The theory provides an elegant proof-of-concept that the attention mechanism can implement a form of optimization, but it does not and cannot prove that this specific linear mechanism is what underlies the sophisticated self-correction abilities observed in models like Qwen2.5-Math. The paper would be stronger if it framed the theory more cautiously as an inspirational, minimal model that demonstrates a core computational capability, rather than a direct explanation of an emergent phenomenon.   \n\n2. The ablation study shows that the minimum-entropy selection criterion is the most critical component of ME-ICPO. However, the paper's justification for this heuristic is purely intuitive, suggesting it avoids \"corrupted\" responses and encourages \"diversified\" ones. This justification is somewhat vague and potentially self-contradictory (low entropy implies low diversity). The success of this heuristic may be domain-specific. For mathematical problems with a single correct reasoning path, low entropy (high agreement among future sampled paths) is likely a strong proxy for correctness. However, for more open-ended or creative tasks, the optimal path might be one that leads to a rich and diverse set of possibilities (high entropy). The paper lacks a more formal justification for this algorithmic choice and does not compare it against more standard selection criteria from RL, such as simply selecting the candidate with the highest self-assessed reward.   \n\n3. Modeling a multi-round reasoning process as a sequence of K-armed bandit pulls is a major simplification. This abstraction ignores the stateful and compositional nature of logical deduction. Each step in a mathematical proof is not an independent choice from a fixed set of K options; rather, it generates a new logical state that constrains all subsequent steps. A more faithful, albeit likely intractable, model would involve a contextual bandit or a full Markov Decision Process (MDP). The paper should explicitly acknowledge and discuss the limitations of this memoryless abstraction and how it impacts the interpretation of the theoretical results.\n\n4. (minor) There is clear over-abuse of spacing in the paper in terms of vspaces. While I realize all authors use this, the authors should not abuse it. Please remove these if your paper is accepted.\n\nTypos and grammatical errors:\n- ...the model's ability to digest the in-context information to improve their response. $\\rightarrow$ ...to improve its response.\n- Such an in-context information can be... $\\rightarrow$ Such in-context information can be...\n- ...without answering why these ability emerge... $\\rightarrow$ ...without answering why this ability emerges... (or ...why these abilities emerge...)\n- ...learn to optimize it's behavior x by optimizing it's policy... $\\rightarrow$ ...learn to optimize its behavior x by optimizing its policy...\n- ...how LLM leverage the in-context actions... $\\rightarrow$ ...how LLMs leverage the in-context actions...\n- ...to improve it's response \\$x\\_\\{t+1\\}\\$... $\\rightarrow$ ...to improve its response \\$x\\_\\{t+1\\}\\$...\n- ...generating it's response \\$x\\_t\\$ and receives... and then improve it's response... $\\rightarrow$ ...generating its response \\$x\\_t\\$ and receives... and then improves its response...\n- ...into it's policy optimization process and to gradually improves its response. $\\rightarrow$ ...into its policy optimization process and to gradually improve its response.\n-  ...where the agent generates and improve it's response... $\\rightarrow$ ...where the agent generates and improves its response...\n- ...denotes its norm \\$l\\_2\\$ For a matrix A. $\\rightarrow$ ...denotes its norm \\$l\\_2\\$. For a matrix A.\n-...during the test-time can improve... $\\rightarrow$ ...during test-time can improve...\n- ...including the Monte-Carol Tree Search... $\\rightarrow$ ...including the Monte-Carlo Tree Search...\n- ...where the LLM evaluate their own response... $\\rightarrow$ ...where the LLM evaluates its own response... (or ...LLMs evaluate their own...)\n-  ...by directly assume the LLM's ability... $\\rightarrow$ ...by directly assuming the LLM's ability...\n-  ...trained linear self attention can implement... $\\rightarrow$ ...trained linear self-attention can implement...\n- ...multi head constructions... $\\rightarrow$ ...multi-head constructions...\n-  ...in which first layer heads preprocess... $\\rightarrow$ ...in which first-layer heads preprocess...\n-  ...rare recent literature have covered... $\\rightarrow$ ...rare recent literature has covered...\n- ...optimize it's policy \\$x\\_t\\$... $\\rightarrow$ ...optimize its policy \\$x\\_t\\$...\n- ...dataset is generating from the policy... $\\rightarrow$ ...dataset is generated from the policy...\n- ...similar with the Follow-the-Regularized Leader... $\\rightarrow$ ...similar to the Follow-the-Regularized Leader...\n- ...defined by \\$s\\propto log~p\\$ In the following... $\\rightarrow$ ...defined by \\$s\\propto log~p\\$. In the following...\n-  ...prefix of trajectory up to... $\\rightarrow$ ...prefix of trajectory \\$\\tau\\$ up to...\n-...exploration parametery is wide... $\\rightarrow$ ...exploration parameter \\$\\gamma\\$ is wide...\n-...and p is a normalization factor... $\\rightarrow$ ...and \\$\\rho\\$ is a normalization factor...\n-  The LSA model parameterized by starts with... $\\rightarrow$ The LSA model parameterized by \\$\\theta\\$ starts with...\n- ...the LSA model updates it's policy... $\\rightarrow$ ...the LSA model updates its policy...\n-  ...corresponding K dimension... $\\rightarrow$ ...corresponding \\$K\\$ dimensions...\n- The expected matrix I is inspired by... $\\rightarrow$ The expected matrix \\$\\Gamma\\$ is inspired by...\n- The Fisher-weighted loss provide new loss... $\\rightarrow$ The Fisher-weighted loss provides a new loss...\n- ...that common KL loss between... $\\rightarrow$ ...that the common KL loss between...\n- ...using the KL loss enable the transformers... $\\rightarrow$ ...using the KL loss enables the transformers...", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)