🤖 AI

MASPRM: Multi-Agent System Process Reward Model

The paper introduces MASPRM, a multi-agent system process reward model trained on terminal outcomes to score intermediate agent messages, which significantly improves inference-time search performance on reasoning benchmarks compared to existing models by enabling more effective step-level beam search and Monte Carlo Tree Search.

Original authors: Milad Yazdani, Mahdi Mostajabdaveh, Zirui Zhou, Ying Xiong

Published 2026-07-16

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Milad Yazdani, Mahdi Mostajabdaveh, Zirui Zhou, Ying Xiong

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a really tricky puzzle, like a complex math problem or a logic riddle. In the world of artificial intelligence, we have "Large Language Models" (LLMs) that are like super-smart, but sometimes scattered, students. If you ask them a hard question, they might try to answer it in one go, but they often get stuck or make a mistake early on and keep going down the wrong path. To fix this, scientists have started using "Multi-Agent Systems." Think of this not as one student working alone, but as a study group. You have a "Reader" who reads the question, a "Planner" who makes a strategy, a "Solver" who does the math, and a "Verifier" who checks the work. They pass notes back and forth in a specific order.

However, there's a catch. If you let this study group run wild, they might waste a lot of time and energy (computing power) arguing about ideas that are already wrong. It's like a group of friends trying to solve a mystery, but they keep chasing a red herring because no one stopped to say, "Wait, this clue doesn't make sense." To stop this waste, researchers use "search algorithms" like Monte Carlo Tree Search (MCTS). This is like a coach who says, "Let's try five different paths for the next step, see which one looks most promising, and then focus on that one." But here is the big problem: How does the coach know which path is promising? Usually, the coach only knows if the final answer is right or wrong. They don't know if the middle steps were good. This paper introduces a new tool to help the coach make better decisions right in the middle of the game.

The Paper: MASPRM

The researchers behind this paper, from the University of British Columbia and Huawei, are tackling exactly that problem. They created something called MASPRM (Multi-Agent System Process Reward Model). You can think of MASPRM as a "Super Coach" or a "Process Judge" that doesn't just wait for the final answer to see if the group succeeded. Instead, it watches the study group's conversation as it happens and scores every single note they pass to each other.

In the past, if an AI team made a mistake in step 3, the system wouldn't know until the very end when the final answer was wrong. By then, the AI had already wasted a ton of energy exploring that bad path. MASPRM changes the game. It looks at the "routed transcript"—which is just a fancy word for the ordered list of messages the agents sent to each other. It scores these messages to say, "Hey, this specific message from the Planner was a great idea, let's keep going," or "This message from the Solver is leading us nowhere, let's cut that branch off immediately."

How They Taught the Coach

One of the coolest parts of this paper is how they trained this Super Coach. Usually, to teach a computer to judge steps, you need a human to sit there and label every single step as "good" or "bad." That is incredibly expensive and slow. The authors found a clever shortcut. They let the AI agents run thousands of simulations using a search method called MCTS. In these simulations, the agents would explore many different paths. At the very end of each path, they knew if the answer was right (+1) or wrong (-1).

Then, they used a mathematical trick called "backpropagation." Imagine a tree where the leaves are the final results. If a leaf is a "win," the value of that win trickles back up the branches, making the earlier steps look valuable. If a leaf is a "loss," the value trickles back up as a negative. MASPRM learned to predict these values just by looking at the conversation history, without a single human ever saying, "This step was good." It learned to spot the difference between a promising conversation and a dead-end argument just by seeing which ones eventually led to the right answer.

What They Found

The team tested this new coach on four famous benchmarks: GSM8K and MATH (for math problems), and MMLU and LOGIQA (for general knowledge and logic). They compared MASPRM against two other methods:

Policy Likelihood: This is like the AI just guessing, "I think this next word is likely," without any real understanding of the goal.
ORM (Outcome Reward Model): This is the old-school coach that only checks the final answer, ignoring the steps in between.

