Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory

Durante más de sesenta años, los científicos han dependido de un poderoso conjunto de reglas para comprender cómo los átomos se unen para formar los materiales que nos rodean. Este marco, conocido como la teoría del funcional de la densidad, trata la nube de electrones que rodea a un átomo no como una colección de partículas individuales, sino como una densidad continua y suave. Este enfoque permite a los investigadores predecir las propiedades de todo, desde nuevas medicinas hasta baterías avanzadas. Sin embargo, hay un inconveniente. Para obtener resultados precisos, la teoría requiere un cálculo complejo y repetitivo que implica resolver el comportamiento de cada uno de los electrones en el sistema. A medida que el sistema crece, el tiempo requerido para estos cálculos se dispara, creciendo tan rápido que simular una pieza de metal de tamaño modesto o una molécula biológica grande se vuelve imposible incluso en las supercomputadoras más potentes. Este cuello de botella ha mantenido muchas preguntas científicas importantes fuera del alcance, obligando a los investigadores a elegir entre precisión o escala.

Un equipo de investigadores del Instituto de Tecnología de California ha encontrado una manera de romper esta barrera enseñando a una computadora a saltarse la parte más costosa del cálculo sin perder precisión. En lugar de intentar predecir la respuesta final directamente o aprender un atajo matemático difícil que a menudo falla, le enseñaron a una inteligencia artificial a realizar el paso específico y repetido que ocurre en cada ciclo de cálculo. Al centrarse en este paso intermedio, crearon un método que puede simular sistemas con decenas de miles de átomos en una sola tarjeta gráfica, una tarea que anteriormente requería miles de procesadores. El resultado es una nueva forma de modelar la materia que es casi tan rápida como las aproximaciones más simples, pero que conserva la alta precisión de los métodos más rigurosos.

El núcleo del problema radica en cómo la teoría maneja los electrones. Para encontrar el estado estable de un material, la computadora debe ajustar repetidamente un mapa de la densidad electrónica hasta que deje de cambiar. En el método tradicional, este ajuste requiere resolver un problema de autovalores masivo, que es esencialmente encontrar los patrones específicos de movimiento de los electrones para cada átomo en el sistema. Esta operación es tan computacionalmente pesada que su costo crece cúbicamente con el número de electrones; duplicar el tamaño del sistema hace que el cálculo sea ocho veces más difícil. Durante décadas, los científicos han intentado evitar esto eliminando la necesidad de rastrear los patrones de electrones individuales por completo, una estrategia conocida como la teoría del funcional de la densidad libre de orbitales. Sin embargo, los intentos previos de aprender este atajo utilizando aprendizaje automático han tenido dificultades. Algunos intentaron aprender la relación entre la densidad electrónica y la energía, pero esta relación es matemáticamente inestable y propensa a errores. Otros intentaron predecir la densidad electrónica final directamente a partir de la disposición de los átomos, pero estos modelos fallaron cuando se les pidió predecir sistemas más grandes o químicamente diferentes de los que fueron entrenados.

Los investigadores se dieron cuenta de que la clave era dejar de intentar predecir el destino final y, en su lugar, aprender el viaje. En cada paso del cálculo tradicional, la computadora toma un mapa de las fuerzas que actúan sobre los electrones y produce un nuevo mapa de dónde es probable que se encuentren los electrones. Este paso hacia adelante es matemáticamente estable y bien definido. El equipo entrenó una red neuronal especializada, un tipo de inteligencia artificial diseñada para comprender patrones espaciales, para aprender exactamente esta transformación. Alimentaron la red con miles de ejemplos de estos mapas de fuerzas y los correspondientes mapas de densidad electrónica que producían, extraídos de un conjunto diverso de 8,504 moléculas y materiales sólidos, incluyendo compuestos orgánicos, aislantes y metales. La red aprendió a predecir la nueva densidad electrónica directamente a partir de las fuerzas, reemplazando efectivamente el cálculo lento y pesado con una predicción rápida y aprendida.

