Resumen Técnico: Autoencoders Dispersos Residualizados para Intervenciones en Transformadores de Múltiples Capas

Declaración del Problema

Los Autoencoders Dispersos (SAEs) son una herramienta principal para la interpretabilidad mecánica, descomponiendo los estados ocultos de los transformadores en características dispersas e interpretables. Sin embargo, la práctica estándar entrena SAEs de forma independiente sobre las activaciones crudas de cada capa. Este enfoque ignora una propiedad estructural fundamental de los transformadores: el flujo residual está fuertemente acoplado a través de la profundidad. Una fracción significativa de la activación de cualquier capa consiste en información transportada hacia adelante o ligeramente transformada desde capas anteriores.

Entrenar SAEs independientes sobre activaciones crudas crea dos problemas prácticos para las intervenciones de múltiples capas:

1. Redundancia de Capacidad: Diferentes diccionarios por capa desperdician capacidad reaprendiendo la misma información predecible y transportada.
2. Interacciones Impredecibles: Cuando múltiples capas son reemplazadas simultáneamente (por ejemplo, mediante ablación o reconstrucción de características), las interacciones entre capas no son aditivas. Si la misma información está representada en múltiples diccionarios, el reemplazo simultáneo puede causar que las características se dupliquen, se cancelen o impulsen las activaciones fuera de la distribución de entrenamiento, lo que lleva a errores compuestos y rendimiento degradado del modelo (entropía cruzada).

Metodología: Autoencoders Dispersos Residualizados (ReSAEs)

Los autores proponen Autoencoders Dispersos Residualizados (ReSAEs) para desacoplar la información específica de la capa de la estructura predecible entre capas. El método implica tres etapas clave:

1. Residualización Afín Pareada

Para un subconjunto seleccionado de capas $L = (\ell_1, \ell_2, \dots, \ell_M)$, el método ajusta un mapa afín entre capas seleccionadas consecutivas. Para un par $(\ell_m, \ell_{m+1})$, se aprende una transformación afín $A_m h^{(\ell_m)} + c_m$ mediante mínimos cuadrados ordinarios regularizados con ridge para predecir la activación de la capa posterior a partir de la anterior.
La activación residualizada $r^{(\ell_{m+1})}$ se define entonces como el componente de la activación de la capa posterior no explicado por este mapa afín:
$$r^{(\ell_{m+1})} = h^{(\ell_{m+1})} - (A_m h^{(\ell_m)} + c_m)$$
La primera capa de anclaje se centra por su media empírica.

2. Entrenamiento Independiente de SAE por Capa

Cada SAE se entrena sobre el residuo en lugar de la activación cruda.

• Entrada: La activación residual se normaliza (mediante normalización RMS por bloque) y se alimenta a un SAE estándar BatchTopK.
• Objetivo: El SAE minimiza la pérdida de reconstrucción del objetivo residual.
• Independencia: Los gradientes no fluyen entre SAEs de diferentes bloques; cada SAE aprende solo la información "nueva" introducida en su capa específica.

3. Reconstrucción en Espacio Original e Intervención en Línea

Para evaluar los ReSAEs utilizando protocolos estándar, el residuo reconstruido se mapea de nuevo al espacio de activación original a través de la cadena afín ajustada.

• Reconstrucción: $\hat{h}^{(\ell_{m+1})} = A_m \hat{h}^{(\ell_m)} + c_m + S^{-1}_{m+1} \hat{z}^{(m+1)}$, donde $\hat{z}$ es la reconstrucción del SAE del residuo.
• Intervención en Línea: Durante un paso hacia adelante con ganchos, la entrada residual para la capa $\ell_m$ se calcula restando la contribución predicha de la capa previamente reconstruida $\hat{h}^{(\ell_{m-1})}$ de la activación entrante actual. Esto asegura que el SAE en la capa $\ell_m$ vea una distribución de entrada similar a su objetivo de entrenamiento, incluso si las capas aguas arriba han sido modificadas.

Contribuciones Clave

