Resumen Técnico: Similitud Tensorial para la Interpretabilidad Mecanística

Enunciado del Problema

La interpretabilidad mecanicista busca descomponer las redes neuronales en partes constituyentes significativas para comprender sus computaciones internas. Un prerequisito crítico para este campo es la capacidad de verificar rigurosamente si dos componentes descompuestos (o dos modelos diferentes) implementan la misma función subyacente. Los métodos de evaluación actuales adolecen de dos limitaciones principales:

1. Similitud Empírica/Conductual: Los métodos que comparan las salidas o activaciones de los modelos en conjuntos de datos específicos son ciegos a los mecanismos fuera de distribución (OOD). No pueden detectar problemas como puertas traseras que se activan solo en entradas específicas o fallar en generalizar si la descomposición se basa en heurísticas aproximadas en lugar de la función real.
2. Similitud Basada en Parámetros: Los métodos que comparan matrices de pesos (por ejemplo, similitud del coseno de matrices) son altamente sensibles a las simetrías del espacio de pesos (como la permutación y la escalación). No logran reconocer la equivalencia funcional cuando los modelos se entrenan de manera diferente o cuando los mecanismos se distribuyen a través de capas en lugar de concentrarse en neuronas individuales.

Existe una falta de una métrica rigurosa e independiente del conjunto de datos que sea invariante a estas simetrías y capaz de capturar la equivalencia funcional global.

Metodología

Los autores proponen la Similitud Tensorial, una métrica basada en pesos diseñada para redes neuronales basadas en tensores (específicamente modelos multilineales como los transformadores bilineales). Este enfoque sustituye el uso de funciones de activación no lineales estándar por arquitecturas polinómicas (multilineales) para permitir un análisis algebraico de los pesos.

Cambio Arquitectónico Central

El método se basa en modelos multilineales donde las funciones de activación son reemplazadas por compuertas multiplicativas (por ejemplo, capas bilineales y atención bilineal).

• Capas Bilineales: Mapean una entrada elevada $\tilde{x} = (1, x)$ a una salida mediante un polinomio de grado 2: $A(x) = D((L\tilde{x}) \odot (R\tilde{x}))$.
• Representación Multilineal: Estas capas pueden expresarse como una contracción de un tensor de pesos $A$ con $n$ copias de la entrada: $A(x) = A \cdot \text{in}^{\otimes n} \tilde{x}$.
• Modelos Profundos: Las redes multilineales profundas se construyen como descomposiciones de estructura arbórea de tensores globales, lo que permite un cálculo eficiente sin materializar el tensor global exponencialmente grande.

La Métrica de Similitud Tensorial

La métrica se define como un producto interno generalizado entre dos tensores de pesos $A$ y $B$, normalizado por sus normas inducidas:
$$\text{sim}_M(A, B) = \frac{\langle A | M | B \rangle}{\|A\|_M \|B\|_M}$$
donde $M$ es un tensor métrico.

Para garantizar que la métrica satisfaga los criterios de invarianza de simetría e independencia del conjunto de datos, los autores derivan una forma específica para $M$:

1. Simetrización: Basándose en el isomorfismo de polarización, los autores establecen que dos modelos multilineales son funcionalmente equivalentes si y solo si sus tensores de pesos tienen el mismo representante en el subespacio simétrico. La métrica proyecta los tensores sobre este subespacio simétrico para eliminar los grados de libertad relacionados con las simetrías de permutación de la entrada.
2. Métrica Gaussiana: Para refinar aún más la métrica y relacionarla con la similitud conductual, los autores definen $M$ utilizando la esperanza del tensor de entrada bajo una distribución gaussiana estándar ($x \sim \mathcal{N}(0, I)$). Usando el teorema de Isserlis, esto resulta en una métrica de forma cerrada que calcula el producto interno de las activaciones sobre todo el espacio de entrada gaussiano, en lugar de un conjunto de datos muestreado.
3. Cálculo Eficiente: Un algoritmo recursivo basado en Gram permite el cálculo de esta similitud para redes profundas procesando tensores locales capa por capa, evitando la necesidad de instanciar el tensor global completo.

Contribuciones Clave

