PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Adaptive Parallel Reasoning with Language Models

Jiayi Pan, Xiuyu Li, Long Lian, Charlie Snell, Yifei Zhou, Adam Yala, Trevor Darrell, Kurt Keutzer, Alane Suhr

arXiv:2504.15466v2cs.AIcs.CL

TL;DR

Existing inference-time reasoning methods face context, latency, and coordination limits. APR combines adaptive parent-child threading with end-to-end reinforcement learning to orchestrate serial and parallel computation. On Countdown, it improves performance within fixed context, scaling budgets, and matched latency, while evaluation remains limited to non-pretrained models and Countdown tasks.

  • Problem

    Serialized reasoning can produce long, latency-heavy outputs that exhaust context windows, while parallel methods may lack coordination and perform redundant computation.

  • Method

    APR lets language models adaptively distribute computation across parent and child threads using spawn() and join(), then optimizes both through end-to-end reinforcement learning.

  • Results

    83.4% vs 60.0% at 4k context, 80.1% vs 66.6% with 20k total tokens, and 75.2% vs 57.3% at around 5,000ms demonstrate APR’s reported advantages on Countdown.

  • Takeaways & Limitations

    APR demonstrates the potential of reasoning systems that dynamically structure inference to improve scalability and efficiency.

  • Takeaways & Limitations

    Experiments are restricted to non-pretrained language models on Countdown tasks, so broader applicability remains to be validated.

Abstract

from arXiv · show

Scaling inference-time computation has substantially improved the reasoning capabilities of language models. However, existing methods have significant limitations: serialized chain-of-thought approaches generate overly long outputs, leading to increased latency and exhausted context windows, while parallel methods such as self-consistency suffer from insufficient coordination, resulting in redundant computations and limited performance gains. To address these shortcomings, we propose Adaptive Parallel Reasoning (APR), a novel reasoning framework that enables language models to orchestrate both serialized and parallel computations end-to-end. APR generalizes existing reasoning methods by enabling adaptive multi-threaded inference using spawn() and join() operations. A key innovation is our end-to-end reinforcement learning strategy, optimizing both parent and child inference threads to enhance task success rate without requiring predefined reasoning structures. Experiments on the Countdown reasoning task demonstrate significant benefits of APR: (1) higher performance within the same context window (83.4% vs. 60.0% at 4k context); (2) superior scalability with increased computation (80.1% vs. 66.6% at 20k total tokens); (3) improved accuracy at equivalent latency (75.2% vs. 57.3% at approximately 5,000ms). APR represents a step towards enabling language models to autonomously optimize their reasoning processes through adaptive allocation of computation.

1 Introduction

Adaptive Parallel Reasoning addresses the latency, context-window, and coordination limits of serialized and independent parallel reasoning by combining parent-child computation with end-to-end optimization. On Countdown, it improves performance under matched context, compute, and latency conditions.

  • Serialized chain-of-thought produces lengthy outputs that increase latency and strain context windows, while parallel methods can lack coordination and waste computation.
  • APR lets parent threads spawn multiple child threads, which independently explore subtasks and return outcomes through join() before parent decoding continues.Child inference is performed simultaneously through batching, reducing real-time latency.
  • 83.4% vs 60.0% at 4k context: APR achieves higher performance than sequential methods within the same context window.
  • 80.1% vs 66.6% with a 20k total-token budget: APR shows superior scaling through wider parallelization alongside longer reasoning chains.
  • 75.2% vs 57.3% at around 5,000ms: APR achieves higher success rates than serialized search at equivalent latency.
  • These results indicate that models can learn when to reason serially and when to branch into parallel computation to allocate inference-time compute more efficiently.

2 Related Work

