Resumen Técnico: Alineación de Políticas de Búsqueda en Árbol con Presupuestos de Tokens Fijos en el Escalado en Tiempo de Prueba de LLMs

Planteamiento del Problema

Los Modelos de Lenguaje de Gran Escala (LLMs) pueden mejorar la calidad de sus respuestas mediante el Escalado en Tiempo de Prueba (Test-Time Scaling), que asigna cómputo adicional durante la inferencia sin actualizar los parámetros del modelo. Si bien la decodificación de búsqueda en árbol (por ejemplo, Monte Carlo Tree Search, MCTS) es una forma efectiva de escalado en tiempo de prueba, las implementaciones en el mundo real suelen imponer un presupuesto de tokens fijo por consulta.

Las políticas de búsqueda en árbol existentes son mayoritariamente ajenas al presupuesto (budget-agnostic). Normalmente tratan el presupuesto de tokens simplemente como una condición de terminación en lugar de una restricción dinámica que guíe el comportamiento de la búsqueda. Este desajuste conduce a dos modos principales de falla:

1. Sobre-ramificación en etapas tardías: La búsqueda continúa expandiendo nodos superficiales tarde en el proceso, agotando el presupuesto antes de que pueda refinar o completar candidatos prometedores.
2. Terminación prematura: La búsqueda se detiene antes de tiempo, dejando una parte del presupuesto asignado sin utilizar.

Las variantes actuales de parada temprana se centran en ahorrar tokens, pero carecen de una política para cambiar dinámicamente la estrategia de búsqueda de una exploración amplia hacia el refinamiento a medida que el presupuesto restante disminuye.

Metodología: BG-MCTS (MCTS Guiado por Presupuesto)

Los autores proponen BG-MCTS (Budget-Guided MCTS), un algoritmo de decodificación de búsqueda en árbol que condiciona explícitamente su política de búsqueda al presupuesto de tokens restante. El mecanismo central se basa en una razón de suficiencia de presupuesto $\rho \in [0, 1]$, definida como:
$$\rho = 1 - \frac{C_{used}}{B}$$
donde $C_{used}$ son los tokens generados acumulados y $B$ es el presupuesto total. A medida que $\rho$ disminuye (acercándose a 0), la política cambia de la exploración a la explotación y la completación.

BG-MCTS introduce dos mecanismos de control condicionados por el presupuesto:

1. Selección Guiada por Presupuesto (BG-PUCT)

La selección PUCT estándar se modifica para amortiguar la bonificación de exploración e introducir un sesgo de completación:
$$\text{BG-PUCT}(p, s, \rho) = \frac{\tilde{W}(s, \rho)}{m_s} + \rho \cdot c \cdot P(s|p) \sqrt{\frac{\ln(m_p)}{m_s}}$$

• Amortiguación de la Exploración: El término de exploración se multiplica por $\rho$. A medida que el presupuesto se agota ($\rho \to 0$), el incentivo para explorar nuevas ramas disminuye.
• Sesgo de Completación: El valor acumulado $\tilde{W}(s, \rho)$ se corrige mediante un término proporcional a la profundidad del nodo $d(x)$ y $(1-\rho)$. Esto sesga la selección hacia nodos más profundos que están más cerca de generar una respuesta final, sin potenciar artificialmente los nodos que ya contienen una respuesta completada.

2. Ensanchamiento de Árbol Guiado por Presupuesto

Para evitar la ramificación de etapa tardía que resulta en un desperdicio, BG-MCTS trata el "ensanchamiento" (generar un nuevo hijo de un padre) como una acción seleccionable junto con los hijos existentes. Se añade un hijo generativo virtual $s_{gen}(p)$ al conjunto de selección. Su puntuación es:
$$E_{gen}(p, \rho) = \mu(p) + \lambda \cdot \rho \cdot \sigma^2(p)$$

• $\mu(p)$ y $\sigma^2(p)$ representan la media y la varianza de las puntuaciones de los hijos existentes.
• El término se escala por $\rho$. Cuando el presupuesto es amplio, se fomenta el ensanchamiento (especialmente si los hijos discrepan). Cuando el presupuesto es escaso, el incentivo para ensanchar se suprime, forzando a la búsqueda a profundizar en las trayectorias prometedoras existentes.

