Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement

El estudio introduce la descomposición de sondas composicionales (CPD) para demostrar que la alineación de la tarea de entrenamiento y la arquitectura equivariante determinan un gradiente de accesibilidad lineal en la información geométrica de los modelos fundamentales atómicos, revelando cómo la información se enruta selectivamente según su tipo de simetría.

Lo esencial

Imagina que tienes una caja de herramientas mágica (un modelo de Inteligencia Artificial) diseñada para predecir propiedades de moléculas, como si fueran pequeños bloques de construcción. Los científicos han creado muchas de estas cajas, pero se preguntaban: ¿Cómo están organizadas las herramientas dentro de la caja?

¿Separa la caja claramente "de qué está hecha la molécula" (su composición, como si fuera una lista de ingredientes) de "cómo están dispuestos esos ingredientes" (su geometría, como si fuera la receta de cómo armarlos)?

Este paper, escrito por Joshua Steier, es como un detective que entra en estas cajas de herramientas para ver cómo están organizadas. Aquí te explico los hallazgos principales con analogías sencillas:

1. El Problema: La Mezcla de Ingredientes y Receta

En el mundo de las moléculas, es difícil separar la "composición" (ej. tiene 3 átomos de carbono y 6 de hidrógeno) de la "geometría" (cómo están conectados esos átomos).

• El error común: Antes, los científicos usaban herramientas muy potentes y complejas (llamadas "árboles de decisión" o Gradient Boosted Trees) para intentar leer la caja. El problema es que estas herramientas eran tan inteligentes que podían "adivinar" la receta incluso si les habían quitado los ingredientes de la lista. Era como si un chef te diera una sopa y, aunque le quitaras la sal, él pudiera decirte exactamente cuánta sal había puesto porque sabía cómo sabe el caldo. Esto les daba resultados falsos y demasiado optimistas.

2. La Solución: El Filtro "CPD"

El autor inventó un nuevo método llamado CPD (Descomposición de la Sonda Composicional).

• La analogía: Imagina que tienes una foto de una casa. Quieres saber cómo es la arquitectura (geometría) sin que la foto te diga de qué color son los ladrillos (composición). El CPD es como un filtro que borra digitalmente el color de los ladrillos. Luego, le pide a una herramienta muy simple y honesta (una regresión lineal, como una regla recta) que intente adivinar la forma de la casa con la foto que queda.
• El hallazgo clave: Si la herramienta simple puede ver bien la forma de la casa después de borrar el color, significa que la caja de herramientas (el modelo) tiene la información de la forma bien organizada y separada. Si no puede verla, es que la información está mezclada y desordenada.

3. El Gran Descubrimiento: Lo que importa es el "Entrenamiento", no la "Caja"

El estudio probó 10 cajas de herramientas diferentes. Esperaban que las cajas más modernas y complejas (con "equivalencia", un término técnico que significa que giran y se adaptan como objetos reales) fueran las mejores. Pero no fue así.

El factor más importante fue qué le enseñaron a la caja:

• La analogía del entrenador: Imagina dos atletas. Uno entrena para correr maratones (geometría) y otro para levantar pesas (energía).
  • Si quieres que el atleta te diga cómo es la forma de una montaña (geometría), el que entrenó para correr (modelo entrenado en "HOMO-LUMO gap") lo hará genial, porque su cerebro está enfocado en la forma.
  • El que entrenó para levantar pesas (modelo entrenado en "energía") tendrá dificultades, aunque sea un atleta más fuerte o tenga mejor equipamiento. Su cerebro está enfocado en el peso, no en la forma.
• Conclusión: No importa cuán sofisticada sea la arquitectura del modelo; si no se le entrenó específicamente para entender la geometría, no tendrá la información bien organizada para que un humano (o una herramienta simple) pueda leerla fácilmente.

4. El Mapa Interno: Las "Autopistas" de la Información

En uno de los modelos más avanzados (llamado MACE), descubrieron algo fascinante: la información viaja por "carriles" separados según su tipo.

• La analogía: Imagina una autopista con carriles.
  • Los carriles rojos (canales escalares) están dedicados exclusivamente a propiedades que no tienen dirección, como la brecha de energía (HOMO-LUMO).
  • Los carriles azules (canales vectoriales) están dedicados a propiedades que sí tienen dirección, como el momento dipolar (como un imán).
• El modelo MACE sabe usar estos carriles perfectamente. Pero otro modelo, ViSNet, aunque también tiene carriles, parece mezclar todo en el carril rojo y deja el azul vacío. Esto muestra que la forma en que se construye el modelo afecta cómo organiza la información.

5. La Lección para el Futuro