1. Arquitectura ReSAE: Introducción de un paradigma de entrenamiento donde los SAEs se entrenan sobre el componente de la activación no predecible linealmente desde una capa seleccionada anterior, con reconstrucciones compuestas mediante una cadena afín.
2. Mejora en Intervenciones de Múltiples Capas: Demostración de que, a pesar de reconstruir menos varianza cruda, los ReSAEs recuperan la entropía cruzada del transformador de manera más fiel bajo reemplazo simultáneo de múltiples capas en comparación con los SAEs crudos.
3. Reducción de Redundancia: Evidencia de que la residualización produce direcciones de decodificador más ortogonales, lo que indica que la capacidad se dirige hacia información local de la capa en lugar de una estructura duplicada entre capas.
4. Evaluación Exhaustiva: Benchmarks extensos en Pythia-1.4B y Gemma-2-9B que cubren calidad de reconstrucción, degradación de entropía cruzada, similitud de decodificador, sondeo disperso, perturbación de sondeo dirigido (TPP) y eliminación de correlaciones espurias (SCR).

Resultados

Los experimentos se realizaron en Pythia-1.4B (capas 0, 6, 12, 18) y Gemma-2-9B (varios huecos entre capas).

• Reconstrucción vs. Funcionalidad: Los ReSAEs exhiben una Varianza Explicada (EV) más baja en el espacio de activación original en comparación con los SAEs crudos, ya que descartan intencionalmente la estructura predecible linealmente. Sin embargo, esta EV más baja se correlaciona con una menor similitud coseno de decodificador (menos redundancia).
• Recuperación de Entropía Cruzada: Bajo reemplazo secuencial en línea de múltiples capas (donde los errores se propagan), los ReSAEs superan significativamente a los SAEs crudos en la recuperación de la entropía cruzada del modelo, particularmente en niveles de dispersión más altos ($k \ge 160$). En configuraciones forzadas por el profesor (sin propagación de errores), los ReSAEs logran consistentemente la mejor recuperación de entropía cruzada.
• Sobreinteracción: Los SAEs crudos muestran una gran "sobreinteracción" negativa (degradación no aditiva) cuando se reemplazan múltiples capas. Los ReSAEs reducen esta métrica hacia cero, lo que indica que sus reemplazos por capa se componen de manera más predecible.
• Calidad de Características:
  • Sondeo Disperso y TPP: Los ReSAEs generalmente superan a los SAEs crudos en precisión de sondeo disperso y perturbación de sondeo dirigido, lo que sugiere que las características aprendidas son más útiles para la lectura lineal y la ablación selectiva.
  • Eliminación de Correlaciones Espurias (SCR): Los ReSAEs no superan uniformemente a los SAEs crudos aquí; en algunos configuraciones (particularmente ablaciones conjuntas por capa en Pythia), los SAEs crudos obtienen puntuaciones más altas. Esto sugiere que las características residualizadas pueden ser menos selectivas para eliminar atributos espurios específicos sin dañar las señales de la tarea.
• Escala y Espaciado: Los beneficios de los ReSAEs persisten en el modelo más grande Gemma-2-9B a través de diferentes huecos entre capas, con reducciones consistentes en la redundancia de decodificador y mejoras en la recuperación de entropía cruzada.

Significado y Afirmaciones

El artículo argumenta que tratar el flujo residual del transformador como un objeto acoplado en lugar de una colección de capas aisladas es crucial para intervenciones efectivas de múltiples capas.

• Asignación de Capacidad: Al eliminar la estructura predecible linealmente entre capas, los ReSAEs obligan al diccionario disperso a centrarse en los componentes de la activación más relevantes para el cálculo aguas abajo del modelo en esa capa específica.
• Estabilidad de la Intervención: El método proporciona una configuración predeterminada más estable para intervenciones de múltiples capas, reduciendo los errores compuestos y las interacciones impredecibles que afectan las implementaciones estándar de SAE.
• Compensación: Los autores reconocen que los ReSAEs no son una mejora universal para todas las tareas de interpretabilidad (por ejemplo, los resultados de SCR son mixtos). Sin embargo, para el objetivo específico de reemplazo de múltiples capas e intervención basada en características, la compensación de menor varianza cruda por mayor fidelidad funcional y menor redundancia es beneficiosa.

El trabajo sugiere que eliminar la estructura predecible mediante regresión afín es un paso de preprocesamiento útil para los SAEs en configuraciones de múltiples capas, aunque la elección óptima del predictor (afín vs. no lineal) y la selección de capas siguen siendo un espacio de diseño abierto.