AtomBench: A Benchmarking Framework for Generative Crystal Reconstruction Models in Conventional Superconductors

Resumen Técnico: AtomBench

Planteamiento del Problema
El artículo aborda una brecha crítica en la evaluación de modelos generativos para la ciencia de materiales: la falta de un marco estandarizado y justo para evaluar las tareas de reconstrucción cristalina (diseño inverso). Mientras que muchos modelos afirman tener éxito en la generación de novo (evaluando estabilidad, unicidad y novedad), la capacidad de reconstruir con precisión una estructura cristalina conocida a partir de información parcial suele pasarse por alto.

Se identifica como problema principal que las comparaciones existentes frecuentemente confunden el rendimiento arquitectónico con la asimetría de información. Los diferentes modelos se evalúan a menudo bajo condiciones en las que poseen cantidades desiguales de información sobre el cristal objetivo (por ejemplo, algunos reciben solo la composición, mientras que otros reciben la composición más un embedding latente de la estructura completa). Esto invalida las comparaciones arquitectónicas directas. Además, las métricas estándar como la "tasa de coincidencia" (match rate) y el "RMSD emparejado" (matched RMSD) dependen de umbrales rígidos (por ejemplo, pymatgen StructureMatcher), lo que crea una evaluación discontinua donde los modelos solo se califican sobre el subconjunto de estructuras que logran coincidir, lo que puede enmascarar fallos sistemáticos en la mayoría de las predicciones.

Metodología
Los autores presentan AtomBench, un marco extensible y agnóstico al modelo diseñado para evaluar modelos generativos en una tarea de reconstrucción de cristales bien definida utilizando superconductores convencionales.

1. Conjuntos de Datos: El estudio utiliza dos conjuntos de datos generados mediante DFT:

   • JARVIS Supercon-3D: 1,058 pares cristal/$T_c$, dominado por O, Al y Ti, con diversos sistemas cristalinos.
   • Alexandria DS-A/B: 8,253 pares cristal/$T_c$, dominado por Ti, Rh y Al, principalmente cúbico y tetragonal.
   • Ambos conjuntos de datos utilizan la Teoría de la Perturbación de la Función de Densidad (DFPT) con protocolos específicos (por ejemplo, valores de $\mu^*$, pseudopotenciales).

2. Modelos Evaluados: Se entrenaron y evaluaron cuatro arquitecturas generativas distintas:

   • AtomGPT: Un modelo de lenguaje de gran tamaño (Mistral-7B) ajustado (fine-tuned) que trata los cristales como texto.
   • CDVAE: Un autoencoder variacional de difusión que utiliza redes neuronales de grafos equivariantes SE(3).
   • FlowMM: Una red de ajuste de flujo (flow-matching) Riemanniana.
   • MatterGen: Un modelo de difusión conjunta con redes de puntuación (score networks) equivariantes.
   • Variantes de Condicionamiento: Para los modelos que lo admiten (AtomGPT, MatterGen), se entrenaron variantes con y sin condicionamiento en la temperatura crítica de superconductividad ($T_c$).

3. Protocolo de Evaluación:

   • División de Datos: Se reservó el 10% de cada conjunto de datos para pruebas.
   • Regímenes de Información: Los modelos se agruparon según la información disponible antes de la predicción:
     • Embedding de estructura completa: CDVAE (recibe un vector latente de la estructura objetivo completa).
     • Composición + $T_c$: AtomGPT/MatterGen (condicionados por $T_c$).
     • Solo Composición: FlowMM, AtomGPT, MatterGen (no condicionados).
   • Métricas:
     • Parámetros de Red: Error Absoluto Medio (MAE) y Divergencia de Kullback-Leibler (KLD) para distribuciones.
     • Coordenadas Atómicas:
       • RMSD Emparejado (Matched RMSD): Error promedio sobre las estructuras emparejadas exitosamente por StructureMatcher.
       • RMSD Corregido Continuo (ccRMSD): Una métrica novedosa introducida en este trabajo. Utiliza el descriptor de Distancia Mínima Promedio (AMD) y la distancia de Chebyshev para medir la fidelidad geométrica local de forma continua a través de todas las estructuras de prueba, independientemente de si se encuentra una coincidencia discreta. Esto elimina el sesgo de descartar estructuras "difíciles de emparejar".

