Disentangling MLP Neuron Weights in Vocabulary Space

1. El Problema: La Polisemanticidad y la Interpretación de Pesos

La interpretación mecánica de los modelos de lenguaje (LM) enfrenta un desafío fundamental: la polisemanticidad. La mayoría de las neuronas en las capas MLP (Perceptrones Multicapa) de los transformadores no representan un solo concepto coherente, sino que codifican múltiples conceptos entrelazados.

• Limitaciones de los enfoques actuales: La mayoría de los métodos de interpretación se basan en los activaciones (forward passes) del modelo utilizando grandes conjuntos de datos. Esto introduce sesgos dependientes de la distribución de datos y conlleva altos costos computacionales.
• El vacío en la investigación: Inspeccionar la información codificada directamente en los pesos del modelo permitiría un análisis independiente de los datos de entrada, pero encontrar la "unidad de análisis correcta" en el espacio de pesos ha sido difícil. Los métodos existentes, como los Autoencoders Escasos (SAEs), requieren datos de activación masivos y a menudo no logran desentrelazar completamente las neuronas polisémicas en componentes monosémicos.

2. Metodología: ROTATE

Los autores proponen ROTATE (Rotation-Optimized Token Alignment in weighT spacE), un método libre de datos (data-free) que no requiere realizar pasadas forward a través del modelo. En su lugar, opera directamente en el espacio de pesos para descomponer las neuronas MLP.

Observación Clave: La Curtosis en el Espacio del Vocabulario

El núcleo de la metodología se basa en una observación estadística:

• Las neuronas que codifican conceptos coherentes y monosémicos exhiben una alta curtosis cuando sus vectores de peso se proyectan en el espacio del vocabulario del modelo (mediante la matriz de desenmascaramiento $U$).
• Una alta curtosis indica una distribución de logits con colas pesadas y un pico agudo, lo que significa que la neurona actúa fuertemente sobre un conjunto escaso de tokens específicos, mientras que tiene poco efecto en la mayoría de los demás.
• Por el contrario, las neuronas polisémicas o aleatorias tienden a tener distribuciones más gaussianas (baja curtosis).

El Algoritmo ROTATE

Dado un vector de peso de una neurona $w$, ROTATE busca una matriz de rotación $R$ que transforme $w$ en un nuevo vector $v = wR$ que maximice la curtosis en el espacio del vocabulario ($z = vU$).

1. Función de Pérdida: Se optimiza una función de pérdida compuesta por dos términos:
   • Pérdida de Curtosis ($L_{kurt}$): Maximiza la curtosis de la proyección de vocabulario para empujar al vector hacia direcciones monosémicas.
   • Regularización ($L_{reg}$): Penaliza la distancia coseno entre el vector rotado $v$ y el peso original $w$. Esto asegura que los canales descubiertos sean fielmente derivados de la neurona original y no direcciones arbitrarias de alta curtosis.
2. Parametrización: Se utilizan matrices de reflexión de Householder para garantizar que la transformación sea ortogonal.
3. Proceso Iterativo con Enmascaramiento: Dado que una neurona puede contener múltiples conceptos, el algoritmo se ejecuta iterativamente. Después de encontrar un canal (vector $v_i$), se enmascaran los tokens que contribuyeron más a la curtosis de ese canal. Esto fuerza a las iteraciones subsiguientes a descubrir nuevas direcciones de alta curtosis que no se solapan con las anteriores.
4. Resultado: El método produce un conjunto de canales de vocabulario (vocabulary channels), que son proyecciones esparsas e interpretables de la neurona original alineadas con el vocabulario del modelo.

3. Contribuciones Clave

1. Observación Teórica: Demostraron que la alta curtosis en las distribuciones de vocabulario es un indicador fiable de direcciones monosémicas en el espacio de pesos de los LM.
2. Método ROTATE: Introdujeron un algoritmo eficiente y libre de datos que descompone los pesos de las neuronas MLP en componentes interpretables sin necesidad de datos de entrenamiento o inferencia.
3. Validación Empírica: Demostraron que los canales recuperados son fieles al comportamiento de la neurona y cubren un espectro de activación más amplio que los métodos basados en SAE.
4. Mejora en Descripciones: Mostraron que agregar las descripciones de los canales individuales genera explicaciones de neuronas más completas y precisas que los enfoques automatizados basados en activaciones.

4. Resultados Experimentales

Los experimentos se realizaron en modelos Gemma-2-2B-it y Llama-3.1-8B-Instruct.

• Fidelidad (Faithfulness):

  • Entrada: Los canales descubiertos predicen con mayor precisión los patrones de activación de la neurona que los baselines basados en SAE (Gemma Scope y Llama Scope). ROTATE obtuvo puntuaciones de fidelidad entre 0.46 y 0.71, superando consistentemente a los SAE (0.33–0.49).
  • Salida: Mediante ablación causal, se demostró que eliminar un canal específico suprime selectivamente la promoción de los conceptos asociados a ese canal, sin afectar a los demás. La ablatión de un canal redujo la frecuencia de sus tokens asociados en un 87-90%, mientras que la ablatión cruzada no tuvo efecto significativo.

• Completitud (Completeness):

  • Se evaluó si el conjunto de canales descubiertos podía explicar la mayoría de las activaciones de la neurona. ROTATE logró una precisión de emparejamiento del 49% al 60% para asignar descripciones de canales a ejemplos de activación, superando a los SAE (36-49%) y muy por encima del nivel aleatorio (20%).

• Generación de Descripciones:

  • Al agregar las descripciones de los canales de vocabulario, ROTATE superó a los baselines de estado del arte (MaxAct+VocabProj y MaxAct++) en comparaciones "cabeza a cabeza" (head-to-head).
  • ROTATE logró tasas de victoria de 2 a 3 veces superiores en la mayoría de las configuraciones, especialmente en activaciones moderadas (rango 100-500) y en distribuciones de datos no vistas (FineWeb), donde los métodos basados en activaciones extremas fallan.

5. Significado e Impacto

• Escalabilidad y Eficiencia: Al ser un método libre de datos, ROTATE elimina la necesidad de recolectar y procesar millones de vectores de activación, lo que lo hace extremadamente escalable y aplicable a modelos grandes sin costo computacional adicional de inferencia.
• Desenredo Fino: Proporciona una unidad de análisis más granular que las neuronas completas, permitiendo a los investigadores entender cómo una sola neurona implementa múltiples funciones distintas (por ejemplo, activarse positivamente para conceptos de negación y negativamente para marcadores temporales).
• Fundamento para Circuitos: Al proporcionar una descomposición fiable de los pesos, ROTATE ofrece un bloque de construcción fundamental para descubrir circuitos neuronales más precisos y estudiar las interacciones entre componentes de la red.
• Superioridad sobre SAE: Los resultados sugieren que la información necesaria para la interpretabilidad monosémica ya está codificada en la estructura geométrica de los pesos, y que los métodos basados en activaciones (SAE) pueden estar perdiendo señal o introduciendo ruido al depender de datos específicos.

En resumen, ROTATE representa un avance significativo en la interpretación mecánica al demostrar que es posible "leer" directamente los conceptos monosémicos dentro de los pesos de los modelos de lenguaje mediante la optimización de la curtosis en el espacio del vocabulario, sin depender de datos externos.