Reinforced Context Order Recovery for Adaptive Reasoning and Planning

Thank you for your review! We address your concerns below.

Q1. Training procedure of ReCOR.

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.

Q2. Generalization to other architectures.

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.

Q3. Visualization in Sudoku.

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.

• The given cells in the prompt are trivial and filled first;
• 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;
• 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.

Puzzle:
 6  0  0  0  0  9  0  0  0 
 4  5  0  0  6  3  0  9  2 
 0  0  2  0  0  0  0  8  0 
 0  0  6  0  0  4  0  0  0 
 0  0  0  0  9  0  0  0  0 
 9  0  0  8  1  0  0  0  4 
 0  8  0  0  4  0  2  0  0 
 0  0  0  5  0  0  0  0  0 
 0  2  0  0  0  0  7  0  5 

Generation order:
15 68 71 35 26  4 59 62 70 
 6 14 25 31  2  9 30  8 13 
51 17  1 34 32 33 56 16 55 
64 81  7 37 75 21 49 58 69 
65 36 67 38 20 73 60 78 57 
22 74 72 10 24 77 54 79 18 
53  3 52 42 12 45 19 46 47 
50 23 48 11 80 76 61 66 63 
40  5 43 41 27 29 28 44 39 

Full response:
 6  1  3  2  8  9  5  4  7 
 4  5  8  7  6  3  1  9  2 
 7  9  2  4  5  1  3  8  6 
 8  7  6  3  2  4  9  5  1 
 2  4  1  6  9  5  8  7  3 
 9  3  5  8  1  7  6  2  4 
 5  8  7  1  4  6  2  3  9 
 3  6  9  5  7  2  4  1  8 
 1  2  4  9  3  8  7  6  5

Q4. Approximation error of $\mathcal{V}$-information.

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.

Q5. Evaluation metrics.

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.

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

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

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

Thank you for your thoughtful review! We are glad to see that you think ReCOR is "elegant and promising". We address your questions below.

Q1. Limitation of training tasks.

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.

Q2. Scaling to larger models and problem domains.

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.

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.

Q1. Comparison with LMs with more layers and reasoning models.

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

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

In conclusion, we show that scale alone is inefficient and insufficient, while our ReCOR solves the problem effectively and efficiently.

Q2. Interaction with pretrained weights.

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.

Q3. Training details of baselines.

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.

Q4. Comparison between ReCOR and MDM.

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.

Q5. Modelling of $H_\mathcal{V}(B \mid \varnothing)$; conditioning of arbitrary contexts.

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.

Q6. Derivation and training of $L_\text{LM}$ (Eq. 13).

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.

Q7. Interpretation of the thresholded reward (Eq. 14-15).

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.

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

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

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

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.

Q1. Scalability of method; potential when compared with existing large models.

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.

Q2. Architecture and training details.

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.

Q3. Comparison with models of different scales.

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

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

In conclusion, we show that scale alone is inefficient and insufficient, while ReCOR solves the problem effectively and efficiently.

Q4. Terminology references.

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.

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

Dear Reviewer,

This is a gentle reminder to carefully read the rebuttal and respond to the authors about whether it addresses the concerns.