Entrenamiento Eficiente de Modelos de Recompensa de Procesos mediante Aprendizaje Activo

Efficient Process Reward Model Training via Active Learning

April 14, 2025

Autores: Keyu Duan, Zichen Liu, Xin Mao, Tianyu Pang, Changyu Chen, Qiguang Chen, Michael Qizhe Shieh, Longxu Dou

cs.AI

Resumen

Los Modelos de Recompensa de Procesos (PRMs, por sus siglas en inglés) proporcionan supervisión a nivel de paso para los modelos de lenguaje grandes (LLMs), pero escalar la anotación de datos de entrenamiento sigue siendo un desafío tanto para humanos como para LLMs. Para abordar esta limitación, proponemos un enfoque de aprendizaje activo, ActPRM, que selecciona proactivamente las muestras más inciertas para el entrenamiento, reduciendo sustancialmente los costos de etiquetado. Durante el entrenamiento, utilizamos el PRM para estimar la incertidumbre después del paso hacia adelante, conservando solo los datos altamente inciertos. Un modelo de razonamiento capaz pero costoso etiqueta entonces estos datos. Luego calculamos la pérdida con respecto a las etiquetas y actualizamos los pesos del PRM. Comparamos ActPRM frente al ajuste fino convencional, en un entorno de aprendizaje activo basado en un conjunto de datos, demostrando que ActPRM reduce un 50% la anotación, pero logrando un rendimiento comparable o incluso mejor. Más allá de la eficiencia en la anotación, avanzamos aún más el PRM entrenado activamente filtrando más de 1 millón de trayectorias de razonamiento matemático con ActPRM, conservando el 60% de los datos. Un entrenamiento posterior en este conjunto de datos seleccionado produce un nuevo PRM de última generación (SOTA) en ProcessBench (75.0%) y PRMBench (65.5%) en comparación con modelos del mismo tamaño.

English

Process Reward Models (PRMs) provide step-level supervision to large language models (LLMs), but scaling up training data annotation remains challenging for both humans and LLMs. To address this limitation, we propose an active learning approach, ActPRM, which proactively selects the most uncertain samples for training, substantially reducing labeling costs. During training, we use the PRM to estimate uncertainty after the forward pass, retaining only highly uncertain data. A capable yet costly reasoning model then labels this data. Then we compute the loss with respect to the labels and update the PRM's weights. We compare ActPRM vs. vanilla fine-tuning, on a pool-based active learning setting, demonstrating that ActPRM reduces 50% annotation, but achieving the comparable or even better performance. Beyond annotation efficiency, we further advance the actively trained PRM by filtering over 1M+ math reasoning trajectories with ActPRM, retaining 60% of the data. A subsequent training on this selected dataset yields a new state-of-the-art (SOTA) PRM on ProcessBench (75.0%) and PRMBench (65.5%) compared with same sized models.