Teaching LLMs for Step-Level Automatic Math Correction via Reinforcement Learning
Junsong Li, Jie Zhou, Yutao Yang, Bihao Zhan, Qianjun Pan, Yuyang Ding, Qin Chen, Jiang Bo, Xin Lin, Liang He
TL;DR
StepAMC tackles the lack of step-level feedback in automatic math correction by using reinforcement learning with a space-constrained policy network and a fine-grained reward network to teach LLMs stepwise reasoning. It reframes the task as a text classification RL problem with per-step rewards and stability constraints, employing PPO-style optimization and a learnable constraint loss. Experiments on the PRM-42K and MSD-22K datasets demonstrate that StepAMC outperforms eleven baselines in F1 and accuracy and achieves better per-class balance. The work highlights the value of structured RL-based feedback for mathematical reasoning and points to scalable future directions across larger datasets and diverse foundation-model backbones.
Abstract
Automatic math correction aims to check students' solutions to mathematical problems via artificial intelligence technologies. Most existing studies focus on judging the final answer at the problem level, while they ignore detailed feedback on each step in a math problem-solving process, which requires abilities of semantic understanding and reasoning. In this paper, we propose a reinforcement learning (RL)-based method to boost large language model (LLM) for step-level automatic math correction, named StepAMC. Particularly, we convert the step-level automatic math correction within the text classification task into an RL problem to enhance the reasoning capabilities of LLMs. Then, we design a space-constrained policy network to improve the stability of RL. Then, we introduce a fine-grained reward network to convert the binary human feedback into a continuous value. We conduct extensive experiments over two benchmark datasets and the results show that our model outperforms the eleven strong baselines.