Learning Adaptive Parallel Reasoning with Language Models

Thank you for your valuable feedback. Here are our responses to the comments:

Task selection and model capacity. In our experiments, we intentionally follow the methodological choices of the Stream of Search (SoS) study (Gandhi et al., 2024; COLM 2024 oral presentation), which centers on the Countdown task due to its structured reasoning challenge, high branching factor, and ease of result verification. Similarly, we train our model from scratch to allow for controlled investigations. As in SoS, focusing on a synthetic yet challenging domain enables rigorous evaluation across accuracy, latency, and scalability, providing clear insights into the effectiveness of our proposed approach.

For model scale, we trained a version that is approximately 3x larger, and the performance trend remains consistent (results shown in our response to Reviewer RZt7), demonstrating that our test-time scaling method is orthogonal to training-time scaling. The size of our pre-trained model (~600M parameters) roughly matches some of the pre-trained models (e.g., Qwen-2.5-0.5B). The fact that APS still significantly outperforms the baseline at this scale confirms that our method remains effective even when applied to models of pretraining-level size.

Regarding generalization, the main innovation of our method (adaptive task decomposition and parallelization optimized through reinforcement learning) is inherently suitable for more domain-general structured reasoning tasks beyond Countdown, e.g., for math reasoning. In this current work, we follow the precedent established by the SoS study of focusing on the well-defined but challenging reasoning task of Countdown. While it is not feasible to include results on benchmarks like GSM8K and HumanEval within the rebuttal period, we plan to explore applying APS to these domains as a next step. We will revise the discussion section to better highlight the generalizability of SoS and outline future directions for extending it to a wider range of tasks.

Symbolic solver reliance. While our experiments with APS use symbolic solvers for generating supervision data for pre-training from scratch, our core method is generalizable to broader reasoning domains. For example, we could use LLMs to decompose existing open-ended reasoning tasks (e.g. math word problems) into subtasks, and identify which of the subtasks can be parallelized; these heuristically-generated decompositions could serve as weak supervision to bootstrap APS training with reinforcement learning.

To clarify a misunderstanding, APS does not require an optimal search path in its pre-training data from symbolic solvers. Our symbolic solver only needs to create viable search paths (i.e., ones that find a solution, not necessarily optimally). The generated demonstrations we pre-train on are suboptimal. We show that we can optimize search paths via end-to-end fine-tuning with reinforcement learning, without requiring an optimal solver.

Complexity and overheads. We appreciate the reviewer raising these important considerations for real-world deployment. Careful system design, including optimized scheduling, efficient KV-cache management, and techniques like block-wise computation (akin to frameworks like SGLang), can significantly mitigate these. For instance, very short child threads could be executed on the main thread, and while uneven threads might wait for the longest to complete, freed resources from shorter, completed threads can be dynamically reallocated. While these are important engineering considerations, our work is a preliminary exploration focused on the novel adaptive parallel reasoning capability. The significant performance and efficiency improvements, along with better leverage of parallel hardware, highlight the strong potential that we believe outweighs these manageable challenges for future deployment-focused work.

