Polar probe linearly decodes semantic structures from LLMs

Resumen Técnico: La Sonda Polar Decodifica Linealmente Estructuras Semánticas de los LLM

Formulación del Problema
El artículo aborda la pregunta fundamental de cómo los Modelos de Lenguaje Grandes (LLM) vinculan conceptos para formar estructuras semánticas complejas. Si bien trabajos anteriores han demostrado que las estructuras sintácticas (por ejemplo, dependencias sujeto-verbo) pueden decodificarse linealmente a partir de las activaciones de los LLM utilizando principios geométricos (Hewitt & Manning, 2019; Diego-Simon et al., 2024), sigue sin estar claro si principios similares se aplican al vínculo semántico. Los autores hipotetizan que los LLM representan la existencia y el tipo de relaciones entre entidades no mediante mecanismos complejos y no lineales, sino a través de un código geométrico simple: la distancia euclidiana entre los incrustamientos de las entidades codifica la existencia de una relación, mientras que la dirección relativa codifica el tipo de relación.

Metodología
Para probar esta hipótesis, los autores introducen una Sonda Polar, una transformación lineal entrenada para mapear las activaciones de las capas de un LLM a un espacio de sonda donde se puede recuperar el grafo relacional.

1. Conjuntos de Datos: El estudio utiliza un conjunto de datos sintético que abarca cinco dominios semánticos distintos:

   • Dominios euclidianos: Ordinalidad variable (variables matemáticas), disposiciones espaciales (arreglos de objetos 2D) y roles temáticos (interacciones sociales).
   • Dominios no euclidianos: Árboles genealógicos (relaciones no conmutativas, de muchos a muchos) y mapas de metro (trayectorias más cortas basadas en redes).
   • Cada muestra consiste en una descripción en lenguaje natural de un grafo relacional. Para evitar que el modelo aprenda correlaciones a nivel superficial, el orden de las relaciones y las entidades se aleatoriza, y se generan múltiples descripciones textuales para el mismo grafo.

2. Arquitectura de la Sonda:

   • Pérdida Estructural ($L_s$): Optimiza la correlación de rango de Spearman entre la matriz de distancias pares predicha y la matriz de distancias de trayectorias más cortas de verdad fundamental del grafo semántico.
   • Pérdida Angular ($L_a$): Optimiza la alineación entre el vector de diferencia de dos incrustamientos de entidades ($\delta_{ij}$) y un vector prototipo aprendido ($p_r$) específico para el tipo de relación $r$.
   • La sonda se entrena de forma independiente en cada capa de varios LLM (Llama3-8B, OLMo-7B y la familia Pythia).

3. Evaluación:

   • Rendimiento de Decodificación: Medido mediante $\rho$ de Spearman para la existencia de la relación (ordenamiento de distancias) y el tipo de relación (alineación de direcciones).
   • Generalización: Probado en nombres de entidades y formas superficiales de relaciones fuera de la distribución (OOD).
   • Intervención Causal: Los autores intervienen en los estados ocultos dirigiendo las activaciones a lo largo de las direcciones de prototipos aprendidas para observar cambios en el rendimiento de tareas de Preguntas y Respuestas (QA) aguas abajo.
   • Correlación: Los errores de la sonda se correlacionan con la probabilidad de logit de respuestas correctas en tareas de QA.

Resultados Clave

• Decodificabilidad Lineal: Las estructuras semánticas reales pueden recuperarse linealmente a partir de las activaciones de los LLM. La Sonda Polar decodifica con éxito tanto la existencia como el tipo de relación en los cinco dominios.
• Emergencia por Capas: El código geométrico es más prominente en las capas intermedias de los LLM preentrenados (alcanzando un pico alrededor de las capas 12–15 para Llama3-8B). En modelos inicializados aleatoriamente, el rendimiento de decodificación se mantiene en niveles de azar, lo que indica que el código emerge durante el preentrenamiento.
• Tamaño del Modelo y Entrenamiento: El rendimiento escala con el tamaño del modelo (de 14M a 6.9B de parámetros) y aumenta con los pasos de preentrenamiento, lo que sugiere que el código polar es una capacidad aprendida dependiente de la capacidad del modelo y la duración del entrenamiento.
• Generalización: Las sondas generalizan a nuevas entidades y formas superficiales de relaciones, aunque el rendimiento disminuye moderadamente con entidades OOD y más significativamente con tipos de relación OOD. El rendimiento también disminuye a medida que aumenta el número de entidades en el grafo, probablemente debido a la expansión del espacio combinatorio.
• Influencia Causal: Dirigir las activaciones del modelo a lo largo de las direcciones de prototipos aprendidas desplaza causalmente la probabilidad de respuestas correctas en tareas de QA, confirmando que estas representaciones geométricas son funcionalmente relevantes para el razonamiento del modelo.
• Disyunción de Subespacios: Los subespacios aprendidos para diferentes dominios semánticos son en gran medida disjuntos en el espacio de activaciones, con la excepción notable de las disposiciones espaciales y la ordinalidad variable, que muestran una superposición significativa (apoyando la hipótesis de superposición).

Significado y Afirmaciones
El artículo afirma demostrar que los LLM aprenden a construir estructuras semánticas complejas vinculando representaciones a través de un principio geométrico simple. Específicamente, extiende el alcance de la sonda lineal más allá de la sintaxis hacia una clase más amplia de grafos dirigidos y etiquetados.

• Más allá de la Sintaxis: A diferencia de trabajos anteriores limitados a árboles sintácticos, este enfoque decodifica con éxito estructuras semánticas que no son necesariamente de estructura arbórea e incluyen topologías no euclidianas (por ejemplo, árboles genealógicos, mapas de metro), aunque con algunas limitaciones en la precisión para relaciones no conmutativas.
• Más allá de la Recuperación de Conocimiento: El estudio apunta a estructuras composicionales que no podrían haberse memorizado durante el entrenamiento (debido a entidades y relaciones aleatorizadas), proporcionando evidencia de cómo los modelos realizan el aprendizaje en contexto y el vínculo.
• Tensión Simbólico-Conexionista: Los hallazgos sugieren una resolución a la tensión entre los enfoques simbólicos y conexionistas, proponiendo que las estructuras simbólicas pueden representarse dentro de sistemas vectoriales mediante subespacios lineales y relaciones geométricas.

Los autores mantienen una postura modesta respecto a las limitaciones del método, señalando que la Sonda Polar depende de distancias lineales y está adaptada para grafos euclidianos. En consecuencia, no puede representar perfectamente relaciones enteras no conmutativas, de muchos a muchos, ni distancias complejas de trayectorias más cortas encontradas en dominios no euclidianos, aunque captura la coherencia local dentro de ellos.