Context Structure Reshapes the Representational Geometry of Language Models

Resumen Técnico: La estructura del contexto reconfigura la geometría representacional de los modelos de lenguaje

Planteamiento del problema
Los Grandes Modelos de Lenguaje (LLMs) demuestran la capacidad de adaptarse a nuevas tareas y aprender estructuras dentro de la ventana de contexto (Aprendizaje en Contexto, o ICL, por sus siglas en inglés). Si bien investigaciones previas han establecido que los LLMs organizan las secuencias de entrada en trayectorias neuronales más "rectas" en las capas profundas para facilitar la predicción del siguiente token mediante la extrapolación lineal, no está claro si este enderezamiento representacional es un mecanismo universal para todas las formas de ICL. Específicamente, la relación entre la estructura del contexto y cómo los LLMs reconfiguran sus representaciones internas para adaptar su comportamiento no se comprende completamente. Este artículo investiga si el enderezamiento representacional ocurre de manera consistente a través de diversas tareas de contexto o si depende de estructuras de tareas específicas.

Metodología
Los autores analizaron el modelo Gemma-2-27B (versión base/preentrenada) a través de tres clases distintas de tareas diseñadas para sondear diferentes demandas cognitivas y estructuras de datos:

1. Dependencias de largo alcance en lenguaje natural: Utilizando el conjunto de datos LAMBADA, el estudio examinó cómo los modelos integran contextos más amplios para predecir la última palabra de un pasaje, contrastándolo con un control de "texto aleatorio" donde se destruyó la estructura semántica.
2. Tareas de Mundo de Rejilla (Grid World): Tareas sintéticas que requieren que el modelo infiera estructuras de grafos latentes a partir de secuencias lineales de tokens. Estas incluyeron:
   • Mapeo Directo: Caminatas aleatorias en una red de 6x6 de palabras únicas.
   • Mapeo Jerárquico: Inferir un grafo latente oculto de 4x4 a partir de observaciones ruidosas de tokens "hijos" asociados a nodos latentes.
   • Transiciones de 0 pasos (0-shot): Probar la capacidad del modelo para inferir transiciones entre nodos latentes que nunca fueron observados explícitamente en el contexto.
3. Tareas de Predicción Estructurada:
   • Aprendizaje de pocos ejemplos (Few-Shot Learning): Benchmarks estándar (por ejemplo, País-Capital, razonamiento algorítmico) formateados como pares alternos de "Q: [Entrada]" y "A: [Salida]".
   • Acertijos (Riddles): Una tarea de desambiguación semántica del dataset BIG-bench, contrastando el rendimiento de 0 pasos frente al de 8 pasos.

Métricas Geométricas y de Comportamiento
Para cuantificar la geometría representacional, los autores emplearon cuatro métricas complementarias a través de las capas del transformador:

• Enderezamiento (Straightening): Medido como la reducción en la curvatura local (ángulo entre vectores de transición consecutivos) desde la primera capa hacia las capas subsiguientes.
• Enderezamiento de Menger: Una medida de curvatura complementaria basada en el radio del círculo circunscrito de tres puntos consecutivos.
• Dimensionalidad Efectiva: Estimada mediante la Razón de Participación de los autovalores de PCA para medir el volumen del espacio de representación.
• Elongación: Medida a través de la anisotropía de la trayectoria (varianza a lo largo del primer frente al segundo componente principal).

El rendimiento conductual se evaluó mediante diferencias de logit (para mundos de rejilla) y precisión de coincidencia exacta (para tareas de pocos ejemplos).

Resultos Clave
El estudio revela una dicotomía fundamental en cómo cambian las representaciones de los LLM durante el ICL, dependiendo de la estructura de la tarea:

1. Tareas de Predicción Continua (Lenguaje Natural y Mundos de Rejilla):

   • En tareas que requieren la integración de dependencias de largo alcance o la inferencia de estructuras de grafos latentes, el aumento de la longitud del contexto conduce a un incremento consistente en el enderezamiento representacional, particularmente en las capas medias (Capas 15–25).
   • Este enderezamiento se correlaciona fuertemente con la mejora conductual (por ejemplo, una mayor separación de logits entre vecinos válidos y no vecinos).
   • En los mundos de rejilla jerárquicos, el enderezamiento persiste incluso en las condiciones de transición de 0 pasos, lo que sugiere que el modelo aprende el grafo latente subyacente en lugar de simplemente memorizar pares de tokens.
   • Geométricamente, estas tareas muestran una contracción de la representación sobre un manifold de menor dimensión y más anisotrópico (elongado).

2. Tareas de Predicción Estructurada (Few-Shot y Acertijos):

   • En tareas que requieren el mapeo de ejemplos discretos a salidas (por ejemplo, País-Capital, Acertijos), el enderezamiento representacional es inconsistente y se disocia del núcleo computacional de la tarea.
   • Disociación de Fase: Al descomponer los ejemplos de pocos pasos en fases de Pregunta, Transición (formato) y Respuesta:
     • La fase de Transición (por ejemplo, "\n A:") exhibe un enderezamiento progresivo a medida que se añaden más ejemplos, comportándose como una continuación predecible.
     • La fase de Pregunta a menudo muestra una caída en el enderezamiento después del primer ejemplo.
     • Crucialmente, la fase de Respuesta (donde se genera la solución real) no muestra un enderezamiento progresivo y, en algunos casos (como en los acertijos), presenta una disminución significativa del enderezamiento con el aumento del contexto.
   • Esto indica que el mecanismo que resuelve la lógica central de la tarea (mapeo de entrada a salida) no depende de la linealización de las trayectorias neuronales observada en la predicción continua.

Significado y Afirmaciones
El artículo argumenta que el Aprendizaje en Contexto no es un proceso monolítico impulsado por un único mecanismo geométrico. En su lugar, los autores proponen que los LLM funcionan como una "navaja suiza", seleccionando dinámicamente estrategias computacionales basadas en la estructura específica de la tarea.

• Para las tareas de predicción continua, el modelo utiliza una estrategia de enderezamiento representacional para linealizar el manifold y facilitar la extrapolación.
• Para las tareas de mapeo estructurado, el modelo emplea mecanismos distintos que no necesariamente se manifiestan como enderezamiento de trayectoria.

Los autores concluyen que, si bien el enderezamiento representacional es una firma robusta del aprendizaje en entornos de predicción continua, no es un indicador universal de la capacidad de ICL. Este hallazgo ofrece una perspectiva cautelosa para la interpretabilidad mecánica: confiar en una única firma geométrica (como el enderezamiento) puede fallar al capturar el espectro completo de cómo los modelos se adaptan a diversos contextos. El estudio enfatiza la necesidad de probar las hipótesis de interpretabilidad a través de un conjunto diverso de estructuras de tareas, en lugar de asumir que una estructura representacional generalizada subyace a todos los comportamientos de ICL.