Comparing with Tree-of-Thought (ToT). We thank the reviewer for raising this important comparison. One key difference, as mentioned by the reviewer, is that ToT relies on a fixed prompt / inference structure to determine branching factors. Our proposed approach contrasts with ToT and many other structured inference methods (e.g., multi-agent debate (https://arxiv.org/abs/2305.14325), graph-of-thoughts (https://arxiv.org/abs/2308.09687), etc.) that require manually defining fixed prompt and inference structures. Instead, APS allows a model to adaptively decide at inference time when and how to branch its search. Another key difference with ToT (and other fixed-structure inference methods) is that such methods cannot be optimized for search beyond their manually-defined inference structures; they are all based on zero-shot prompting with no model training. Instead, APS initially uses supervised training to bootstrap branching behavior, and then applies reinforcement learning to perform end-to-end optimization of search strategies.

We sincerely appreciate the reviewer's comments and recognition of our approach.

Larger models. Thank you for the suggestion. Following this, we additionally trained APS and the baseline with a 600M-parameter model during the rebuttal. We kindly refer you to the general response for the results. In summary, APS continues to significantly outperform the SoS+ baseline with the 600M-parameter model, especially as total compute increases.

Performance vs GPUs. APS’s accuracy remains unchanged with different GPU counts, as the reasoning process itself is unaffected. We refer the reviewer to the general response for variations in latency when we have varying number of GPUs available.

Performance with SOTA-sized models. Our rebuttal experiments with the 600M-parameter model, detailed in response to the previous comment, show that APS scales positively with model size and continues to outperform the SoS+ baseline. Note that this is of similar size to representative pretrained models such as Qwen-2.5-0.5B.

This suggests its benefits are likely to persist at larger scales. Due to compute constraints, we couldn’t evaluate billion-parameter models during the rebuttal phase, but as noted in the paper, generalizing APS to larger pretrained LLMs and across broader applications remains an important and promising direction for future work.

Potential challenges in adapting to large pre-trained models with larger context. We thank the reviewer for raising this important point. Adapting APS to large pre-trained models with extended context windows introduces two key challenges. First, the increased context length leads to larger KV caches, and when combined with APS’s multi-threaded execution, this results in substantial GPU memory overhead. Efficient resource management becomes critical to ensure that memory constraints do not hinder performance.

Second, integrating APS into distributed serving systems is non-trivial. Large models are typically deployed with tensor or pipeline parallelism, where parameters are split across multiple GPUs or servers. In contrast, APS leverages data parallelism to balance workloads across decoding threads. Combining APS’s dynamic multi-threaded inference with existing parallelism strategies requires careful coordination of tensor, pipeline, and data parallelism, along with thoughtful placement of model parameters across devices to maintain optimal performance. As a promising future direction, we aim to leverage powerful modern serving frameworks such as SGLang and vLLM to address these integration challenges and enable scalable deployment.

We would like to thank you for the insightful feedback on our work. Below are our clarifications and responses to your comments:

Empirical experiments and generalization. In our experiments, we intentionally follow the methodological choices of the Stream of Search (SoS) study (Gandhi et al., 2024; COLM 2024 oral presentation), which centers on the Countdown task due to its structured reasoning challenge, high branching factor, and ease of solution verification. As in SoS, focusing on a controlled yet challenging domain enables us to rigorously evaluate across accuracy, latency, and scalability, providing clear insights into the effectiveness of our proposed approach.

Regarding generalization, the core innovation of our method—adaptive task decomposition and parallelization optimized via reinforcement learning—is naturally suited to more domain-general structured reasoning tasks beyond Countdown, such as puzzle solving and math reasoning. While it is not feasible to include results on benchmarks like GSM8K and HumanEval within the rebuttal period, we plan to explore applying APS to these domains as a next step. We will revise the discussion section to better highlight the generalizability of SoS and outline future directions for extending it to a wider range of tasks.

Prompt template. Since we performed supervised training on our model from scratch, there is no prompt in the form of textual descriptions for the task. During inference time, the model is asked to perform completion following the prefix for both APS and the baseline. We list the prefix for APS and the baseline below, using the task “constructing 79 from [45, 41, 94, 82]” as an example:

Prefix (APS): Moving to Node #0\nCurrent State: 79:[45, 41, 94, 82], Operations: []
Prefix (Baseline): Current State: 79:[45, 41, 94, 82], Operations: []

Budget conditioning. For the baseline, length-controlled models were trained by conditioning on context window size. Training samples are binned by length in 512-token increments, and this bin size is used as a conditioning signal. During inference, specifying a bin size controls the generated token count. For example, if the sample length or expected output is between 1537 and 2048 tokens, we prepend "Token Budget: 2048" to the prompt during training and inference. No other additional information, such as increased shot numbers, is included as context length increases. For instance, the SoS+ model in Fig. 3(a) uses a 4096 token budget conditioning, and Fig. 3(b) demonstrates the effect of token budget conditioning of 1024 and 4096.

The approach is straightforward and standard. We see the simplicity of our approach as a strength rather than a limitation. The use of 'spawn' and 'join' operations, inspired by multithreading in operating systems, is intentionally straightforward and forms a clear foundation for scalable parallel reasoning.

While there are prior works exploring parallelized decoding (e.g., self-consistency sampling), we are among the first to adapt and apply multithreading principles to enable language models to orchestrate parallel reasoning end-to-end. This approach offers an elegant and effective solution to the limitations of serial decoding in LMs (e.g., latency, context window bottlenecks). We demonstrate that LMs can learn this complex orchestration through a combination of supervised and reinforcement learning, leading to significant improvements in reasoning efficiency and scalability.

We would greatly appreciate citations for specific prior work that renders our approach “standard.” To the best of our knowledge, few existing work has demonstrated this kind of dynamic, learned task decomposition combined with parallel execution in language models. We believe our method is a pioneering step toward enabling language models to perform structured reasoning in a more scalable and efficient manner.

More powerful models & longer context. We trained a more powerful model with ~600M parameters and found that our method can continue to scale with increased model size, as detailed in our general response. This model size is already larger than some of the pretrained models (e.g., Qwen-2.5-0.5B). APS continues to significantly outperform the SoS+ baseline in this regime.

Note that one strength of APS is that it could reduce the requirements on the context length by exploring multiple directions in parallel. This makes the used context much less compared to the baseline. We demonstrated that the preferred way (as chosen by our end-to-end RL) to scale APS is through increasing the number of child threads, rather than the sequential search. This is shown in Fig. 6: For APS, RL scales the number of total tokens from 10k to 19k, achieving 7% performance boost. In contrast for our baseline, RL only scales the number of tokens from 2.2k to 2.7k, achieving 2.5% performance boost.