Visualizing LLM Latent Space Geometry Through Dimensionality Reduction

Resumen Técnico: Visualización de la Geometría del Espacio Latente de los LLM mediante la Reducción de la Dimensionalidad

Planteamiento del Problema
Aunque los Modelos de Lenguaje de Gran Escala (LLM) alcanzan un rendimiento de vanguardia, sus procesos internos de toma de decisiones siguen siendo opacos. La interpretabilidad mecánica busca la ingeniería inversa de estos componentes internos para lograr una comprensión causal y granular. Una vía prometedora para esto es la geometría de características: analizar la estructura y organización de las representaciones dentro de espacios latentes de alta dimensionalidad. Sin embargo, estos espacios son difíciles de interpretar directamente debido a su extrema dimensionalidad. El desafío radica en extraer, procesar y visualizar estos estados latentes para descubrir cómo se codifican, transforman y organizan las características abstractas a través de las capas del modelo sin perder información estructural relevante.

Metodología
Los autores proponen un flujo de trabajo sistemático para capturar y visualizar las geometrías de los estados latentes en modelos Transformer de solo decodificador (específicamente GPT-2 y LLaMa). La metodología consta de tres etapas principales:

1. Pipeline de Extracción de Latentes:

   • El pipeline procesa datos de texto (del conjunto de datos PG-19) o tokens individuales a través de la arquitectura Transformer.
   • Captura activaciones en seis puntos distintos dentro de cada bloque Transformer: las salidas de las capas de normalización pre-atención y pre-MLP, las salidas puras de los módulos de atención y MLP, y los estados del flujo residual después de la adición de estos componentes (pre-adición y post-adición).
   • Los autores adoptan una visión "centrada en el flujo residual", tratando al modelo como un canal de comunicación compartido donde los componentes leen de y escriben en un espacio común.

2. Procesamiento y Preprocesamiento de Datos:

   • Para gestionar la alta dimensionalidad y las restricciones de memoria, los estados latentes pueden reducirse opcionalmente (por ejemplo, de 4096 a 512 dimensiones) mediante PCA.
   • Los pasos específicos de preprocesamiento incluyen la conversión de vectores a longitud unitaria para enfatizar la estructura direccional sobre la magnitud, y el promedio de vectores a través de dimensiones de muestra o secuencia para aislar fuentes específicas de variabilidad (por ejemplo, la evolución capa por capa frente al ruido a nivel de token).
   • Los autores excluyen explícitamente la posición de la secuencia 0 (token inicial) en muchas visualizaciones debido a su norma anomalamente alta, que puede oscurecer otros patrones geométricos.

3. Reducción de la Dimensionalidad y Visualización:

   • Análisis de Componentes Principales (PCA): Utilizado como una línea base lineal para capturar direcciones de máxima varianza. Se aplica sin centrado de media para preservar la posición del origen.
   • Aproximación y Proyección de Manifolds Uniformes (UMAP): Utilizado como una alternativa no lineal para preservar vecindades locales y la estructura del manifold global. Los autores utilizan la distancia de coseno para UMAP con el fin de alinearse con la Hipótesis de Representación Lineal, que postula las características como direcciones en el espacio latente.
   • Las visualizaciones están codificadas por colores según la capa (azul para atención, rojo para MLP) y la posición de la secuencia (gradiente de oscuro a claro).

Contribuciones Clave y Resultados

• Separación Geométrica Atención vs. MLP:
  El hallazgo más significativo es una separación geométrica clara y persistente entre las salidas de los componentes de atención y los componentes de Perceptrón Multicapa (MLP) en las capas intermedias. En las visualizaciones de dimensionalidad reducida (tanto PCA como UMAP), las salidas de atención ocupan consistentemente regiones distintas de las salidas de MLP. Los autores señalan que este patrón no ha sido documentado en trabajos previos.

• Norma Alta de la Posición Inicial de la Secuencia:
  El estudio confirma y expande las observaciones de que la norma del estado latente en la posición 0 de la secuencia es significativamente mayor que en las posiciones subsiguientes.

  • En GPT-2 (usando embeddings posicionales aprendidos), esto es esperado.
  • En LLaMa (usando Codificaciones Posicionales Rotatorias/RoPE), los autores encuentran que la norma alta no se limita al token <BOS>, sino que se aplica a la gran mayoría de los tokens del vocabulario cuando aparecen como el token inicial. Esto sugiere que LLaMa ha desarrollado un mecanismo complejo para detectar el inicio de la secuencia, distinto de los embeddings aprendidos triviales de GPT-2.

• Evolución por Capas y Dinámica de la Norma:
  Las visualizaciones de las normas de los estados latentes a través de las capas revelan que las normas generalmente aumentan a través de las capas medias y caen después de la última capa de normalización. El token inicial exhibe un "pico" único de norma que emerge después de algunas capas y se disipa en las capas finales.

• Geometrías de los Embeddings Posicionales:

  • GPT-2: Las visualizaciones confirman la estructura "helicoidal" de los embeddings posicionales aprendidos de GPT-2 en el espacio de alta dimensión, visible a través de varios pares de dimensiones de PCA.
  • LLaMa: Los patrones geométricos para RoPE son menos distintivos. En las primeras dimensiones de PCA, una "cola" de posiciones iniciales de la secuencia conduce hacia una nube densa de posiciones posteriores. En dimensiones posteriores, este patrón se disuelve en una nube uniforme, lo que sugiere que la geometría posicional en los modelos RoPE puede ser de baja dimensionalidad.

Significancia y Reivindicaciones
El artículo afirma proporcionar un marco robusto para visualizar la geometría interna de los LLM, haciendo accesible el concepto abstracto de "geometría de características". Al demostrar la separación de las salidas de atención y MLP, el trabajo ofrece una nueva perspectiva sobre cómo los diferentes componentes Transformer contribuyen al flujo residual compartido. Los autores posicionan este trabajo como un paso hacia la interpretabilidad basada en principios, sugiriendo que comprender estas estructuras geométricas es esencial para analizar cómo los modelos generalizan, razonan y abstraen. Los autores declaran explícitamente que sus hallazgos respecto a la separación atención/MLP son novedosos según su conocimiento y que los patrones observados en las normas del token inicial de LLaMa ameritan una mayor investigación sobre los mecanismos de las codificaciones posicionales relativas.

El trabajo no pretende resolver el problema completo de la interpretabilidad, sino que proporciona un "trasfondo y exposición sólidos" y un conjunto de herramientas (el pipeline de visualización) para facilitar la investigación posterior sobre la dinámica de los estados latentes a través de capas y modelos.