Learning to Reason Across Parallel Samples for LLM Reasoning

TL;DR

This paper introduces a compact LLM called Sample Set Aggregator (SSA) that learns to effectively combine multiple reasoning samples to improve accuracy in large language model reasoning tasks, outperforming traditional aggregation methods.

Contribution

The paper presents SSA, a novel reinforcement learning-trained model that aggregates multiple samples for reasoning, enhancing performance and efficiency over naive methods and larger models.

Findings

SSA improves pass@5 by 8% over majority voting on MATH.

SSA surpasses larger model-based re-ranking methods.

SSA generalizes well across datasets, sample sizes, and model scales.

Abstract

Scaling test-time compute brings substantial performance gains for large language models (LLMs). By sampling multiple answers and heuristically aggregate their answers (e.g., either through majority voting or using verifiers to rank the answers), one can achieve consistent performance gains in math domains. In this paper, we propose a new way to leverage such multiple sample set. We train a compact LLM, called Sample Set Aggregator (SSA), that takes a concatenated sequence of multiple samples and output the final answer, optimizing it for the answer accuracy with reinforcement learning. Experiments on five reasoning datasets demonstrate both the efficacy and efficiency of SSA. Notably, SSA improves over naive majority voting by 8% pass@5 on MATH. Furthermore, our 3B SSA surpasses model-based re-ranking with a much larger 72B process reward model. Our analysis also shows promising generalization ability of SSA, across sample set sizes, base model families and scales, and tasks. By separating LLMs to generate answers and LLMs to analyze and aggregate sampled answers, our approach can work with the outputs from premier black box models easily and efficiently.