Constructing Interpretable Features from Compositional Neuron Groups

Resumen Técnico: Construcción de Características Interpretables a partir de Grupos de Neuronas Composicionales

Enunciado del Problema
Un desafío central en la interpretabilidad mecánica es identificar las unidades de análisis apropiadas dentro de los Modelos de Lenguaje Grandes (LLM) que expliquen causalmente las salidas del modelo. Aunque la investigación temprana se centró en neuronas individuales, la evidencia sugiere que las neuronas a menudo codifican múltiples conceptos (polisemicidad), lo que motiva un cambio hacia el análisis de direcciones en el espacio de activación. Los enfoques no supervisados actuales dependen de Autoencoders Dispersos (SAE) entrenados en activaciones de la corriente residual para aprender direcciones de características desde cero. Sin embargo, los SAE enfrentan limitaciones significativas: a menudo luchan en evaluaciones causales (dirección), carecen de interpretabilidad intrínseca porque sus características aprendidas no están explícitamente vinculadas a los cálculos internos del modelo, y sus características no están restringidas al espacio de representación del modelo. Existe la necesidad de un método que descubra características interpretables de manera no supervisada, manteniendo al mismo tiempo un vínculo causal directo con los mecanismos del modelo.

Metodología
Los autores proponen descomponer las activaciones de la Red Neuronal de Perceptrón Multicapa (MLP) directamente utilizando Factorización de Matriz Semi-No Negativa (SNMF). A diferencia de los SAE, que aprenden nuevas direcciones, este método representa las características como combinaciones lineales dispersas de neuronas co-activadas existentes dentro de la MLP.

1. Formulación: Dada una matriz $A \in \mathbb{R}^{d_a \times n}$ de activaciones de la MLP (donde $d_a$ es la dimensión interna y $n$ es el número de tokens), el método factoriza $A$ en dos matrices:

   • $Z \in \mathbb{R}^{d_a \times k}$: Una matriz de $k$ características de la MLP, donde cada columna define una combinación lineal dispersa de neuronas.
   • $Y \in \mathbb{R}^{k \times n}_{\geq 0}$: Una matriz de coeficientes no negativa que mapea las entradas a las características.
     La descomposición sigue $A \approx ZY$.

2. Mecanismo de Interpretabilidad:

   • Dispersión: Las columnas de $Z$ están restringidas para ser dispersas utilizando un operador rígido de "ganador se lleva todo", asegurando que cada característica active solo un pequeño subconjunto de neuronas.
   • Fundamentación Causal: La matriz no negativa $Y$ codifica explícitamente qué tokens contribuyeron a la creación de una característica, permitiendo la asignación directa de etiquetas de conceptos textuales basadas en contextos de alta activación.
   • Mapeo a la Corriente Residual: Cada vector de característica $z_i$ en el espacio de activación se mapea a la corriente residual a través de los pesos de salida de la MLP ($W_V$), definiendo una característica de corriente residual $f_i = W_V z_i$. Esto asegura que las características estén fundamentadas en las vías computacionales reales del modelo.

3. Optimización: El método utiliza actualizaciones multiplicativas (Ding et al., 2010) para minimizar la norma de Frobenius entre $A$ y $ZY$, con la dispersión aplicada en $Z$.

Contribuciones Clave

• SNMF para Descomposición de MLP: El artículo introduce un método no supervisado para descomponer las activaciones de la MLP en características interpretables definidas por grupos de neuronas co-activadas, en lugar de aprender direcciones arbitrarias desde cero.
• Superioridad en la Dirección Causal: Los autores demuestran que las características derivadas de SNMF superan a los SAEs de última generación (GemmaScope, Llamascope) y a una línea base supervisada sólida (Diferencia de Medias) en tareas de dirección causal, mientras que desempeñan un rendimiento comparable en la detección de conceptos.
• Revelación de Composición Jerárquica: Mediante SNMF recursiva, los autores muestran que las características de la MLP son composicionales. Los conceptos finos (por ejemplo, días de la semana específicos) se forman combinando neuronas "núcleo" compartidas con neuronas exclusivas. Esto revela una estructura jerárquica donde conceptos más amplios (por ejemplo, "día de la semana") están compuestos por subconceptos más específicos, invirtiendo efectivamente el fenómeno de "división de características" a menudo observado en los SAE.
• Reutilización de Neuronas: El análisis confirma que las características semánticamente relacionadas comparten un conjunto central de neuronas, con neuronas exclusivas que controlan la promoción o supresión de subconceptos específicos.

Resultados
Se realizaron experimentos en Llama 3.1-8B, Gemma 2-2B y GPT-2 Small.

• Detección de Conceptos: Las características de SNMF lograron puntuaciones de detección positivas para la mayoría de las características (>75%), lo que indica que se activan para entradas relacionadas con conceptos y permanecen inactivas para entradas neutrales. El rendimiento fue comparable o superó al de los SAE entrenados en salidas de MLP.
• Dirección de Conceptos: En las evaluaciones de dirección causal (midiendo la capacidad de dirigir la generación del modelo hacia un concepto mientras se preserva la fluidez), SNMF superó consistentemente a los SAE en todas las capas. También coincidió o superó frecuentemente el rendimiento de la línea base supervisada de Diferencia de Medias.
• Análisis Jerárquico: La aplicación recursiva de SNMF fusionó con éxito características finas (por ejemplo, "lunes", "martes") en conceptos de nivel superior (por ejemplo, "día de la semana", "fin de semana"). Las intervenciones causales confirmaron que amplificar las neuronas "núcleo" promovió todos los tokens relacionados, mientras que amplificar las neuronas "exclusivas" promovió tokens específicos y suprimió otros.

Significado y Afirmaciones
El artículo afirma que SNMF proporciona una herramienta simple y efectiva para identificar características interpretables en LLM. Al descomponer directamente las activaciones de la MLP, el método vincula las características a los componentes específicos del modelo responsables de su formación. Los resultados sugieren que la MLP opera mediante actualizaciones aditivas compuestas por conjuntos de neuronas interpretables. Este trabajo posiciona a SNMF como una alternativa viable a los SAE, ofreciendo un mejor rendimiento causal e interpretabilidad intrínseca sin requerir supervisión. Además, el descubrimiento de la composición jerárquica de neuronas proporciona una nueva lente para comprender cómo los LLM construyen conceptos complejos a partir de bloques de construcción neuronales reutilizables, extendiendo la función funcional de las neuronas más allá de la codificación de un solo concepto.

Los autores reconocen limitaciones, incluido el enfoque actual en conteos de características $k < 500$ y la dependencia de LLM como jueces para la evaluación, aunque reportan un fuerte acuerdo con los juicios humanos. Sugieren que el trabajo futuro podría explorar la escalabilidad a valores de $k$ más grandes y utilizar el descenso de gradiente proyectado para la optimización.