The results were quite clear. When the researchers used MASPRM to guide the search (specifically using MCTS and a method called Step-Level Beam Search), the AI got significantly better at solving problems.

On the 1.5 billion parameter model (a smaller, faster AI), MASPRM improved the success rate by 2.0 to 3.0 points compared to the old Outcome Reward Model.
On the 7 billion parameter model (a larger, smarter AI), the improvement was massive, jumping by 4.1 to 14.5 points.
In the hardest test (GSM8K), the combination of the 7B model and MASPRM reached a Hit@1 of 82.9%, meaning it got the right answer on the first try nearly 83% of the time.

The paper suggests that MASPRM is especially good at "stepwise search." This is when the AI has to make a series of decisions, like choosing which agent speaks next or what the next sentence should be. Because MASPRM can judge these intermediate steps, it prevents the AI from wasting time on bad ideas. It also made the AI's choices more reliable; the gap between the "best" guess and the "fifth-best" guess shrank, meaning the AI was more confident in its top choices.

The Limits and the Future

The authors are careful to note that this isn't a magic wand that fixes everything. Their system works best when the agents have a fixed schedule (like a strict order of who speaks when) and when the final answer can be clearly checked (like a math problem with a specific number). They admit that for open-ended tasks where there is no single "correct" answer, or for systems where the agents change their roles dynamically on the fly, this method is still being figured out. They also found that if the coach only sees what one agent sees (instead of the whole group's conversation), it doesn't work as well, suggesting that having a "global view" of the team's chat is crucial for the best performance.

In short, MASPRM is a breakthrough in teaching AI teams how to self-correct in real-time. Instead of waiting until the end to realize they were wrong, they now have a coach that can spot a bad path early and steer the team back on track, saving time and getting better results. It suggests that for complex reasoning tasks, the secret isn't just having a bigger brain, but having a better way to guide the thinking process as it happens.

Technical Summary: MASPRM (Multi-Agent System Process Reward Model)

Problem Statement

Inference-time search over Multi-Agent Systems (MAS) often wastes computational resources because existing methods struggle to identify which specific agent's intermediate message advanced progress toward a solution. While answer-level methods (e.g., majority voting, verifier reranking) can select among completed solutions, they fail to provide step-level guidance. Consequently, search procedures like Monte Carlo Tree Search (MCTS) or Beam Search may continue to expand prefixes that are likely under the policy but misaligned with eventual correctness, leading to unproductive exploration.

Extending single-agent Process Reward Models (PRMs) to MAS is non-trivial due to four specific challenges:

Score Granularity: A "step" in MAS is an agent action (planning, routing, tool use) rather than a token, requiring scoring of routed transcript prefixes.
Schedule and Graph: The value of a transcript depends on the specific order of agent actions and message routing.
Heterogeneity: Agents have distinct roles, tool access, and prompts, meaning intermediate prefixes originate from different distributions.
Partial Observability: Each agent only observes its local view of the global transcript, yet a global scorer is needed for consistent guidance.

Methodology

The authors propose MASPRM, a process reward model designed to score routed transcript prefixes and agent-proposed actions, acting as an inference-time controller for step-level Beam Search (SBS) and MCTS.

Core Architecture

State Representation: A state $s_t$ is defined as $(x, H_{t-1}, i_t)$ , where $x$ is the input, $H_{t-1}$ is the full routed transcript prefix, and $i_t$ is the acting agent.
Observability: While acting agents generate actions based on their local views ( $H^{(i)}_{t-1}$ ), the MASPRM scorer conditions on the full routed transcript prefix ( $H_{t-1}$ ) and the acting agent's identity. This allows the scorer to handle heterogeneity and scheduling implicitly within the input representation.
Training Data Generation (MAS-MCTS): The model is trained without human step-level annotations. Instead, the authors use MAS-MCTS rollouts labeled only with terminal outcome rewards ( $R(s_T) \in \{+1, -1\}$ $R (s_{T}) \in {+ 1, - 1}$ ).
- Terminal rewards are backed up to edges to estimate action values $\hat{Q}(s, a)$ .
- Preference pairs are constructed by comparing sibling successor states (children of the same parent) where one has a higher backed-up value than the other.
Training Objective: MASPRM is trained using a Bradley-Terry ranking loss on these preference pairs. This is chosen because the scorer's primary role at inference is to rank candidate successor states during search, rather than perform pointwise regression.
Inference Guidance:
- MCTS: MASPRM scores newly expanded edges to initialize virtual visits and guides the Upper Confidence Bound (UCT) selection.
- SBS: MASPRM scores all sampled successors at each beam depth, allowing the algorithm to prune unpromising branches and retain the top candidates.

