TriForces: Augmenting Atomistic GNNs for Transferable Representations

TriForces es un marco de tres flujos agnóstico al modelo que combina el aprendizaje auto-supervisado con representaciones de composición y estructura separadas para mejorar significativamente la transferibilidad y la eficiencia de datos de las redes neuronales gráficas atómicas para potenciales interatómicos de aprendizaje automático.

Lo esencial

Imagina que estás intentando enseñar a un chef robot cómo cocinar.

El Problema: El Chef "Valido para Todos"
Actualmente, los científicos utilizan modelos de IA potentes (llamados MLIPs) para predecir cómo se comportan los átomos, como la cantidad de energía que tiene un material o qué tan difícil es mover los átomos. Estos modelos se entrenan con grandes cantidades de datos provenientes de supercomputadoras (DFT).

Sin embargo, estos modelos tienen un defecto: son como un chef que memorizó el sabor exacto de un plato específico pero olvidó por qué sabía así. Si les pides que cocinen un plato ligeramente diferente (un nuevo tipo de material), tienen dificultades. Confunden los ingredientes (composición) con la forma de la olla (estructura). Si cambias los ingredientes, se confunden sobre la forma, y viceversa. Esto los hace malos aprendiendo nuevas tareas rápidamente, especialmente cuando no tienes muchos datos para enseñarles.

La Solución: TriForces (La Cocina de Tres Corrientes)
Los autores presentan TriForces, una nueva forma de construir estos chefs de IA. En lugar de un solo cerebro gigante intentando recordar todo a la vez, dividen el cerebro en tres "corrientes" o departamentos especializados:

1. La Corriente de Ingredientes (Composición): Este departamento solo mira qué hay en la olla (por ejemplo, "Tenemos 2 Hidrógenos y 1 Oxígeno"). Ignora por completo la forma. Aprende la química.
2. La Corriente de Forma (Estructura): Este departamento solo mira cómo están dispuestos los átomos en el espacio (por ejemplo, "Están en un triángulo"). Ignora qué son realmente los átomos. Aprende la geometría.
3. La Corriente de Interacción: Este es el chef principal que toma las notas de los departamentos de Ingredientes y Forma y las combina para predecir el resultado final (energía o fuerza).

La Salsa Secreta: Aprendizaje Auto-supervisado
Antes de que se le pida al modelo predecir una propiedad específica, los autores lo entrenan usando un juego llamado "Aprendizaje Auto-supervisado". Piensa en esto como una sesión de práctica donde la IA tiene que:

• Denoising (Eliminar ruido): Mirar una imagen ligeramente rota o ruidosa de una molécula y arreglarla.
• Enmascaramiento: Cubrir un ingrediente y adivinar qué era basándose en los vecinos.
• Emparejamiento: Mirar dos versiones ligeramente diferentes de la misma molécula y darse cuenta de que son lo mismo.

Este entrenamiento obliga a la IA a organizar su conocimiento ordenadamente. Aprende que los "ingredientes" pertenecen a una carpeta y las "formas" a otra, en lugar de mezclarlos todos juntos.

Por Qué Esto Importa (Los Resultados)
El artículo muestra que esta nueva cocina de "Tres Corrientes" funciona mucho mejor que las antiguas cocinas de "Un Solo Cerebro":

• Aprendizaje Más Rápido: Cuando se le da una pequeña cantidad de datos nuevos (como 20,000 ejemplos en lugar de millones), TriForces aprende mucho más rápido y comete menos errores. Es como un chef que puede aprender una nueva receta después de probarla una vez, en lugar de necesitar cocinarla mil veces.
• Mejor Memoria: La IA no olvida lo que aprendió. Puede transferir su conocimiento de un tipo de material a otro sin confundirse.
• Conocimiento Buscable: Como la IA mantiene separados los "ingredientes" y las "formas", puedes pedirle que encuentre materiales que se vean iguales pero tengan ingredientes diferentes, o materiales con los mismos ingredientes pero formas diferentes. Los modelos antiguos no podían hacer esto porque su conocimiento estaba demasiado mezclado.

En Resumen
TriForces es un marco que descompone el trabajo complejo de entender los átomos en tres trabajos más simples: conocer los ingredientes, conocer la forma y conocer cómo funcionan juntos. Al entrenar a la IA para mantener estos trabajos separados y practicar con "juegos de adivinanza" (aprendizaje auto-supervisado), el modelo se convierte en una herramienta mucho más flexible, eficiente y precisa para descubrir nuevos materiales.