Prior inference-time scaling methods improve reasoning but face sequential latency and context limits, while independent parallel calls can duplicate computation. Related work explores coordination, serving efficiencies, and learned control of inference.

  • Inference-time scaling: Increasing test-time compute can improve language-model performance, but longer autoregressive outputs increase latency and strain single-context-window inference.
  • Parallelization in language model inference: Independent parallel inference can enhance performance, yet limited coordination across threads causes redundant computation and suboptimal resource utilization.
  • Parallelization in language model inference: Shared prefixes enable prefix caching and reuse, and APR’s parent-child structure creates opportunities to reduce parallelization overhead.
  • Training language models to control their own outputs: Other approaches train models to discard context or insert pause tokens to control inference for efficiency and performance.

3 Adaptive Parallel Reasoning

APR adaptively orchestrates parent and child threads to traverse reasoning trees through serialized and parallel inference. It uses model-generated spawn() and join() operations, supervised initialization, and end-to-end reinforcement learning, while current experiments remain limited to from-scratch Countdown settings.

  • Adaptive Parallel Reasoning: APR distributes inference across an adaptively orchestrated hierarchy of parent and child threads to parallelize reasoning at inference time.
  • 3.2 Multi-Threading at Inference Time: Unlike methods with developer-provided search structures, APR lets the model determine search structure during inference through spawn(msgs).The operation passes distinct contexts to multiple child threads executed in parallel.
  • 3.2 Multi-Threading at Inference Time: Child threads use the same language model to execute distinct subtasks simultaneously, with contexts limited to tokens passed by the parent.
  • 3.2 Multi-Threading at Inference Time: Child threads terminate with join(msg), returning selected tokens to the parent; retaining intermediate traces in children reduces parent token usage and inference cost.
  • 3.3 Training Models to Adaptively Parallelize their Reasoning: APR uses supervised demonstrations generated by a symbolic solver to initialize models for generating spawn() and join() reasoning traces.
  • 3.3 Training Models to Adaptively Parallelize their Reasoning: Hybrid search combines breadth-first and depth-first examples to diversify trajectories and support spontaneous parallelization behaviors.
  • 3.3 Training Models to Adaptively Parallelize their Reasoning: APR demonstrations reduce context-window bottlenecks because each decomposed demonstration contains only part of the search, enabling more extensive searches.
  • 3.3 Training Models to Adaptively Parallelize their Reasoning: End-to-end reinforcement learning further optimizes reasoning effectiveness and computational efficiency beyond supervised imitation.

4 Experiments

Experiments evaluate APR against serialized and parallel baselines on Countdown across context limits, compute budgets, and efficiency measures. APR benefits from reinforcement learning that increases both reasoning length and child-thread allocation, improving accuracy and scaling.

  • Experiment setup: The experiments compare APR with SoS+ and self-consistency baselines using accuracy, total tokens, sequential tokens, and latency.The evaluation uses Countdown problems and measures both task performance and computational efficiency.
  • Scaling with higher compute: 13.5% absolute improvement (66.6% →80.1%) over SoS+ cons@7 at 20k tokens shows APR scales more effectively with increased compute.APR initially under-performs below 4k tokens because some computation is spent orchestrating threads, but parallel exploration becomes advantageous as compute increases.
  • Scaling with context window size: APR evaluates performance under 1k-to-4k context-window constraints using 3, 6, and 10 child threads, while SoS+ uses context-conditioned models.Cumulative accuracy counts only traces whose lengths remain within each specified window.
  • Reinforcement learning: 83.4% accuracy after reinforcement learning, up from 75.5%, accompanies increases in average sequence length from 1,471 to 1,796 tokens and child threads from 6.1 to 8.2.The reported behavior indicates broader search was more optimal than deeper search for the Countdown task.
  • Efficiency of APR: APR improves efficiency by boosting accuracy with minimal additional sequential tokens beyond 2,048, while SoS+ approaches 3,000 tokens for marginal gains.Sequential token usage is defined as the average length of the longest non-parallelizable component across parent and child threads.
  • Ablation study: APR’s accuracy gain from reinforcement learning is 7.9%, compared with 2.7% for SoS+, alongside a 52.5% increase in APR total tokens versus 20.9% for SoS+.With ten child threads enforced before and after training, accuracy changes only from 83.2% to 83.3%, indicating the main gain comes from scaling test-time compute.

5 Conclusions, Limitations, and Future Work