• Para los científicos: Si quieres usar un modelo de IA para predecir algo relacionado con la forma de una molécula, no elijas el modelo más "moderno" o "grande". Elige el que haya sido entrenado específicamente para entender la geometría.
• Advertencia: Si usas herramientas de análisis demasiado complejas (como los árboles de decisión mencionados al principio) para leer estos modelos, podrías engañarte a ti mismo y pensar que el modelo sabe más de lo que realmente sabe.

En resumen:
Este paper nos dice que la "organización" de la inteligencia artificial no depende solo de qué tan complejo sea su diseño, sino de qué le enseñaron a hacer. Si quieres que la IA entienda la forma de las cosas, debes entrenarla pensando en la forma, no solo en el peso o la energía. Además, nos enseña a ser humildes con nuestras herramientas de medición: a veces, la herramienta más simple es la que nos dice la verdad.

Resumen técnico

Aquí tienes un resumen técnico detallado del artículo en español, estructurado según los puntos solicitados:

Título: Enrutamiento de Información en Modelos Fundamentales Atomísticos: Cómo la Alineación de Tareas y la Equivarianza Dan Forma al Desentrelazamiento Lineal

1. El Problema

Los modelos fundamentales atomísticos (como MACE, SchNet, PaiNN, ViSNet) han logrado predecir propiedades moleculares con una precisión cercana a la Teoría del Funcional de la Densidad (DFT) a un costo computacional reducido. Sin embargo, surge una pregunta fundamental: ¿cómo se organizan sus representaciones intermedias?

Específicamente, ¿separan limpiamente la composición (qué elementos están presentes y en qué proporción) de la geometría (cómo están dispuestos los átomos en el espacio)?

• El desafío: Las propiedades moleculares dependen de ambos factores, los cuales están correlacionados. Las sondas (probes) tradicionales que se entrenan directamente sobre las representaciones brutas no pueden distinguir si el modelo ha aprendido la geometría real o si está aprovechando "atajos" basados en la composición.
• El fallo metodológico previo: Los autores identificaron que las sondas no lineales (como los árboles de decisión o Gradient Boosted Trees - GBT) aplicadas a representaciones donde se ha eliminado la señal de composición mediante regresión lineal, producen resultados engañosamente altos. Estas sondas no lineales son capaces de reconstruir la señal de composición proyectada a través de interacciones de características de alta dimensión, inflando artificialmente las métricas de rendimiento.

2. Metodología: Descomposición de Sondas Composicionales (CPD)

Para abordar estos problemas, los autores introducen Compositional Probe Decomposition (CPD), un protocolo riguroso para aislar y medir la información geométrica linealmente accesible.

• Proceso de CPD:

  1. Extracción de Representaciones: Se extraen las activaciones finales de modelos preentrenados y congelados para un conjunto fijo de moléculas (2,000 de QM9).
  2. Proyección OLS (Mínimos Cuadrados Ordinarios): Se ajusta un modelo de regresión lineal para predecir las representaciones ($X$) a partir de las características de composición ($Z$, fracciones elementales y conteo de átomos). Se calculan los residuos: $X_{geom} = X - Z\hat{\beta}$. Estos residuos representan la información que es linealmente ortogonal a la composición.
  3. Probing Lineal: Se entrena una regresión de Ridge (lineal) sobre los residuos $X_{geom}$ para predecir propiedades objetivo (ej. brecha HOMO-LUMO). La métrica clave es $R^2_{geom}$.
  4. Validación Fold-wise: La proyección se realiza dentro de cada pliegue de validación cruzada para evitar fugas de información del conjunto de prueba.

• Validación Crítica:

  • Sondas Lineales vs. No Lineales: Se demuestra que las sondas GBT sobre residuos recuperan $R^2$ de 0.68–0.95 en objetivos puramente composicionales (como la masa atómica promedio), lo cual es teóricamente imposible si la composición se ha eliminado correctamente. Esto confirma que solo las sondas lineales ofrecen una medida fiel de la accesibilidad lineal.
  • Benchmarks de Isómeros Estructurales: Se utiliza un conjunto de isómeros (misma composición, diferente geometría). La componente composicional debe tener un rendimiento al azar (50%), mientras que el residuo geométrico debe permitir clasificar correctamente la diferencia estructural.

3. Contribuciones Clave

