Reinforced Context Order Recovery for Adaptive Reasoning and Planning
Authors: Long Ma
Fangwei Zhong
Yizhou Wang
Venue: NeurIPS 2025 poster ()
Keywords: Adaptive Token Order
Autoregressive Models
Diffusion Models
Reinforcement Learning
Primary area: deep_learning
TL;DR: We propose an algorithm to recover the correct generation order from textual data and solve reasoning and planning problems adaptively without order annotations.

Abstract: Modern causal language models, followed by rapid developments in discrete diffusion models, can now produce a wide variety of interesting and useful content. However, these families of models are predominantly trained to output tokens with a fixed (left-to-right) or random order, which may deviate from the logical order in which tokens are generated originally. In this paper, we observe that current causal and diffusion models encounter difficulties in problems that require adaptive token generation orders to solve tractably, which we characterize with the $\mathcal{V}$-information framework. Motivated by this, we propose Reinforced Context Order Recovery (ReCOR), a reinforcement-learning-based framework to extract adaptive, data-dependent token generation orders from text data without annotations. Self-supervised by token prediction statistics, ReCOR estimates the hardness of predicting every unfilled token and adaptively selects the next token during both training and inference. Experiments on challenging reasoning and planning datasets demonstrate the superior performance of ReCOR compared with baselines, sometimes outperforming oracle models supervised with the ground-truth order.

Reinforced Context Order Recovery for Adaptive Reasoning and Planning
Long Ma Academy for Advanced Interdisciplinary Studies, Peking University, Beijing, China malong@pku. edu.
Fangwei Zhong
School of Artificial Intelligence,
Beijing Normal University, Beijing, China
fangweizhong@bnu. wtweeiizhon@bbueed..ccc edu. cn
Yizhou Wang
School of Computer Science, Institute for Artificial Intelligence, State Key Laboratory of General Artificial Intelligence, Peking University, Beijing, China yizhou. wang0pku. edu. cn
Abstract
Modern causal language models, followed by rapid developments discrete dif- fusion models, now produce wide variety of interesting and useful content. However, these families models are predominantly trained to output tokens with fixed (left-to-right) random order, which may deviate from the logical order in which tokens are generated originally. In this paper, we observe that current causal and diffusion models encounter difficulties problems that require adaptive token generation orders to solve tractably. which we characterize with the V-information framework. Motivated by this, we propose Reinforced Context framework to extract
Order Recovery (ReCOR), Seinforremeen--eaaminn--- adaptive, data-dependent token generation orders from text data without annota- tions. Self-supervised by token prediction statistics, ReCOR estimates the hardness of predicting every unfilled token and adaptively selects the next token during both training and inference. Experiments on challenging reasoning and planning datasets demonstrate the superior performance of ReCOR compared with baselines. sometimes outperforming oracle models supervised with the ground-truth order
Introduction
Text generation models have seen remarkable advancements in the past few years, with causal language models (CLMs) trained by next-token prediction taking the lead 22 followed by more recent discrete diffusion models 4. 5 61. These classes of models have demonstrated impressive capabilities across wide range of tasks, from writing code to executing tasks as agents 7 g 9].
Despite the developments, current models still encounter significant challenges when faced with complex reasoning and planning problems that require a horizon and flexible decision-making process. A crucial part of this challenge lies in the token generation order [0]. As illustrated in Fig. adaptive reasoning tasks like Sudoku, some cells could be hard to predict instantly, depending on other cells to be filled first and provide additional constraints to eliminate candidates. However, CLMs always follow a rigid left-to-right generation paradigm, encountering many intractable tokens along the way. In contrast, humans rarely solve complex reasoning problems in strictly linear fashion; we tackle the easiest parts first and use those insights to address progressively more challenging ones. Taking note of this issue, existing works either explicitly train the model to predict the easiest next
Project page: httpe////sstee.gogllll.o.//////e
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

