🤖 machine learning

FunPRM: Function-as-Step Process Reward Model with Meta Reward Correction for Code Generation

FunPRM enhances code generation by treating modular functions as reasoning steps for Process Reward Models and employing a meta-learning mechanism to correct noisy partial-solution rewards using unit-test-based final evaluations, thereby achieving state-of-the-art performance and more readable code.

Original authors: Ruiyi Zhang, Peijia Qin, Qi Cao, Eric Xue, Pengtao Xie

Published 2026-02-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ruiyi Zhang, Peijia Qin, Qi Cao, Eric Xue, Pengtao Xie

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 teach a very talented but sometimes overconfident robot chef how to cook a complex, multi-course meal. The robot is great at following instructions, but when asked to make a complicated dish, it often rushes, mixes up steps, or adds the wrong ingredients halfway through, resulting in a burnt meal.

This paper introduces FunPRM, a new "tasting coach" designed to help these AI chefs (Large Language Models) cook better code. Here is how it works, broken down into simple concepts:

1. The Problem: The "Line-by-Line" Confusion

Previously, when trying to grade a robot's cooking process, coaches would look at the recipe one sentence at a time.

The Issue: In math, steps are clear (Step 1: Add numbers, Step 2: Divide). But in coding, a "step" is hard to define. Is a step a single line of code? If so, a complex dish might have 500 steps. Grading 500 tiny steps is slow, confusing, and often wrong because a single line might look okay but be part of a broken plan.
The Analogy: It's like a teacher grading a student's essay by checking every single letter. If the student writes "T-h-e," the teacher says "Good!" even if the next word is "b-a-d." You miss the big picture.

2. The First Innovation: "Chain-of-Function" (The Recipe Book)

FunPRM changes the rules. Instead of looking at every single line, it tells the robot chef: "Break your recipe into distinct, named chapters."

How it works: The AI is prompted to write code where every major task is its own separate "function" (a mini-program with its own name and description).
The Analogy: Instead of a giant wall of text, the robot now writes a recipe book with clear chapters: Chapter 1: Chop the Vegetables, Chapter 2: Sauté the Onions, Chapter 3: Simmer the Sauce.
The Benefit: The coach (FunPRM) can now grade the whole "Chop Vegetables" chapter as one step. If the chopping is messy, the coach knows immediately. This makes the code easier for humans to read and easier for the AI to learn from.

3. The Second Innovation: The "Meta-Coach" (Cleaning the Noise)

Even with clear chapters, the coach still has a problem: How do we know if a half-finished dish is good?

The Issue: To train the coach, we usually guess if a partial dish is good by simulating "what if we finished it?" (a method called Monte Carlo sampling). This is like guessing if a half-baked cake will rise by shaking the pan. It's fast, but the guess is often "noisy" (noisy = full of static or errors).
The Solution: FunPRM uses a "Meta-Coach" system.
1. It knows for sure if the final dish is good because it can run the code against a strict test (like a taste test).
2. It uses this "clean" final score to fix the "noisy" guesses about the earlier steps.
The Analogy: Imagine a sports coach who isn't sure if a player's warm-up was good. But, the coach does know the player won the final game. The Meta-Coach looks at the win and says, "Since they won the game, that warm-up must have been decent, even if my initial guess was shaky." It uses the known truth of the finish to clean up the confusion of the start.

4. The Results: A Better Chef

The researchers tested this new system on two major "cooking competitions" (datasets called LiveCodeBench and BigCodeBench).

