¿Los modelos representan conocimiento fáctico internamente? Esta publicación presenta BizzaroWorld, un estudio de interpretabilidad mecanicista que intenta localizar circuitos de recuerdo de hechos en la familia de modelos Gemma utilizando parches de activación en 60 pares de indicaciones y 20 categorías de conocimiento. El trabajo técnico aquí está muy influenciado por el trabajo realizado por Prakash et al.¹, quienes analizaron el seguimiento de entidades dentro de la serie de modelos LLaMa.
El objetivo: localizar dónde se encuentra el conocimiento real dentro de un transformador y si esa ubicación es consistente en toda la escala del modelo. El código base completo está disponible aquí.
La configuración experimental
En primer lugar, quería saber por qué las diferencias logit a través de pares de mensajes limpios y corruptos eran la forma ideal de solucionar este problema. Encontrar hechos parecía similar a encontrar entidades como Prakash et al. Lo hice, así que me concentré en la identificación indirecta de objetos (IoI). Esto queda más claro si presento tres indicaciones objetivas y sus objetivos claros.
“Cuando la pintura roja se mezcla con pintura amarilla, el resultado es” -> “ Naranja” “El infierno épico fue escrito por” -> “ Dante” “El dios romano Mercurio coincide con el dios griego” -> “ Hermes”
Me pareció que, para responder preguntas como estas, los LLM necesitarían encontrar entidades dentro de su representación. Entonces, busqué más respuestas sobre cómo funcionarían las diferencias logit y el IoI en este contexto y, para ello, el curso ARENA² fue de gran ayuda. Descubrí que las diferencias logit entre pares de mensajes limpios y corruptos proporcionan una señal escalar clara y adecuada para medir el efecto causal de las intervenciones de parcheo, así que eso es lo que medí.
Por lo tanto, diseñé una batería de hechos de 60 pares de mensajes limpios/corruptos, en 20 categorías diferentes de hechos.
Antes de comenzar cualquier experimento de parcheo, quería identificar los pares de indicaciones de señal más alta; aquellos en los que la diferencia logit entre ejecuciones limpias y corruptas fue más pronunciada e informativa para la aplicación de parches de activación. Así que creé mi propia métrica para medir esto, TotalSwing, que calcula el efecto neto de aplicar parches en estos pares de indicaciones. Aquí hay un ejemplo de lo que estoy hablando.
Mi intuición fue que calcular las diferencias logit en ambos lados y restarlas sería la señal más limpia, ya que el lado derecho (aplicar el objetivo limpio al mensaje corrupto) generalmente resulta en una diferencia logit negativa como se muestra arriba. Eso significa que ambos valores terminan sumados, lo que da como resultado un escalar positivo y agradable a través del cual puedo ordenar los 60 pares de indicaciones que había creado.
Funcionó bien y con esto se me ocurrió un archivo CSV con todas las indicaciones, ordenadas por TotalSwing. Los llamé pares de indicaciones doradas y, usándolos, creé tres modos experimentales para cada experimento posterior que realizaría.
Y ahora estaba listo para comenzar la experimentación.
Aislar los componentes de Gemma-2B
Los LLM son estructuras enormes. Para saber dónde está sucediendo algo, necesitamos un aislamiento despiadado. TransformerLens³ de Neel Nanda fue fundamental para lograr precisamente eso. Quería conectar todas las diferentes piezas del rompecabezas, es decir, los flujos residuales antes y después de todos los componentes relevantes (las cabezas de atención y las subcapas de MLP) en todas las capas. Y eso es precisamente lo que hice. Realicé cuatro experimentos, en los que progresivamente reduje el modelo a sus piezas.
Experimento 1 = aplicar parches en la posición final del token Experimento 2 = aplicar parches antes y después de cada una de las subcapas Experimento 3 = aplicar parches en la posición del token de entidad Experimento 4 = aplicar parches antes de cada cabeza de atención
Los números que generaron estos experimentos resaltaron un hallazgo claro: existe un circuito de recuperación factual de tres fases dentro de la familia de modelos Gemma.
Fase 1: Almacenamiento (capas 0 a 14, posición del token de entidad): los hechos se codifican como direcciones en el flujo residual en el token de entidad. El flujo residual domina causalmente aquí, contribuyendo 40 veces más que los resultados de atención y 18 veces más que los resultados de MLP. El 86,7% de los 15 pares de indicaciones principales liberaron su señal almacenada en las capas 13 a 15, con una peor capa media de 16,3 en todos los modos experimentales (Pearson r = -0,83 entre la confianza del modelo y la puntuación de daño).
Fase 2: Enrutamiento (cabezas de atención distribuidas): la señal se mueve desde la posición del token de entidad hasta la posición de predicción final a través de cabezas de atención colectivamente. Ningún jefe fue el único responsable, aunque el jefe 2 fue desproporcionadamente activo; por ejemplo, estuvo activo en el 40% de los pares rápidos en el modo experimental A. Sin embargo, el daño individual a la cabeza (ΔLD = -0,68) fue insignificante en comparación con el daño total del flujo residual (ΔLD = -11,47).
Fase 3: Lectura (capas 15 a 17, posición final de la ficha): la respuesta se recupera, no se calcula. Los bloques tardíos son de paso, es decir, la señal ya está codificada y simplemente se lee. Este hallazgo fue unánime en los tres modos experimentales y 20 categorías de conocimiento.
El circuito trifásico a escala: Gemma-12B-IT
Buscaba un resultado generalizable, así que el siguiente paso fue ver si este patrón se aplicaba también al modelo más grande, Gemma-12B-IT. Aunque quería probarlo con modelos Gemma aún más grandes, como Gemma-31B o Gemma-27B, estaba en deuda con la limitación de espacio en disco HPC de mi universidad, de la que hablaré más adelante. Aún así, pude replicar todo el conjunto para el modelo 12B, incluidos todos los modos de experimento A, B y C.
Encontré algunos resultados interesantes al hacerlo, pero antes que nada, revisemos cómo y dónde se diferencian estas dos arquitecturas.
Además de estas diferencias arquitectónicas, todo lo demás, como los tokenizadores utilizados⁴, son iguales entre los dos modelos. Sin embargo, vi algo que era bastante diferente para Gemma-12B-IT con respecto al comportamiento del tokenizador que influyó en la rápida selección de pares.
Cuando hice el pase de clasificación inicial, como describí anteriormente para clasificar mis 60 pares de mensajes, este modelo Gemma más grande eliminó tres pares de mensajes dorados, a pesar de que ambos usan el mismo tokenizador.
Esta exclusión ocurre porque, durante el paso hacia adelante, el modelo asigna tokens individuales a matrices de ID de tokens. Para que estos pases funcionen, las formas de la matriz deben coincidir con todo lo demás; de lo contrario, la multiplicación de matrices no funciona. Observé esto cuando hice este proceso con Gemma-2B, donde vi un comportamiento extraño, como la unidad física “hercios” asignada a dos tokens⁵. Muy poco intuitivo. Esperaba que mis 60 mensajes pasaran por Gemma-12B-IT sin problemas, pero me equivoqué aquí. Este efecto fue obviamente más pronunciado cuando hice experimentos iniciales con LLaMa-70B⁶, que detallaré en la sección de trabajo futuro a continuación, pero me tomó por sorpresa.
La consecuencia de esto es que las comparaciones mecanicistas entre modelos están parcialmente limitadas por la deriva del conjunto de datos inducida por el tokenizador, y las diferencias reportadas deben interpretarse teniendo esa advertencia en mente.
Por lo tanto, antes de diseñar cualquier batería de hechos para tales experimentos, debemos analizar los hechos en todos los modelos que se están probando, de modo que cualquier desviación del tokenizador se informe de inmediato y los pares de avisos afectados puedan reemplazarse antes de comenzar a realizar cualquier experimentación.
Después de notar estas anomalías, realicé los experimentos del 1 al 4, para todos los modos experimentales A, B y C.
Descubrí que este circuito trifásico se replicaba a escala: el almacenamiento se desplazaba a las capas 0 a 27, el enrutamiento permanecía distribuido sin un cabezal significativamente dominante y la lectura se concentraba en las capas finales, estructuralmente idénticas a Gemma-2B, escaladas proporcionalmente.
He aquí algunas cifras que lo demuestran.
Cada categoría de hecho muestra su propio patrón de comportamiento, lo cual era consistente con mi hipótesis.
La otra parte interesante que vi fue que el efecto de los cabezales de atención de Gemma-12B-IT parecía incluso más distribuido y más apagado que lo que había visto en Gemma-2B. Este hallazgo se destaca mediante mapas de calor medios ld_delta para los dos modelos anteriores, para los cabezales de atención.
Para Gemma-12B-IT, el mapa de calor está casi completamente vacío, excepto pequeñas celdas de colores alrededor de las capas 20 y 28⁷. Una vez más, la magnitud de los valores de diferencia logit fue mucho mayor para la corriente residual que para estos componentes de la subcapa.
Problemas de cuota de disco, experimentos futuros y conclusión
A lo largo de este experimento, encontré muchas técnicas interesantes que estaba ansioso por probar, como la aplicación de parches de ruta (con DCM) y CMAP. También quería comprobar si la cuantificación o la ejecución de todo el proceso a través de variantes ajustadas cambiaría este circuito de recuperación de hechos de tres fases. Sin embargo, me encontré con un cuello de botella debido a una restricción de cuota de disco de 30 GB. De hecho, como se indicó anteriormente, ya había preparado una versión cuantificada de 8 bits de LLaMa-70B y me había ocupado de la exclusión de 22 pares de avisos dorados, debido a diferencias en el tokenizador. La inferencia estaba funcionando bien, pero no pude expandir el proceso más allá debido a las limitaciones del disco, así que dejé la idea para más adelante.
En mi opinión, ampliar esta línea de trabajo con parches de ruta es el siguiente paso natural, ya que mostraría mucho más. El parcheo de ruta, formalizado por Goldowsky-Dill et al⁸., agudiza el parcheo de activación desde el nivel de nodo hasta la precisión del nivel de borde. El parche de activación estándar mide el efecto causal total de un nodo reemplazando su salida y observando todas las consecuencias posteriores; En cambio, el parche de ruta aísla los bordes individuales en el gráfico computacional, revelando con precisión qué componentes se comunican con cuál.
Estos hallazgos sientan las bases para una intervención dirigida: saber dónde reside el recuerdo de los hechos es un requisito previo para saber dónde intervenir cuando falla. Para llevar estas ideas más allá, me gustaría continuar con mis planes iniciales y también ver cómo los jefes de atención colaboran mediante SAE. Sí, el flujo residual está haciendo el trabajo pesado, pero ¿qué significa eso? Necesito más detalles.
En resumen, el siguiente paso lógico para este trabajo es la replicación entre arquitecturas en LLaMA y otras variantes⁹. Además, el hallazgo de enrutamiento distribuido en el Experimento 4 justifica experimentos de parcheo de ruta para establecer relaciones causales dirigidas entre los componentes.
Referencias
Nikhil Prakash, Tamar Rott Shaham, Tal Haklay, Yonatan Belinkov y David Bau. El ajuste mejora los mecanismos existentes: un estudio de caso sobre seguimiento de entidades, 2024 https://github.com/callummcdougall/ARENA_3.0 https://github.com/TransformerLensOrg/TransformerLens Gemma usa el tokenizador SentencePieza Mientras que algo así como la unidad física “Watt” era un token, como era de esperar, lógicamente LLaMa usa el tokenizador Tiktoken o SentencePieza, dependiendo de la versión del modelo. También se resaltan otras células, pero son mucho más apagadas en comparación con Gemma-2B Nicholas Goldowsky-Dill, Chris MacLeod, Lucas Sato y Aryaman Arora. Localización del comportamiento del modelo con parcheo de ruta, 2023. Sería especialmente interesante incluir aquí un modelo de lenguaje de difusión como LLaDA-8B, ya que el mecanismo de atención difiere fundamentalmente de los transformadores autorregresivos estándar, lo que requiere una infraestructura de enlace personalizada más allá de lo que TransformerLens admite actualmente.