3
5
Let's solve this puzzle one by one... But I'm stuck at the first!
7
Be flexible: we can always fill the easy ones first!
6
8
5
3
8
8
9
6
2
6
2
7
5
7
17
1
2
9
2
6
2
4
2
Figure Illustration of ReCOR (right) compared with standard causal language modeling (left). While causal language modeling always tries to produce tokens left-to-right, ReCOR estimates the hardness of each token and adaptively prioritizes the easy ones without external supervision.
token 10 or leverage the properties of masked diffusion models (MDMs) 12 to adaptively decode easy tokens during inference. However, these methods require additional annotations or introduce large distribution shifts between training and inference, hurting performance.
Facing these challenges, in this paper, we introduce Reinforced Context Order Recovery (ReCOR), self-supervised framework that learns to adaptively determine the optimal token generation order without explicit order annotations. Our approach is motivated by the insight that the hardness of predicting different tokens varies dramatically conditioned on the current context, which we characterize using the framework of predictive V-inffrmation.. To operationalize the objective derived under this framework, ReCOR casts order prediction as decision-making problem and trains policy that adaptively selects the next token (Fig. D. ReCOR jointly optimizes the token prediction model and the order prediction policy, generating rewards with the former as self-supervision for the latter. Furthermore, unlike previous approaches that apply adaptive strategies for inference only, ReCOR follows the same distribution of order during both training and inference. This ensures that the model not only adapts its generation behavior during inference but also benefits from learning informative, tractable token prediction tasks during training. the experiments, we demonstrate the effectiveness of ReCOR and logic puzzles, where ReCOR consistently outperforms the state-of-the teefthh-arttt methods.
variety of reasoning and planning tasks, including arithmetic problems
The contributions of our paper are fourfold: 1) We observe and characterize the token ordering problem using the framework predictive Y-informmtionn and propose a corresponding objective; 2) We propose an RL--based formulation and algorithm for optimizing the objective and obtaining an adaptive order predictor; 3) We empirically validate ReCOR on multiple reasoning and planning benchmarks, demonstrating state-of-the-art performance that is competitive even better than oracle models with access to ground-truth orders; 4) We offer an analysis on specific failure modes of inference-only adaptive methods, showing the necessity of our training-time adaptive designs.
Related Work
(Any-order) Autoregressive Models and (Discrete) Diffuusion Models. Autoregressive and diffusion models are currently the two dominant families of generative models in various domains [3] 14 15).
For textual modeling, recent years, there has been a seismic rise popularity for large language models 2 of which the vast majority are trained using the next-token prediction objective, or
causal language modeling, c.g. GPT [16] [7]. Alternatively, discrete variants [5] 4 of diffusion models [18
19 20 like masked diffusion models (MDMs) apply denoising objective to randomly corrupted data to learn the underlying structures. Furthermore, any-order autoregressive models generalize the causal modeling paradigm by allowing arbitrary orders masks during training and inference. Compared with these approaches, ReCOR uses a learned adaptive order during both training and inference, avoiding issues with fixed or random orders. We primarily compare to adaptive inference variants of MDMs [2 since AO-ARMs are trained similarly.
(AO-ARMs) P 2
Token Ordering and Reflections on Causal I Language Modeling, Within the active research field of language models, there have long been criticisms and attempts at improvements regarding various aspects of the causal modeling paradigm, including autoregressive inference (23). teacher-forcing

---

training [24], and the left-to-right order [25] In particular, current LLMs have been found to struggle
with many complex reasoning and planning problems, especially adaptive ones [26] T0 27 28].
We focus token ordering, motivated by the fact that these reasoning problems often require intricate, data-dependent orders be tractably solved. Recently, adaptive inference methods based on MDMs [2 were proposed to address this problem, combining random masking at training time with selective token inference. We note that adaptive orders are needed during both training and inference, and design ReCOR to automatically recover the correct order with self-supervision.
3 Preliminary
3.1 Problem Setup
We focus our attention the classic sequence-to-scquencc setting, where the goal is to learn conditional distribution p(y x) with x E X' _ T as the prompt (space), E 3 S T as the response (space), and T as the token vocabulary. For simplicity,
assume prompts of length N and responses of length M: x = (I1, T2.... .y) E X' = = . = (V1- V2. yM E y TM which can be achieved without loss of generality through padding. A training dataset D [(x is available with i.i.d. samples y) ~ p(x)p(y x).
3.2 Existing Approaches
In recent years, causal language models (CLMs) have witnessed giant wave of interest, trained using the prevalent next-token prediction (NTP) objective:
M
min (xy))
Σ
log popo yxy
with po(- ys)) parameterized as an autoregressive sequence model (e.g. GPT 16
As a generalization of CLMs, any-order autoregressive models (AO-ARMs) 21 22 perform
"next-token" predictions under an arbitrary generation order p SM instead of the causal order:
min (x..))
Po x. x.Y Pi)
where Up is the distribution of training orders, usually taken as the uniform distribution over all M- permutations SM (21] or further canonicalized using simple rules (22]. Note that p the generation order and distinct from the raw content order, which is still retained; i.e. for y abc and p1 = 3. the partially generated response at the first step should be **C instead of c**.
Alternatively, masked diffusion models (MDMs) 55 special [MASK] token forward process Grjo (y' y) that masks each token independently at random, then learn to predict the masked tokens with denoising objective (constants omitted):
replace some response tokens with
min
Σ lb [MASK] log Po(s: xxy',tt,)
(3)
which is also often implemented with a transformer.
Problem Analysis and Method
This section is structured as follows:
We first motivate our focus on generation order by characterizing
token hardness [Sec./4.1 and propose to cast order recovery as a decision- problem (Sec. (.2).
The core training algorithm for ReCOR is presented in Sec. 4.3 with architectural details in Sec. 4.4
4.1 Generation Order and Token Hardness
We begin our discussion about the order problem by noting that prior works mostly handle generation orders passively during training, delegating the issue to fixed data-independent strategy (Sec. 3.2
3

---

always left-to-right for CLM and uniformly random for AO-ARM and MDM). This is justified from probabilistic standpoint, where any joint distribution over a series of random variables can be decomposed arbitrarily into the product of series of conditional distributions via the chain rule: P(11, U2.-.., yas) I P(Up; Sp...... Dop......p.....W.. Vp SM- Furthermore, for problems with a unique solution, the prompt can completely determine each response token, and we have H(Y,, X, (Y,, b<<)) HYp,, X) 0. I(Yp,: (Y hes X) 0, V. i, predicted response tokens do not bring additional information, also implying that order is irrelevant.
[1] for hard reasoning and planning problems, plain causal modelling frequently encounters computationally intractable intermediate steps and fails dramatically; tailored generation order is required for each instance to make the problem tractable in practice. Consequently, we choose to explicitly model the generation order p as an unobserved variable to be recovered from the plain-text-only training data.
However, as illustrated in Fig. and observed by prior works 10 12
Under such formulation, the problem then turns to defining and finding good p. Intuitively, we would like to start with the easy parts of the response and work our way to the harder parts step-by-step. utilizing the partially generated solution as a form of chain-of-thougee to guide later generations. We formalize this intuition using the framework of predictive Y-information (30]:
Definition 1 (Predictive V-information (30). Let V S { A [Ø} + As} be a predictive function class that predicts the distribution of random variable taking values over B, optionally given the value of another random variable over A under some regularity conditions.
Define the predictive conditional V-entropy of random variables A, B as
Hy(B [ A) inf B.,b-A log f[a] feV
(4)
(y Ø) inf
E log [b] JEV
(5)
Then the predictive V-information from random variable to B is defined as
lv(A + B3 = Hy(B Ø) Hv(B A).
(6)
Def. T captures the hardness of a token given the current context under computational constraints, which evades standard probabilistic and information-theoreiic arguments. Given prompt X and set of already predicted response tokens| {Y,; Ji<s, an easy token Y,, would have a large [v(X, (Y hiss + Y,,) whereas a hard token would have a small V-information. Building on this characterization, we set our overall objective maximize the following cumulative predictive V-information using learned autoregressive p p-generator parameterized by A:
max
Eprrpp((1X..tX bed be (Yp + Y,,
Note that Def. [ of V-information contains an optimization problem. Naively training a new predictor to solve the optimization problem within (x(X, (Y,, Jiاز Y),) for every Ps: would require an exponential number of training instances. For tractability, we set the predictive function class V to the family of autoregressive token generators conditioned on arbitrary contexts and parameterized by sb and share it across inner subproblems for all p. yielding the following objective up to constants:
max E 0.6
(x,))~D
Σ E pppopp log x, Kpyye Di)
(8)
4.2 Order Recovery as Decision-making Problem
We now proceed to solve Eq. for 0 and the associated v.. Although [30) proposed a method for structure learning based on the Chu-Liu algorithm (3I, it assumes that the edge weight depends only on the two vertices connected by the edge, while under our setting, the likelihood depends on all random variables in the context with an exponential number of combinations. Furthermore, structure learning requires a uniform structure over the whole dataset, while the correct generation order can be data-dependent and vary between instances (e.g. Sudoku (10 in Fig. [].
Facing these challenges, we formulate the recovery of p as a decision- problem and employ reinforcement learning (RL) techniques 332
33 34
to train the p-generator. We construct the

---

Acting in the Environment
Training with Actions
yp
x
3
Dataset D
3
2
5
SelfAttn Qa
3
P.
CrossAttn
2
5
3
Pv
p
Policy 7Eg
by,
Select
R
Lqu
L
3
5
0
1
yp
0 6 1
Token
Position
Vector
Scalar
Figure 2: Illustration of the training procedure ReCOR. ReCOR first rolls out the order prediction policy Te autoregressively
sampled data (x,y) y) to obtain the actions p, (left), then performs parallelized training on the sampled actions (right). The token predictor py is trained to answer queries generated by Ty using LLM while Qo is supervised by reward signals from py using LsQL.
following Markov decision process (MDP) P, where S == X' UICIMIT' is the state space with the prompt and a partially generated response; A := [MM] is the action space corresponding the position of the next token to generate; / is the transition that adds the newly generated token to the state; Rt is the reward function as in Eq. 14 and 15 P1 p(x) x 8(2) is the initial state distribution consisting of a sampled prompt and an empty response; and 7 [0. 1] the = Pt. which transitions to (x, (x.Ypsa yp.)) with Up. taken from the training data or sampled from P x,Y Pt) during inference.
discount factor. Concretely, at time step S t < M,
(x, yp..), ag
Under this formulation, the objective is to train policy me S + A satisfying
Σ Ex'
max Es, ap (st Rs (144 (e) 0
(e)
(9)
which recovers Eq. with (- se) = po x. x.,, R(st, a.) log Pai (Upp x.ypec Pt) and
y 1. This formulation also allows more room for design choices, e.g. RL algorithms, discount factor, and alternative reward functions; see Sec. 4.3 below for details.
Equipped with ma for order prediction and py for token prediction, we now perform inference by sampling ~ (- x, yyype )c:) and Hr. ~ (- x, 8 autoregressively (see Alg. 2b.
4.3 Reinforced Training for Adaptive Order Prediction
With an RL formulation in place solve Eq. as illustrated in Fig. 2 soft Q-learning 333 that optimizes the following entropy-regularized discounted-retumn objective:
use a discrete version of
max Exa 6
Er [x)))
(10)
where a 2 is the entropy coefficient. This objective balances exploration and exploitation under maximum entropy RL framework by requiring the policy to act as stochastically as possible while optimizing returns. In addition, since our action space A = [M] discrete and tractable, efficiently compute wa and Vo from Qo without separately learning the policy and value function through function approximation. To optimize this objective, we use soft Bellman update 33 implemented with the following mean squared error loss:
LSOL-MSE (0) E.a.,re' [[Qo(s, a) -
a)(s,
(1I)
where Q0(s, a) = r Vo') r + ya log E. (xp(Qo(s,a is the Q-value target with gradients detached. Note that we do not target network techniques.
5

---

Alternatively. we can also use the binary cross entropy loss as value loss:
ESQL-BCE( := Br,,r [- exp Qo(s, a) Qo(s, a) - (1 cxp (o(,aa))))) log(1 - exp Qo(s,
(12)
which may deliver better robustness to outliers for 0 v exp Qo(s, a) - and is compatible with sparse reward functions, as discussed below.
Finally, we discuss the optimization of the token predictor parameters vp and the choice of reward function. We train the token predictor Pv online along with the policy me with the following (permuted) language modelling loss, echoing Eq.
g
M
(bb == (xxx))DD
Σ
log Pub (Ups x. X.y pi)
(13)
Note that Pe is trained on P sampled from the current policy me ensure that it always stays on-policy. As LIM is supervised, Pe learns much faster than the RL-trained Te. and we treat ei as always being near the optimality and approximate (y,X, {Y + yx)) with the current log (yppp x. yes Ps). As result, we may set the reward function to the (negated) perplexity:
R(*,, a+) = (x. pe) == log X, Yp
(14)
Empirically, we found the following thresholded, sparse reward with a better performance:
R(st, a+) = Y ).)))))))))) px) == log (pp Xxype 2 v/]
(15)
where 7, E (0, 1) is the probability threshold. Note, however, that the logarithm is undefined when (Yp, x, y. Pi) < 77. we use this reward in conjunction with LSQL-BEE, which only requires exp Rep(((x. (px. Pt) == (Bp. y.
2 17] that is well-defined.
Training pseudocode is presented in Alg. [ Furthermore, since the dynamics of the MDP we defined fully known, we can sample potentially multiple actions Pt.1... at the same state (x. Xpe)) and optimize fIM and Lsor for all of the K actions to improve learning quality. See Sec. 5.6 for details.
4.4 Multi-stream Architecture for Order and Token Predictions
With the algorithm of ReCOR described above, we now instantiate the p-generator To(- x, Y pcc) and token predictor (- X. y p Pt). For our experiments, we use GPT-2 []6] as the backbone, which could be replaced with other transformer variants. To encode a permuted response, we add the absolute positional embedding of GPT-2 at position P; to token Up.- while the prompt x is encoded as a regular left-to-right sequence before the response tokens. We use a full mask on the prompt x and a causal mask on the response yo to enable parallelized training and KV-cached inference.
This suffices to handle the sequential inputs x, ypp The token predictor P (- x, Spcco Pt) takes an additional scalar query position as input at every time step l. Ideally, we'd 'dd like the queries at every time step not to interfere with each other; consequently, we adopt multi- conceptually similar to XLNet (21) where main stream takes x. X. as input and performs self- attention as in regular transformers. and token query stream takes p as queries and cross-attends onto the main stream hidden states as keys and values without self-attention between queries. This guarantees that every query would be independent of others. We produce token predictions via a token head on the outputs of the token query stream. For order predictions the main stream outputs since order predictions do not depend on additional positions. To further enhance expressiveness, we can optionally use an order query stream instead of the main stream to perform order predictions. With C' learnable order query embeddings QI... as inputs, we forward the order query stream for C' times and concatenate the outputs before projecting to Q values. This design allows us to scale compute for better performance; see Sec. 5.6 Sec.5566 for details.
architecture
we can directly use
Experiments
In this section, we aim to answer the following questions with our experiments: 1) Can ReCOR solve arithmetic problems without special data preprocessing? 2) Can ReCOR solve reasoning and planning problems adaptively without annotations? 3) Do we need adaptive orders during training or for inference only? 4) How does ReCOR compare with the state-of-the-ar methods under fair inference compute settings? 5) Can the performance of ReCOR scale with more compute?

---

Table Performance of ReCOR and baselines on arithmetic datasets. ReCOR outperforms baselines and is competitive with the oracle.
Task
AR-GT
ReCOR
CLM
MDM
AdaMDM
ARG
0.994 = 0.003
0.987 + 0.007
0.017 + 0.002
0.035 + 0.010
0.174 = 0.036
MUL
0.999 = 0.001
0.964 0.007
0.594 + 0.018
0.943 + 0.006
0.951 = 0.038
Table 2:
Performance of ReCOR and baselines on puzzle datasets. ReCOR outperforms all ap- proaches, even including the oracle AR-GT supervised with ground-truth orders.
Task
AR-GT
ReCOR
CLM
MDM
AdaMDM
Sudoku
0,8718
0.9017 + 0.0004
0.0973
0.0688
0.8949
Zebra 0.9117
0.9905 + 0.0021
0.8031
0.769
0.985
5.1 Experiment Setup
We chose the following datasets to validate ReCOR's ' performance at adaptive reasoning: 1) Arith- metic datasets, including synthetic autoregression and multiplication, which are originally generated in reverse order during data generation. 2) Puzzle datasets, including Sudoku and 7ebra (10), which may require an adaptive, data-dependent order to solve. We use accuracy (full response match) as the main evaluation metric and report standard deviations over 3 training seeds.
We compare with the following representative baselines: 1) Causal language models (CLM) [16] that always follow left-to-right order. This baseline serves to demonstrate the necessity of non- left-to-right generation orders. 2) Autoregressive models supervised with ground-truth generation orders (AR-GT) I0 that is an oracle for ReCOR, as ReCOR does not have access to the ground truth order during both training and inference. 3) Masked diffusion models (MDM) [5] 4 that randomly mask the input and perform denoising, in effect trying to learn arbitrary orders. 4) Adaptive masked diffusion models (AdaMDM) [12 that use recent state-of-the-art adaptive inference methods.
5.2 Can ReCOR solve arithmetic problems without special data preprocessing?
We start with arithmetic problems generated using fixed, non-causal order. We employ a synthetic Autoregression ( (ARG) task where each response token depends on the prompt and all response tokens after it, and a multiplication (MUL) task between 20-digit and 2-digit numbers. It has been shown that CLMs often need manually reversed training data additional CoT annotations to solve certain arithmetic problems [35| 36 37 38) due to the reverse dependencies between tokens brought by carry digits. In this work, we are interested in whether ReCOR can automatically recover the correct order without manual preprocessing or additional annotations. Consequently, we apply the reversed order only to AR-GT and keep the natural order for the rest of the approaches.
As shown in Tab. I ReCOR achieves competitive performance compared with all baselines and outperforms CL.Ms by a large margin, showing that the plain next-token prediction objective is not sufficient for solving these problems. Notably, ReCOR competitive with AR-GT, which uses the ground-trute truth order during both training and inference. This demonstrates its ability recover order in a self-supervised manner. Furthermore, ReCOR outperforms (Ada)MDM significantly in Autoregression; we compare both methods in more detail in Sec. 5.4
5.3 Can ReCOR solve reasoning and planning problems adaptively without annotations?
In this section, study classic logic puzzles Sudoku and Zebra 100 The two datasets are examples of problems that require adaptive, data-dependent reasoning and planning, rendering them difficult for approaches that generate solutions with a fixed order. For example, in Sudoku, we may need to fill certain cells first in order to eliminate candidates for other cells and arrive at the correct answer, as illustrated in Fig. The cell dependency structure varies between instances, posing great challenge.
We show results on Sudoku and Zebra in Tab. 2 12]. ReCOR is the best among all approaches, outperforming even the oracle AR-GT. We remark
Values without standard deviations are reported by
7

---

Table Impact of suffix unmasking for MDMs on Autoregression. Suffix unmasking dramatically improves the performance of AdaSufMDM over AdaMDM, but still lags behind ReCOR.
ReCOR
SufMDM
AdaSufMDM
MDM
AdaMDM
0.987 = 0.007 0.040 + 0.005 0.536 + 0.051 0.035 + 0.010 0.174 + 0.036
Table 4: Performance of ReCOR and Ada(Suf)MDM under different inference compute settings. MDM underperforms ReCOR even with 10x compute, and is much worse with the same compute. We use AdaMDM on MUL and AdASufMDM on ARG, as AdaMDM fails on ARG.
Method
ReCOR
Ada(Suf)MDM
FLOPs
1x
Ix(T' 2)
10x = 20)
ARG
0.987 + 0.007
0.028 + 0.009
0.536 + 0.051
MUL
0.964 + 0.007
0.875 + 0.043
0.951 + 0.038
that at every time step, ReCOR can generate estimated reward for every unfilled cell, which is denser and more diverse signal than the single next position offered by the ground truth, potentially contributing to its superior performance. CLM and vanilla MDM perform similarly and are both worse than ReCOR and AdaMDM, showing that an adaptive generation order is indeed necessary.
5.4 Do we need adaptive orders during training or for inference only?
We look closer into the difference between ReCOR and AdaMDM [2 which are recent, state- of-the-art methods also focused on the order problem. Both works showed that MDMs encounter intractable sub-problems due to random masking during training, and proposed to use adaptive decoding strategies during inference to prioritize easy tokens and circumvent these hard scenarios. One of the core differences between ReCOR and AdaMDM that AdaMDM trained under random masking while ReCOR follows the recovered order during both training and inference. Thus problem arises: do we really need adaptive orders during training?
We answer the question in the affirmative with the Autoregression task (Tab. I. In Autoregression, each response token is generated conditioned on all response tokens behind it, and consequently, can only be tractably predicted when all of these tokens are present in the context. Similar properties hold for many long-horizon reasoning problems with multiple steps and dense dependencies between the steps. However, under random masking, the probability of such an event grows exponentially smaller with the length of the context, and MDMs rarely see any long, complete contexts during training. As result, MDMs are bad at the corresponding sub-problems, even though they are supposedly easy and tractable, leading to bad overall performance as evidenced in Tab. [ In contrast, ReCOR automatically recognizes and frequently visits these good sub- -robllmss during training, making sure that the on-policy token predictor pa can reliably solve them.
To further support this analysis, we propose (Ada)SufMDM, variant of (Ada)MDM that unmasks a suffix of length 1 24jr 1 during training. This operation dramatically boosts the probability of each "correct" sub-problem from (ceM) to 1(M-1)). Note that this fix applics only to ARG, where the correct order is known to be reversed; applying the operation in general would require ground- -ound--ruuhh order annotations. The updated results are shown in Tab. B SufAdaMDM obtains a large performance improvement and still outperforms SufMDM. demonstrating that correct orders important during both training and inference. However, there remains a gap between AdaSufMDM and ReCOR, which is likely due to the additional intractable sub- probleessem brought by random masking, as argued [1).. Our analysis complements their findings in that random masking leads to not only many intractable sub-problems, but also dearth of tractable ones.
5.5 How does ReCOR compare with the sfate-of-the-art methods under fair inference compute settings?
We show another advantage of our autoregressive design by comparing ReCOR with MDMs under different inference compute settings. The standard MDM inference setting used in our primary
8

---

0.8-
6.8
Wildelton Accuncty
Accunry
0.6
6.6-
64
C=B
6.2
Km
6.2
Eed
Cel
Km1 0.0
Main Stream
0.0
20000
40000
G0000
B0000
100000
120000
140000
20000
40000
60000
80000
100000
120000
140000
Training Gradient Step
Training Gradient Step
(a) Token query scaling with C 8.
(b) Order query scaling with K 8.
Figure 3: ReCOR's performance when scaling the number of token queries K (a) and order queries C' (b). Main Stream in (b) denotes using the main stream outputs without separate order query stream. ReCOR can improve its performance with more computation during training and inference.
experiments employs a large number of diffusion steps; consequently, many transformer forward passes are needed during inference. We note that for inherently serial problems where tokens need to be generated one-by onc-by-one by-one, MDMs require diffusion steps to the order of (MM)), ), bringing the overall time complexity to ((MM)). In comparison, since ReCOR is autoregressive and uses a causal mask for the response, we can use K KV-caching during inference, reducing the total number of complete transformer forward passes to
where is the number of order stream queries and a constant with respect to the context length, maintaining an O(M²) time complexity.
Empirically, we infer MDMs on the arithmetic datasets under alternative inference compute settings in Tab. 4 ReCOR uses C = 0 for these datasets, leading to an inference FLOPS equivalent to complete transformer forward passes. When restricting the FLOPs of MDMs the same level, the number of diffusion steps is limited to T = 2, dramatically reducing performance.
5.6 Can the performance of ReCOR scale with more compute?
In this section, we demonstrate scaling properties arising from the designs of ReCOR. As described Sec. 4 in contrast to the standard RL setting, ReCOR can take K > 1 actions at the same state and learn from all of these actions, since the MDP we defined can be easily simulated perfectly. Furthermore, we can use C' queries to the order query stream to obtain more expressive Q function. These designs allow us to scale compute to improve performance even with the underlying backbone parameter count fixed.
To demonstrate the scaling properties, we perform ablation experiments on Sudoku that vary one of K C from the primary setting K C 8. We show training curves of validation accuracy in Fig. 3 It can be seen that in both groups of experiments, performance improves monotonically with the varying parameter. As a special case, directly using the main stream outputs to compute Q values (Main Stream in Fig. 3b) achieves decent performance between C = 1 and 2 without using separate order queries. This makes a good, lightweight choice that we adopt for the arithmetic datasets.
5.7 How does ReCOR compare with large pretrained models?
To further demonstrate the performance of ReCOR, we compare with o3, state-of-the-art pretrained large reasoning model by OpenAI, our ARG dataset. Since we could fine-tune o3, we take the following two setups to evaluate its performance:
1. Prompting o3 with in-context examples and a detailed description of the ground-truth data generation algorithm.
2. Prompting o3 with in-context examples only.
We evaluate random samples from the test set and allocate token budget of 10000 tokens each. Note that in main experiments, only the generated training data is available to the models without

---

the ground-truth generation rule, corresponding to the harder setup 2 above. As result, under setup I with the ground-truth algorithm annotation, o3 achieves a success rate 0.3, where the majority of failures come from the depletion of token budgets. Even considering only the completely generated solutions, the success rate of o3 0.938 and still lower than ReCOR's s 0.987. Furthermore, o3 takes about 4 minutes on average to solve each problem instance, while ReCOR solves the entire 1000-sample test set within 10 seconds. Under setup with no access to ground-truth rules, o3 fails to recover the correct data generation algorithm and has success rate of O. We stress again that ReCOR operates under this harder setting without any additional annotations. The results show that even large pretrained reasoning models may struggle reason about token ordering effectively and efficiently, especially when relevant annotations are missing from the training data.
6 Conclusion and Limitations
In this paper, we introduce Reinforced Context Order Recovery (ReCOR), a ceinfrreeeentllerriinc based algorithm for automatically recovering the correct generation order from textual data without annotations. While modern text generation models are predominantly causal, they have been shown to fail adaptive reasoning problems due the presence of intractable tokens. We characterize this phenomenon with the framework of V-information and propose to optimize the corresponding objective to learn suitable data-dependent generation order. Subsequently, we operationalize this objective by introducing ReCOR with reinforcement learning setups, recovering unobserved genera- tion order from purely textual data in a self-supervised manner. Empirically, ReCOR outperforms strong baselines on various datasets, including recent adaptive inference approaches using masked diffusion models and oracle models supervised with the ground-trutt order.
There are certain limitations and future works for ReCOR. We primarily run experiments on reasoning and planning-relates related datasets, echoing the setups of our baselines. Going forward, we are excited to further scale up ReCOR to larger parameter counts and more diverse problems and datasets with more compute available. Applying ReCOR to modern LLMs and multi-modal models is the path forward to fully realizing its potential. The RL-based formulation also opens up vast space to integrate with other RL techniques and further improve the performance of ReCOR, which we did not exhaust due to limited resources. Furthermore, there could be more unobserved variables driving the generation of textual data beyond the generation order. The recovery of these variables presents exciting opportunities for future explorations.
Acknowledgements
This work was supported by the National Science and Technology Major Project (2022ZD01 14904), USFC-6247070125, NSFC-62406010, the State Key Lab of General Artificial Intelligence at Peking University, the Fundamental Research Funds for the Central Universities, and Qualcomm University Research Grant.
References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0877 2023.
[2] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
[3] Gemini Team, Rohan Anil, Sebastian Borgeaud, Jean-Bapteeeeee Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, Katie Millican, et al. Gemini: a family of highly capable multimodal models. arXiv preprint 2312. 2023.
[4] Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. Advances neural information processing systems, 34:12454-12465, 2021.
10

---

[5] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete stacesspaces. Advances in neural information processing systems, 34:1790 7791-17993,, 2021.
[6] Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint arXiv:2502.09992 2025.
[7] Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Advances in neural information processing systems, 35:24824 2282377--4837, 2022.
prompting elicits reasoning in large language models.
[8] Lei Wang, Chen Ma, Xueyang Feng, Zeyu Zhang, Hao Yang, Jingsen Zhang. Zhiyuan Chen, Jiakai Tang, Xu Chen, Yankai Lin, et al. A survey on large language model based autonomous agents. Frontiers of Computer Science, 18(6):186345, 2024.
[9] Daya Guo, Qihao Zhu, Dejian Yang. Zhenda Xie, Kai Dong, Wentao Zhang, Guanting Chen, Xiao Bi, Yu Wu, YK Li, et al. Deepseck-coder. When the large language model meets programming-the rise of code intelligence. arXiv preprint arXiv:2401.14196 2024.
[10] Kulin Shah, Nishanth Dikkala, Xin Wang, and Rina Panigrahy. Causal language modeling can elicit search and reasoning capabilities on logic puzzles. In The Thirty-eigtth Annual Conference on Neural Information Processing Systems, 2024.
[11] Jiacheng Ye, Jiahui Gao, Shansan Gong, Lin Zheng, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Beyond autoregression: Discrete diffusion for complex reasoning and planning. In The Thirteenth International Conference on Learning Representations. 2025.
[12] Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. arXiv preprint arXiv:2502.06768. 2025.
[13] Ling Yang. Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: comprehensive survey of methods and applications. ACM Computing Surveys, 56(4): 2023.
[14] Hanqun Cao, Cheng Tan, Zhangyang Gao, Yilun Xu, Guangyong Chen, Pheng-Am Heng, and Stan Z Li. A survey on generative diffusion models. IEEE Transactions on Knowledge and Data Engineering. 2024.
[15] Jing Xiong, Gongye Liu, Lun Huang, Chengyue Wu, Taiqiang Wu, Yao Mu, Yuan Yao, Hui Shen, Zhongwei Wan, Jinfa Huang, et al. Autoregressive models in vision: survey. arXiv preprint arXiv:2411.05900 2024.
[16) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[17] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 33:1877-1901, 2020.
[18] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[19] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[20] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[21] Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoe V
Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.

---

[22] Andy Shih, Dorsa Sadigh, and Stefano Ermon. Training and inference on any-order autoregres- sive models the right way. Advances in Neural Information Processing Systems, 35:2762-2775 2022.
[23] Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, et al. Faith and fate: Limits of transformers compositionality Advances in Neural Information Processing Systems, 36:70293-70332, 2023.
[24] Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. In Interna- tional Conference on Machine Learning, pages 2296-2318 PMLR, 2024.
[25] Li Du, Hongyuan Mei, and Jason Eisner. Autoregressive modeling with lookahead attention. arXiv preprint arXiv:2305. 12272, 2023.
[26] Karthik Valmeekam, Matthew Marquez, Sarath Sreedharan, and Subbarao Kambhampati. On the planning abilities of large language models-a critical investigation. Advances in Neural Information Processing Systems, 36:75993-76005, 2023.
[27] Subbarao Kambhampati, Karthik Valmeekam, Lin Guan, Mudit Verma, Kaya Stechly, Siddhant Bhambri, Lucas Paul Saldyt, and Anil Murthy. Position: LLMs can't t plan, but can help planning in LLM-modulo frameworks. In Forty-first International Conference Machine Learning 2024.
[28] Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. In The Twelfth International Conference on Learning Representations, 2024.
[29] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In International Conference on Machine Learning, pages 328 12849-3 PMLR, 2024.
[30] Yilun Xu, Shengjia Zhao, Jiaming Song, Russell Stewart, and Stefano Ermon. A theory of usable information under computational constraints. In International Conference on Learning Representations, 2020.
[31] Yoeng-Jin Chu. On the shortest arborescence of directed graph. Scientia Sinica, 14:1396-1400, 1965.
[32] Richard S Sutton. Reinforcement learning: An introduction. A Bradford Book, 2018.
[33] Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy- based policies. In International conference on machine learning, pages 1361- PMLR, 2017.
[34] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal 2017.
policy optimization algorithms. arXiv preprint
[35] Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. In The Twelfth International Conference on Learning Representations, 2024.
[36] Ruoqi Shen, Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, Yuanzhi Li. and Yi Zhang. Positional description matters for transformers arithmetic. arXiv preprint arXiv:2311.14733 2023.
[37] Daniel Zhang-Li, Li, Nianyi Lin, Jifan Yu, Zheyuan Zhang, Zijun Yao, Xiaokang Zhang, Lei Hou, Jing Zhang, and Juanzi Li. Reverse that number! decoding order matters in arithmetic learning. arXiv preprint arXiv:2403.05845 2024.
[38] Sean McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild, et al. Transformers can do arithmetic with the right embeddings. Advances in Neural Information Processing Systems, 37:108012-1080 2024.
12

---

[39] Shengyi Huang and Santiago Ontañón. A closer look at invalid action masking in policy gradient algorithms, arXiv preprint arXiv:2006.14177, 2020.
[40] Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In Proceedings of the IEEE international conference on computer vision, pages 2980-2988. 2988, 2017.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We support our main claims with extensive analysis in Sec. and experiments
in Sec.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in Sec. 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
13

---

judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification:
We do not include theorems.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We describe our method and experiment setup in Sec. 4 and 5.1 while additional
details (e.g. hyperparameters) are provided the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
14

---

(d)
We report standard deviations for experiments we ran in the tables in the main We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
We will release our code and data upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7 http:/scccpbbllcuuuddoeoooo
Coweeoomissiieieeooeeeeeeoo Pollyyy))))) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide experimental details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
text.
Guidelines:
The answer NA means that the paper does not include experiments.
15

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide relevant details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurrps.. http://aarrrpsc.c.ppp..t.htaaaaa
Answer:
[Yes]
Justification:
We confirm compliance with the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss potential societal impacts in Sec.
Guidelines:
16

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not produce any high-risk assets.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the original papers in the main text for the code and data used with details the supplementary material.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
17

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thood cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[Yes]
We provide details in the supplementary material for the synthetic datasets we constructed.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do perform crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not perform crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
18

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: We do not use LLMs for important, original, or non-standard components of the methods in this research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https://neurl)) aattpp///eearrrp..ccc///offeeeer /tt///////oeeeeeee////////../. for what should or should not be described.
19