1. Metodología CPD Validada: Un nuevo estándar para el probing en química computacional que elimina sistemáticamente el sesgo de composición, junto con la demostración de que las sondas no lineales inflan los resultados en representaciones residualizadas.
2. Gradiente de Accesibilidad Lineal: Identificación de una variación de 6.6x en la información geométrica accesible entre diferentes modelos, explicada por tres factores interactivos.
3. Descubrimiento de Factores Determinantes: Demostración de que la alineación de la tarea (el objetivo de entrenamiento) es el factor dominante, superando a la arquitectura y la diversidad de datos.
4. Enrutamiento de Información por Simetría: Evidencia de que en arquitecturas como MACE, los canales de representación (escalares $L=0$ vs. vectoriales $L=1$) enrutan preferentemente propiedades que coinciden con su tipo de simetría.

4. Resultados Principales

A. El Gradiente de Accesibilidad Lineal
Al evaluar 10 modelos de 5 familias arquitectónicas en el dataset QM9, se observa un gradiente claro en $R^2_{geom}$ para la brecha HOMO-LUMO:

• Rango: Desde $R^2_{geom} = 0.081$ (MACE entrenado solo en energía en QM9) hasta $0.533$ (PaiNN entrenado en brecha HOMO-LUMO).

B. Los Tres Factores Explicativos

1. Alineación de Tarea (Dominante): Es el factor más importante. Los modelos entrenados directamente en la propiedad objetivo (brecha HOMO-LUMO) superan a los entrenados en energía por un margen de $\sim 0.25$ en $R^2$, independientemente de la arquitectura.
   • Ejemplo: PaiNN (entrenado en HL) tiene $0.533$, mientras que PaiNN-Energy (misma arquitectura, entrenado en energía) cae a$0.310$.
   • Interpretación: Si el objetivo de entrenamiento es sensible a la geometría, el modelo se ve forzado a codificarla de manera accesible. Si el objetivo es la energía total (dominada por composición), la geometría se codifica de forma menos accesible linealmente.
2. Equivarianza (Amplifica pero no sustituye): La arquitectura equivariante por sí sola no garantiza un buen desentrelazamiento. MACE (equivariante) entrenado en energía tiene un rendimiento peor que SchNet (invariante) entrenado en energía. La combinación de equivarianza + objetivo alineado es la que produce los mejores resultados.
3. Diversidad de Datos (Compensación Parcial): El preentrenamiento en datos masivos y diversos (MPTraj) mejora la accesibilidad geométrica incluso si el objetivo no está alineado (MACE preentrenado: $0.364$vs MACE QM9:$0.081$), pero no llega a igualar a los modelos con alineación de tarea.

C. Enrutamiento por Representación Irreducible (MACE)
En MACE, la información se enruta según el tipo de simetría:

• Canales Escalares ($L=0$): Codifican fuertemente propiedades escalares como la brecha HOMO-LUMO ($R^2 = 0.76$).
• Canales Vectoriales ($L=1$): Codifican fuertemente propiedades vectoriales como el momento dipolar ($R^2 = 0.59$).
• Contraste: ViSNet, aunque también tiene flujos escalares y vectoriales, concentra casi toda la información en el flujo escalar, sugiriendo que su enrutamiento es específico de la arquitectura de productos tensoriales de MACE.

D. Eficiencia de Muestra
Las representaciones linealmente desentrelazadas (como las de PaiNN) requieren significativamente menos datos para extraer señal geométrica. PaiNN con solo 50 moléculas supera el rendimiento de SchNet con 2,000 moléculas.

5. Significado e Impacto

• Para la Selección de Modelos: La elección del objetivo de entrenamiento es más crítica que la elección de la arquitectura. Para tareas sensibles a la geometría (propiedades electrónicas, reactividad), es preferible usar un modelo preentrenado en una tarea geométrica alineada, incluso si es una arquitectura "más simple" (invariante), que un modelo equivariante preentrenado solo en energía.
• Implicaciones Metodológicas: El hallazgo de la inflación por sondas no lineales es crucial para todo el campo del probing (no solo en química, sino también en NLP y visión por computadora). Sugiere que las conclusiones sobre la "ausencia" de información en representaciones residualizadas basadas en sondas complejas (GBT, MLP) pueden ser erróneas.
• Teoría de Representación: Los resultados desafían la hipótesis de que la inductividad arquitectónica (equivarianza) es suficiente para organizar representaciones. Demuestran que el sesgo supervisado (el objetivo de pérdida) es el motor principal para la organización lineal de la información, superando a los sesgos arquitectónicos.
• Generalización: Los hallazgos se mantienen al extenderse a cristales inorgánicos (Materials Project), validando que la separación composición-geometría y el gradiente de accesibilidad son propiedades generales de los modelos fundamentales moleculares.

En resumen, el paper establece que la "calidad" de una representación para tareas de transferencia no se trata solo de cuánta información contiene, sino de cómo está organizada. La alineación de la tarea es la clave para que la información geométrica sea linealmente accesible y utilizable por cabezas de predicción simples.