Uso de autoencoders (SAE) para interpretabilidad mecanicista

En un post pasado vimos una introducción a la interpretabilidad mecanicista. Vimos que una sola neurona de un modelo de lenguaje (LLM) se activa ante conceptos completamente diferentes y que los SAEs (Sparse Autoencoders) descomponen estas activaciones complejas en características o conceptos (features) dispersos y monosemánticos (que significan una sola cosa). También referenciamos tutoriales de introducción, recursos y librerías como TransformerLens.

El proceso para que el SAE descomponga estas activaciones complejas en conceptos (features) dispersos y monosemánticos es el siguiente:

EXTRACCIÓN DE ACTIVACIONES

En esta fase se elige un corpus de texto amplio y diverso para asegurar que el modelo exhibe una gran variedad de conceptos. Utilizando librerías de interpretabilidad (como TransformerLens o NNsight), se intercepta el flujo de ejecución del LLM, colocando un hook en la capa de interés.

Se pasan los tokens de cada secuencia del dataset por el LLM generando un tensor de activaciones de tamaño (BxT, D) donde B es el número de secuencias, T el número de tokens y D la dimensión de la capa seleccionada.

ENTRENAMIENTO DEL SAE

El SAE tiene la siguiente arquitectura:

• Capa de Entrada: Recibe un vector de activación del LLM, .
• Capa Oculta (Diccionario): Codifica la entrada. Tiene una dimensión , donde F es el factor de expansión, típicamente entre 8 y 64. Contiene los conceptos dispersos.
• Capa de Salida: Reconstruye o decodifica la entrada, .

El SAE se entrena minimizando una pérdida combinada:

• El primer término (Error Cuadrático Medio) fuerza a que la reconstrucción sea fiel.
• El segundo término penaliza la suma de las activaciones de la capa oculta, forzando a que la gran mayoría de los conceptos tengan un valor de cero para cualquier token dado.

ETIQUETADO DE CONCEPTOS

Una vez entrenado, hay que asociar una etiqueta o explicación semántica a cada neurona (concepto) del SAE. Se pasa un nuevo dataset por el LLM y el SAE. Para cada concepto (feature) i del SAE, se registran los tokens y fragmentos que causaron las activaciones más altas (máximos activadores) y aquellos que no la activaron.

Posteriormente, los fragmentos registrados para las activaciones de cada concepto se pasan a un LLM evaluador, que devuelve una explicación conceptual para cada concepto i del SAE, del estilo “Este concepto se activa en verbos en pasado relacionados con la informática”).

INTERVENCIÓN DEL SAE EN EL LLM

Se utiliza el SAE para sustituir dinámicamente a la capa original del LLM para realizar experimentos.

• Clamping o sobreescritura (añadir dirección): En el momento del cálculo en una capa específica, el vector decodificado de un concepto objetivo del SAE se suma a la activación natural del LLM para añadir artificialmente ese concepto.
• Ablación o supresión (restar dirección): Se puede forzar que el vector de activaciones de un concepto de comportamiento no deseado (como “vulnerabilidades de código”) sea igual a cero. El modelo pierde la capacidad de evocar ese concepto, alterando su respuesta.

Hemos visto el proceso para el uso de SAEs en interpretabilidad mecanicista. Parece un enfoque útil pero está sujeto a las objeciones que comentamos en la introducción a la interpretabilidad mecanicista. En los últimos meses ha habido noticias de que DeepMind ha quitado prioridad al uso de SAEs con fines de interpretabilidad.