Algoritmo Unificado

El algoritmo sigue un bucle estándar de MCTS, pero hace que tanto la decisión de selección como la de ensanchamiento dependan de $\rho$. Esto induce un esquema de ancho-a-profundidad (wide-to-deep schedule): una exploración amplia al inicio de la búsqueda, seguida de un refinamiento enfocado y la completación de la respuesta a medida que el presupuesto se agota.

Principales Contribuciones

1. Diseño de Política Consciente del Presupuesto: El artículo introduce el primer marco de decodificación de búsqueda en árbol que alinea dinámicamente el comportamiento de búsqueda (exploración vs. explotación vs. ensanchamiento) con el presupuesto de tokens restante, en lugar de tratar el presupuesto meramente como una señal de parada estática.
2. Mecanismos Novedosos: La propuesta de BG-PUCT (con exploración condicionada al presupuesto y corrección de valor sesgada por profundidad) y Hijos Generativos Virtuales para un ensanchamiento consciente del presupuesto.
3. Validación Empírica: Evaluación exhaustiva en tareas de razonamiento matemático (MATH500, AIME24/25) y un benchmark de razonamiento de física (UGPhysics) utilizando LLMs de pesos abiertos (Llama-3.1-8B, Qwen2.5-7B, Qwen3-32B).

Resultados Experimentales

Los autores evaluaron BG-MCTS frente a líneas base ajenas al presupuesto (MCTS estándar, AB-MCTS-M, LiteSearch) y líneas base de escalado (Muestreo Repetido, Refinamiento Secuencial) bajo presupuestos fijos de 10k, 20k y 30k tokens.

• Rendimiento: BG-MCTS superó consistentemente a todas las líneas base en todos los modelos y presupuestos. Por ejemplo, en MATH500 con Qwen2.5-7B, BG-MCTS alcanzó un 71.1% de precisión a 30k tokens, comparado con el 65.9% de la siguiente mejor línea base de búsqueda en árbol (MCTS).
• Utilización del Presupuesto: A diferencia de LiteSearch, que a menudo se detiene prematuramente, BG-MCTS continúa mejorando la precisión a medida que se acerca al agotamiento del presupuesto.
• Comportamiento Cualitativo: El análisis de los árboles de búsqueda reveló que BG-MCTS mantiene un ancho amplio al principio, pero cambia hacia la profundización de ramas específicas al final de la búsqueda. En contraste, el MCTS estándar suele continuar ensanchando nodos superficiales hasta que el presupuesto se agota, fallando en completar los caminos prometedores.
• Razonamiento de Física: La efectividad del método se extendió a UGPhysics, demostando que el comportamiento de ancho-a-profundidad alineado con el presupuesto es beneficioso más allá del razonamiento puramente matemático.

Significancia y Reivindicaciones

El artículo sostiene que condicionar las decisiones de la búsqueda en el presupuesto restante es más efectivo que tratar el presupuesto meramente como una condición de parada.

• Escalado Predecible: BG-MCTS ofrece una interfaz sencilla para un escalado en tiempo de prueba predecible, donde los usuarios especifican un presupuesto de tokens y la política adapta automáticamente su programa de exploración a completación dentro de ese margen de costo.
• Implicaciones para la Síntesis de Datos: Los autores señalan que BG-MCTS tiende a producir menos nodos respondidos por instancia de problema, pero con una precisión (exactitud) significativamente mayor. Este régimen de "menos instancias, mejores respuestas" se destaca como particularmente valioso para la síntesis de datos, donde las trazas de razonamiento de alta calidad son más críticas que la cobertura bruta.
• Robustez: El método sigue siendo efectivo incluso cuando se utilizan modelos de recompensa de proceso (PRM) imperfectos que proporcionan señales débiles o saturadas, lo que sugiere que los cambios estructurales conscientes del presupuesto pueden compensar las señales de evaluación ruidosas.

El trabajo concluye que alinear las políticas de búsqueda con las restricciones de tokens fijos es un paso crucial para desplegar LLMs en entornos del mundo real con recursos limitados.