SPIRAL: Learning to Search and Aggregate

Abstract: LLM reasoning can be substantially improved at test time via scaffolds that scale inference compute across different primitives -- sequential reasoning within a trace, independently sampled parallel traces, and aggregation of multiple reasoning traces into a final response. During post-training, however, LLMs are optimized only for sequential reasoning within a single trace. We introduce Sequential-Parallel-Aggregative Reinforcement Learning (SPIRAL), a framework in which a LLM is trained to use all three primitives, as part of a unified inference compute pipeline. Concretely, the LLM first samples a set of independent traces in parallel, each produced through sequential chain-of-thought reasoning, and then generates a final aggregation trace conditioned on those traces; all components are optimized end-to-end against the reward of the final aggregated response. To train this system, SPIRAL uses set reinforcement learning to teach models to produce a set of traces that are collectively useful for an aggregator and standard reinforcement learning to teach models to aggregate the set into improved final responses. Our experiments on reasoning tasks show that SPIRAL effectively scales with inference compute, outperforming GRPO by up to 11$\times$ scaling efficiency and 15% higher performance when all three compute primitives are scaled.

• The paper introduces Spiral, a framework that integrates sequential, parallel, and aggregative reinforcement learning to optimize language model inference.
• It leverages set reinforcement learning with a marginal advantage estimator to assign collective rewards to sets of generated responses.
• Empirical results demonstrate significant gains in pass@k performance and recursive self-aggregation accuracy, highlighting improved search diversity and scalability.

Spiral: Learning to Search and Aggregate

Motivation and Background

Current reinforcement learning (RL) pipelines for LMs overwhelmingly focus on sequential inference—training the model to optimize a single chain-of-thought trace end-to-end. However, RL fine-tuning practices diverge substantially from test-time deployment, where practitioners increasingly rely on scaffolds that scale inference across three primitives: sequential compute (deep tokens per trace), parallel compute (multiple traces for the same prompt), and aggregative compute (combining multiple traces into a synthesized output). The lack of explicit training to coordinate these primitives creates an alignment gap, hampering models' ability to utilize available compute effectively and limiting the scalability of LMs on tasks requiring complex reasoning.

The Spiral framework, Sequential-Parallel-Aggregative Reinforcement Learning, addresses this limitation by exposing the model to all three inference compute primitives during RL fine-tuning. The central objective is to enable LMs to co-evolve search procedures that natively balance exploration (through parallel traces), exploitation (via sequential chain-of-thought), and synthesis (via aggregation), thereby internalizing strategies that surpass the rigid, hand-designed scaffolds in prior works.

Figure 1: Spiral trains LMs end-to-end over sequential, parallel, and aggregative compute; empirically, Spiral demonstrates superior scaling efficiency compared to GRPO across compute axes, particularly in hybrid and recursive self-aggregation regimes.

Formalization of Set RL and Inference Compute Coupling

Spiral's foundation is the extension of set reinforcement learning (set RL) to handle LM fine-tuning with set-valued credit assignment. In standard RL, trajectories are individually rewarded; by contrast, set RL assigns a collective reward to sets of generations sampled conditionally independent from the same input. The policy is updated with a low-variance, marginal advantage estimator reflecting the average utility of each sampled generation across all sets in which it appears. This formulation explicitly couples the search process across parallel traces, rewarding their collective usefulness for downstream aggregation rather than individual correctness.

The training objective for Spiral is:

$$\max_{\theta, \phi} \mathbb{E}_{y_{1:n} \sim \pi_\theta(\cdot \mid x)}\left[ \mathbb{E}_{y_* \sim \pi_\phi(\cdot \mid x, y_{1:n})} [r(x, y_*)] \right]$$

where $\pi_\theta$ parameterizes the search trace policy and $\pi_\phi$ (often $\pi_\theta$ itself) parameterizes the aggregation trace. $r(x, y_*)$ is the final scalar reward for the aggregated response, strictly aligning credit assignment with utility induced at the level of the synthesized output.

Spiral Algorithmic Structure

The Spiral algorithm implements a two-level data collection and optimization pipeline. For each training prompt, the model samples a set of $N_1$ search traces. From these, it constructs $K$ sets of size $n$ by random sampling, and for each set, it samples $N_2$ aggregation traces conditioned jointly on the problem and the selected search traces. Rewards are only provided for final aggregation outputs. The set RL marginal advantage estimator is applied to assign credit to search traces (optimized for "aggregatability"), while standard RL is applied to aggregation traces. Gradient estimators are derived to handle both symmetric and asymmetric set objectives—with explicit theoretical treatment for ordered tuples and the effect of ordering bias.

(Figure 2)