The Outcome: FunPRM consistently helped the AI chefs produce better code than other methods.
The Record: When paired with a top-tier AI (OpenAI's O4-mini), FunPRM achieved the highest score ever recorded on the LiveCodeBench leaderboard.
Human Feedback: When human developers looked at the code, they preferred the FunPRM version. They found it easier to read and reuse, much like preferring a recipe with clear chapters over a messy paragraph of instructions.

Summary

FunPRM is a system that teaches AI to write code in organized "chapters" (functions) rather than a messy stream of text. It then uses a smart "cleaning" trick to fix its own grading mistakes by looking at the final result to understand the intermediate steps. The result is code that is not only more likely to work but is also easier for humans to understand.

Technical Summary: FunPRM

Problem Statement

Code generation remains a core application of Large Language Models (LLMs), yet even state-of-the-art models frequently fail on complex programming tasks due to multi-step reasoning requirements and the tendency to hallucinate. While test-time scaling approaches like Process Reward Models (PRMs) have shown success in mathematical reasoning by evaluating intermediate steps, their application to code generation faces two significant hurdles:

Lack of Meaningful Step Decomposition: Unlike mathematical problems which naturally decompose into explicit reasoning steps under Chain-of-Thought (CoT) prompting, code lacks an obvious definition of "steps." Existing methods often treat individual lines of code as steps, leading to hundreds of steps per program and prohibitive computational costs.
Noisy Partial-Solution Rewards: Training PRMs requires ground-truth correctness scores for partial solutions. While human annotation is too costly for large-scale use, automated Monte Carlo (MC) estimation methods produce noisy rewards. These noisy signals, derived from sampling whether a partial solution leads to a correct final answer, degrade PRM performance.

Methodology: FunPRM

The authors propose FunPRM, a Process Reward Model tailored specifically for code generation that addresses the challenges of step definition and reward noise through two primary mechanisms.

1. Chain-of-Function (CoF) Prompting

To solve the step decomposition problem, FunPRM employs a novel prompting strategy that encourages LLMs to generate modular code organized into functions.

Function as Step: Instead of treating lines of code as steps, FunPRM treats each function as a reasoning step.
Prompting Strategy: The system prompt guides the LLM to extract logically independent operations into separate functions, organize them top-down (high-level functions first), and include descriptive docstrings.
Benefit: This creates a clear, semantically meaningful sequence of reasoning steps for the PRM to evaluate, while simultaneously improving code readability and reusability for human developers.

2. Meta-Learning-Based Reward Correction

To address the noise in Monte Carlo-estimated partial-solution rewards, FunPRM introduces a bi-level meta-learning framework that leverages the unique property of coding tasks: the ability to obtain clean, reliable correctness scores for final solutions via unit-test evaluation systems.

Noisy vs. Clean Signals: The framework utilizes a noisy dataset of partial solutions (estimated via MC sampling) and a clean meta-dataset of final solutions (evaluated via unit tests).
Reward Correction Table: A lightweight, trainable reward-correction table ( $g_\theta$ ) is introduced to add a learnable residual to the noisy partial-solution rewards.
Bi-Level Optimization:
1. Inner Loop: The PRM ( $f_\phi$ ) is trained on the noisy partial solutions using the current corrected rewards. The PRM parameters are updated via a single gradient descent step.
2. Outer Loop: The updated PRM is evaluated on the clean final-solution dataset. The resulting loss (meta-loss) is used to compute gradients with respect to the reward-correction parameters $\theta$ .
Goal: This process optimizes the correction table to "purify" the noisy partial-solution rewards, ensuring that the PRM trained on them generalizes well to correctly identifying successful final solutions.

Key Contributions

FunPRM Framework: A PRM tailored for code generation that redefines reasoning steps at the function level via Chain-of-Function prompting, enabling effective step-level evaluation without the computational overhead of line-level granularity.
Meta Reward Correction: A novel mechanism that uses clean final-solution rewards from unit tests to denoise Monte Carlo-estimated partial-solution rewards, improving the quality of training signals without requiring human annotation.
Empirical Superiority: Extensive experiments demonstrating that FunPRM consistently outperforms existing test-time scaling baselines (including Self-Certainty, Outcome Reward Models, and Skywork-PRM) across multiple base LLMs and benchmarks.

Experimental Results

The authors evaluated FunPRM on LiveCodeBench and BigCodeBench, as well as for domain generalization on HumanEval+ and MBPP+.

State-of-the-Art Performance: When combined with OpenAI's O4-mini (High), FunPRM achieved a pass@1 of 80.9% on LiveCodeBench, surpassing all other models on the leaderboard, including O3 (High), DeepSeek-R1, and Gemini-2.5-Pro.
Consistency Across Models: FunPRM outperformed all baseline test-time scaling methods across five different base LLMs (Qwen3-30B, GPT-4o-mini, DeepSeek-Coder, Qwen2.5-7B, and O4-mini) on both LiveCodeBench and BigCodeBench.
Ablation Studies: Removing either the Chain-of-Function prompting or the Meta Reward Correction mechanism resulted in performance drops, confirming the necessity of both components.
Human Evaluation: In a study involving three researchers, FunPRM-generated code was consistently preferred over baseline code for both readability (in correct solutions) and reusability/ease of fixing (in incorrect solutions).
Domain Generalization: FunPRM trained on LiveCodeBench and BigCodeBench successfully improved performance on HumanEval+ and MBPP+ without additional training, indicating robust generalization across different coding styles and difficulty levels.

Significance

The paper claims that FunPRM represents a practical and effective approach to improving test-time scaling in LLM-based code generation. By aligning the definition of reasoning steps with the natural modularity of software engineering (functions) and leveraging the deterministic nature of unit tests to refine training signals, FunPRM overcomes the specific limitations that have hindered PRM adoption in coding. The results suggest that this approach can substantially push the performance frontier of code generation, even for already high-performing models, while producing code that is more readable and reusable for developers.