---

A Broader Impacts
ReCOR aims to enhance the adaptive reasoning and planning capabiliiies of machine learning models, with potential applications in many fields, e.g. code writing, assisting scientific discoveries, and general LLM agents etc. Embodied agents with real-world tasks also require these capabilities. We regard ReCOR as mostly foundational research while noting that advancements in such capabilities have both positive and negative potential consequences if subject to misuse.
B Pseudocode of Training and Inference Algorithms for ReCOR
Algorithm 1 Training of ReCOR.
Require: Dataset D
Ensure:
Trained parameters 0, Initialize 8
while Not converged do
Sample mini-batch B {(x, D for / = 1... .M do
Rollout Ty on B
Sample Pt, Pt,[K] TT( X
end for
Compute rewards r using d on B, p.
Update d with LLM and 0 with LsQi on B, p. p.r end while
Algorithm 2 Inference of ReCOR.
Require: Eyaluation prompt parameters 0", t/*
Ensure: Generated response y
for t I...M do
Sample Pt ~ TT- x, Xx)))))
Sample úp. ~ p (: x, Pece
end for
C Generation Examples
Here we show generation examples of ReCOR and AdaMDM on Autoregression. The prompt is
64610246433563440133 while the correct response is 1544311665421664555, Grey denotes unfilled positions, green denotes correct tokens, and red ones incorrect tokens.
As observed below, ReCOR (left) recovers the reversed generation order and generates the correct tokens, while AdaMDM (right) encounters difficulties and incorrectly generates later tokens. Note that even though AdaMDM uses adaptive inference techniques, it is confused about which token can actually be predicted correctly, and hallucinates wrong tokens. We analyze and explain this phenomenon in the experiment section in the main text.
20

---

2
3
5
5
4
4
5
2
5
4
6
6
5
4
2
5
4
3
5
4
5
5
5
4
2
5
4
3
5
6
5
4
2
5
4
3
5
5
5
5
6
5
4 2
6
5
5
3
5
5
5
6
5
5
5
6
5
4 2
5
5
3
5
6
5
5
5
5
6
5
4 2
6
4
3
5
2
6
5
5
3
5
5 6
5 4 2
6
4
4
5
2
5
5
4
3
5
5 6
5 4 2
6
5
4
5
4
5
5
5
*
4
4
3
5
5
6
5 4 2
6
5
4
5
5
5
5
*
5
4
4
3
5
5
6
5
4
2
5
4
1 5
5
5
5
4
4
3
5
5 6 5 2
3 5
4
5
D
Experiment Details
D.1
Evaluation Datasets
Here describe details about the datasets used in experiments. For all datasets, randomly split small validation set of size 448 from the training set for validation purposes.
D.1.1 Synthetic Autoregression
We generate synthetic autoregression dataset (ARG) via the following procedure: For length L and prime modulus p, we first generate the prompt x; ~ Up,, S < L, the generate response satisfying the following autoregression formula:
UL ØL
(16)
8/
II +2,%%(p 1) + 11 Sp.1 J>1
(17)
This results in every depending on all リj.j resulting in =
> i. We use L = 20, p = 7 for our experiments, M 20. The size of the training set is 10 while the size of the test set 103.
D.1.2 Multiplication
We generate multiplication dataset with prompts of the form x*y= and responses z, where [1019 1020) 20-digit positive integer, y E [2, 100) is positive integer, and z = x y. This leads to N = 24, M = 22. Both x and y are uniformly at random within the respective range. We generate a training set of size 10a and a test set size 10
D.1.3 Sudoku
We use the Sudoku dataset from [10]. The prompt is string of length 81 that contains the flattened
initial configuration, where each element is either in [1,9] 9) denoting given cell, or indicating that
the value in this cell is missing. The response is also string of length 81 that contains the full
solution. N M 81. The training set contains approximately 1.8 106 samples while the test set contains 1 instances.
21

