💬 NLP

Enhanced LLM Reasoning by Optimizing Reward Functions with Search-Driven Reinforcement Learning

This paper proposes a search-driven reinforcement learning framework that iteratively optimizes reward function specifications through automated generation, validation, and GRPO-based screening, demonstrating that a ranked-feedback loop significantly enhances mathematical reasoning performance on GSM8K compared to static or random reward ensembles.

Original authors: Arash Ahmadi (Mike), Sarah Sharif (Mike), Yaser (Mike), Banad

Published 2026-05-05

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Arash Ahmadi (Mike), Sarah Sharif (Mike), Yaser (Mike), Banad

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 have a very smart student (a Large Language Model) who is great at talking but sometimes struggles with math. You want to teach them to solve math problems better. Usually, you'd hire a teacher to grade their homework. But in this paper, the authors ask a different question: "What if we don't know exactly how to grade the homework perfectly? What if we let an AI help us invent the best grading rules?"

Here is the story of how they did it, explained simply.

1. The Problem: The "Grading Rubric" Mystery

In the world of AI, to teach a model to think better, you use a system called Reinforcement Learning. Think of this like training a dog.

The Dog: The AI model.
The Treat: A "reward" (a positive score).
The Trick: Solving a math problem correctly.

The problem is that designing the "treat" (the reward function) is hard. If you tell the AI, "Good job if you get the right answer," it might cheat by guessing or copying patterns without actually thinking. If you try to reward the steps it takes, you have to write complex rules, and it's easy to make a mistake in those rules.

2. The Solution: The "Grading Rule Factory"

The authors built a system that treats the grading rules themselves as something to be optimized. They didn't write the rules by hand; they set up a factory where an AI (a "frontier" model named Kimi K2) acts as a Rule Inventor.

Here is how the factory works, step-by-step:

Round 1: The Rule Inventor AI is given 20 math problems and asked to write 10 new ways to grade them. It writes these rules as simple computer code (like a recipe).
The Safety Check: Before these rules are used, they go through a "security scanner" to make sure they aren't dangerous or broken.
The Test Drive: The researchers take a small math student (a 3-billion-parameter AI model) and let it practice solving problems using one of these new rules for a short time (500 steps).
The Scorecard: They see how well the student did on a test. If the new rule helped the student get better scores, the rule gets a high ranking. If the rule confused the student, it gets a low ranking.
The Feedback Loop: The top-performing rules are summarized and fed back to the Rule Inventor AI. The AI is told, "Hey, rules that counted the number of thinking steps worked well. Rules that just looked at the final answer didn't work as well. Try to invent new rules based on that."

They repeated this cycle 5 times, generating 50 different candidate rules.

3. The Results: Finding the "Golden Rules"

After 5 rounds, they found some winners.

The Best Single Rule: One rule, called thinking_steps_count, was a champion. It gave extra points if the AI wrote out its thinking process in at least three distinct lines. It didn't check if the answer was right (the base system did that); it just encouraged the AI to show its work.
The Power of Teamwork (Ensembles): They realized that one rule isn't enough. So, they took the best 5 rules and combined them into a "super-grading team."
- The Result: This team of rules helped the AI jump from a 60% success rate (using only basic rules) to a 79.5% success rate.
- The "Magic" Check: To prove it wasn't just luck or the sheer number of rules, they tried a "random team" of 5 rules picked from the pile. That random team crashed, with the AI failing almost everything. This proved that the feedback loop (learning which rules worked and feeding that back to the inventor) was the secret sauce, not just having more rules.

4. Did the AI Cheat? (The "Reward Hacking" Audit)

A common fear is that if you tell an AI "give me more lines of thinking," it will just fill the page with nonsense to get points.

The Audit: The authors checked the AI's answers. They found that when the AI got the answer right, it actually wrote more lines and used more math symbols than when it got it wrong.
The Conclusion: The AI wasn't cheating by padding empty space; it was genuinely thinking harder because the rules encouraged it to.

5. Why This Matters

It's Accessible: You don't need a supercomputer. They ran this whole experiment on a single, high-end consumer graphics card (like the kind gamers use) in about 40 hours.
It's Automated: Instead of a human expert spending weeks guessing what the perfect grading rubric looks like, this system automates the discovery of those rules.
It's Transparent: The rules aren't a black box; they are simple Python code that humans can read, understand, and tweak.

In a nutshell: The paper shows that if you let an AI help you design the grading system for another AI, and you keep refining those rules based on how well the student performs, you can teach the student to think much better than if you just gave it a static, human-written set of rules.

Technical Summary: Enhanced LLM Reasoning by Optimizing Reward Functions with Search-Driven Reinforcement Learning

Problem Statement

Mathematical reasoning serves as a critical benchmark for Large Language Models (LLMs), requiring not just fluent text generation but intermediate consistency, quantity carrying, and precise final answers. While Reinforcement Learning (RL) has emerged as a standard post-training mechanism to align LLMs with desired behaviors, its efficacy is highly sensitive to the design of the reward function. Current approaches often rely on manual, expert-driven reward specification, which is prone to error, requires significant trial-and-error, and may inadvertently induce "reward hacking" or brittle behaviors. Furthermore, existing literature largely treats the reward function as fixed infrastructure (e.g., handcrafted rules or trained reward models) rather than an object of optimization itself. This paper addresses the question of whether the reward specification used during RL post-training can be systematically searched and improved, particularly in a setting where the base model is held fixed and the reward function is the primary design variable.

Methodology

