Search-Based Reward Function Optimization for LLM Reasoning via Reinforcement Learning

Abstract

Reinforcement learning has become a central post-training mechanism for improving the reasoning capabilities of large language models (LLMs), yet performance remains sensitive to the design of the reward function that drives policy optimization. This thesis introduces a search-driven reward discovery framework that treats the reward specification itself as an object of optimization rather than a fixed design choice. Candidate reward functions are generated by a frontier language model (Kimi K2), validated automatically, screened through short Group Relative Policy Optimization (GRPO) training runs on a Llama-3.2-3B-Instruct base model with Low-Rank Adaptation (LoRA), and ranked by downstream mathematical reasoning performance on the GSM8K benchmark. The search produces 50 candidate reward functions over five iterative rounds, with each round receiving feedback on the performance of all previously evaluated rewards. The search completes in approximately 40 hours on a single NVIDIA RTX 5090 GPU. Results show a progressive improvement in reward quality over rounds: the mean F1 score increases from 0.596 in Round 1 to 0.632 in Round 5, and the top-performing individual rewards, thinking_steps_count (F1 = 0.7871) and thinking_has_calc (F1 = 0.7860), emerge from the later rounds. Only the top 12% of generated rewards exceed 0.60 accuracy, so screening separates effective reward signals from weak ones before any reward enters ensemble training. Seven ensemble configurations combining top-ranked rewards are evaluated. The best-performing ensemble, ens_diverse_rounds, selects five rewards from five different rounds and achieves an F1 of 0.7954 and accuracy of 0.6603, which surpasses the best individual reward by +0.0083 in F1 and +0.0113 in accuracy. This diversity-focused ensemble numerically outperforms the ranking-based ensembles in the default run, suggesting that reward diversity over rounds is a useful selection signal alongside raw ranking position. Comparison against five off-the-shelf models of comparable size (2B-4B parameters) shows that the GRPO-trained model achieves competitive F1 and strong precision (0.7181), which indicates that the discovered rewards improve the quality of the model's reasoning process. These results demonstrate that systematic search over reward functions is a viable and interpretable approach to improving reinforcement learning for LLM mathematical reasoning.