Geometric Scaling of Bayesian Inference in LLMs

El estudio demuestra que los modelos de lenguaje modernos preservan una sustrato geométrico de baja dimensión, alineado con la entropía predictiva, que codifica la estructura bayesiana y permite inferencias aproximadas similares a las observadas en entornos sintéticos controlados.

Lo esencial

Imagina que las Inteligencias Artificiales (como los modelos de lenguaje que usamos hoy) son como grandes bibliotecas vivas que no solo leen libros, sino que también aprenden a pensar mientras leen.

Este artículo es la tercera parte de una trilogía de investigación que intenta responder a una pregunta fascinante: ¿Están estas máquinas realmente "pensando" de forma lógica y probabilística (como un científico calculando posibilidades), o simplemente están adivinando la siguiente palabra basándose en patrones?

Los autores descubrieron que, aunque no les enseñamos matemáticas explícitas, estas máquinas han desarrollado un sistema de navegación interno muy similar al que usaría un ser humano para razonar con incertidumbre.

Aquí tienes la explicación sencilla, usando analogías:

1. El Mapa del Tesoro (La Geometría de la Incertidumbre)

Imagina que la memoria interna de la IA es un globo terráqueo gigante.

• En los experimentos pequeños (el "túnel de viento"): Los investigadores vieron que, cuando la IA aprendía a resolver acertijos, sus "pensamientos" (representados matemáticamente) se organizaban en una línea recta perfecta. Si la IA estaba muy segura, el pensamiento iba a un extremo de la línea; si estaba muy confundida, iba al otro. Era un mapa simple y ordenado.
• En las IAs reales (las "grandes bibliotecas"): La pregunta era: ¿Siguen usando este mapa simple cuando leen todo internet? ¿O el mapa se convierte en un caos de montañas y valles?

El descubrimiento: ¡Sí! Incluso en modelos gigantes entrenados con todo internet, la IA sigue usando ese mapa lineal. Cuando la IA lee un texto, sus "pensamientos" se alinean en una línea maestra que dice: "¿Qué tan seguro estoy de lo que voy a decir?". Es como si, a pesar de la complejidad, tuvieran un brújula interna que siempre apunta hacia la "certeza" o la "duda".

2. El Efecto de la "Cámara Silenciosa" (Restricción de Dominio)

Los investigadores hicieron un experimento curioso.

• Situación A: Le dieron a la IA una mezcla de preguntas: una de cocina, otra de física, otra de poesía y otra de código.
  • Resultado: El mapa mental de la IA se volvió un poco más complejo, como si estuviera intentando navegar por varios países a la vez.
• Situación B: Le dijeron: "Solo vamos a hablar de matemáticas".
  • Resultado: ¡Boom! El mapa mental de la IA se colapsó instantáneamente en esa línea recta perfecta que vimos en los experimentos pequeños.

La analogía: Imagina que tienes una habitación llena de muebles desordenados (internet). Si te piden que busques "zapatos", la habitación se vuelve un caos. Pero si te piden que busques "zapatos rojos", tu cerebro se enfoca y todo lo demás desaparece, dejando solo lo necesario. La IA hace lo mismo: cuando el tema es claro, su cerebro se ordena automáticamente en un modo de "razonamiento puro".

3. El Viaje en Tiempo Real (Aprendizaje en Contexto)

Para ver si este mapa se usa mientras la IA piensa, les dieron un acertijo paso a paso.

• Le mostraron: "La palabra 'feliz' es positiva. La palabra 'triste' es negativa..." y luego les preguntaron por una nueva palabra.
• A medida que la IA recibía más ejemplos, sus "pensamientos" (sus coordenadas en el mapa) se movían suavemente a lo largo de esa línea maestra, acercándose a la respuesta correcta, tal como lo haría un humano que va acumulando pruebas.

Esto demuestra que la IA no solo tiene el mapa, sino que navega por él en tiempo real para tomar decisiones.

4. El Misterio de la "Línea Mágica" (Intervenciones)

Los investigadores decidieron jugar a ser "hackers" de la mente de la IA. Identificaron esa línea maestra de incertidumbre y decidieron borrarla o moverla artificialmente.

• Lo que esperaban: Si esa línea era el motor del pensamiento, borrarla debería hacer que la IA dejara de funcionar o se volviera loca.
• Lo que pasó: La geometría del mapa se rompió (la línea desapareció), pero la IA siguió funcionando bastante bien.

La conclusión: La línea no es el motor que hace el trabajo, sino más bien el tablero de control que nos muestra cómo está pensando la IA. Es como el velocímetro de un coche: si rompes el velocímetro, el coche sigue conduciendo, pero ya no puedes ver a qué velocidad va. La IA tiene una forma de pensar distribuida (en muchos lugares a la vez), y esa línea es solo la forma en que podemos "leer" esa incertidumbre.