Los autores han publicado su código y modelos pre-entrenados para que otros científicos puedan usar esta "cocina de tres corrientes" para construir mejores IA para la ciencia de materiales.

Resumen técnico

Resumen Técnico: TriForces

Enunciado del Problema

Los Potenciales Interatómicos de Aprendizaje Automático (MLIPs) basados en Redes Neuronales de Grafos Geométricas (GNN) han logrado alta precisión en la predicción de propiedades atómicas cuando se entrenan con conjuntos de datos de Teoría del Funcional de la Densidad (DFT) a gran escala. Sin embargo, su utilidad práctica se ve obstaculizada por dos limitaciones principales:

1. Transferibilidad Inconsistente: Los MLIPs a menudo fallan al transferirse efectivamente a nuevas químicas o tareas, particularmente al ajustar finamente sobre conjuntos de datos pequeños, costosos y específicos de la tarea. Los modelos preentrenados frecuentemente luchan con tareas diagnósticas simples (por ejemplo, identificar sistemas cristalinos o elementos mayoritarios) y exhiben "olvido catastrófico" o un rendimiento de transferencia inestable a través de dominios.
2. Mala Reutilización de Representaciones: Los MLIPs actuales optimizan representaciones para objetivos de predicción específicos (energía y fuerzas) en lugar de para una reutilización general. En consecuencia, estas representaciones a menudo entrelazan la composición (química) y la estructura (geometría), haciéndolas inadecuadas para análisis exploratorio, recuperación de vecinos más cercanos o búsqueda de similitud descompuesta. Los objetivos supervisados estándar fomentan representaciones suficientes para la regresión pero no organizadas para preservar información accesible de composición y estructura.

Metodología: El Marco TriForces

Los autores proponen TriForces, un marco agnóstico al modelo que aumenta las GNNs geométricas existentes con una arquitectura de tres flujos y aprendizaje auto-supervisado (SSL) multi-objetivo. La innovación central es la factorización explícita de las representaciones atómicas en tres componentes distintos:

1. Arquitectura de Tres Flujos

En lugar de un único vector latente, TriForces descompone la representación a nivel de nodo $h_i$ en tres flujos concatenados:

• Flujo de Composición ($h^{comp}$): Codifica información química sin coordenadas. Procesa el conjunto de elementos atómicos únicos y sus conteos estequiométricos utilizando un Transformer con atención ponderada por conteo. Este flujo preserva los conteos absolutos de elementos para codificar el tamaño del sistema y la escala de energía, independientemente de la geometría.
• Flujo de Estructura ($h^{struct}$): Codifica la geometría sin identidad de elemento (agnóstico al tipo). Construye descriptores locales invariantes a la rotación inspirados en SOAP (Superposición Suave de Posiciones Atómicas), utilizando funciones de base radial, armónicos esféricos y cortes multi-escala. Este flujo captura motivos geométricos reutilizables y patrones topológicos, potenciados por paso de mensajes invariante.
• Flujo de Interacción ($h^{int}$): Una GNN geométrica base estándar (por ejemplo, MACE, eSEN, Orb-v3) que captura el acoplamiento entre composición y geometría, preservando la expresividad de la arquitectura original.

2. Preentrenamiento Auto-Supervisado

Para organizar el espacio latente y mejorar la transferibilidad, TriForces emplea una estrategia de preentrenamiento SSL multi-objetivo utilizando aumentaciones estocásticas (ruido de posición, enmascaramiento de elementos, variación de grafos y rotaciones). El marco combina tres objetivos complementarios:

• No-Reconstrucción (LeJEPA): Alinea incrustaciones desde dos vistas aumentadas de la misma estructura a nivel de nodo y de grafo. Esto impone invarianza a las aumentaciones y organiza el espacio latente global sin requerir gradientes de parada o codificadores de momento.
• Denoising (Eliminación de Ruido): Entrena al modelo para recuperar posiciones atómicas limpias a partir de entradas ruidosas. Esto estabiliza las representaciones geométricas y proporciona implícitamente aumentación de rotación.
• Enmascaramiento: Predice tipos atómicos enmascarados basándose en la geometría y composición circundantes. Esto fomenta que el modelo aprenda patrones composicionales y contexto.

