Treinamento Eficiente de Modelos de Recompensa de Processo por meio de Aprendizado Ativo

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

Resumo

Os Modelos de Recompensa de Processo (PRMs) fornecem supervisão em nível de etapa para grandes modelos de linguagem (LLMs), mas a escalabilidade da anotação de dados de treinamento continua sendo um desafio tanto para humanos quanto para LLMs. Para abordar essa limitação, propomos uma abordagem de aprendizado ativo, ActPRM, que seleciona proativamente as amostras mais incertas para treinamento, reduzindo substancialmente os custos de rotulagem. Durante o treinamento, usamos o PRM para estimar a incerteza após a passagem direta, retendo apenas os dados altamente incertos. Um modelo de raciocínio capaz, porém custoso, então rotula esses dados. Em seguida, calculamos a perda em relação aos rótulos e atualizamos os pesos do PRM. Comparamos o ActPRM com o ajuste fino tradicional, em um cenário de aprendizado ativo baseado em pool, demonstrando que o ActPRM reduz 50% da anotação, mas alcança desempenho comparável ou até melhor. Além da eficiência de anotação, avançamos ainda mais o PRM treinado ativamente ao filtrar mais de 1 milhão de trajetórias de raciocínio matemático com o ActPRM, retendo 60% dos dados. Um treinamento subsequente nesse conjunto de dados selecionado resulta em um novo PRM de estado da arte (SOTA) no ProcessBench (75,0%) e no PRMBench (65,5%) em comparação com modelos de mesmo tamanho.

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.