Lo que hace que este enfoque sea único es cómo maneja la complejidad de diferentes sistemas. Los investigadores diseñaron la red para que sea invariante al dominio, lo que significa que puede aplicar las mismas reglas aprendidas a una molécula diminuta y a un cristal masivo sin necesidad de ser reentrenada para cada nuevo tamaño. También se aseguraron de que la red respetara las simetrías fundamentales del espacio, de modo que entiende que rotar o desplazar una molécula no cambia sus propiedades físicas. Al ser probada en sistemas que nunca había visto, el modelo demostró ser notablemente robusto. En un conjunto de moléculas de tipo fármaco que contenían elementos y tamaños no presentes en los datos de entrenamiento, el nuevo método mantuvo una alta precisión, mientras que los modelos de predicción directa anteriores vieron cómo sus errores se disparaban a medida que las moléculas crecían. La naturaleza autoconsistente del método le permite corregir sus propios errores en cada paso, de forma muy similar a un navegante que comprueba su posición con frecuencia en lugar de adivinar toda la ruta de una vez.

El poder de este método quedó claro cuando el equipo lo aplicó a sistemas periódicos, como los cristales, y a defectos extendidos en metales. En los cálculos estándar para cristales, la computadora debe repetir el cálculo pesado para muchos puntos diferentes en el espacio matemático que describe el patrón repetitivo del cristal. El nuevo método se salta esta repetición por completo, evolucionando la densidad electrónica en la celda unidad sin necesidad de muestrear estos puntos individualmente. Esto permitió a los investigadores converger cálculos para sistemas metálicos y semiconductores con un solo modelo, reproduciendo los espectros electrónicos y las propiedades estructurales con la misma precisión que el método tradicional. Además, el modelo mostró una capacidad sorprendente para transferirse entre diferentes aproximaciones teóricas. Un modelo entrenado en un tipo de cálculo podía utilizarse con una aproximación diferente y más precisa para las fuerzas sin necesidad de reentrenamiento, simplemente intercambiando las fuerzas de entrada.

La prueba definitiva llegó cuando los investigadores aplicaron el método a una dislocación de magnesio, un tipo de defecto en un cristal metálico que es crítico para entender cómo los materiales se deforman y se rompen. Estos defectos crean campos elásticos de largo alcance que requieren simulaciones de miles de átomos para modelarlos con precisión. Los intentos previos de simular tal sistema con alta precisión requerían una supercomputadora con miles de unidades de procesamiento gráfico y tardaban días en ejecutarse. El nuevo método, ejecutándose en una sola tarjeta gráfica moderna, logró converger el cálculo para un sistema que contenía 8,250 átomos de magnesio y 82,500 electrones de valencia. El tiempo que tomó completar el cálculo creció de forma casi lineal con el tamaño del sistema, confirmando que el cuello de botella computacional había sido eliminado. Las densidades electrónicas resultantes fueron indistinguibles de la referencia de alta precisión, lo que permitió a los investigadores resolver las sutiles diferencias de energía entre diferentes tipos de defectos que determinan la resistencia del metal.

Este trabajo demuestra que el aprendizaje automático puede extender el alcance de los cálculos de estructura electrónica no reemplando la teoría física completa, sino acelerando su operación más costosa y repetida. Al aprender el mapa directo de las fuerzas a la densidad, los investigadores han creado una herramienta que es estable, transferible a diferentes tipos de materia y capaz de manejar sistemas a una escala que antes era inaccesible. El método proporciona un diagnóstico directo de su propia fiabilidad; si el cálculo no logra converger, señala que el sistema está fuera del rango de confianza del modelo, una característica que carecen los modelos de predicción directa. Aunque el modelo actual predice la densidad electrónica, la energía total y otras propiedades aún pueden obtenerse con un único paso final, preservando la precisión de la teoría completa. Este enfoque abre la puerta a la simulación de materiales complejos, sistemas biológicos y reacciones químicas a un nivel de detalle y escala que antes era el dominio de la imposibilidad teórica.