La pérdida final de preentrenamiento es una suma ponderada de estos tres objetivos.

Contribuciones Clave

1. Descomposición Arquitectónica: Un diseño de tres flujos que separa explícitamente composición, estructura e interacción, asegurando que ambos factores se preserven por diseño en lugar de estar entrelazados.
2. Estrategia de Preentrenamiento Híbrida: Un enfoque auto-supervisado que combina objetivos basados en reconstrucción (denoising, enmascaramiento) con aprendizaje de predicción latente (LeJEPA) para estructurar el espacio de incrustaciones para una mejor transferencia aguas abajo.
3. Recuperación Interpretable: La capacidad de realizar búsqueda de similitud dirigida en los espacios de incrustación composicional, estructural o conjunta, permitiendo la comparación de materiales basada en criterios específicos (por ejemplo, solo química o solo estructura).
4. Validación Empírica: Experimentos extensos a través de múltiples arquitecturas (Orb-v3, eSEN, MACE) y puntos de referencia (OMat24, MatBench, QM9) que demuestran una mayor eficiencia de datos, rendimiento de transferencia y calidad de representación.

Resultados

• Rendimiento de Transferencia (OMat24): En regímenes de datos limitados, TriForces supera significativamente a las líneas base. Con 20K muestras, reduce el Error Absoluto Medio (MAE) de energía en un 57% en comparación con el modelo base. Mejora el MAE de fuerza en todos los tamaños de muestra y reduce los errores de tensión.
• Eficiencia de Datos: TriForces logra errores más bajos en cada tamaño de conjunto de datos (de 20K a 2M de muestras), con las ganancias más significativas observadas en configuraciones de bajos datos.
• Rendimiento en Puntos de Referencia:
  • MatBench: Las variantes de TriForces logran los mejores resultados generales en 6 de 8 tareas, superando tanto a las líneas base de preentrenamiento auto-supervisado como a las etiquetadas con DFT. Por ejemplo, el MAE de Fonones mejoró de 57.8 a 19.5 cm$^{-1}$.
  • MatBench Discovery: TriForces eSEN-sm logra un MAE de energía comparable a un modelo eSEN-30M-OAM mucho más grande mientras utiliza un 60% menos de parámetros y entrena hasta 5$\times$ más rápido.
  • QM9: El preentrenamiento sobre entradas químicas diversas (bulk + moléculas) reduce consistentemente el MAE en comparación con las líneas base de solo bulk o sin SSL.
• Calidad de Representación: La sonda lineal sobre incrustaciones congeladas muestra que TriForces preserva información fundamental (sistema cristalino, elemento mayoritario, número de coordinación) que los MLIPs estándar pierden. TriForces alcanza una precisión del 96–100% en la clasificación de sistemas cristalinos y elementos mayoritarios, mientras que las líneas base luchan (55–73%).
• Recuperación: El marco permite una recuperación efectiva de k-NN donde el flujo de composición sobresale en la recuperación de conjuntos de elementos y el flujo de estructura sobresale en la recuperación de grupos espaciales, una capacidad ausente en modelos de flujo único.

Significado y Afirmaciones

El artículo posiciona a TriForces no meramente como un método auto-supervisado, sino como un marco arquitectónico cuyas representaciones se ven potenciadas adicionalmente por SSL.

• Dependencia del Régimen: Los autores afirman que la factorización de flujos proporciona las ganancias dominantes en configuraciones supervisadas a gran escala, mientras que SSL es más valiosa para la transferencia de bajos datos, la organización de representaciones y las tareas de recuperación.
• Desacoplamiento: Al separar composición y estructura, TriForces aborda los problemas de "transferencia frágil" y "difícil de reutilizar" de los modelos fundacionales actuales. Permite que los modelos aprendan representaciones organizadas para el análisis (por ejemplo, sondas, recuperación) además de la predicción.
• Practicidad: El marco es agnóstico al modelo y de tipo enchufable, haciéndolo inmediatamente aplicable a arquitecturas atómicas existentes o nuevas. Los autores liberan checkpoints preentrenados y código para facilitar la reutilización en el modelado atómico aguas abajo.

El trabajo sugiere que los futuros modelos fundacionales atómicos deben ir más allá de las arquitecturas de predicción de flujo único hacia representaciones factorizadas que preserven explícitamente los factores físicos distintos de los sistemas químicos.