Contribuciones Clave

1. Marco AtomBench: Un paquete de Python de código abierto (atombench) que estandariza la evaluación de modelos de diseño inverso, permitiendo la sumisión directa al JARVIS-Leaderboard.
2. Métrica ccRMSD: La introducción de una métrica continua, independiente de la tasa de coincidencia, que proporciona una evaluación de la fidelidad geométrica local más robusta que el RMSD emparejado tradicional.
3. Evaluación de Información-Teórica: Un agrupamiento riguroso de los modelos por regímenes de "entropía condicional" (cantidad de información conocida) para asegurar comparaciones arquitectónicas justas.
4. Análisis Sistemático: La primera comparación a gran escala de estas arquitecturas específicas en la reconstrucción de superconductores, analizando el impacto del condicionamiento de $T_c$ y la simetría del sistema cristalino.

Resultos

• Rendimiento por Métrica:
  • Coordenadas Atómicas: MatterGen logró la mejor fidelidad de reconstrucción (menor ccRMSD y RMSD emparejado), seguido de AtomGPT. CDVAE y FlowMM mostraron errores significativamente mayores.
  • Parámetros de Red: CDVAE fue el más preciso en la reconstrucción de parámetros de red (menor MAE y KLD), seguido de AtomGPT y MatterGen. FlowMM fue el menos preciso.
  • Velocidad vs. Precisión: FlowMM fue el más rápido en inferencia pero el menos preciso. MatterGen ofreció la mayor precisión de coordenadas pero con un mayor coste computacional.
• Impacto del Condicionamiento de $T_c$: El condicionamiento en la temperatura crítica ($T_c$) no mejoró consistentemente la fidelidad de reconstrucción para AtomGPT o MatterGen, lo que sugiere que un escalar único aporta poca información estructural más allá de la composición en este contexto.
• Tendencias del Sistema Cristalino: El error de reconstrucción varió según el sistema cristalino. Los sistemas cúbicos (alta simetría, menos grados de libertad) exhibieron los errores más bajos, mientras que los sistemas monoclínicos y hexagonales mostraron errores más altos, particularmente en el parámetro de red $c$.
• Discrepancia de Métricas: El estudio destacó que el RMSD emparejado puede ser engañoso; por ejemplo, FlowMM tuvo una tasa de coincidencia muy baja (<10%), lo que hace que su RMSD emparejado sea estadísticamente poco fiable, mientras que el ccRMSD proporcionó una estimación de error significativa para todo el conjunto de datos.

Significancia y Reivindicaciones
El artículo afirma que AtomBench proporciona una corrección necesaria al campo al exponer cómo la asimetría de información puede inflar artificialmente las métricas de rendimiento. Al introducir el ccRMSD, los autores abogan por un cambio desde la evaluación de "coincidir o descartar" hacia medidas de fidelidad continua que tengan en cuenta la distribución completa de las predicciones.

Los autores declaran modestamente que sus hallazgos no clasifican universalmente las arquitecturas, sino que ilustran cómo la fidelidad de reconstrucción varía bajo diferentes regímenes de información de inferencia. Enfatizan que la elección del modelo debe depender de la aplicación específica:

• FlowMM para muestreo de alto rendimiento donde la velocidad es primordial.
• CDVAE para tareas que requieren alta precisión en los parámetros de red.
• MatterGen para tareas que requieren una colocación de coordenadas atómicas de alta fidelidad.

El trabajo concluye que, si bien los modelos actuales (CDVAE, AtomGPT, MatterGen) alcanzan niveles de precisión adecuados para inicializar cálculos de DFT, el campo requiere métricas continuas como el ccRMSD para evaluar y mejorar adecuadamente los modelos generativos para el descubrimiento de materiales.