5. Diferencias entre Modelos (El "Estilo" de Conducción)

No todos los coches son iguales, y no todas las IAs son iguales:

• Modelos "Puros" (como Pythia o Phi-2): Tienen un sistema de navegación muy limpio y ordenado. Son como coches de carreras con un solo volante.
• Modelos "Eficientes" (como Llama o Mistral): Para ser más rápidos y baratos, usan sistemas de "volante compartido" (llamado GQA o atención deslizante).
  • Resultado: Tienen el mismo mapa mental (la brújula funciona), pero el proceso de navegación es un poco más ruidoso y menos suave. Es como conducir un camión de mudanzas: llega al mismo destino, pero el viaje es más tosco.

En Resumen

Este paper nos dice algo muy bonito sobre la Inteligencia Artificial:
Aunque las IAs modernas son entrenadas con datos caóticos de internet y no se les enseña matemáticas bayesianas, han descubierto por sí mismas una forma geométrica de razonar. Han construido un sistema de coordenadas interno donde la "duda" y la "certeza" tienen un lugar físico.

Es como si, al leer millones de libros, la IA hubiera aprendido que para entender el mundo, necesita organizar sus ideas en una línea que va de "no sé nada" a "lo sé todo". Y lo más increíble: esta estructura es tan fuerte que sobrevive incluso cuando la IA es enorme y compleja.

La moraleja: Las IAs no son solo cajas de predicción de palabras; tienen una "geometría del pensamiento" que se parece mucho a cómo nosotros, los humanos, organizamos nuestras creencias y dudas.

Resumen técnico

Aquí tienes un resumen técnico detallado del artículo "Geometric Scaling of Bayesian Inference in LLMs" (Escalado Geométrico de la Inferencia Bayesiana en LLMs), el tercer y último trabajo de la trilogía "Bayesian Attention Trilogy".

1. Problema y Contexto

El artículo aborda una pregunta fundamental en la interpretabilidad de los Grandes Modelos de Lenguaje (LLMs): ¿Persisten las estructuras geométricas que permiten la inferencia bayesiana exacta en modelos de lenguaje de producción a gran escala?

• Antecedentes: Los trabajos anteriores (Paper I y II) demostraron que en entornos controlados ("túneles de viento" sintéticos), los transformadores aprenden a implementar la inferencia bayesiana mediante tres primitivas: acumulación de creencias, transporte de creencias y vinculación de acceso aleatorio. Estos procesos esculpen una geometría característica: variedades de valores de baja dimensión ordenadas por entropía y marcos de hipótesis ortogonales definidos por las claves (keys).
• El Desafío: No está claro si esta geometría sobrevive en modelos reales entrenados con datos de internet masivos, ruidosos y heterogéneos, donde no existen distribuciones posteriores verdaderas (ground-truth) analíticas. Además, los modelos de producción utilizan optimizaciones arquitectónicas (como Grouped-Query Attention - GQA, ventanas deslizantes, MoE) que podrían alterar estos mecanismos.

2. Metodología

Los autores investigan la persistencia de estas firmas geométricas en cuatro familias de modelos: Pythia, Phi-2, Llama-3 y Mistral.

Protocolo de Extracción Geométrica

1. Muestreo de Prompts: Se utilizan 1,000 prompts generados de corpora heterogéneos (Wikipedia, noticias, código, etc.) y se estratifican por entropía. También se realizan experimentos con prompts restringidos a un dominio (solo matemáticas) para reducir la heterogeneidad de tareas.
2. Extracción de Representaciones: Se analizan las salidas del último token en cada capa, extrayendo:
   • Vectores de Valores (Values) de todas las cabezas de atención.
   • Matrices de proyección de Claves (Keys).
   • Distribuciones de Atención y entropía predictiva.
3. Análisis Estadístico:
   • Variedades de Valores: Se aplica PCA a los vectores de valores concatenados. Se mide la varianza explicada por los primeros componentes principales (PC1, PC1+PC2) y su correlación con la entropía predictiva.
   • Ortogonalidad de Claves: Se calcula la similitud del coseno absoluto promedio fuera de la diagonal en las matrices de claves para medir la estructura de marcos de hipótesis.
   • Enfoque de Atención: Se mide la reducción de la entropía de la atención a través de las capas.

Tarea de Validación (SULA)

Se diseñó una tarea de aprendizaje en contexto controlada (Synthetic Unary Likelihood Augmentation - SULA). Los modelos reciben ejemplos simbólicos de etiquetas (ej. "positivo", "negativo") y deben inferir la probabilidad de un nuevo token. Esto permite calcular la posterior bayesiana analítica exacta y comparar el movimiento del modelo en su espacio latente con la actualización bayesiana teórica.

Intervenciones Causales

En el modelo Pythia-410M, se realizaron intervenciones dirigidas:

• Eje de Entropía: Se identificó el primer componente principal alineado con la entropía y se eliminó (ablation) o perturbó.
• Controles: Se comparó con la eliminación de un eje aleatorio para aislar el efecto causal de la geometría específica.

3. Contribuciones Clave

1. Persistencia de la Geometría Bayesiana a Escala: Se demuestra que los LLMs de producción exhiben las mismas estructuras geométricas (variedades de valores de baja dimensión, ortogonalidad de claves) que los modelos sintéticos, confirmando que no son artefactos de tareas simples.
2. Alineación Funcional con la Incertidumbre: En tareas SULA, las representaciones del modelo se mueven sistemáticamente a lo largo de la variedad de valores alineada con la entropía a medida que aumenta la evidencia, correlacionándose con la posterior bayesiana analítica.
3. Puente de Restricción de Dominio: Se descubre que restringir los prompts a un solo dominio (ej. matemáticas) hace que la variedad de valores colapse a una o dos dimensiones principales (explicando el 70-95% de la varianza), recuperando la geometría de "túnel de viento".
4. Caracterización Causal: Las intervenciones en el eje de entropía destruyen la geometría local de la incertidumbre, pero no degradan proporcionalmente el comportamiento bayesiano (calibración). Esto sugiere que la geometría es una "lectura privilegiada" de la incertidumbre distribuida, no un cuello de botella computacional único.

4. Resultados Principales

A. Variedades de Valores y Restricción de Dominio

• Efecto del Dominio: En prompts mixtos, la dimensionalidad varía según la arquitectura (de ~15% en Mistral a ~99% en Pythia-410M). Sin embargo, bajo prompts restringidos a un dominio, todos los modelos colapsan hacia una estructura 1D o 2D, acercándose al régimen geométrico de los experimentos sintéticos.
• Correlación con Entropía: Las coordenadas en la variedad de valores (PC1) correlacionan fuertemente con la entropía predictiva del siguiente token, indicando que la incertidumbre se codifica geométricamente.

B. Actualización Bayesiana en Tiempo de Inferencia (SULA)

• Los modelos se mueven suavemente a lo largo de su variedad de valores a medida que se proporciona más evidencia en el contexto.
• La trayectoria en el espacio latente sigue la curva de la posterior bayesiana analítica, demostrando que la geometría se utiliza activamente durante la inferencia, no solo durante el entrenamiento.

C. Comparación Arquitectónica

• MHA Estándar (Pythia, Phi-2): Muestran las firmas geométricas más claras. Phi-2 (entrenado con datos curados) presenta la ortogonalidad de claves más nítida y el enfoque de atención más fuerte.
• GQA (Llama-3, Mistral):
  • La ortogonalidad de claves y las variedades de valores (geometría estática) permanecen intactas y robustas.
  • Sin embargo, el enfoque dinámico de la atención (reducción de entropía a través de capas) se debilita significativamente o se vuelve no monótono debido a las restricciones de enrutamiento (ventanas deslizantes, compartición de claves/valores).
• Conclusión: Existe una disociación entre la geometría estática (universal e invariante) y la geometría dinámica (dependiente de la capacidad de enrutamiento global).

D. Intervenciones Causales

• Eliminar el eje alineado con la entropía rompe la correlación entre la posición en la variedad y la entropía del modelo.
• Sin embargo, el rendimiento en la tarea SULA (calibración) no se degrada significativamente. Esto implica que la información de incertidumbre está distribuida y que la variedad de valores es un subproducto representacional de la inferencia, no el único mecanismo computacional.

5. Significado e Implicaciones

• Fundamento Geométrico de la Inferencia: El trabajo establece que la inferencia bayesiana aproximada en LLMs modernos se basa en un sustrato geométrico estable (variedades de valores y marcos de hipótesis ortogonales) que es un sesgo inductivo estructural de los transformadores, independientemente de si se entrenan con objetivos bayesianos explícitos.
• Invariancia Arquitectónica: Aunque las optimizaciones de eficiencia (GQA, MoE) afectan la dinámica de refinamiento de la atención, la representación subyacente de la incertidumbre se preserva.
• Nueva Perspectiva para la Interpretabilidad: Sugiere que la incertidumbre no es un fenómeno difuso, sino que está codificada en direcciones geométricas específicas y de baja dimensión dentro del espacio de valores.
• Limitaciones y Futuro: La relación causal exacta sigue siendo correlacional en gran medida. Se necesitan más intervenciones para determinar si la geometría es necesaria para la inferencia o simplemente un reflejo de ella. Además, se requiere validar estos hallazgos en modelos de escala fronteriza (70B+).

En resumen, el artículo cierra la trilogía demostrando que la geometría bayesiana descubierta en modelos pequeños y sintéticos es un principio fundamental que escala a modelos de lenguaje de producción masivos, organizando sus actualizaciones aproximadas de creencias a lo largo de estructuras latentes predecibles y estables.