1. Una Métrica Basada en Pesos Principista: El artículo introduce una métrica que es invariante a las simetrías del espacio de pesos (permutaciones, escalado) e independiente del conjunto de datos de entrada. Verifica formalmente la equivalencia funcional comprobando si dos modelos implementan la misma lógica polinómica.
2. Análisis de Forma Cerrada y Libre de Datos: A diferencia de las métricas conductuales, la similitud tensorial no requiere muestrear datos. Integra sobre todo el espacio de entrada gaussiano, lo que la hace capaz de detectar mecanismos OOD (como puertas traseras) que son invisibles para las métricas de rendimiento estándar.
3. Capacidades de Localización: La métrica puede "cortarse" para enfocarse en dimensiones de salida específicas o "diferenciarse" de tensores para aislar cambios funcionales entre puntos de control, permitiendo la atribución de mecanismos específicos a salidas específicas.
4. Escalabilidad: El algoritmo recursivo asegura que la métrica sea computacionalmente tratable para redes profundas, evitando los límites de convergencia exponenciales o cuadráticos asociados con el muestreo de Monte Carlo.

Resultados

Los autores validan la similitud tensorial en cuatro dominios distintos, demostrando su superioridad sobre las métricas existentes (CKA, coseno de matrices, similitud conductual):

1. Olvido Catastrófico (Visión): En una configuración de entrenamiento progresivo en SVHN, la similitud tensorial localiza con precisión la pérdida de conocimiento para dígitos específicos (por ejemplo, el dígito '9') cuando se reduce el conjunto de datos. Muestra una caída brusca en la similitud para la clase afectada mientras permanece estable para las demás, whereas las métricas conductuales muestran señales ambiguas.
2. Grokking (Aritmética Modular): En tareas de suma modular, la similitud tensorial rastrea la reorganización continua del modelo hacia frecuencias de Fourier específicas. Revela que la reorganización estructural continúa incluso después de que la precisión salta a la perfección, proporcionando una visión más legible de la dinámica de entrenamiento que las curvas de precisión por sí solas.
3. Detección de Puertas Traseras (OOD): Cuando se inyecta una puerta trasera en SVHN (activándose en un parche de diamante negro), las métricas conductuales estándar en datos limpios no logran detectar el cambio. La similitud tensorial, sin embargo, señala la puerta trasera solo con datos limpios y localiza el cambio en la clase atacada específica ('9'), demostrando su capacidad para detectar mecanismos OOD.
4. Modelado del Lenguaje: Aplicada a un transformador de atención bilineal de dos capas entrenado en The Pile, la similitud tensorial revela estructuras de bloques claras y transiciones agudas en los regímenes de entrenamiento que las métricas existentes difuminan. Rastrea la formación de estadísticas complejas de n-gramas de manera más efectiva que los proxies conductuales.

Significado y Afirmaciones

El artículo posiciona la similitud tensorial como una herramienta fundamental para la interpretabilidad mecanicista. Su significado principal radica en reducir el problema de medir la similitud y verificar la fidelidad de un problema de aproximación empírica a un problema algebraico resuelto.

• Fidelidad: Al depender de las propiedades algebraicas de los modelos multilineales, la métrica asegura que las descomposiciones capturen la función subyacente real en lugar de heurísticas aproximadas.
• Robustez: Proporciona afirmaciones de robustez más fuertes al aplicarse a toda la distribución, descartando efectivamente los mecanismos que solo se activan en entradas específicas (como puertas traseras).
• Limitaciones: Los autores reconocen modestamente que el método depende actualmente de modelos multilineales, lo que excluye la mayoría de las arquitecturas más avanzadas que utilizan no linealidades no polinómicas. Sin embargo, argumentan que los modelos multilineales sirven como un campo de pruebas vital para la interpretabilidad y que la evidencia sugiere que pueden igualar el rendimiento de las arquitecturas estándar.
• Perspectiva Futura: El trabajo sugiere que la restricción a modelos multilineales es un compromiso necesario para un análisis algebraico tratable, pero allana el camino para futuros métodos de atribución (por ejemplo, calcular la similitud entre caminos en transformadores profundos) y respalda la visión de que las arquitecturas multilineales son más transparentes que las redes profundas ordinarias.