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

Por mais de sessenta anos, os cientistas têm confiado em um poderoso conjunto de regras para entender como os átomos se unem para formar os materiais ao nosso redor. Essa estrutura, conhecida como teoria do funcional da densidade, trata a nuvem eletrônica que envolve um átomo não como uma coleção de partículas individuais, mas como uma densidade contínua e suave. Essa abordagem permite que pesquisadores prevejam as propriedades de tudo, desde novos medicamentos até baterias avançadas. No entanto, há um problema. Para obter resultados precisos, a teoria exige um cálculo complexo e repetitivo que envolve resolver o comportamento de cada único elétron no sistema. À medida que o sistema cresce, o tempo necessário para esses cálculos explode, crescendo tão rápido que simular um pedaço modestamente dimensionado de metal ou uma grande molécula biológica torna-se impossível mesmo nos supercomputadores mais poderosos. Esse gargalo manteve muitas questões científicas importantes fora de alcance, forçando os pesquisadores a escolher entre precisão ou escala.

Uma equipe de pesquisadores do Instituto de Tecnologia da Califórnia encontrou uma maneira de romper essa barreira, ensinando um computador a pular a parte mais cara do cálculo sem perder a precisão. Em vez de tentar prever a resposta final diretamente ou aprender um atalho matemático difícil que frequentemente falha, eles ensinaram uma inteligência artificial a realizar a etapa específica e repetida que ocorre em cada ciclo de cálculo. Ao focar nesta etapa intermediária, criaram um método que pode simular sistemas com dezenas de milhares de átomos em uma única placa de vídeo, uma tarefa que anteriormente exigia milhares de processadores. O resultado é uma nova forma de modelar a matéria que é quase tão rápida quanto as aproximações mais simples, mas mantém a alta precisão dos métodos mais rigorosos.

O cerne do problema reside em como a teoria lida com os elétrons. Para encontrar o estado estável de um material, o computador deve ajustar repetidamente um mapa da densidade eletrônica até que ela pare de mudar. No método tradicional, esse ajuste requer a resolução de um enorme problema de autovalores, que é essencialmente encontrar os padrões específicos de movimento dos elétrons para cada átomo no sistema. Essa operação é tão computacionalmente pesada que seu custo cresce cubicamente com o número de elétrons; dobrar o tamanho do sistema torna o cálculo oito vezes mais difícil. Por décadas, os cientistas tentaram contornar isso eliminando a necessidade de rastrear padrões de elétrons individuais inteiramente, uma estratégia conhecida como teoria do funcional da densidade livre de orbitais. No entanto, tentativas anteriores de aprender esse atalho usando aprendizado de máquina tiveram dificuldades. Alguns tentaram aprender a relação entre a densidade eletrônica e a energia, mas essa relação é matematicamente instável e propensa a erros. Outros tentaram prever a densidade eletrônica final diretamente a partir do arranjo dos átomos, mas esses modelos falharam quando confrontados com sistemas maiores ou quimicamente diferentes daqueles para os quais foram treinados.

Os pesquisadores perceberam que a chave era parar de tentar prever o destino final e, em vez disso, aprender a jornada. Em cada etapa do cálculo tradicional, o computador pega um mapa das forças que atuam sobre os elétrons e produz um novo mapa de onde os elétrons provavelmente serão encontrados. Este passo à frente é matematicamente estável e bem definido. A equipe treinou uma rede neural especializada, um tipo de inteligência artificial projetada para entender padrões espaciais, para aprender exatamente essa transformação. Eles alimentaram a rede com milhares de exemplos desses mapas de forças e os correspondentes mapas de densidade eletrônica que eles produziam, extraídos de um conjunto diversificado de 8.504 moléculas e materiais sólidos, incluindo compostos orgânicos, isolantes e metais. A rede aprendeu a prever a nova densidade eletrônica diretamente das forças, substituindo efetivamente o cálculo lento e pesado por uma previsão rápida e aprendida.

