Resumen Técnico: Escalamiento de Test-Time Eficiente y Entrenable para Modelos de Lenguaje mediante el Enrutamiento de Ramas Locales (Local Branch Routing)

Declaración del Problema

Los avances recientes en el razonamiento de modelos de lenguaje (LLM) demuestran que asignar computación adicional durante la inferencia (escalamiento de test-time) puede mejorar significamente el rendimiento. Sin embargo, los enfoques existentes enfrentan un compromiso crítico:

• Muestreo de Cadena de Pensamiento (CoT - Chain-of-Thought): Aunque es efectivo, el CoT estándar sigue siendo de un solo hilo (single-threaded), comprometiéndose con un solo token antes de observar las consecuencias de las alternativas.
• Métodos Basados en Búsqueda: Los enfoques como el Árbol de Pensamientos (Tree of Thoughts) o la Búsqueda en Árbol de Monte Carlo (MCTS) introducen amplitud mediante la exploración de árboles de pensamientos intermedios. Sin embargo, estos suelen ser de grano grueso, computacionalmente costosos y difíciles de optimizar como una política de decodificación única y entrenable de extremo a extremo.
• Métodos de Tokens Suaves (Soft-Token): Técnicas como Soft Thinking o Multiplex Thinking ofrecen amplitud a nivel de token al fusionar las representaciones (embeddings) de los candidatos en mezclas continuas. Aunque son entrenables, estos métodos desdibujan la identidad discreta de las ramas candidatas, lo que podría oscurecer la información específica de cada concepto.

El artículo postula que los modelos de lenguaje podrían beneficiarse de un mecanismo de escalamiento de test-time a nivel de token y ligero, que permita que las decisiones se informen mediante los estados ocultos de candidatos locales futuros sin incurrir en el costo de una búsqueda completa de la solución ni en la ambigüedad de las mezclas suaves.

Metodología: Enrutamiento de Ramas Locales (LBR)

Los autores proponen el Enrutamiento de Ramas Locales (LBR), un marco de decodificación de ramificación y enrutamiento a nivel de token que es entrenable. El LBR aumenta la generación autoregresiva estándar con un árbol de mirada hacia adelante (lookahead) local y rodante.

Pipeline de Decodificación

En cada paso de decodificación $t$, dado un prefijo comprometido $x_{<t}$, el LBR ejecuta un proceso de cuatro etapas de "poda-desplazamiento-crecimiento" (prune–shift–grow):

1. Crecer (Grow): En lugar de muestrear un único siguiente token, el LBR expande un árbol local $\mathcal{A}_t$ de ancho $K$ y profundidad $L$ enraizado en $x_{<t}$. Cada nodo muestreado es reenviado (forwarded) a través del modelo de lenguaje para calcular su estado oculto $h_t(v)$. Esto crea un árbol de estados ocultos locales en lugar de solo un conjunto de cadenas candidatas.
2. Enrutar (Route): Un enrutador ligero $\rho_\phi$ observa el árbol reenviado y selecciona un subárbol de profundidad 1 para comprometerse. El enrutador utiliza un mecanismo de atención de conjuntos (set-attention) para codificar cada subárbol candidato (por ejemplo, un solo token para $L=1$, o un token más sus continuaciones para $L>1$) en un vector. Estos vectores se comparan mediante atención entre subárboles para producir una puntuación relativa, permitiendo que el enrutador tome una decisión contrastiva entre futuros hermanos.
3. Podar (Prune): El enrutador selecciona el mejor subárbol (índice $k^*_t$). Todos los demás subárboles se descartan. El token raíz del subárbol seleccionado se convierte en el siguiente token comprometido $x_t$.
4. Desplazar y Regrow (Shift and Regrow): El subárbol seleccionado se desplaza hacia adelante (su raíz se elimina y sus descendientes se convierten en el nuevo prefijo para el siguiente paso). Se cultiva una nueva capa para restaurar la profundidad del árbol $L$, y el ciclo se repite.

Objetivo de Entrenamiento

El LBR define una verosimilitud de trayectoria de árbol (tree-trajectory likelihood) que permite el aprendizaje por refuerzo de extremo a extremo con recompensas verificables (RLVR).

• Estocasticidad: Solo dos operaciones son estocásticas: el muestreo de los nuevos nodos crecidos $G_t$ del modelo base y el muestreo de la decisión del enrutador $k^*_t$. La poda, el desplazamiento y la reutilización de nodos son deterministas.
• Verosimilitud (Likelihood): La log-probabilidad de una trayectoria $F$ es la suma de las log-probabilidades de los nuevos tokens crecidos y las probabilidades de selección del enrutador:
  $$\log p_{\theta,\phi}(F | q) = \sum_{t=1}^T \sum_{v \in G_t} \log \tilde{\pi}_\theta(\text{tok}(v) | \text{ctx}(v)) + \sum_{t=1}^T \log \rho_\phi(k^*_t | x_{<t}, \mathcal{A}_t; \theta)$$
