Resumen Técnico: VASAE – Asignación de Nombres al Diccionario de Direcciones de SAE mediante el Anclaje Alineado con el Vocabulario

1. Declaración del Problema

Los Autoencoders Dispersos (SAEs) se utilizan cada vez más para descomponer los flujos residuales de alta dimensión de los modelos Transformer en códigos dispersos sobre diccionarios aprendidos, lo que ayuda en la interpretabilidad mecánica. Sin embargo, una limitación significativa de los SAE estándar es que sus características aprendidas (direcciones del diccionario) no están intrínsecamente conectadas con el vocabulario de tokens del modelo.

En la práctica estándar, el nombramiento de características es un proceso post hoc. Los investigadores deben inspeccionar los contextos donde se activan características específicas o ejecutar procedimientos de interpretación automatizados por separado para asignar etiquetas semánticas. Esta desconexión surge porque la geometría del diccionario aprendido y la geometría de los embeddings de los tokens son típicamente ajenas; una característica puede reconstruir eficazmente las activaciones mientras se encuentra lejos de cualquier vector de embedding de token, lo que la deja sin un nombre "intrínseco" a nivel de vocabulario.

El artículo plantea: ¿Podemos entrenar las características de un SAE de modo que sigan siendo lo suficientemente flexibles para una reconstrucción de alta calidad pero estén geométricamente lo suficientemente ancladas cerca de los embeddings de los tokens como para adquirir nombres de tokens intrínsecos durante el entrenamiento, sin sacrificar el rendimiento de la reconstrucción?

2. Metodología: Autoencoder Disperso Alineado con el Vocabulario (VASAE)

Los autores proponen VASAE, una modificación al objetivo de entrenamiento del SAE estándar que introduce un mecanismo de anclaje alineado con el vocabulario.

Arquitectura Central

• Entrada: El VASAE opera sobre los vectores del flujo residual post-procesados ($h^{(\ell)}$) de una capa $\ell$ de un Transformer.
• Codificador (Encoder): Utiliza un codificador disperso top-k estándar (mapa lineal + ReLU + selección Top-K) para producir un código disperso $z$.
• Decodificador (Decoder): Utiliza una matriz de diccionario aprendible $W_D$ (donde las columnas son vectores de características $d_i$) para reconstruir el flujo residual.
• Anclas (Anchors): El método utiliza los embeddings de tokens de entrada fijos del modelo ($W_E$) como direcciones de referencia. Estos embeddings comparten la misma dimensionalidad que el flujo residual (o son proyectados para hacerlo).

El Objetivo de Anclaje

La innovación clave es la adición de un término de pérdida de anclaje al error de reconstrucción estándar.

1. Alineación con el Token más Cercano: Para cada característica aprendida $d_i$, el método calcula la similitud de coseno con todos los embeddings de tokens fijos $w_v$.
   $$s_i = \max_{v \in V} \cos(d_i, w_v)$$
   El token $v^*_i$ que logra este máximo define el nombre de token intrínseco para la característica $i$.
2. Pérdida de Anclaje (Anchor Loss): El objetivo de entrenamiento fomenta que las características permanezcan cerca de su embedding de token más cercano:
   $$L_{anchor} = -\frac{1}{d_{sparse}} \sum_{i=1}^{d_{sparse}} s_i$$
   Minimizar esta pérdida maximiza la similitud de coseno entre las características y sus tokens de vocabulario más cercanos.
3. Pérdida Total: El objetivo combinado es:
   $$L_{vasae} = L_{recon} + \lambda_{anchor} L_{anchor}$$
   donde $\lambda_{anchor}$ es un hiperparámetro que controla la fuerza de la restricción de alineación.

Distinción del Acoplamiento Estricto (Hard-Tying)

Los autores distinguen explícitamente el VASAE de una línea base de decodificador con "acoplamiento estricto" (hard-tied), donde los pesos del decodificador son exactamente iguales a la matriz de embeddings de tokens ($W_D = W_E$). El acoplamiento estricto obliga a las características a ser exactamente los embeddings de los tokens, lo cual, según la hipótesis del artículo, es insuficiente para una reconstrucción de alta calidad. El VASAE mantiene el diccionario aprendible, pero lo atrae hacia la geometría del vocabulario mediante la pérdida de anclaje.

3. Configuración Experimental

El método fue evaluado en dos modelos Transformer de solo decodificador:

• GPT-2-small: 12 capas, $d_{model}=768$, tamaño del vocabulario $|V|=50,257$.
• Llama-3.1-8B: 32 capas, $d_{model}=4096$, tamaño del vocabulario $|V|=128,256$.

Líneas Base (Baselines):

1. SAE Estándar: Entrenado únicamente con la pérdida de reconstrucción.
2. Decodificador con Acoplamiento Estricto (Hard-Tied Decoder): Pesos del decodificador fijados a los embeddings de los tokens; solo se entrena el codificador.