O que torna essa abordagem única é como ela lida com a complexidade de diferentes sistemas. Os pesquisadores projetaram a rede para ser invariante ao domínio, o que significa que ela pode aplicar as mesmas regras aprendidas a uma molécula minúscula e a um cristal massivo sem a necessidade de ser treinada novamente para cada novo tamanho. Eles também garantiram que a rede respeitasse as simetrias fundamentais do espaço, de modo que ela entenda que rotacionar ou deslocar uma molécula não altera suas propriedades físicas. Quando testado em sistemas que nunca havia visto antes, o modelo provou ser notavelmente robusto. Em um conjunto de moléculas do tipo fármaco contendo elementos e tamanhos não presentes nos dados de treinamento, o novo método manteve alta precisão, enquanto os modelos de previsão direta anteriores viram seus erros dispararem conforme as moléculas cresciam. A natureza autoconsequente do método permite que ele corrija seus próprios erros em cada etapa, tal como um navegador que verifica sua posição frequentemente, em vez de apenas adivinhar toda a rota de uma só vez.

O poder deste método ficou claro quando a equipe o aplicou a sistemas periódicos, como cristais, e a defeitos estendidos em metais. Nos cálculos padrão para cristais, o computador deve repetir o cálculo pesado para muitos pontos diferentes no espaço matemático que descreve o padrão repetitivo do cristal. O novo método pula essa repetição inteiramente, evoluindo a densidade eletrônica na célula unitária sem a necessidade de amostrar esses pontos individualmente. Isso permitiu que os pesquisadores convergissem cálculos para sistemas metálicos e semicondutores com um único modelo, reproduzindo os espectros eletrônicos e as propriedades estruturais com a mesma precisão do método tradicional. Além disso, o modelo mostrou uma capacidade surpreendente de transferência entre diferentes aproximações teóricas. Um modelo treinado em um tipo de cálculo poderia ser usado com uma aproximação diferente e mais precisa para as forças sem qualquer retreinamento, simplesmente trocando as forças de entrada.

O teste definitivo veio quando os pesquisadores aplicaram o método a uma dislocação de magnésio, um tipo de defeito em um cristal metálico que é crucial para entender como os materiais se deformam e quebram. Esses defeitos criam campos elásticos de longo alcance que exigem simulações de milhares de átomos para serem modelados com precisão. Tentativas anteriores de simular tal sistema com alta precisão exigiam um supercomputador com milhares de unidades de processamento gráfico e levavam dias para rodar. O novo método, rodando em uma única placa de vídeo moderna, convergiu com sucesso o cálculo para um sistema contendo 8.250 átomos de magnésio e 82.500 elétrons de valência. O tempo para completar o cálculo cresceu quase linearmente com o tamanho do sistema, confirmando que o gargalo computacional havia sido removido. As densidades eletrônicas resultantes foram indistinguíveis da referência de alta precisão, permitindo que os pesquisadores resolvessem as sutis diferenças de energia entre diferentes tipos de defeitos que determinam a resistência do metal.

Este trabalho demonstra que o aprendizado de máquina pode estender o alcance dos cálculos de estrutura eletrônica não substituindo toda a teoria física, mas acelerando sua operação mais cara e repetida. Ao aprender o mapa direto das forças para a densidade, os pesquisadores criaram uma ferramenta que é estável, transferível entre diferentes tipos de matéria e capaz de lidar com sistemas em uma escala que antes era inacessível. O método fornece um diagnóstico direto de sua própria confiabilidade; se o cálculo falhar em convergir, isso sinaliza que o sistema está fora do intervalo de confiança do modelo, um recurso que os modelos de previsão direta não possuem. Embora o modelo atual preveja a densidade eletrônica, a energia total e outras propriedades ainda podem ser obtidas com um único passo final, preservando a precisão da teoria completa. Esta abordagem abre as portas para simular materiais complexos, sistemas biológicos e reações químicas em um nível de detalhe e escala que antes era o domínio da impossibilidade teórica.