Adaptive Parallel Reasoning distributes computation across serial and parallel reasoning paths, using parent-child threading and supervised plus end-to-end reinforcement learning. On Countdown, it improves performance under context, compute, and latency constraints, while current experiments remain limited to non-pretrained models and Countdown tasks.

  • Conclusions: Adaptive Parallel Reasoning adaptively distributes computation across serial and parallel reasoning paths through a parent-child threading mechanism.The method combines supervised training with fully end-to-end reinforcement learning to learn parallelization strategies without manually designed structures.
  • Conclusions: 83.4% vs. 60.0% at 4k context demonstrates higher performance within the same context window on Countdown.APR performs more effective search within fixed context constraints than sequential methods.
  • Conclusions: 80.1% vs. 66.6% using 20k tokens demonstrates superior scaling behavior as compute budgets increase.The improvement comes from combining wider parallelization with longer individual reasoning chains.
  • Conclusions: 75.2% vs. 57.3% at around 5,000 ms demonstrates higher success rates than serialized search methods at equivalent latency.This result indicates improved performance under the reported latency constraint.
  • Limitations and Future Work: Current experiments are restricted to non-pretrained language models on Countdown tasks, leaving broader applicability for future validation.The authors plan to extend APR to pretrained models and general reasoning tasks.
  • Limitations and Future Work: The current setup requires supervised bootstrapping by mimicking a symbolic solver, and future work aims to reduce this reliance.The authors propose exploring direct reinforcement learning from strong pretrained checkpoints.
  • Limitations and Future Work: APR currently uses fork and join operations for inter-thread communication, with other orchestration protocols left for future exploration.Proposed alternatives include any-to-any messaging, all-to-all communication, and subscription-based methods.

A.1 Additional Related Works

Related reasoning methods use test-time computation, independent parallel inference, structured reasoning paths, and multi-agent interaction to improve performance. Their limitations include long serialized outputs, redundant uncoordinated computation, fixed structures, and prompting-based rather than end-to-end optimization.

  • Language model reasoning algorithms: Chain-of-thought prompting established that additional test-time compute can improve language-model performance on end tasks.Later approaches optimize reasoning traces through supervised fine-tuning or reinforcement learning.
  • Parallel inference: Independent parallel inference improves task performance but lacks coordination between threads, causing redundant computation and suboptimal resource utilization.Self-consistency generates multiple candidate solutions independently and selects the most common solution.
  • Structured parallel reasoning: Tree-of-Thought, Graph-of-Thought, Skeleton-of-Thought, Atom-of-Thought, and Self-Ask structure reasoning exploration into multiple parallelizable calls.These methods build on simple ensembling approaches by organizing reasoning paths into structured calls.
  • Related optimization: Zhuge et al. optimize prompts and agent connectivity with reinforcement learning, but their underlying language models remain fixed.This differs from approaches that update the reasoning model itself.
  • Concurrent approaches: PASTA decomposes sequential tasks into parallel subtasks that merge into one main thread, whereas APR spawns exploratory threads and selectively incorporates successful outcomes.The comparison distinguishes APR’s selective integration from PASTA’s decomposition-and-merge design.
  • Concurrent approaches: Hogwild! Inference remains a prompting method, while APR optimizes collaborative reasoning end-to-end through supervised and reinforcement learning.The cited comparison contrasts prompting-based coordination with APR’s model optimization.
  • Controlling inference: PENCIL trains models to discard context, and pause-token methods add tokens to increase inference-time computation for efficiency or performance.These approaches also train language models to control aspects of their inference process.
  • Language model serving systems: Modern serving systems such as vLLM and SGLang improve token and total throughput through batched processing and related optimizations.The passage frames serving systems as infrastructure for efficient language-model inference.

A.2 Implementation Details

