Resumen Técnico: Transporte Óptimo Semántico para la Correspondencia de Características de Autoencoders Dispersos y la Compresión de Circuitos

1. Definición del Problema

Los Autoencoders Dispersos (SAE) han surgido como una herramienta principal para interpretar los Modelos de Lenguaje Grande (LLM) al descomponer los espacios latentes en características interpretables. Sin embargo, persisten dos desafíos críticos en la escalabilidad del análisis de SAE:

1. Correspondencia de Características entre Capas: Identificar características semánticamente similares a través de diferentes capas de un LLM. Los métodos existentes dependen típicamente de comparar vectores decodificadores. Este enfoque falla al emparejar características entre capas no consecutivas porque los vectores decodificadores se entrenan para reconstruir distribuciones de activación distintas, lo que hace que la comparación directa de vectores sea inexacta.
2. Compresión de Circuitos de Características: Los circuitos de características (grafos causales que conectan características entre capas) a menudo contienen cientos de nodos, lo que hace que la inspección manual sea impráctica. Los métodos actuales de compresión dependen de agrupar manualmente nodos semánticamente similares en "supernodos", un proceso que no escala.

El problema central identificado por los autores es la estimación de la distancia semántica entre características que residen en diferentes variedades de activación. Los enfoques existentes tratan las características como vectores individuales (por ejemplo, pesos del decodificador), lo que descarta el rico contexto de dónde y con qué fuerza se activa una característica, lo que conduce a una pérdida de información y a la incapacidad de manejar comparaciones multi-variedad.

2. Metodología

Los autores proponen un marco distribucional que reformula la comparación de características como un problema de Transporte Óptimo (OT).

2.1 Representación de Características como Distribuciones

En lugar de representar una característica mediante un único vector decodificador, el método representa cada característica $i$ en la capa $\ell$ como una distribución empírica ponderada por activación sobre los estados ocultos donde se activa.

• Dado un corpus de posiciones de tokens $T$, el método selecciona las $K$ posiciones de tokens superiores donde la característica $i$ se activa con mayor fuerza.
• Cada token seleccionado $t$ se asigna un peso $w_{i,t}$ proporcional a su magnitud de activación.
• La característica se representa como una distribución $\hat{\mu}_{i,T} = \sum w_{i,t} \delta_{x_t}$, donde $x_t$ es el estado oculto en el token $t$.

2.2 Espacio de Referencia Compartido

Para comparar características entre diferentes capas (que residen en variedades diferentes), el método proyecta estas distribuciones específicas de la capa en un espacio de referencia compartido $\mathcal{R}$.

• Para la correspondencia entre capas, el espacio de referencia es típicamente el espacio de estados ocultos de la capa objetivo. Las activaciones de la capa fuente se proyectan a las posiciones de tokens correspondientes en la capa objetivo.
• Para la compresión de circuitos, el espacio de referencia es la concatenación de los estados ocultos a través de todas las capas.

2.3 Distancia Semántica mediante Distancia de Wasserstein

Una vez proyectadas en el espacio compartido, la distancia semántica entre dos características se calcula utilizando la distancia de Wasserstein (Distancia del Movilizador de Tierra).

