SPIRAL: Learning to Search and Aggregate

27 July 2026

Paper authors: Jubayer Ibn Hamid, Ifdita Hasan Orney, Michael Y. Li, Omar Shaikh, Yoonho Lee, Dorsa Sadigh, Chelsea Finn, and Noah Goodman.

Idea

Mismatch between how reasoning models are trained and how people actually spend extra inference compute at test time.

GRPO trains the model to produce a strong single reasoning trace, but with a lot of compute people typically try parallel strategies and then combine useful ones and then repeat.

Thus, SPIRAL trains on that second procedure. Objective:

J (θ) = 𝔼_{y_{1 : n} ~ π_{θ} (\cdot ∣ x)} [𝔼_{y_{*} ~ π_{θ} (\cdot ∣ x, y_{1 : n})} [r (x, y_{*})]] .

Here $x$ is the original problem, $y_{1}, \dots, y_{n}$ are independent search traces, and $y_{*}$ is the later aggregation trace. $r (x, y_{*})$ scores only the final aggregation answer. Same model is used for both stages.

Key point is that the search traces are rewarded for helping the later aggregator produce a correct answer (not just for being correct).

Three types of inference compute: sequential + parallel + aggregative. It jointly learns all three.

Basically, the issue with final-answer-only GRPO is you can’t directly reward exploration and partial progress → too sparse rewards. Uses set RL.

Set reinforcement learning

Scores a collection of samples jointly, $f (x, y_{1 : n})$ . In SPIRAL, the set score is:

f_{S P I R A L} (x, y_{1 : n}) = 𝔼_{y_{*} ~ π_{θ} (\cdot ∣ x, y_{1 : n})} [r (x, y_{*})] .

Feed $y_{1 : n}$ to the aggregator. Measure the expected reward of what the aggregator produces.

Then the simplest policy gradient estimator would give every trace in the set the same advantage:

\nabla_{θ} J_{s e a r c h} \approx A^{♯} (x, y_{1 : n}) \sum_{i = 1}^{n} \nabla_{θ} \log π_{θ} (y_{i} ∣ x) .

$A^{♯}$ is the advantage for a set. This is noisy though, so they sample a large pool of traces and construct overlapping subsets.

Their implementation

For each maths problem:

1. Sample eight search traces

y_{1}, \dots, y_{8} \overset{i . i . d .}{~} π_{θ} (\cdot ∣ x) .

Ordinary CoT from the prompt.

2. Construct four sets of four traces

From the eight traces, construct $G_{1}, G_{2}, G_{3}, G_{4}$ , where each $G_{i}$ contains four distinct search traces. Overlap lowers noise on whether a certain trace is helpful.

These are picked uniformly at random, without replacement, from the $(\binom{8}{4})$ available sets.

3. Sample four aggregation traces for each set

y_{1}^{G_{i}}, \dots, y_{4}^{G_{i}} ~ π_{θ} (\cdot ∣ x, G_{i}) .

So 16 aggregation traces. The aggregation prompt tells the model to audit each candidate trace, identify useful supported ideas, identify errors and unsupported leaps, verify the useful components, and synthesise a final solution.

4. Reward the final aggregation

16 rewards.

5. Advantage for the aggregation traces

The aggregations get ordinary group-relative advantages with respect to the other aggregations from the same $G_{i}$ :

A (x, y_{j}^{G_{i}}) = r (x, y_{j}^{G_{i}}) - {\bar{r}}_{G_{i}},

{\bar{r}}_{G_{i}} = \frac{1}{4} \sum_{j = 1}^{4} r (x, y_{j}^{G_{i}}) .

The model gets better at aggregating what it gets given.

6. Advantage for each set

For a set $G_{i}$ , SPIRAL averages the reward of its four aggregations ( ${\bar{r}}_{G_{i}}$ ), then takes the average across all four sets to get group-relative advantages for each $G_{i}$ :

A^{♯} (x, G_{i}) = {\bar{r}}_{G_{i}} - \frac{1}{4} \sum_{k = 1}^{4} {\bar{r}}_{G_{k}} .

7. Advantage for each search trace

Suppose search trace $y_{1}$ appeared in $G_{1}, G_{2}, G_{3}$ . Then SPIRAL gives $y_{1}$ the average advantage of the sets containing it:

A_{m a r g}^{♯} (x, y_{1}) = \frac{A^{♯} (x, G_{1}) + A^{♯} (x, G_{2}) + A^{♯} (x, G_{3})}{3} .

Thus:

\nabla_{θ} J_{s e a r c h} \approx \sum_{i = 1}^{8} A_{m a r g}^{♯} (x, y_{i}) \nabla_{θ} \log π_{θ} (y_{i} ∣ x) .

They can apply it recursively too.

Issue: the sets $G_{i}$ are not ordered, but the attention is. Therefore you would actually have to sample from $_{8} P_{4}$ rather than $(\binom{8}{4})$ to avoid bias.

Gradient

\nabla_{θ} J (θ) = \nabla_{θ} 𝔼_{y_{1 : n} ~ π_{θ} (\cdot ∣ x)} [𝔼_{y_{*} ~ π_{θ} (\cdot ∣ x, y_{1 : n})} [r (x, y_{*})]] .

\nabla_{θ} J (θ) = 𝔼_{y_{1 : n} ~ π_{θ} (\cdot ∣ x)} 𝔼_{y_{*} ~ π_{θ} (\cdot ∣ x, y_{1 : n})} [r (x, y_{*}) \nabla_{θ} \log (π_{θ} (y_{1 : n} ∣ x) π_{θ} (y_{*} ∣ x, y_{1 : n}))] .

= 𝔼_{y_{1 : n} ~ π_{θ} (\cdot ∣ x)} 𝔼_{y_{*} ~ π_{θ} (\cdot ∣ x, y_{1 : n})} [r (x, y_{*}) (\nabla_{θ} \log π_{θ} (y_{1 : n} ∣ x) + \nabla_{θ} \log π_{θ} (y_{*} ∣ x, y_{1 : n}))] .

Set RL gradient

= 𝔼_{y_{1 : n} ~ π_{θ} (\cdot ∣ x)} [\nabla_{θ} \log π_{θ} (y_{1 : n} ∣ x) 𝔼_{y_{*} ~ π_{θ} (\cdot ∣ x, y_{1 : n})} [r (x, y_{*})]]

+ Standard RL gradient

+ 𝔼_{y_{1 : n} ~ π_{θ} (\cdot ∣ x)} [𝔼_{y_{*} ~ π_{θ} (\cdot ∣ x, y_{1 : n})} [\nabla_{θ} \log π_{θ} (y_{*} ∣ x, y_{1 : n}) r (x, y_{*})]] .

Thoughts

I like this, but I wonder whether it’s really needed anymore with agentic systems that can launch sub-agents during RL, since this seems kind of overlapping. Or maybe we could fuse the two together and do this within the large-scale agentic stuff somehow.