• Optimización: El marco se entrena utilizando un objetivo de razón de verosimilitud (por ejemplo, GRPO) donde la ventaja $A(F)$ de un verificador recompensa tanto la generación de tokens del modelo base como las decisiones de selección del enrutador.

Contribuciones Clave

1. Ramificación a Nivel de Token con Identidad Discreta: A diferencia de los métodos de tokens suaves que fusionan los candidatos en mezclas continuas, el LBR preserva los candidatos como ramas discretas reenviadas. Esto permite que el enrutador aproveche los estados ocultos post-candidato—información que se revela solo después de que un token candidato es generado y procesado por el modelo.
2. Entrenamiento de Extremo a Extremo Tratable: Al definir una verosimilitud específica para la trayectoria del árbol, el LBR permite que el modelo base y el enrutador se optimicen conjuntamente bajo el mismo principio de RLVR utilizado en sistemas de razonamiento de tokens discretos.
3. Escalamiento de Test-Time Eficiente: El marco introduce amplitud a nivel de token con una mirada hacia adelante superficial ($L$ es pequeño, por ejemplo, 1 o 2), evitando la sobrecarga computacional de una búsqueda de profundidad completa, pero aprovechando aun así la evidencia más allá de la distribución del siguiente token raíz.
4. Enrutador de Atención de Conjuntos: Los autores introducen una arquitectura de enrutador que utiliza atención entre subárboles para comparar candidatos hermanos, demostrando que el enrutamiento es una decisión contrastiva en lugar de un re-ranking independiente.

Resultados Experimentales

Planificación Jerárquica Sintética

En una tarea sintética de alcanzabilidad de grafos con traslación de radio:

• Rendimiento: El LBR alcanzó la mayor precisión objetivo (0.967), superando al CoT Discreto (0.914) y al Soft Thinking (0.845).
• Análisis de Mecanismo: Los experimentos de sondeo revelaron que los estados ocultos post-candidato en LBR son significativamente más predictivos del objetivo final que los estados previos a la ramificación. Esto confirma que el reenvío de candidatos permite al modelo acceder a información descendente que no está disponible para la decodificación autoregresiva estándar.
• Representación: A diferencia de Soft Thinking, que desdibujó las identidades de los conceptos mediante mezclas continuas, el LBR preservó las identidades discretas de las ramas, haciendo que el camino correcto fuera más recuperable.

Benchmarks de Razonamiento Matemático

Evaluado en modelos DeepSeek-R1-Distill-Qwen (1.5B y 7B) a través de seis benchmarks (Minerva, AIME, MATH500, etc.):

• Ganancias de Rendimiento: El LBR mejoró tanto el Pass@1 como el Pass@32 sobre:
  • Cadena de Pensamiento Discreta (Discrete CoT).
  • RLVR de token discreto vanilla.
  • Líneas base de ramificación de tokens suaves compatibles con RL (por ejemplo, Multiplex Thinking).
• Ablación de Profundidad: Aumentar la profundidad de la mirada hacia adelante de $L=1$ a $L=2$ mejoró aún más el rendimiento, particularmente en el backbone de 7B, lo que sugiere que las continuaciones locales más profundas proporcionan más evidencia de enrutamiento.
• Ablación del Enrutador: Eliminar la atención entre subárboles (usando un enrutador independiente) resultó en un menor rendimiento, validando que la comparación contrastiva de las ramas hermanas es crucial.
• Dinámicas de Entrenamiento: Durante el entrenamiento RLVR, la entropía del enrutador disminuyó a medida que el enrutador se volvía más decisivo, mientras que la entropía del modelo base en las posiciones enrutadas también disminuyó, indicando una co-adaptación donde el modelo aprende a generar ramas locales más distinguibles.

Significado y Reivindicaciones

El artículo afirma que el Enrutamiento de Ramas Locales ofrece una forma eficiente, entrenable y discreta de escalamiento de test-time para modelos de lenguaje. Su significado principal radica en demostrar que:

1. La evidencia post-candidato es valiosa: Las decisiones de enrutamiento se benefician de los estados ocultos inducidos por los tokens candidatos, no solo de la distribución previa a la ramificación.
2. La ramificación discreta es superior a la mezcla suave: Preservar las identidades de las ramas discretas permite una toma de decisiones más clara en comparación con las mezclas de tokens suaves, que pueden oscurecer la identidad de los conceptos.
3. Escalabilidad: El marco proporciona un punto medio entre el CoT de un solo hilo y la búsqueda de árbol completa y costosa, permitiendo un escalamiento efectivo mediante el ramificado local ligero.

Los autores reconocen la limitación de la localidad: el LBR está limitado por la ventana de mirada hacia adelante superficial y no realiza una búsqueda global sobre soluciones completas. Se sugiere como trabajo futuro combinar LBR con profundidad adaptativa o búsqueda global selectiva.