The authors propose a search-driven framework that treats reward function specification as an iterative optimization problem. The pipeline operates over five rounds, generating and evaluating candidate reward functions using a fixed base model and a frontier LLM as a generator.

Core Components

Base Model and Training Setup:
- Model: Meta's Llama-3.2-3B-Instruct, fine-tuned using Low-Rank Adaptation (LoRA) with a rank of 32.
- Algorithm: Group Relative Policy Optimization (GRPO), which replaces the separate critic network in PPO with group-relative normalization of sampled outputs, reducing memory overhead and improving mathematical reasoning.
- Dataset: GSM8K (7,473 training examples, 1,319 test examples).
- Output Format: Structured output with <thinking> and <solution> tags to separate reasoning from the final answer.
Fixed Base Rewards:
Three protected base rewards are included in every training run and cannot be overridden by generated candidates:
- match_format_exactly: Rewards correct pairing of <thinking> and <solution> tags.
- match_format_approximately: Provides graduated signals for tag presence/absence.
- check_answer: Evaluates the extracted numeric answer against ground truth with tiered scoring (exact match, normalization, ratio checks).
Search-Driven Generation Loop:
- Generator: A frontier LLM (Kimi K2) generates 10 candidate reward functions per round (50 total over 5 rounds).
- Prompting: The generator receives 20 random GSM8K examples, a signature specification, and a ranked summary of previous rounds' performance to guide generation and avoid semantic duplication.
- Validation: Candidates undergo a four-stage sandbox validation (AST import check, locked-down execution, restricted namespace, and probe call) to ensure safety and executability.
- Screening: Valid candidates are trained for 500 GRPO steps on the GSM8K training split alongside the three base rewards.
- Ranking: Models are evaluated on the GSM8K test set using F1 score (harmonic mean of precision and recall) as the primary ranking metric. The ranked summary is fed back into the next generation round.
Ensemble Construction:
After the search, top-performing rewards are combined into ensembles. Seven configurations were tested, varying in size (3 to 10 discovered rewards) and selection strategy (top-ranked vs. diversity-based across rounds). The ensemble reward is the element-wise sum of the three base rewards and the selected discovered rewards.

Key Contributions

Formulation of Reward Search: The paper formulates reward specification for LLM reasoning as an iterative generate–validate–train–rank loop over executable Python code, treating the reward itself as a search object.
GRPO Integration: The framework applies this search loop specifically under GRPO, a computationally efficient and reasoning-relevant policy optimizer, demonstrating that reward discovery is feasible within parameter-efficient fine-tuning constraints.
Ensemble Analysis: The study examines whether combining multiple discovered rewards yields stronger signals than individual rewards, finding that ensembles of 5 or more rewards significantly outperform single rewards or base-only configurations.
Empirical Rigor: The paper provides a comprehensive analysis including bootstrap confidence intervals, pairwise McNemar tests with Bonferroni correction, and specific audits for reward hacking (e.g., line-padding) and format compliance.

Results

Search Trajectory: Over five rounds, the mean F1 score of generated rewards increased from 0.596 (Round 1) to 0.632 (Round 5). The top individual reward (thinking_steps_count) achieved an F1 of 0.787.
Ensemble Performance: The best ensemble configuration (ens_diverse_rounds, combining 3 base rewards + 5 diverse discovered rewards) achieved an F1 of 0.795 and an accuracy of 0.660 (95% CI [0.635, 0.686]).
Statistical Significance: Pairwise McNemar tests with Bonferroni correction ( $\alpha \approx 0.0024$ ) showed that all ensembles with 5 or more discovered rewards were statistically indistinguishable from each other but significantly better than the base-only baseline (F1 = 0.609) and smaller ensembles.
Control Experiments:
- Random Baseline: A randomly drawn 5-reward ensemble collapsed to an F1 of 0.047, demonstrating that the gain is driven by the ranked-feedback loop, not merely the additive signal of having more rewards.
- Seed Robustness: Retraining the best ensemble with three different seeds yielded F1 scores of $0.785 \pm 0.013$ , confirming the qualitative ordering of results is robust to training variance.
Reward Hacking Audit: Analysis of the top structural reward (thinking_steps_count) showed that correct responses naturally contained more reasoning lines and higher arithmetic density than incorrect ones, suggesting the model did not exploit the reward via empty line padding.
Baseline Comparison: The GRPO-trained Llama-3.2-3B model outperformed five off-the-shelf models in the 2B–4B parameter range (including Qwen and Phi variants) on both strict accuracy and F1.

Significance and Claims

The paper claims that the reward specification in RL post-training can be treated as a reliably searchable object, automating a process typically reliant on human intuition.

Accessibility: The entire pipeline runs on a single consumer-grade GPU (NVIDIA RTX 5090) in approximately 40 hours, making automated reward discovery accessible to academic and small-industry budgets.
Auditability: Unlike black-box reward models, the discovered rewards are readable Python code that practitioners can audit, modify, and combine.
Generalizability: While evaluated on GSM8K, the authors posit that the "recipe" transfers to other reasoning domains (e.g., code generation, scientific problem solving) and offers a practical alternative to expensive reward-model training.
Modesty: The authors explicitly state that the contribution is methodological rather than a claim to a new state-of-the-art on GSM8K itself. The primary evidence is internal to the pipeline: the iterative improvement of rewards, the collapse of random controls, and the robustness of the ensemble ordering across seeds. The paper does not claim the discovered rewards are universally optimal but rather that the search process effectively identifies high-performing structural signals for a fixed base model.