---

D.1.4 Zebra
We also use the Zebra dataset from 10]. The prompt for Zebra puzzles consists of a set of clues; we tokenize the special words in the clues with corresponding special tokens. After tokenization, N = 455, M 42. The training set contains about 1.5 10 puzzles while the test set contains 109.
D.2
Implementation Details
D.2.1
Details of ReCOR
The core algorithm of ReCOR is described in the method section of the main text. We document additional details below.
The ReCOR-related hyperparameters are listed in Tab. 6 For the token query stream and (optional) order query stream, we reuse all parameters of the main stream (including attention projection layers and MLP layers) to ensure that ReCOR always has the same smaller parameter count than baselines. During training, for the sampling of p, we always greedily sample the position with minimal hardness as predicted by Qe. while the token queries p sampled using Exploration mode in Tab. O without duplications. Uniform means uniform sampling p ~ UA while a = means using the corresponding entropy coefficient for exploration. For inference, always greedily sample both Ps and Vp.- We use action masking [39] and never pick actions at already filled positions. We employ y across the board, in effect solving a contextual bandit problem. Since our evaluation datasets require an exact match of all response tokens, the correctness of every token counts, and we do not need the trade-off between short-term and long-term rewards introduced by 7- y simplifies the learning procedure by eliminating additional value function evaluations while still achieving superior performance, as evidenced in our experiments. Furthermore, we use the thresholded, sparse reward since this choice slightly improves performance (see additional ablations below), which we also attribute the exact-match evaluation metric. Note that the binary cross-entropy loss associated with this reward actually computes an upper bound of the true Q function via Jensen's 's inequality. However, since we use y 0, the expectation for the Q value concentrates at single point, rendering the bound tight.
focal-like techniques [40 with focal y listed in Tab. as Focal coefficients We use focal loss for the RI. value loss LsO instead of the token (language modeling) loss LIM- differing from similar techniques used in [11]. We remark that, argued in the main text, certain token prediction sub-problems are inherently hard, and we should not expect the model to solve them or penalize it for these failures; rather, we simply want the model to recognize that these sub-problems are bad through the value loss.
Optionally,
D.2.2 Hyperparameters
We list hyperparameters on Autoregression and Multiplication in Tab. while describing ReCOR- related ones all datasets in Tab. 6 Baseline performances for Sudoku and Zebra are reported by [12. In Autoregression and Multiplication, compared with ReCOR, we double the batch size and number of epochs for baselines to match the amount of compute per iteration and number of gradient steps of ReCOR to ensure a fair comparison.
D.2.3 Other Training Details
For all experiments, we use mixed- preccssoe training with bfloat16. The model parameters are kept in full precision (float32). On Autoregression, Multiplication, and Sudoku, we use the tiny version of GPT-2 with layers, 384 hidden dimensions, and 12 attention heads, resulting approximately 6M parameters. On Zebra, we use the nano version with double the depth and parameter count. Fach experiment takes couple of hours using a single NVIDIA RTX4090 on Autoregression and Multiplication, and less than 2 days using single NVIDIA AL00-80G on Sudoku and Zebra.
22

---

Table 5: Hyperparameters for ReCOR and baselines on Autoregression and Multiplication.
Hyperparameter
ReCOR
CLM
AR-GT
(Ada)MDM
Model parameters
6M
6M
6M
6M
Batch size
1024
2048
2048
2048
Epochs
100
200
200
200
Optimizer
A
AdamW
AdamW
Adam W
Learning rate
10-3
10-3
10-3
10-3
LR scheduler
Cosine
Cosine
Cosine
Cosine
Weight decay
0.1
0.1
0.1
0.1
Diffusion steps
N/A
N/A
N/A
20
Table 6: ReCOR-related hyperparameters evaluation datasets.
Hyperparameter
ARG
MUL.
Sudoku
Zebra
# Model parameters
6M
6M
6M
11M
Batch size
1024
1024
512
512
# Epochs Optimizer
100
100
40
50
AdamW
A
AdamW
AdamW
Learning rate
10
10-3
3 * 10-4
3 * 10-4
LR scheduler
Cosine
Cosine
Cosine
Cosine
Weight decay
0.1
0.1
0.1
0.1
Discount factor
0.0
0.0
0.0
0.0
Focal coefficient
2.0
2.0
0.0
0.0
Threshold !)
0.8
0.8
0.8
0.8
# Token queries K
# Order queries C'
0
8
2
Exploration mode
Uniform
Uniform
a 0.1
E
Additional Experiments
E.1
Exploration Strategies
We present additional ablation experiments regarding alternative exploration strategies for token queries p on Sudoku in Fig. 4a| alpha=... corresponds an entropy-regularized exploration with the designated entropy coefhcient, Uniform denotes uniform sampling, and TopK denotes sampling the K positions with the largest Q-values. We find ReCOR pretty robust to alternative exploration strategies and opt for 0.1 in our main experiments.
E.2 Reward Functions and RL algorithms
Here we study the effects of alternative reward functions and RL algorithms on Sudoku. We compare our primary setting (Soft Q-learning sparse reward with threshold 7) 0.8) alternative settings with negated perplexity as rewards and policy-gradien class of RL algorithms. Note that the sparse reward cannot used with policy-gradient algorithms. We use the following loss to implement PPO 34 llike algorithm, replacing Lsqi.:
Ta s) A (To( s))
(0) = E.,a,,,.
(18)
Te (a
where is the and does not reuse actions, Gol is the detached version of the current f, and the ratio is always so we do not need the clipping in standard PPO.
advantage. Note that since ReCOR always stays exactly on-policy
As shown in Fig. 4b our value-based choice of RI. algorithm with sparse reward achieves the best overall performance. Switching to dense rewards still yields decent performance, validating motivation. We found value-based algorithm to slightly outperform
methods,
23

---

6.8
Valcantoon Accurncy
6.6
6.4
0.4
0.2
apha=8.01
0.2-
#pham]
SQL. sparse
Tapk
SQL. I91
Uniform
PVG, IPL
20000
40000
60000
100000
120000
140000
20000
40000
60000
8:0000
100000
120000
140000
Thsining Gradiert Shep
Thiining Gradlent Step
(a) Exploration strategy ablation with sparse rewards.
(b) Reward function and RL algorithm ablation.
Figure 4: Additional ablation experiments regarding the exploration strategy of p (a) and the choice of reward function and RL algorithm (b).
Table 7: Performance of ReCOR and baselines on the mixed dataset between ARG and MUL. ReCOR can perform uniformly well when trained on both datasets simultaneously.
ReCOR
CLM
MDM
AdaMDM
0.945
0.265
0.398
0.558
potentially due to the fact that it can make full use of return signals while policy-based methods only model the relative difference between actions.
E.3 Handling Multiple Tasks
To validate the generalization capability of ReCOR across different types of tasks, we conduct mixed-dataset experiment that trains ReCOR and baselines on ARG and MUL simultaneously. The average evaluation accuracies on both test sets are shown in Tab. We can see that ReCOR learns from both datasets and achieves uniformly good performance, while baselines still struggle.
24


Reviews:
[{"review_id": "4SfGPUEaVI", "reviewer": "Reviewer_UpeJ", "summary": "The paper proposed Reinforced Context Order Recovery (ReCOR), a self-supervised framework that learns to adaptively determine the optimal token generation order without explicit order annotations. ReCOR estimates the hardness of predicting every unfilled token and adaptively selects the next token during both training and inference. They showed their method outperforms baselines on arithmetic problems and c logic puzzles Sudoku and Zebra.", "questions": "Please refer to weaknesses above.", "limitations": "yes", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The problem is well motivated and automatically recovering the correct generation order from textual data is very important in some planning tasks. \n\n- The proposed method jointly optimizes the token prediction model and the order prediction policy, generating rewards with the former as self-supervision for the latter.\n\n- Propose method solves arithmetic problems without special data preprocessing\n\nWeaknesses:\n\n- The paper formulates automatically recovering the correct generation order from textual data as a decision making problem and uses RL-techniques to solve this. Even though the problem is well motivated and seems interesting, I am not sure how this method can scale. Specially how this can be compared against existing large models? What is the potential here? It might be the case that it is able to recover the correct generation order but other capabilities might be sacrificed. \n\n- In experiment the details about model architecture and details of training are missing.\n\n- I think since the proposed approach changes the training of the model, the limitations need to be extended to cover the scalability, comparing other capabilities, and even comparing with other existing models in different scales.\n   \n- Some terms or abbreviations seem to be unconventional and there is no reference for that: for example L_SQL (I think it meant soft Q-learning) ,  \\nu-information framework.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our paper! We are glad to see that you regard our problem as ·\"well motivated and interesting\"·. We address your concerns below.\n\n> *Q1. Scalability of method; potential when compared with existing large models.*\n\n**R1:** We have demonstrated the scalability of ReCOR in Section 5.6 with compute scaling experiments. ReCOR can scale along multiple axes, improving with more training and test-time compute. We note that ReCOR is highly scalable due to its **self-supervised nature**; ReCOR requires only purely textual data to train without any additional annotations, making it as general as standard pretraining objectives like next-token prediction. Compared with existing large models that only learn to predict the next token, ReCOR-trained models can predict orders in addition to tokens, capturing structures in the training data, which in turn enhances the token prediction capability, as shown in our experiments.\n\n> *Q2. Architecture and training details.*\n\n**R2:** We have included the architecture and training details in our supplementary material PDF with pseudocode in the appendix of the main text. To summarize, we largely followed the training setups of our baselines, in particular [1], and built upon their released codebase to ensure a fair comparison. All methods are trained from scratch and use the same GPT-2 architecture, training data (except for the oracle AR-GT, which has additional access to privileged order information), and weight decay regularization of 0.1.\n\n> *Q3. Comparison with models of different scales.*\n\n**R3:** We note that simply scaling up under the next-token prediction paradigm is **not** sufficient for handling the order problem; furthermore, reasoning models require suitable CoT annotations to warm up, and it is extremely challenging to solve novel, unfamiliar problems.\nTo demonstrate the ineffectiveness of existing state-of-the-art large reasoning models, we perform an experiment with o3 on ARG. We design two prompts where prompt A includes the actual generation rule of the dataset and examples, and prompt B includes in-context examples only. Note that in our main experiments, only the generated data is available to the models without the ground-truth generation rule, corresponding to the harder scenario B. We evaluate 50 random samples from the test set and allocate a token budget of 10000 tokens each. As a result, \n- Under prompt A with ground-truth rules, o3 achieves only a 0.3 success rate, where the majority of failures come from depletion of token budgets. Even considering only the completely generated solutions, the success rate of o3 is 0.938 and is still lower than ReCOR's 0.987. Furthermore, o3 takes ~4min on average to solve each individual question, while ReCOR solves the entire 1000-sample test set within 10 seconds.\n- Under prompt B without access to ground-truth rules, o3 fails to recover the correct data generation algorithm and has 0 success rate. We stress again that ReCOR operates under this harder setting without any additional annotations.\n\nIn conclusion, we show that scale alone is inefficient and insufficient, while ReCOR solves the problem effectively and efficiently.\n\n> *Q4. Terminology references.*\n\n**R4:** We apologize for any confusions; $L_\\text{SQL}$ indeed means soft Q-learning, which we cite as [33] in Sec. 4.3, line 161 and 166 before we define the soft Q-learning loss function Eq. 11. $\\mathcal{V}$-information is cited as [30] in Sec. 4.1, line 120, where we invoke the theoretical formulation from existing literature to characterize the token hardness problem. We will revise the manuscript for better clarity in the revision.\n\n[1] Ye, Jiacheng, Jiahui Gao, Shansan Gong, Lin Zheng, Xin Jiang, Zhenguo Li, and Lingpeng Kong. \"Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning.\" In The Thirteenth International Conference on Learning Representations (ICLR 2025)."}, {"review_id": "031ebksA74", "reviewer": "Reviewer_MWbF", "summary": "This paper studies the problem that for certain tasks (such as arithmetic calculation and sudoku), solving the task from a easy-to-hard order is more preferred than the traditional left-to-right order. The paper proposes Reinforced Context Order Recovery (ReCOR), which models token order prediction, before predicting the token at each predicted position. Specifically, ReCOR leverages V-information to recover token orders as a markov decision making problem, and applies Q-learning to predict tokens. At inference time, a multi-stream architecture is used for order and token predictions. Evaluated on various tasks (including synthetic autoregression, multiplication, Sudoku, and Zebra) results show that a GPT-2 based ReCOR performs better than autoregressive methods (CLM) and diffusion methods (MDM, AdaMDM) and is close to an oracle-based AR-GT.", "questions": "1. Equation 4 - 7 defines how order might be defined from entropy to V-information, but how is 8 derived from 7? i.e., how do you model Hv (B | ∅)? And how do you condition on arbitrary contexts as you mentioned that the naive solution is not tractable?\n2. In equation 13, where is the language modeling objective from? What is considered as the ground truth?\n3. Why is the threshold required from equation 14 - 15? What is the interpretation of adding this threshold?", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. This paper is well motivated and targets at an interesting problem with autoregressive models in terms of decoding orders.\n2. The proposed ReCOR is intuitive and straightforward, and achieves better performance on various tasks.\n\nWeaknesses:\n1. The scope and generalization of the proposed method is limited. Typically I wouldn't criticize on the base model used (in this case GPT-2) because I understand resource limitations. However, in this particular case, on the one hand, GPT-2 is a very shallow and small LM that is not representative of more recent LMs where more layers may already take ordering into considering (even though still decoded in a left-to-right order), especially when easy-to-hard prediction is easily targeted with reasoning models. On the other hand, GPT-2 is trained in an autoregressive way, which raises the question on how ReCOR, a framework taking into arbitrary order interacts with the pretrained weights. Furthermore, it is not clear how the baselines are implemented (e.g., how is the CLM and MDM trained, including the base model, training examples, architecture, regularization). Moreover, although MDM is similar to ReCOR in the sense that it is not autoregressive, diffusion methods are designed from a totally different perspective (and is not a meaningful baseline). Therefore, it is not convincing to me that this work would inspire a wide range of audience interested in similar topics.\n2. Some details are missing. See details below.", "was_revised": true, "final_justification": "Because of the concerns I raised including the generalization of the proposed method which require more careful analysis and changes of narratives, I believe a major revision is required to make the paper beneficial to broader audiences.", "rebuttal": "Thank you for your review! We are happy to see that you think ReCOR is \"intuitive\", \"well motivated\", and \"targets an interesting problem.\" We address your concerns below.\n\n> *Q1. Comparison with LMs with more layers and reasoning models.*\n\n**R1:** We note that simply scaling up under the next-token prediction paradigm is **insufficient** for handling the order problem; furthermore, reasoning models require suitable CoT annotations to warm up, and it is extremely challenging to solve novel, unfamiliar problems.\nTo demonstrate the ineffectiveness of existing state-of-the-art large reasoning models, we perform an experiment with o3 on ARG. We design two prompts where prompt A includes the actual generation rule of the dataset and examples, and prompt B includes in-context examples only. Note that in our main experiments, only the generated data is available to the models without the ground-truth generation rule, corresponding to the harder scenario B. We evaluate 50 random samples from the test set and allocate a token budget of 10000 tokens each. As a result, \n- Under prompt A with ground-truth rules, o3 achieves only 0.3 success rate, where the majority of failures come from depletion of token budgets. Even considering only the completely generated solutions, the success rate of o3 is 0.938 and still lower than ReCOR's 0.987. Furthermore, o3 takes ~4min on average to solve each individual question, while ReCOR's inference procedure solves the entire 1000-sample test set within 10 seconds.\n- Under prompt B without access to ground-truth rules, o3 fails to recover the correct data generation algorithm and has 0 success rate. We stress again that ReCOR operates under this harder setting without any additional annotations.\n\nIn conclusion, we show that scale alone is inefficient and insufficient, while our ReCOR solves the problem effectively and efficiently.\n\n> *Q2. Interaction with pretrained weights.*\n\n**R2:** We note that ReCOR is a framework for pre-training from scratch without relying on any pretrained weights. We consider fine-tuning from existing checkpoints as an important direction for future work, where we can adapt left-to-right pretrained weights into an adaptive-order model. This has been shown to be feasible by recent literature, e.g., [1] adapts a LLaMA checkpoint to be a masked diffusion model, but doing so qualifies as an independent contribution and is out of scope for the current work.\n\n> *Q3. Training details of baselines.*\n\n**R3:** We have included training and architecture details of all methods in the supplementary material. To summarize here, we largely followed the training setups of our baselines, in particular [2,3], and build upon their released codebase to ensure a fair comparison. All methods (including CLM and MDM) are trained from scratch and use the same GPT-2 backbone architecture, training data (except for the oracle AR-GT which has additional access to privileged order information), and weight decay regularization of 0.1.\n\n> *Q4. Comparison between ReCOR and MDM.*\n\n**R4:** We'd like to clarify that ReCOR ***is*** autoregressive; it is just not left-to-right, differing from standard causal language models. For a complete view of related works, we compare ReCOR with diverse baselines, including both causal language models and masked diffusion models, and show that ReCOR outperforms all of these approaches. We additionally provide detailed analysis and comparison between ReCOR and MDM in Sec. 5.4 and 5.5, extending the results of recent high-profile works in diffusion language modelling [2, 3]. We believe our algorithms and findings are of value to these communities and beyond.\n\n> *Q5. Modelling of $H_\\mathcal{V}(B \\mid \\varnothing)$; conditioning of arbitrary contexts.*\n\n**R5:** Since $H_\\mathcal{V}(B \\mid \\varnothing)$ conditions on an empty set, it remains a constant with respect to the generation order $\\rho$, and consequently our optimization parameter $\\theta$ in Eq. 7. Thus we can remove the term from our objective, leading to the simplification from Eq. 7 to 8. For the conditioning of arbitrary contexts, as shown in Fig. 2, we feed the context consisting of a set of tokens at arbitrary positions into the transformer. The token is represented by vocabulary embeddings while the position is represented by additive absolute positional embeddings, as stipulated by GPT-2. This design allows us to tractably encode and condition on an arbitrary context with the same set of parameters, without having to instantiate a separate model for each set of positions.\n\n> *Q6. Derivation and training of $L_\\text{LM}$ (Eq. 13).*\n\n**R6:** The language modelling objective $L_\\text{LM}$ is directly derived from the overall objective Eq. 8; in particular, we optimize Eq. 8 with respect to $\\psi$ using $L_\\text{LM}$. The ground-truth response token sequence $\\mathbf{y}$ comes from the training data, which is the textual corpus we're training on.\n\n> *Q7. Interpretation of the thresholded reward (Eq. 14-15).*\n\n**R7:** We can interpret the threshold as a form of *reward shaping*, a commonly used technique in the field of reinforcement learning to improve learning efficiency. Intuitively, in ReCOR, we use the threshold to provide a sharper, binary distinction between \"easy\" and \"hard\" tokens to assist policy learning.\n\n[1] Gong, Shansan, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An et al. \"Scaling Diffusion Language Models via Adaptation from Autoregressive Models.\" In The Thirteenth International Conference on Learning Representations.\n\n[2] Ye, Jiacheng, Jiahui Gao, Shansan Gong, Lin Zheng, Xin Jiang, Zhenguo Li, and Lingpeng Kong. \"Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning.\" In The Thirteenth International Conference on Learning Representations.\n\n[3] Kim, Jaeyeon, Kulin Shah, Vasilis Kontonis, Sham M. Kakade, and Sitan Chen. \"Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions.\" In Forty-second International Conference on Machine Learning."}, {"review_id": "jSRKBIxekz", "reviewer": "Reviewer_jkyP", "summary": "(Motivation) Language models are predominantly trained to output tokens with a fixed (left-to-right) or random order, which may deviate from the logical order in which tokens should ideally be used. Causal LMs always follow a rigid left-to-right generation paradigm, encountering many intractable tokens along the way, and random order or masked diffusion-based models struggle with performance from introducing large distribution shifts between training and inference or attending to sub-problems (the latter of which is also shown as a result in this work).\n\nThe authors create an objective based on \\nu-information to turn token order selection into a decision-making problem that can be solved with an MDP and soft Q learning. The resulting algorithm is called ReCOR (“Reinforced Context Order Recovery”). The submitted work demonstrates that ReCOR outperforms or performs similarly to baselines and oracles on simpel arithmetic tasks, synthetic autoregression, and puzzle datasets (Sudoko, Zebra), evaluating using GPT2 as backbone for all algorithms. The authors use their results to discuss algorithmic performance, balanced compute requirements and scaling, and the impact of adaptive ordering during inference and training.", "questions": "- I assume the results for puzzle benchmarks and arithmetic benchmarks were from different training runs and models. I would like to understand if ReCOR limits a trained model to a very specific narrow task or not. Would ReCOR perform better across the tested tasks than the baselines if trained on both? Can a ReCOR-trained GPT2 still perform general language generation? \n\n- Have the baselines been evaluated for harder tasks (e.g., coding) with larger models (larger than 1B parameters)? Based on your results, how would you expect ReCOR to perform in comparison at these scales for these tasks?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Reasons to accept:\n\n- The paper is very well written and, to the best of my knowledge, original work (Clarity). As I am not familiar with most of the token-reordering literature, I will reduce my confidence to 3.\n\n- The authors provide good theoretical description and background for the introduction of ReCOR and conduct a series of experiments to support all their claims. The proposed algorithm is of high quality and optimized for parallel computation (Quality).\n\n- The idea of using a learned and adaptive re-ordering for reasoning and planning problems is elegant and promising. The submission has the potential to advance the field of (low-compute, task-specific) reasoning and planning work significantly. (Significance).\n\nReasons to reject:\n\n- It is unclear to me how well the learned adaptive reordering generalizes across tasks or to more complex tasks. It would have been ideal to see ReCOR tested for something like (simple) coding benchmarks with smaller models in the 1.5-3B parameter regime.", "was_revised": true, "final_justification": "I had little concerns regarding the quality of this submission and only asked for clarifications on cross-task performance and scaling the proposed algorithm to larger models.\n\nI still believe this is an elegant and promising approach that can be the foundation of crucial future work and should be featured at NeurIPS 2025.", "rebuttal": "Thank you for your thoughtful review! We are glad to see that you think ReCOR is `\"elegant and promising\"`. We address your questions below.\n\n> *Q1. Limitation of training tasks.*\n\n**R1:** We note that ReCOR is *not* limited to a narrow task. Similar to standard causal language models, ReCOR can learn to model all of the training data provided to it `without any annotation requirements`. We support this statement with a mixed data experiment where we mix the training sets of ARG and MUL and retrain ReCOR and baselines. The average evaluation accuracies on both test sets are shown below. We can see that ReCOR learns from both datasets and achieves uniformly good performance, while baselines still struggle. We also note that ReCOR, like its baselines, is trained from scratch without initializing from pretrained GPT-2 checkpoints and never trained on general language data.\n\n| Method | ReCOR |  CLM   |  MDM   |  AdaMDM   |\n| -------- | -------- | --- | --- | --- |\n| Average Accuracy         |  **0.945**        |  0.265   |  0.398   | 0.558    |\n\n\n> *Q2. Scaling to larger models and problem domains.*\n\n**R2:** Our baselines primarily focus on the same reasoning and planning domains and have similar parameter scales to ReCOR. Scaling up model size and problem domains is a key future work that we are actively pursuing to further demonstrate the potential of ReCOR. For example, as mentioned in your review, we believe ReCOR can be applied to coding domains to automatically recover the hierarchical structure of code and enable better modelling of the training data with its order recovery capabilities, while our baselines are limited to fixed linear (causal language models) or random (masked diffusion models) structures."}, {"review_id": "9pdark0NgH", "reviewer": "Reviewer_81xH", "summary": "This paper proposes a ReCOR based on reinforcement learning, which aims to solve the performance bottleneck caused by fixed or random generation order in traditional causal language models (CLMs) and diffusion models (MDMs) in complex reasoning tasks. By modeling the generation order as a reinforcement learning problem and combining it with a multi-stream Transformer architecture, ReCOR achieves self-supervised learning of the optimal generation order without labeled data. Experiments have verified the effectiveness of the method on tasks such as arithmetic, Sudoku, and logic puzzles. Some tasks even surpass the baseline model based on real order annotation.", "questions": "1) Training complexity. RL relies in part on online policy sampling, which may require a large amount of interaction data. The authors mention using a pre-trained model for warm-starting, but do not specify a specific strategy (such as imitation learning or curriculum learning)\n\n2)  ​​Generalization risk​​. The modification of the backbone model (such as GPT-2) by the multi-stream architecture may limit its transferability. If it is replaced with another architecture (such as LLaMA), do I need to adjust the hyperparameters?", "limitations": "1) ​Interpretability of generation order: We recommend visualizing the generation order ρ predicted by ReCOR in the Sudoku task (as shown in the right side of Figure 1) to verify whether it conforms to the intuition of “from easy to difficult”.\n\n2) The calculation of V-information in Formula 7 depends on the shared token predictor pψ, and the impact of its approximate error with the ideal V-information on policy learning needs to be further discussed.\n\n3) ​​Evaluation limitations​​:\nThe experiment only uses accuracy (complete match) as an indicator. For partially correct generation results (such as partially correct Sudoku), more fine-grained evaluation (such as cell-level accuracy) can be supplemented.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "1） For the first time, the generation sequence problem is formalized as a reinforcement learning task of maximizing the prediction V-information.\n2） Significantly outperforms baselines on non-causally dependent arithmetic tasks (ARG, MUL) and complex reasoning tasks (Sudoku, Zebra)", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! We address your concerns below.\n\n> *Q1. Training procedure of ReCOR.*\n\n**R1:** We clarify that we do not use a pretrained model for initialization and train all methods from scratch. In contrast to other domains like robotics or games, the RL interaction procedure required by ReCOR is equivalent to transformer inferences of the order policy, which can be implemented highly efficiently.\n\n> *Q2. Generalization to other architectures.*\n\n**R2:** We note that ReCOR is specifically designed to be highly generalizable. The multi-stream design modifies minimally upon the original transformer architecture, introduces little hyperparameter, and is compatible with any variant of transformer, including LLaMA.\n\n> *Q3. Visualization in Sudoku.*\n\n**R3:** Unfortunately, due to restrictions by the NeurIPS rebuttal policy, we cannot display images or videos in the rebuttal. We show here an example of ReCOR in the Sudoku test set; it can be seen that ReCOR follows the easy-to-hard order, e.g.\n- The given cells in the prompt are trivial and filled first;\n- The first non-trivial cell to be filled is the 17th one at (3, 2), which can be deduced from the prompt alone by elimination on number 9;\n- The last cells to be filled, e.g., the 81st at (4, 2), have very few clues initially, both in the row, column, and subgrid.\n\n```\nPuzzle:\n 6  0  0  0  0  9  0  0  0 \n 4  5  0  0  6  3  0  9  2 \n 0  0  2  0  0  0  0  8  0 \n 0  0  6  0  0  4  0  0  0 \n 0  0  0  0  9  0  0  0  0 \n 9  0  0  8  1  0  0  0  4 \n 0  8  0  0  4  0  2  0  0 \n 0  0  0  5  0  0  0  0  0 \n 0  2  0  0  0  0  7  0  5 \n\nGeneration order:\n15 68 71 35 26  4 59 62 70 \n 6 14 25 31  2  9 30  8 13 \n51 17  1 34 32 33 56 16 55 \n64 81  7 37 75 21 49 58 69 \n65 36 67 38 20 73 60 78 57 \n22 74 72 10 24 77 54 79 18 \n53  3 52 42 12 45 19 46 47 \n50 23 48 11 80 76 61 66 63 \n40  5 43 41 27 29 28 44 39 \n\nFull response:\n 6  1  3  2  8  9  5  4  7 \n 4  5  8  7  6  3  1  9  2 \n 7  9  2  4  5  1  3  8  6 \n 8  7  6  3  2  4  9  5  1 \n 2  4  1  6  9  5  8  7  3 \n 9  3  5  8  1  7  6  2  4 \n 5  8  7  1  4  6  2  3  9 \n 3  6  9  5  7  2  4  1  8 \n 1  2  4  9  3  8  7  6  5\n```\n\n> *Q4. Approximation error of $\\mathcal{V}$-information.*\n\n**R4:** We note that our $\\mathcal{V}$-information estimator, the token predictor $p_\\psi$, is trained in a supervised manner and learns much faster than the RL-trained order policy $\\pi_\\theta$. As a result, its estimates of $\\mathcal{V}$-information is very accurate, as indicated by the superior performance of ReCOR. We can further improve the performance of ReCOR by giving a more accurate estimate. For example, we can sample multiple $K>1$ actions at every state and supervise $p_\\psi$ at all of these positions, reducing the approximation error and further improving performance, as shown in Sec. 5.6.\n\n> *Q5. Evaluation metrics.*\n\n**R5:** For reasoning and planning problems, it is often the complete correctness of the solution that matters; for example, any single error in a long mathematical proof could invalidate the entire proof. We thus use full solution match as the primary evaluation metric, following the common practice of our baselines. To address your concerns more thoroughly, we provide a cell-level analysis on ARG, as shown in the table below. Note that this is a significantly easier setup than a full solution match. While the performance of baselines improves, ReCOR achieves a near-perfect result and still significantly outperforms the baselines.\n\n| Method | ReCOR |  CLM[1]   |  MDM[2,3]   |  AdaMDM[2,3]   |\n| -------- | -------- | --- | --- | --- |\n| Cell Accuracy         |  **0.999**        |  0.505   |  0.868   | 0.912    |\n\n[1] Radford, Alec, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. \"Language models are unsupervised multitask learners.\" OpenAI blog 1, no. 8 (2019): 9.\n\n[2] Ye, Jiacheng, Jiahui Gao, Shansan Gong, Lin Zheng, Xin Jiang, Zhenguo Li, and Lingpeng Kong. \"Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning.\" In The Thirteenth International Conference on Learning Representations (ICLR 2025).\n\n[3] Kim, Jaeyeon, Kulin Shah, Vasilis Kontonis, Sham M. Kakade, and Sitan Chen. \"Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions.\" In Forty-second International Conference on Machine Learning (ICML 2025)."}]

Author rebuttal:

Meta-review: A. For certain tasks (such as arithmetic calculation and sudoku), solving the task from a easy-to-hard order is more preferred than the traditional left-to-right order. The paper uses RL to predict the token order before filling in the token at each predicted position. Results are shown on tasks (e.g., multiplication, Sudoku, Zebra), comparing to autoregressive methods and diffusion models.

B. Reviewers found the paper well written and the method a well motivated approach to an important problem. They also appreciated the strong results, the theoretical justification, and the fact that the method could be parallelized.

C. Reviewers raised questions about scalability to larger models, missing experimental details, scope and generalization of the method (only shown on GPT-2), some missing details

D. I think that the reviewers did a sufficient job addressing the reviewer concerns, and that the paper should be accepted. Reviewer UpeJ did not engage in the discussion; both myself and another reviewer believe that their concerns have been addressed, so we're treating the actual scores as 5/4/"4"/3. I also feel like the primary concerns in Reviewer MWbF's review (score=3) have been addressed in the rebuttal. I agree with all reviewers that the empirical results are good, but I feel like they are actually rather marginal (e.g., compare with AdaMDM in Table 2). However, the paper has a great discussion of the comparison with AdaMDM, including some nice conceptual experiments.

E. During the rebuttal, the authors answered reviewer questions, provided results with the o3 model, and demonstrated that the method can generalize across multiple tasks and datasets. The reviewers that engaged in the discussion said that some of their concerns had been addressed. However, one reviewer said that concerns about the generalization of the approach and correctness of the analysis remained. The author's contested reviewer MWbF's understanding of the paper.
Decision: Accept (poster)