Figure 2: Spiral’s two-level data collection structure, highlighting independent sampling of search traces, batched set construction for aggregation, and credit assignment pipelines for both search and aggregation generations.

This approach ensures that the LM does not collapse entropy in its search outputs; traces that appear incorrect in isolation but contribute constructively upon aggregation are positively reinforced, resulting in diverse and complementary candidate pools.

Empirical Results

Parallel Scaling (Pass@$k$)

Spiral demonstrates a significant scaling advantage over strong RL baselines such as GRPO when evaluated under parallel compute scaling (pass@$\pi_\theta$0 rate). When $\pi_\theta$1 independently sampled search traces are drawn per problem, Spiral achieves up to 11-fold improvements in scaling efficiency on mathematical reasoning evaluation, indicating a qualitative difference in how the model utilizes expanded compute for search and exploration.

Figure 3: Pass@$\pi_\theta$2 as a function of sampled parallel attempts; Spiral achieves much higher test-set coverage than GRPO as $\pi_\theta$3 increases, reflecting effective search diversity.

Recursive Self-Aggregation

The framework's most crucial benefit emerges in regimes where both parallel and aggregative compute are scaled, i.e., recursive self-aggregation (RSA). Here, Spiral’s pass@1 rate under recursive aggregation (conditioning on sets of traces and re-aggregating at each level) shows a relative gain of up to 13.5% versus conventional RL training. These gains persist across aggregation depths, demonstrating Spiral’s ability to produce traces that can be verified, refined, and synthesized more effectively.

Figure 4: Spiral’s pass@1 accuracy under multiple steps of RSA, showing stable and superior improvements over baselines.

Diversity and Entropy Analysis

Token-level entropy over the course of training reveals that Spiral does not prematurely converge to low-diversity search behaviors—a common collapse observed with single-trace reward alignment methods like GRPO. Instead, Spiral maintains higher entropy in its search process, allowing for robust population-based exploration.

Figure 5: Token-level entropy comparison during training; Spiral maintains higher entropy for search traces compared to GRPO, confirming reduced collapse and greater search diversity.

Comparison with Rule-Based Aggregation

Spiral-mediated aggregation (self-aggregation) outperforms rule-based schemes (e.g., majority voting over search traces) in the high-compute regime. While both Spiral and GRPO perform similarly with majority voting, Spiral’s pass@1 under learned (model-based) aggregation scales more favorably, demonstrating that credit assignment over candidate search trace populations and synergy with learned aggregation dominate over hand-crafted selection heuristics.

Figure 6: Aggregation scheme comparison: model-based Spiral aggregation scales better than majority voting and outperforms GRPO across the number of parallel traces.

Sequential Compute Scaling

When evaluating pure sequential compute scaling (increasing chain-of-thought length per single trace), all compared models saturate at much lower overall coverage than via parallel or aggregative scaling, highlighting the limitations of further budget increases in a single trajectory versus coordinated parallel-aggregate exploration.

Theoretical and Practical Implications

Spiral brings two main theoretical advances: (1) It bridges the disconnect between RL fine-tuning and test-time harnesses by aligning training credit assignment with complex, multi-primitives inference pipelines; and (2) it introduces scalable, unbiased set RL gradient estimators capable of supporting both symmetric and asymmetric context objectives, confirmed with explicit analytic treatment of ordering effects.

Practically, Spiral allows LMs to internalize cooperative, search-and-aggregate strategies that enable effective utilization of large inference budgets. As a result, models trained with Spiral can outperform equivalent-size models or even larger models in compute-constrained or high-precision domains (e.g., mathematical reasoning, code synthesis, or multi-step theorem proving) where diverse, aggregated reasoning steps are key.

Future Directions

The next steps for Spiral research include scaling up model and dataset size (noting the current results are on 4b-parameter scale), more in-depth analysis of diversity behaviors, independent verification task evaluation to isolate aggregation improvements, and combinations with alternative baselines (e.g., GRPO-trained search with Spiral aggregation). There is also opportunity for extending Spiral to modular settings, where different search and aggregation policies can be trained and composed, and for integrating meta-learning harnesses that adaptively select primitives at test time.

Conclusion

Spiral provides an explicit algorithmic and theoretical framework for training LLMs to effectively utilize, coordinate, and scale sequential, parallel, and aggregative inference compute primitives. Empirical evidence substantiates strong scaling improvements and superior search-aggregate coordination over rigid hand-crafted pipelines and standard RL training. These results reframe the landscape of fine-tuning for advanced LMs, providing a principled foundation for end-to-end learning over compositional inference pipelines and multi-agent scaffolds.

---

Reference: "SPIRAL: Learning to Search and Aggregate" (2606.23595)