Resumen Técnico: Tree SAE: Aprendizaje de Estructuras Jerárquicas de Características en Autoencoders Dispersos

Declaración del Problema

Los Autoencoders Dispersos (SAE) se utilizan ampliamente para extraer características interpretables por humanos de modelos de lenguaje opacos. Sin embargo, los SAE estándar tratan las características como un conjunto independiente, descuidando la estructura jerárquica inherente de los conceptos del mundo real. Esta omisión conduce a fenómenos indeseables como la absorción de características (donde las características de granularidad fina inhiben a las de granularidad gruesa), la división de características (donde un solo concepto se fragmenta en múltiples características) y la composición de características (donde los modelos aprenden amalgamas "polisémicas" fragmentadas o excesivamente complejas).

Las arquitecturas recientes, como los SAE Matryoshka y los SAE de Búsqueda de Coincidencias (Matching Pursuit), intentan abordar estos problemas introduciendo reconstrucción multinivel u ortogonalidad condicional. Se basan en la Similitud Coseno Enmascarada (MCS) para identificar relaciones jerárquicas, la cual se fundamenta en la condición de cobertura de activación: una característica hija se considera descendiente de una padre si se activa únicamente cuando la característica padre se activa.

Los autores demuestran que la cobertura de activación por sí sola es insuficiente para identificar relaciones jerárquicas verdaderas. Proporcionan evidencia de que las características padre e hija pueden satisfacer la condición de cobertura de activación mientras permanecen semánticamente no relacionadas (por ejemplo, una característica padre de "tiempo" emparejada incorrectamente con una característica hija de "política", o una "característica PCA" densa actuando como padre de conceptos no relacionados). Esto conduce a falsos positivos donde los pares jerárquicos identificados carecen de coherencia conceptual.

Metodología

1. La Condición de Reconstrucción

Para abordar las limitaciones de la cobertura de activación, los autores introducen una nueva condición de reconstrucción. Esta condición impone un vínculo funcional más profundo entre los niveles jerárquicos.

• Definición: Para una característica padre $f_p$ y una característica hija $f_c$, la puntuación de reconstrucción $S_{res}$ se define basándose en la alineación de sus vectores decodificadores ($d_p, d_c$) con el vector de concepto verdadero de la hija ($d^*_c$).
• Mecanismo: La condición requiere que tanto los vectores decodificadores del padre como de la hija contribuyan constructivamente a la reconstrucción del concepto de la hija. Específicamente, asegura que el padre capture una dirección conceptual amplia que la hija refina.
• Validación: Los autores muestran que, aunque la cobertura de activación puede ser satisfecha por pares espurios, la puntuación de reconstrucción filtra eficazmente estos casos. Los pares jerárquicos verdaderos exhiben puntuaciones de reconstrucción altas, mientras que los pares no relacionados (incluso con cobertura de activación perfecta) exhiben puntuaciones bajas.

2. Arquitectura Tree SAE

Los autores proponen Tree SAE, un modelo diseñado para aprender estructuras jerárquicas directamente dentro del conjunto de características, en lugar de inferirlas a posteriori.

• Estructura de Árbol: El conjunto de características se divide en múltiples "capas de privilegio". Una característica en una capa superior puede asignarse como descendiente de una característica en cualquier capa inferior. Un nodo raíz imaginario en la capa 0 sirve como línea base para las características sin padres explícitos.
• Vector de Asignación: Un vector de asignación $a_l$ asigna características padre a características hija en cada capa $l$.
• Activación Regularizada: Para imponer la estructura de árbol, la activación de una característica hija $f_i$ se regulariza para ser no nula solo si su característica padre asignada está activa. Esto sigue estrictamente la condición de cobertura de activación durante el entrenamiento.
• Pérdida de Reconstrucción Multinivel: El objetivo de entrenamiento incluye una pérdida de reconstrucción multinivel. Para una capa $l$, la pérdida restringe la reconstrucción acumulada de todas las capas anteriores (incluida $l$) para recuperar con precisión la entrada. Esto asegura que los vectores padre capturen conceptos generales y los vectores hija los refinen en conceptos especializados.
• Asignación Dinámica: Para mitigar el problema de las "características muertas" (características que nunca se activan), los autores introducen un mecanismo de asignación dinámica. Este algoritmo reasigna características hija muertas a características padre basándose en una métrica de "ganancia" (derivada de la pérdida de entrenamiento acumulada), asegurando que las características hija se asignen a padres que puedan apoyar eficazmente su aprendizaje.

