Entraînement efficace des modèles de récompense de processus par apprentissage actif

Efficient Process Reward Model Training via Active Learning

April 14, 2025

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

cs.AI

Résumé

Les modèles de récompense de processus (PRM) fournissent une supervision au niveau des étapes pour les grands modèles de langage (LLM), mais l'augmentation de l'annotation des données d'entraînement reste un défi tant pour les humains que pour les LLM. Pour surmonter cette limitation, nous proposons une approche d'apprentissage actif, ActPRM, qui sélectionne de manière proactive les échantillons les plus incertains pour l'entraînement, réduisant ainsi considérablement les coûts d'étiquetage. Pendant l'entraînement, nous utilisons le PRM pour estimer l'incertitude après la passe avant, en conservant uniquement les données très incertaines. Un modèle de raisonnement performant mais coûteux étiquette ensuite ces données. Nous calculons ensuite la perte par rapport aux étiquettes et mettons à jour les poids du PRM. Nous comparons ActPRM à un réglage fin classique, dans un cadre d'apprentissage actif basé sur un pool de données, démontrant qu'ActPRM réduit de 50 % l'annotation tout en atteignant des performances comparables, voire supérieures. Au-delà de l'efficacité de l'annotation, nous améliorons davantage le PRM entraîné activement en filtrant plus de 1 million de trajectoires de raisonnement mathématique avec ActPRM, en conservant 60 % des données. Un entraînement ultérieur sur cet ensemble de données sélectionné permet d'obtenir un nouveau PRM de pointe (SOTA) sur ProcessBench (75,0 %) et PRMBench (65,5 %) par rapport à des modèles de même taille.

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.