Métricas:

• Reconstrucción: Error Cuadrático Medio (MSE), Varianza Explicada (VE) y recuperación de la Entropía Cruzada (CE) tras la sustitución del flujo residual.
• Alineación: Tasa de alineación geométrica de las características (porcentaje de características con un puntaje de alineación $s_i \ge 0.8$) y cobertura del vocabulario.
• Estudios de Caso: Inspección cualitativa de los nombres de tokens intrínsecos asignados a las características activas en contextos específicos.

4. Resultados Clave

Calidad de la Reconstrucción

El VASAE logra métricas de reconstrucción comparables al SAE Estándar, superando significativamente a la línea base del Decodificador con Acoplamiento Estricto.

• GPT-2-small: Tanto el VASAE como el SAE Estándar logran una Varianza Explicada (VE) de 0.965 y una recuperación de CE de 0.975.
• Llama-3.1-8B: Tanto el VASAE como el SAE Estándar logran una VE de 0.931 y una recuperación de CE de 0.906.
• Línea Base de Acoplamiento Estricto: Tiene un desempeño pobre (por ejemplo, VE negativa en GPT-2-small), confirmando que forzar las características a ser exactamente los embeddings de los tokens destruye la capacidad de reconstrucción.

Alineación Geométrica

• GPT-2-small: Con $\lambda_{anchor} = 10^{-4}$, entre el 89–94% de las características en las capas 0–10 exceden el umbral de alineación ($s_i \ge 0.8$). La alineación es fuerte y consistente en las capas iniciales y medias.
• Llama-3.1-8B: La alineación depende de la profundidad. Con un coeficiente más alto ($\lambda_{anchor} = 5 \times 10^{-3}$), la capa inicial (L0) logra un 92.8% de alineación. Sin embargo, la alineación cae significativamente en las capas más profundas (ej. L31), sugiriendo que el anclaje al vocabulario es más compatible con los estados residuales cercanos al espacio de embeddings.

Estudios de Caso

El análisis cualitativo de los nombres de tokens intrínsecos revela coherencia semántica:

• En GPT-2-small, las características activas en una oración sobre "Baker Street" reciben nombres intrínsecos como located (ubicado), Street (calle), corner (esquina) y around (alrededor).
• Ejemplos adicionales muestran la alineación con adjetivos/adverbios (ej. perfect, feasible), términos relacionados con premios (ej. winner, praise) y términos de auto-presentación (ej. fan, name).
• En Llama-3.1-8B, la capa inicial muestra una alineación similar relacionada con ubicaciones, mientras que las capas profundas muestran menos consistencia, reflejando las métricas de alineación cuantitativas.

5. Contribuciones y Significado

El artículo realiza tres contribuciones principales:

1. Introducción de VASAE: Un objetivo de anclaje de vocabulario suave que asigna nombres de tokens más cercanos a las direcciones del decodificador aprendidas durante el entrenamiento, en lugar de hacerlo post-hoc.
2. Evaluación del Anclaje: Demostrar que el anclaje alineado con el vocabulario preserva una calidad de reconstrucción comparable a los SAE estándar, mientras prueba que el acoplamiento estricto del decodificador a los embeddings de los tokens es insuficiente para esta tarea.
3. Análisis Geométrico: Analizar la alineación de las características en el espacio de embeddings y la cobertura del vocabulario, proporcionando evidencia de que los nombres de tokens intrínsecos actúan como identificadores derivados de la geometría.

Significado:
Los autores argumentan que el VASAE complementa los métodos de interpretación post-hoc existentes al conectar las características aprendidas con los nombres de tokens intrínsecos durante el entrenamiento. Esto proporciona un "mango indexado por vocabulario" para las características, ofreciendo un punto de partida para la inspección.

Crucialmente, el artículo mantiene una postura modesta sobre las implicaciones de estos nombres:

• No son Explicaciones: Un nombre de token intrínseco no es una explicación completa de la característica. No explica por qué una característica se activa ni su papel causal en la computación del modelo.
• Herramienta de Indexación: Los nombres sirven como herramientas preliminares de indexación. La verdadera comprensión requiere un análisis posterior de cómo las características nombradas participan en los circuitos, cómo se conectan con otros componentes y cómo las intervenciones sobre ellas afectan el comportamiento.
• Limitaciones: La alineación es más robusta en GPT-2-small y en las capas iniciales/medias de Llama-3.1-8B. El método actualmente solo se ancla a los embeddings de entrada; para modelos con embeddings de entrada/salida no unidos (untied), la elección de la matriz de anclaje sigue siendo una pregunta abierta.

En resumen, el VASAE logra cerrar la breancia entre el espacio geométrico de las características aprendidas y el espacio del vocabulario del modelo, permitiendo que las características porten nombres de tokens intrínsecos sin comprometer su utilidad para la reconstrucción.