Búsqueda Guiada por Valores para un Razonamiento Eficiente en Cadena de Pensamiento

Value-Guided Search for Efficient Chain-of-Thought Reasoning

May 23, 2025

Autores: Kaiwen Wang, Jin Peng Zhou, Jonathan Chang, Zhaolin Gao, Nathan Kallus, Kianté Brantley, Wen Sun

cs.AI

Resumen

En este artículo, proponemos un método simple y eficiente para el entrenamiento de modelos de valor en trazas de razonamiento de contexto largo. En comparación con los modelos de recompensa de proceso (PRMs) existentes, nuestro método no requiere una noción detallada de "paso", que es difícil de definir para modelos de razonamiento de contexto largo. Al recopilar un conjunto de datos de 2.5 millones de trazas de razonamiento, entrenamos un modelo de valor a nivel de tokens de 1.5B y lo aplicamos a los modelos DeepSeek para mejorar el rendimiento con escalado de cómputo en tiempo de prueba. Descubrimos que la búsqueda guiada por valor (VGS) por bloques con una votación mayoritaria ponderada final logra un mejor escalado en tiempo de prueba que métodos estándar como la votación mayoritaria o best-of-n. Con un presupuesto de inferencia de 64 generaciones, VGS con DeepSeek-R1-Distill-1.5B alcanza una precisión promedio del 45.7% en cuatro benchmarks de matemáticas competitivas (AIME 2024 & 2025, HMMT Feb 2024 & 2025), alcanzando paridad con o3-mini-medium. Además, VGS reduce significativamente los FLOPs de inferencia necesarios para alcanzar el mismo rendimiento que la votación mayoritaria. Nuestro conjunto de datos, modelo y código base son de código abierto.

English

In this paper, we propose a simple and efficient method for value model training on long-context reasoning traces. Compared to existing process reward models (PRMs), our method does not require a fine-grained notion of "step," which is difficult to define for long-context reasoning models. By collecting a dataset of 2.5 million reasoning traces, we train a 1.5B token-level value model and apply it to DeepSeek models for improved performance with test-time compute scaling. We find that block-wise value-guided search (VGS) with a final weighted majority vote achieves better test-time scaling than standard methods such as majority voting or best-of-n. With an inference budget of 64 generations, VGS with DeepSeek-R1-Distill-1.5B achieves an average accuracy of 45.7% across four competition math benchmarks (AIME 2024 & 2025, HMMT Feb 2024 & 2025), reaching parity with o3-mini-medium. Moreover, VGS significantly reduces the inference FLOPs required to achieve the same performance of majority voting. Our dataset, model and codebase are open-sourced.