Key Contributions

Formulation of MASPRM: A process reward model specifically designed to score agent actions and routed transcript prefixes in collaborative multi-agent settings, enabling agent-aware credit assignment.
MAS-MCTS Supervision: A training paradigm that backs up terminal rewards to agent-action values to train the process model without requiring human step-level annotations.
Guided Search Algorithms: The introduction of MASPRM-guided SBS and MCTS over agent rollouts, which improves decision quality and effectively prunes unproductive branches in predefined MAS configurations.

Experimental Results

The authors evaluated MASPRM on four benchmarks: GSM8K and MATH (math reasoning), and MMLU and LOGIQA (knowledge/reading comprehension). The experiments utilized Qwen2.5-1.5B and Qwen2.5-7B as base models for agents and scorers.

Performance vs. Outcome Reward Models (ORM): Under matched scorer sizes and comparable MCTS budgets, MASPRM consistently outperformed size-matched ORMs.
- At 1.5B scorer size: MASPRM exceeded ORM by +2.0 to +3.0 points across all benchmarks.
- At 7B scorer size: MASPRM exceeded ORM by +4.1 to +14.5 points.
Performance vs. Policy Likelihood: MASPRM significantly improved Hit@1 over policy likelihood guidance, particularly in stepwise search methods.
- With a 7B scorer and 1.5B agents under MCTS, MASPRM achieved an average gain of +13.4 points over policy likelihood.
- On GSM8K, the 7B MASPRM + MCTS configuration reached 82.9% Hit@1, a substantial improvement over the single-pass baseline (47.2%) and the 1.5B ORM (66.8%).
Ranking Quality: MASPRM improved the quality of ranked solutions, reducing the gap between Hit@1 and Hit@5 by up to 10.3 points, indicating better identification of correct trajectories early in the search.
Generalization:
- Cross-Dataset: MASPRM trained on one dataset (e.g., GSM8K) transferred zero-shot to related datasets (e.g., MATH) with gains of +5.4 to +11.1 points over policy likelihood, though in-domain training remained superior.
- Cross-Agent: A MASPRM trained on Qwen2.5-1.5B rollouts successfully guided Llama3.2 3B agents without retraining, improving Hit@1 by +8.7 to +13.0 points over policy likelihood.

Significance and Claims

The paper claims that inference-time multi-agent search becomes substantially more reliable when guided by a learned process value model trained solely from terminal rewards.

Key claims regarding the significance of the work include:

Efficiency of Small Agents: The approach demonstrates that smaller LLMs (e.g., 1.5B) can achieve high performance on complex reasoning tasks when paired with a learned process reward model and structured search, offering a cost-effective alternative to scaling up model parameters.
Step-Level Guidance: The results suggest that process guidance is most beneficial for algorithms that make repeated intermediate search decisions (like MCTS), rather than just final answer selection.
Scalability: The method scales effectively with scorer size; moving from a 1.5B to a 7B scorer yielded significant additional gains, particularly in MCTS settings.
Limitations: The authors modestly note that while MASPRM complements large-model scaling, it does not yet substitute for it on the hardest reasoning tasks (e.g., LOGIQA, where it remained below GPT-4o-mini and Mistral 24B). The current scope is limited to predefined, role-pipelined MAS configurations on verifiable terminal-answer tasks.

The work concludes that learned process-level value signals, derived without human step-level annotations, are a useful and effective component for guiding multi-agent search.