Resumen Técnico: Hidden States Know Where Reasoning Diverges

Título original: Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance

1. El Problema: Asignación de Crédito en RLVR

En el aprendizaje por refuerzo con recompensas verificables (RLVR), como el algoritmo GRPO (Group Relative Policy Optimization), un desafío crítico es la asignación de crédito (credit assignment). Actualmente, cuando un modelo genera una cadena de razonamiento larga, el algoritmo asigna la misma ventaja (advantage) a todos los tokens de la trayectoria, ya sea que el resultado final sea correcto o incorrecto.

Esto es ineficiente porque:

• Un solo error en un paso intermedio puede invalidar todo un razonamiento, pero los tokens previos al error reciben la misma penalización que los posteriores.
• Los modelos de recompensa de proceso (PRM) pueden solucionar esto, pero requieren anotaciones costosas a nivel de paso o el entrenamiento de un modelo de recompensa adicional, lo cual es difícil de escalar y propenso a errores de calibración.

2. Metodología: El método SHEAR

Los autores proponen SHEAR (Span-level Hidden state Enabled Advantage Reweighting), un método que extrae señales de crédito de forma autosupervisada utilizando únicamente las representaciones internas (estados ocultos) del propio modelo.

A. Observación Empírica

El estudio demuestra que las distribuciones de los estados ocultos contienen una señal de la calidad del razonamiento local. Al comparar las trayectorias correctas con las incorrectas dentro de un mismo grupo de GRPO, descubrieron que la distancia de Wasserstein entre sus distribuciones de estados ocultos aumenta significativamente en los puntos donde el razonamiento diverge (donde una trayectoria se vuelve correcta y la otra falla).

B. Implementación Técnica

El proceso de SHEAR consta de cuatro pasos principales:

1. Descomposición en fragmentos (Spans): Se divide la trayectoria en fragmentos superpuestos mediante una ventana deslizante.
2. Cálculo de la Distancia de Wasserstein: Para cada fragmento, se calcula la distancia de Wasserstein (usando el algoritmo de Sinkhorn para eficiencia) hacia el conjunto de fragmentos del grupo opuesto (si el fragmento es de una trayectoria correcta, se compara con las incorrectas, y viceversa).
3. Agregación por Max-Pooling: Se proyecta la señal de los fragmentos a nivel de token individual, asignando a cada token el valor máximo de la distancia de los fragmentos que lo contienen. Esto enfatiza las regiones de mayor divergencia.
4. Re-ponderación de la Ventaja: La ventaja de GRPO se escala por este peso derivado de la distancia. Los tokens en regiones donde la trayectoria es muy distinta a la del grupo opuesto reciben actualizaciones más fuertes (positivas si la trayectoria es correcta, negativas si es incorrecta).

C. Fundamento Teórico

El artículo incluye un teorema de separación que demuestra matemáticamente que, bajo ciertas condiciones de ruido de muestra finita, los fragmentos posteriores a un error exhibirán distancias de Wasserstein significativamente mayores que los fragmentos previos al error.

3. Contribuciones Clave

1. Descubrimiento Empírico: Demostración de que la divergencia en los estados ocultos es un proxy fiable de la calidad del razonamiento local sin necesidad de etiquetas externas.
2. Base Teórica: Un teorema que justifica el uso de la distancia de Wasserstein para la asignación de crédito.
3. Algoritmo SHEAR: Una modificación mínima al pipeline de entrenamiento de GRPO que no requiere modelos adicionales ni anotaciones de pasos intermedios.

4. Resultados Experimentales

Los autores evaluaron el método en 5 benchmarks de razonamiento matemático y 5 de generación de código utilizando diversos modelos base (Qwen2.5, Llama 3.1).

• Razonamiento Matemático: SHEAR superó consistentemente a GRPO estándar y, lo más importante, superó a los modelos basados en PRM (modelos de recompensa de proceso). Esto se debe a que SHEAR evita el desajuste de distribución que ocurre cuando un PRM externo se descalibra respecto al modelo que está aprendiendo.
• Generación de Código: Se observaron mejoras consistentes, especialmente en tareas de mayor complejidad y longitud (como LiveCodeBench).
• Eficiencia: El sobrecosto computacional es bajo, situándose entre el 7% y el 16% del tiempo de entrenamiento total.

5. Significado e Impacto

La importancia de este trabajo radica en su capacidad para democratizar la asignación de crédito fina (fine-grained credit assignment). Antes de SHEAR, mejorar el razonamiento paso a paso requería un esfuerzo humano masivo para etiquetar cada paso. SHEAR permite que el modelo "se dé cuenta" de dónde se equivoca simplemente analizando cómo cambian sus propias representaciones internas, permitiendo un entrenamiento de razonamiento mucho más profundo, escalable y eficiente.