Resumen Técnico: Rastreo de Circuitos Markovianos para la Dinámica de Estados en Transformadores

Enunciado del Problema

La interpretabilidad mecanicista se ha centrado en gran medida en analizar circuitos locales, características y cabezas de atención para explicar comportamientos aislados en transformadores. Sin embargo, las tareas de secuencia parcialmente observadas requieren un paradigma computacional diferente: el modelo debe mantener y actualizar un resumen interno (estado) a lo largo del tiempo para informar predicciones futuras. Aunque los Modelos Ocultos de Markov (HMM) proporcionan un marco matemáticamente riguroso para tales tareas, donde la solución óptima es un vector de creencia bayesiana sobre estados latentes, existe una falta de benchmarks controlados para probar si las activaciones de los transformadores codifican realmente estas dinámicas de transición de estados. El trabajo existente a menudo estudia componentes (sondeo, parcheo, extracción de estados) de forma aislada, sin acceso a la verdad fundamental exacta para estados latentes, vectores de creencia y objetivos contrafactuales simultáneamente.

Metodología: Rastreo de Circuitos Markovianos (MCT)

El artículo introduce el Rastreo de Circuitos Markovianos (MCT), un pipeline de diagnóstico diseñado para probar si las activaciones de los transformadores contienen una estructura de transición de estados gruesa. La metodología opera dentro de un benchmark sintético controlado utilizando seis familias distintas de HMM con complejidades variables de transición y emisión.

El Benchmark

El benchmark entrena transformadores causales diminutos (2 capas, ancho 128, ~275k parámetros) sobre emisiones visibles de HMM donde los estados latentes, las matrices de transición y las probabilidades de emisión son conocidos exactamente. Esta configuración proporciona:

1. Verdad Fundamental Exacta: Secuencias de estados latentes, vectores de creencia bayesiana y matrices de transición.
2. Objetivos Exactos: Predicciones del siguiente token óptimas al Bayes y objetivos contrafactuales de estado forzado (la distribución del siguiente token si el estado latente se forzara a un valor específico).

El Pipeline MCT

Dado un transformador entrenado sobre emisiones de HMM, MCT realiza los siguientes pasos:

1. Extracción de Estados: Un extractor de estados $g$ mapea los vectores de activación $a_{\ell,t}$ a estados candidatos discretos $z_t$. El extractor de referencia utiliza agrupamiento K-means sobre las activaciones residuales.
2. Recuperación de Creencias: Una sonda de ridge lineal mapea las activaciones al vector de creencia bayesiana exacto para probar si la información de creencia está presente.
3. Recuperación de Transiciones: La secuencia de estados extraída se utiliza para estimar una matriz de transición, que se compara con la matriz de transición real del HMM utilizando la divergencia KL por filas.
4. Prueba de Orden Markoviano: Se ajustan modelos empíricos de orden 0, 1 y 2 a los estados extraídos para evaluar la fuerza de la señal markoviana de primer orden.
5. Forzamiento de Estados (Intervención): Esta es la prueba causal central. El centroide de un cluster de estado recuperado se parchea en el modelo en una posición de token específica. La salida del modelo se compara luego con el objetivo contrafactual exacto del HMM (la distribución resultante de forzar el estado latente). El rendimiento se mide mediante la divergencia KL entre la salida del modelo parcheado y el objetivo exacto.

Resultados Clave

Aprendizaje de Tareas

En 18 ejecuciones (6 familias de HMM $\times$ 3 semillas), los transformadores aprendieron la tarea de predicción de manera efectiva, logrando una pérdida excedente media sobre el predictor óptimo al Bayes de solo 0.0138.

Recuperación de Estados y Transiciones

• Información de Creencia: Las sondas lineales confirmaron que la información de creencia es recuperable a partir de las activaciones en todas las familias, aunque la calidad varía significativamente según la familia de HMM.
• Estructura de Transición: Las abstracciones de estado basadas en residuos y PCA recuperaron una estructura de transición parcial. La señal fue más fuerte en regímenes persistentes (transiciones adhesivas) y fácilmente separables (baja entropía de emisión).
• Limitaciones: La recuperación fue más débil en regímenes de emisión ambigua y de seis estados. El agrupamiento exacto de creencias permaneció como un límite superior más fuerte que el agrupamiento derivado de activaciones, lo que sugiere que el modelo representa un objeto continuo similar a una creencia en lugar de una copia directa de la etiqueta latente discreta.

Forzamiento de Estados (Evidencia Causal)

La intervención de forzamiento de estados proporcionó la evidencia más clara de utilidad causal:

• Rendimiento: Parchear un centroide de estado recuperado redujo la divergencia KL hacia el objetivo contrafactual exacto de 0.1957 (sin parchear) a 0.0532 en promedio.
• Controles: Esta mejora superó significativamente a los controles, incluidos centros de estados incorrectos, activaciones medias, activaciones aleatorias, centros de etiquetas mezcladas e incluso centros de estados verdaderos en algunas métricas agregadas.
• Implicación: La abstracción de estado recuperada es causalmente significativa y está más cerca del estado predictivo interno del modelo que la etiqueta latente muestreada en ciertos contextos.

Contribuciones

1. Benchmark Controlado: Un marco unificado que proporciona etiquetas latentes exactas, creencias bayesianas, matrices de transición y objetivos contrafactuales para probar métodos de interpretabilidad de transformadores.
2. Pipeline de Diagnóstico (MCT): Un pipeline de referencia para mapear activaciones a abstracciones de estado, probar el orden markoviano y validar la relevancia causal mediante el forzamiento de estados.
3. Hallazgos Empíricos: Evidencia de que los transformadores pequeños aprenden predictores cercanos al Bayes y que sus activaciones residuales contienen una estructura de transición parcial y gruesa, particularmente en regímenes con señales de transición fuertes.

Significado y Afirmaciones

El artículo posiciona a MCT no como un algoritmo completo para recuperar estados ocultos, sino como un marco de diagnóstico para la dinámica de estados de los transformadores. Los autores afirman que:

• MCT detecta con éxito una estructura de transición predictiva parcial en entornos sintéticos.
• Identifica dónde las señales de interpretabilidad son fuertes (regímenes persistentes/fáciles) y dónde son débiles (regímenes ambiguos/de alta entropía).
• Las pruebas de forzamiento de estados que utilizan contrafactuales exactos de HMM proporcionan un método riguroso para evaluar la relevancia causal de las abstracciones de estado extraídas.

El trabajo es modesto en alcance, reconociendo que las tareas son sintéticas, los modelos son diminutos y el extractor de referencia es simple. No afirma que el razonamiento en lenguaje natural posea estructuras de estado idénticas, sino que establece una línea base para probar si las abstracciones derivadas de activaciones pueden exponer dinámicas predictivas gruesas frente a una verdad fundamental exacta.