Score-based diffusion models for accurate crystal-structure inpainting and reconstruction of hydrogen positions

Resumen Técnico: Modelos de difusión basados en puntuación para la reconstrucción e inpainting preciso de la estructura cristalina de posiciones de hidrógeno

Planteamiento del Problema
La determinación precisa de las posiciones atómicas de los átomos de hidrógeno en materiales cristalinos sigue siendo un desafío significativo en la ciencia de materiales computacional. Aunque el hidrógeno es abundante y crítico para aplicaciones que van desde el almacenamiento de energía hasta la industria de semiconductores, su bajo poder de dispersión de rayos X y electrones dificulta su localización mediante la difracción de rayos X estándar. En consecuencia, las bases de datos cristalográficas históricas (como MC3D, derivada de COD, ICSD y MPDS) suelen contener estructuras donde las posiciones del hidrógeno se omiten, se insertan mediante heurísticas o se determinan con una incertidumbre significativa. Aunque los avances experimentales (por ejemplo, la cristalografía cuántica) y la costosa dispersión de neutrones pueden resolver estas posiciones, las bases de datos existentes aún dependen en gran medida de datos incompletos o aproximados. El problema central abordado es la reconstrucción de estas posiciones de hidrógeno faltantes dado un "huésped" de estructura cristalina conocido (el armazón no hidrogenado) para generar estructuras cristalinas precisas, estables y completas.

Metodología
Los autores proponen un enfoque generativo que combina modelos de difusión basados en puntuación con técnicas adaptadas del inpainting de imágenes en visión por computadora.

1. Arquitectura del Modelo: El estudio utiliza MatterGen de Microsoft, un modelo de difusión basado en puntuación preentrenado diseñado originalmente para generar nuevas estructuras cristalinas. Los autores reentrenan y adaptan este modelo para la tarea específica de inpainting (completar datos faltantes) en lugar de la generación completa.
2. Estrategias de Inpainting: Se comparan varios enfoques:
   • Línea base de MatterGen: El modelo original aplicado a la tarea con la red y los tipos atómicos fijos.
   • pos-only: Un modelo reentrenado para eliminar el ruido solo en las posiciones atómicas, manteniendo la red y los tipos fijos.
   • pos-only-RePaint: El modelo reentrenado combinado con el algoritmo RePaint, que resamplea iterativamente las regiones conocidas para armonizarlas con las regiones completadas (inpainted).
   • pos-only-TD (Dependiente del Tiempo): Una adaptación novedosa del concepto TD-Paint de la visión por computadora. Este modelo está entrenado para manejar niveles de ruido variables por sitio dentro de una sola estructura. Condiciona el proceso de eliminación de ruido basándose en posiciones conocidas "limpias" (ruido cero) mientras predice iterativamente los sitios de hidrógeno faltantes (ruido alto). Esto permite al modelo utilizar directamente la estructura huésped conocida sin añadir ruido a la misma.
3. Flujo de Trabajo:
   • Partir de una estructura huésped (sitios de hidrógeno eliminados).
   • Inicializar aleatoriamente $N_H$ sitios de hidrógeno en la celda unidad.
   • Aplicar el modelo de difusión pos-only-TD durante $N_{steps}$ iteraciones para predecir las posiciones del hidrógeno.
   • Realizar una relajación restringida de los sitios de hidrógeno utilizando un Potencial Interatómico de Aprendizaje Automático (MLIP), específicamente NequIP, manteniendo fijos los átomos del huésped y la celda.
   • Realizar una relajación completa de todas las posiciones (incluyendo los átomos del huésped) utilizando NequIP.
4. Métricas de Evaluación: El éxito se define de dos maneras:
   • Coincidencia Estructural: La estructura predicha coincide con la referencia (de la base de datos) dentro de las tolerancias de StructureMatcher de pymatgen.
   • Coincidencia de Menor Energía o Estructural (LES): La predicción coincide con la estructura o encuentra una configuración que es energéticamente más estable que la referencia (según DFT o MLIP).

