Resumen Técnico: Descubriendo Estrategias de Razonamiento Latentes en Modelos de Lenguaje

Planteamiento del Problema

Los modelos de lenguaje (LM) preentrenados entrenados en tareas de razonamiento suelen resolver problemas utilizando múltiples estrategias distintas y de alto nivel (por ejemplo, prueba por contradicción frente a inducción, o diferentes algoritmos de ordenación). Sin embargo, estas estrategias son implícitas y están entrelazadas dentro de la distribución de respuesta del modelo $p_\theta(y|x)$. No existe una variable explícita para identificar o controlar qué estrategia se está utilizando durante la generación.

El artículo formula el problema de descomponer esta distribución de respuesta en una representación condicionada por la estrategia estructurada. El objetivo es aprender una factorización de enrutador-generador $(r_\phi(z|x), g_\phi(y|x, z))$ tal que:

1. Fidelidad Marginal: La distribución integrada $\int r_\phi(z|x)g_\phi(y|x, z) dz$ aproxima la distribución del modelo base original $p_\theta(y|x)$.
2. Alineación de Estrategia: La variable latente $z$ es informativa y está alineada con la estrategia, lo que significa que variar $z$ cambia la estrategia de alto nivel utilizada, y el mismo valor latente conlleva un significado de estrategia consistente a través de diferentes entradas.

El Desafío Central: A diferencia del modelado de variables latentes estándar donde un modelo se entrena desde cero, el modelo base $p_\theta$ ya captura la distribución objetivo sin ninguna variable latente. Por consiguiente, una factorización que ignore $z$ (donde $g_\phi(y|x, z) = p_\theta(y|x)$ para todo $z$) está disponible en la inicialización y es globalmente óptima para los objetivos que solo miden el ajuste marginal. La inferencia variacional estándar (por ejemplo, el ELBO condicional) sufre de colapso del posterior en este escenario, fallando en incentivar a $z$ para codificar la variación a nivel de estrategia.

Metodología

Arquitectura

Los autores adaptan un Transformer preentrenado a un marco de enrutador-generador:

• Enrutador ($r_\phi$): Lee la entrada $x$ y devuelve una distribución sobre una estrategia latente continua $z$ (modelada como una Gaussiana diagonal).
• Generador ($g_\phi$): Un modelo autorregresivo que condiciona en $x$ y un embedding proyectado $E(z)$ del latente muestreado.
• Posterior ($q_\xi$): Un modelo separado utilizado únicamente durante el entrenamiento para inferir $z$ a partir del par completo entrada-respuesta $(x, y)$.
• Parametrización: El enrutador y el generador comparten parámetros inicializados a partir del modelo base $\theta$, actualizados mediante una adaptación de bajo rango tipo LoRA.

Objetivo de Reconstrucción Dirigido por el Modelo

Para superar la subespecificación del ELBO estándar, los autores proponen un objetivo de reconstrucción dirigido por el modelo. Este objetivo utiliza el modelo base fijo $p_\theta$ como referencia para dirigir dónde la variable latente $z$ debe explicar la variación.

La función objetivo $J_\theta$ consiste en un término de reconstrucción $R_\theta$ y una penalización KL:
$$J_\theta(x, y; \phi, \xi) = R_\theta(x, y; \phi, \xi) + \beta \text{KL}(q_\xi(\cdot|x, y) \parallel r_\phi(\cdot|x))$$

Componentes clave del término de reconstrucción:

1. Normalización del Modelo Base: La pérdida de reconstrucción se normaliza por la pérdida de respuesta esperada del modelo base ($c_\theta$), expresando la pérdida en "unidades de referencia del modelo base".
2. Ponderación Dirigida por la Sorpresa: En lugar de un peso uniforme por token, la pérdida se pondera por la sorpresa a nivel de token del modelo base ($b_{\theta,t}$).
   • $b_{\theta,t}(x, y) = -\log p_\theta(y_t | x, y_{<t})$
   • Los tokens con alta sorpresa (puntos de ramificación donde el modelo base es incierto) reciben una mayor presión de reconstrucción.
   • Esto fomenta que $z$ codifique la variación en las posiciones donde el modelo base deja plausibles múltiples continuaciones, lo que típicamente corresponde a los puntos de elección de estrategia.