• La distancia $W_c(\hat{\mu}, \hat{\nu})$ mide el costo mínimo para transportar la masa de probabilidad de la distribución de activación de una característica a la otra.
• El costo base $c(z, z')$ mide el gasto de emparejar dos representaciones de contexto en el espacio compartido.
• Este enfoque captura la estructura geométrica de los patrones de activación en lugar de simplemente la similitud vectorial punto a punto.

2.4 Tareas Derivadas

• Correspondencia de Características: Para una característica objetivo, el método encuentra la característica fuente con la distancia de Wasserstein mínima (vecino más cercano).
• Compresión de Circuitos: El método trata los nodos del circuito como puntos en el espacio distribucional y aplica agrupamiento (por ejemplo, agrupamiento aglomerativo) para agrupar nodos con pequeñas distancias de Wasserstein pareadas en supernodos, minimizando la distancia semántica dentro del grupo.

3. Contribuciones Clave y Garantías Teóricas

El artículo proporciona tres garantías teóricas principales para el marco propuesto:

1. Invarianza de Redimensionamiento de Activación: La distribución de características proyectada es invariante a la escala global de los valores de activación. Dado que la representación depende de la clasificación relativa y la distribución proporcional de las activaciones (selección Top-K y pesos normalizados), los desplazamientos escalares arbitrarios en las magnitudes de activación no alteran la distancia semántica.
2. Estabilidad bajo Perturbaciones: La distancia de Wasserstein entre distribuciones de características es estable bajo perturbaciones suaves de las distribuciones. El cambio en la distancia está acotado por las constantes de Lipschitz del costo base y la distancia de Wasserstein entre las distribuciones perturbadas y originales. Esto asegura robustez frente al ruido en la selección de tokens o el reentrenamiento de SAE.
3. Recuperación Basada en Margen: Los autores establecen condiciones bajo las cuales la correspondencia de vecinos más cercanos empírica y la agrupación semántica recuperan perfectamente las coincidencias verdaderas (oráculo). Si el "margen" (la brecha entre la mejor coincidencia y la segunda mejor coincidencia) supera el doble del error de estimación, se garantiza que la decisión empírica sea correcta.

4. Resultados Experimentales

El método se evaluó en SAE entrenados en GPT-2 y Gemma-2-2B, comparando contra líneas base que incluyen SAE Match (L2 de vector decodificador), Feature Flow (coseno de vector decodificador), Parcheo de Atribución y selección basada en LLM.

• Correspondencia de Características: El método propuesto superó a todas las líneas base en Evaluación de LLM (calificaciones de similitud semántica) y Precisión de Correspondencia en configuraciones de capas consecutivas, de medio a última y de primera a última. Crucialmente, mantuvo un rendimiento competitivo en métricas derivadas (pérdida de Entropía Cruzada y Varianza Explicada), demostrando que identifica características con roles funcionales similares, no solo tokens de activación similares. Destacó notablemente al emparejar sutiles características de "suma de dígitos" donde otros métodos fallaron.
• Compresión de Circuitos: El método logró una mayor precisión al comprimir circuitos de características en supernodos intuitivos en comparación con la agrupación manual y otras líneas base. Identificó con éxito grupos modulares de características (por ejemplo, "cabezas de inducción" y "detección de tokens") que consistentemente desempeñan el mismo rol funcional en diferentes circuitos.
• Estudio de Ablación: Los experimentos confirmaron que representar características como distribuciones preserva más información semántica que colapsarlas en un único vector centrinoide, incluso cuando se proyectan sobre la misma variedad.

5. Significado y Afirmaciones

El artículo afirma proporcionar una métrica semántica unificada que resuelve tanto la correspondencia de características entre capas como la compresión de circuitos como instancias del mismo problema fundamental: estimar la distancia semántica a través de variedades.

• Escalabilidad: Al automatizar la agrupación de características en supernodos, el método elimina la dependencia del etiquetado humano, haciendo viable el análisis de circuitos de características a gran escala.
• Robustez: El enfoque distribucional captura el "contexto" de la activación de la característica, permitiendo una correspondencia precisa incluso cuando las características evolucionan significativamente a través de las capas (distribuciones disímiles), un escenario donde los métodos basados en vectores fallan.
• Interpretabilidad: La capacidad de comprimir automáticamente circuitos complejos en supernodos interpretables mejora la comprensión de los mecanismos del modelo, como se demuestra en estudios de caso sobre circuitos de "Dallas" (geografía), "Rima" y "Oposición".

Los autores posicionan este trabajo como un paso fundamental hacia la interpretabilidad mecánica escalable y automatizada, superando las limitaciones de la comparación de características centrada en vectores.