Resultados Clave
El método fue evaluado en dos conjuntos de datos derivados de la base de datos MC3D:

• Conjunto de Datos DFT: Estructuras con posiciones de hidrógeno conocidas que han sido relajadas mediante la Teoría del Funcional de la Densidad (DFT).

• Conjunto de Datos EXP: Las estructuras experimentales originales antes de la curación por DFT.

• Comparación de Rendimiento: El modelo pos-only-TD superó significativamente al MatterGen base, al modelo pos-only y al enfoque pos-only-RePaint.

  • RePaint mejoró la varianza pero aumentó el costo computacional (requiriendo ~1000 pasos frente a 300) sin aumentar significativamente la tasa de coincidencia.
  • El modelo pos-only-TD logró una tasa de coincidencia estructural del 88% en un subconjunto donde un método de "pinball" basado en DFT alcanzó aproximadamente el 87%, pero con un enfoque de optimización global más eficiente.

• Tasas de Éxito:

  • Al generar una sola muestra por estructura, el método logró altas tasas de coincidencia.
  • Al generar múltiples muestras (por ejemplo, $k=10$ o $k=30$) y seleccionar la configuración más estable, la tasa de coincidencia LES superó el 97% para el conjunto de datos DFT+MLIP y el 97.2% para el conjunto de datos EXP+MLIP.
  • Al excluir las estructuras marcadas como "teóricas" en las bases de datos de origen, la tasa de éxito en el conjunto de datos EXP superó el 99%.

• Estabilidad y Generalización:

  • Muchas predicciones "no coincidentes" resultaron ser energéticamente más estables que las estructuras de referencia originales, lo que sugiere que las referencias experimentales en las bases de datos eran subóptimas o incorrectas.
  • El modelo demostró generalización zero-shot a estructuras con hasta 40 átomos por celda unidad, a pesar de haber sido entrenado solo con estructuras de hasta 20 átomos. Aunque el rendimiento disminuyó ligeramente para sistemas más grandes, la tasa de coincidencia LES se mantuvo por encima del 90%.

• Eficiencia Computacional: El enfoque de difusión es significamente más rápido que el método de pinball basado en DFT, el cual sufre de explosión combinatoria y altos costos computacionales.

Significancia y Reivindicaciones
El artículo sostiene que la transferencia de técnicas de inpainting de imágenes (específicamente TD-Paint) a la ciencia de materiales permite completar las estructuras huésped de manera más rápida y precisa en comparación con los modelos de difusión no condicionados o los enfoques tradicionales basados en DFT.

• Curación de Bases de Datos: El método proporciona una herramienta robusta para curar y extender las bases de datos de materiales (como MC3D) mediante la reconstrucción fiable de las posiciones de hidrógeno faltantes, permitiendo así el cribado de alto rendimiento de materiales que contienen hidrógeno.
• Descubrimiento de Configuraciones Estables: El enfoque no solo reconstruye estructuras conocidas; a menudo identifica configuraciones que son más estables que las reportadas en las bases de datos experimentales, corrigiendo potencialmente errores históricos o inserciones heurísticas.
• Transferibilidad Agnóstica al Hidrógeno: Los modelos se entrenan de manera "agnóstica al hidrógeno". Los autores afirman que esta arquitectura puede transferirse fácilmente a otras tareas de inpainting, como la intercalación de iones (por ejemplo, Li, Na) en materiales de cátodo o la predicción de sitios de parada de muones, sin necesidad de reentrenamiento.
• Base para Trabajos Futuros: El artículo posiciona este método como una base que puede ajustarse para estructuras más grandes o aplicaciones específicas, en lugar de requerir un reentrenamiento completo.

Los autores enfatizan que su definición de éxito incluye tanto la recuperación estructural como el descubrimiento de configuraciones más estables, enmarcando esto último no como un fallo de reconstrucción, sino como la predicción exitosa de un estado físicamente más preciso.