El término de reconstrucción se define como una fracción ponderada y normalizada de la pérdida del modelo base:
$$R_\theta = \frac{1}{c_\theta} \mathbb{E}_{z \sim q_\xi} \left[ \sum_{t} a_{\theta,t}(x, y) [-\log g_\phi(y_t | x, z, y_{<t})] \right]$$
donde los pesos $a_{\theta,t}$ interpolan entre la presión uniforme y la presión dirigida por la sorpresa.

Configuración Experimental y Benchmark

Los autores introducen un benchmark controlado de tareas algorítmicas de múltiples estrategias para hacer medible la alineación de la estrategia.

• Tareas: Suma de listas, algoritmos de ordenación, búsqueda de caminos en rejilla, resolución de ecuaciones lineales, conversión de bases y suma de múltiples dígitos.
• Diseño: Cada tarea permite múltiples trazas de solución válidas (estrategias) para la misma instancia del problema. Las etiquetas de estrategia de verdad absoluta (ground-truth) se generan pero se mantienen ocultas durante el entrenamiento, utilizándose solo para la evaluación.
• Métricas de Evaluación:
  1. Fidelidad Distribucional: La fracción de trazas generadas que siguen siendo soluciones válidas para la tarea.
  2. Alineación de Estrategia (Consistencia Analógica): La probabilidad de que el muestreo de un latente $z$ de una entrada fuente $X$ y su aplicación a una entrada objetivo relacionada $X'$ resulte en una traza de solución $Y'$ que utilice la misma estrategia que la $Y$ generada a partir de $X$.

Resultados Clave

1. La Información Implícita Existe: Sondas lineales en los estados ocultos del modelo base pueden decodificar las etiquetas de estrategia, confirmando que la estructura objetivo existe implícitamente antes de la factorización.
2. Fracaso del ELBO Estándar: Los objetivos variacionales estándar preservan la fidelidad de la distribución pero resultan en un colapso latente, donde $z$ permanece inerte y no logra alinearse con las estrategias.
3. Éxito de la Reconstrucción Dirigida por el Modelo: El objetivo propuesto recupera con éxito códigos latentes que:
   • Mantienen una alta fidelidad distribucional (preservando las salidas válidas del modelo base).
   • Logran una alta Consistencia Analógica, demostrando que $z$ transporta semántica de estrategia reutilizable a través de diferentes entradas.
   • Se separan en regiones distintas específicas de la tarea y de la estrategia en el espacio latente durante el entrenamiento.
4. Estudios de Ablación:
   • Eliminar la escala global (normalización) o la ponderación de tokens (dirección de la sorpresa) degrada la alineación de la estrategia.
   • Usar una ponderación de sorpresa inversa (enfatizando los tokens de baja sorpresa) debilita la alineación, confirmando que la presión debe aplicarse a los tokens inciertos o de ramificación para recuperar la variación de la estrategia.

Análisis Teórico

El artículo inicia un análisis teórico del problema de recuperación. Encuadra el objetivo como la maximización de un puntaje perfilado que involucra la información mutua $I(Y; Z|X)$ menos una penalización de complejidad de aproximación. Los autores proporcionan condiciones suficientes bajo las cuales un objetivo variacional penalizado por complejidad selecciona una factorización latente alineada con una estructura de estrategia objetivo $S^*$, siempre que la "discrepancia de recuperación" entre la solución óptima y las subóptimas sea positiva.

Significado y Reivindicaciones

El artículo afirma abordar una forma patológica específica de colapso del posterior que es única para la adaptación de modelos preentrenados en lugar de entrenarlos desde cero. Al introducir un objetivo de reconstrucción dirigido por el modelo, los autores demuestran un método para:

• Desentrelazar las estrategias de razonamiento implícitas en una variable latente explícita y controlable.
• Preservar la distribución de respuesta del modelo base mientras se hace que la variable latente sea informativa.
• Permitir la Intervención: La estructura latente recuperada permite dirigir la generación hacia estrategias de alto nivel específicas (por ejemplo, forzar un método de prueba específico) sin reentrenar el modelo base.

Los autores señalan que la metodología es validada en tareas algorítmicas sintéticas con estrategias de verdad absoluta. Establecen explícitamente que extender esto a tareas de razonamiento de código abierto y caracterizar completamente la conexión entre el criterio de selección del oráculo teórico y la dinámica de entrenamiento real siguen siendo preguntas abiertas para trabajos futuros. Las aplicaciones potenciales mencionadas incluyen la interpretabilidad del modelo, la controlabilidad y la exploración dirigida en el aprendizaje por refuerzo post-entrenamiento.