Title: Reinforced Context Order Recovery for Adaptive Reasoning and Planning
Venue: NeurIPS 2025
Accepted: yes
Keywords: Adaptive Token Order

Autoregressive Models

Diffusion Models

Reinforcement Learning
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n1） 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)\n\nQUESTIONS:\n1) 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?\n\nLIMITATIONS:\n1) ​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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n(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.\n\nSTRENGTHS AND WEAKNESSES:\nReasons 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\nPlease refer to weaknesses above.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-12"}