Contribuciones Clave

1. Identificación de las Limitaciones de la Cobertura de Activación: El artículo proporciona evidencia concreta de que la cobertura de activación por sí sola identifica pares jerárquicos semánticamente incoherentes y propone una nueva definición que combina la cobertura de activación con la condición de reconstrucción.
2. Modelo Tree SAE: La introducción de una arquitectura SAE novedosa que codifica una estructura de árbol jerárquica directamente dentro del conjunto de características, la cual se aprende dinámicamente durante el entrenamiento.
3. Aprendizaje Jerárquico Superior: Demostración de que los Tree SAE aprenden pares padre-hija significativamente más coherentes que los SAE existentes (incluidos los SAE Matryoshka y Matching Pursuit), evitando al mismo tiempo las patologías de división, absorción y composición de características.
4. Análisis Geométrico: Mostrar la utilidad de la estructura de árbol explícita para mapear la geometría de los subespacios de características hija y comprender cómo los modelos de lenguaje descomponen los conceptos en niveles de granularidad gruesa y fina.

Resultados Experimentales

Los autores evaluaron los Tree SAE en GPT2-small a través de varios niveles de dispersión ($L_0$) y tamaños de diccionario, comparándolos con SAE Top-k, SAE Matryoshka y SAE de Búsqueda de Coincidencias.

• Métricas Jerárquicas: Los Tree SAE descubrieron sustancialmente más pares jerárquicos que satisficieron tanto la cobertura de activación como las condiciones de reconstrucción en comparación con otros modelos. Se encontró que la estructura de árbol interna era ligeramente más efectiva para identificar estos pares que el análisis MCS a posteriori.
• Patologías de Características: Los Tree SAE rindieron de manera competitiva o superaron a los SAE Matryoshka de última generación (SOTA) en la reducción de la división, absorción y composición de características.
• Reconstrucción e Interpretabilidad: Los Tree SAE mantuvieron una pérdida de reconstrucción competitiva y varianza explicada, reconstruyendo con precisión los latentes originales. En términos de interpretabilidad (medida mediante AutoInterp), los Tree SAE igualaron el rendimiento de los SAE Matryoshka.
• Escalabilidad: Los Tree SAE escalaron bien con los tamaños de diccionario (de 6k a 49k), a menudo superando a los SAE Matryoshka en métricas jerárquicas a escalas más pequeñas y manteniendo un rendimiento comparable a escalas más grandes.
• Diversidad de Características Hija: Las características hija aprendidas exhibieron tasas de co-activación bajas (<3.0%), lo que indica que representan aspectos distintos del concepto padre en lugar de información redundante.

Significado y Afirmaciones

El artículo afirma que Tree SAE representa un paso significativo hacia adelante en la comprensión de las representaciones internas de los grandes modelos de lenguaje. Al ir más allá de los conjuntos de características independientes para modelar explícitamente estructuras jerárquicas, los autores argumentan que los Tree SAE:

• Mitigan las patologías conocidas de características jerárquicas (división, absorción, composición).
• Proporcionan un marco más preciso para identificar relaciones verdaderas de conceptos padre-hija, evitando los falsos positivos comunes en métodos anteriores.
• Permiten el análisis geométrico de los subespacios de conceptos jerárquicos, ofreciendo nuevas perspectivas sobre cómo los modelos descomponen conceptos generales en subconceptos especializados.

Los autores señalan limitaciones, específicamente que la estructura de árbol rígida puede conducir a un mayor número de características muertas en comparación con otros SAE, aunque la asignación dinámica ayuda a mitigar esto. Sugieren que el trabajo futuro podría explorar mecanismos de puerta suave (soft-gating) para mejorar aún más la utilización de características.