The implementation uses a Llama2-following model and separates supervised and reinforcement-learning training configurations. Supervised training uses TPUv3 cores, while reinforcement learning uses GRPO with SGLang rollouts on GPUs.

  • Model architecture: The model follows Llama2 and has 228M non-embedding parameters, 293M total parameters, 18 layers, 1024 hidden dimensions, and 16 attention heads.Its context window is 4096 tokens.
  • Supervised training: Supervised training uses 128 TPUv3 cores, batch size 256, learning rate 5e-5, and 19,000 training steps.The training runs for approximately 10 epochs following the cited prior setup.
  • Reinforcement learning: Reinforcement learning uses GRPO on 2 Nvidia GPUs with batch size 64, five rollouts per sample, and SGLang for rollout.Evaluation uses greedy sampling unless otherwise stated because it performs best for both SoS+ and APR.
  • Reinforcement learning: Reinforcement-learning training uses learning rate 1 × 10^-5, PPO clip ratio 0.2, 150 total steps, and two inner PPO optimization steps per step.Gradient clipping and separate KL divergence factors are used for training stability.

A.3 Additional Results with Larger Models

Additional larger-model experiments compare APR and SoS+ across compute levels for 200M and 600M parameter models. APR consistently leads SoS+ and scales more strongly with model size.

  • Scaling with model size: Both methods improve with increased model capacity, while APR maintains a substantial lead over SoS+ at every compute level.The reported comparison uses accuracy over average total compute.
  • Scaling with model size: APR consistently outperforms SoS+ across average total compute levels for 200M and 600M parameter models.Figure 8 plots accuracy against average total compute measured in thousands of tokens.

A.4 Additional Results with Pretrained Models

APR’s gains extend to a pretrained Qwen2.5 1.5B model, where it significantly outperforms SoS+ and mirrors results observed with Llama2.

  • APR significantly outperforms SoS+ when fine-tuned on a pretrained Qwen2.5 1.5B model.The comparison uses approximately 4k supervised training steps on the same SoS+ and APR demonstration data.
  • The pretrained-Qwen results mirror the performance trend observed with Llama2.The paper describes this consistency as evidence that APR’s gains extend beyond models trained from scratch.
  • Table 1 compares SoS+ and APR across pretrained Qwen2.5 1.5B and Llama2 200M models.

A.5 Extended Context Window Experiments

On the five-number Countdown task, APR continues improving with larger context budgets and surpasses SoS+ beyond 3.5k tokens.

  • APR continues to improve up to about 6k tokens and outperforms SoS+ beyond 3.5k tokens.
  • APR achieves gains of 7% and 11% over SoS+ at an 8k-token budget.
  • The evaluation uses a five-number Countdown variant with a search space 40× larger than the referenced task.Context budgets extend up to 8k tokens.

A.6 Symbolic Search Algorithm

The symbolic search implementation combines state expansion with selective parallel exploration, while experiments examine baselines, supervision, temperature, compute scaling, and hardware-related timing limits.

  • Symbolic Search Algorithm: APR spawns parallel sub-searches when the model selects DFS on a promising node, exploring the resulting nodes concurrently.The parallel search version follows the same broad algorithm while replacing selected DFS exploration with parallel sub-searches.
  • Symbolic Search Algorithm: APR’s symbolic solver expands states, checks goal attainment, and either explores promising states in parallel or queues states for further exploration.Algorithm 2 returns when a parallel result reaches the goal; otherwise, states are added to the deque.
  • Symbolic Search Algorithm: SE selects next operations using the top K candidates ranked by a multiply heuristic based on the target number and its factors.The promising function determines whether the current node receives the parallel exploration treatment.
  • Experimental Comparisons: SoS+ consistently outperforms original SoS across sampling temperatures, while enhanced SoS+ supervision remains insufficient to close its gap to APR.Increasing beam size can produce trajectories exceeding context length; rejection sampling provides only modest gains.
  • Experimental Comparisons: APR outperforms SoS+ across temperatures and shows more reliable, efficient scaling as total compute increases.The relative advantages remain consistent before and after reinforcement learning.
  • Timing Considerations: Hardware constraints create a slight mismatch between sequential-token counts and wall-clock time when multiple child threads share GPUs.Additional GPUs and improved load balancing can mitigate uneven workloads and increased device load.
Loading 2504.15466v2…