Expander attention as exchange-correlation

Resumo Técnico: Expander Attention como Troca-Correlação

Enunciado do Problema
A teoria do funcional da densidade de Kohn-Sham (DFT) é o padrão para cálculos de estrutura eletrônica devido ao seu equilíbrio entre precisão e custo computacional. No entanto, sua utilidade prática depende de aproximações para o funcional de troca-correlação (XC) desconhecido. Embora existam muitas Aproximações de Funcional de Densidade (DFAs), elas lutam com sistemas fortemente correlacionados, como a curva de dissociação do hidrogênio ou H4 planar, frequentemente falhando em capturar a energetica correta. DFAs aprendidas por máquina (ML) emergiram como uma alternativa promissora para abordar essas limitações, aprendendo interações não locais. No entanto, um gargalo persistente permanece: funcionais ML de alta precisão capazes de capturar fortes correlações tipicamente sofrem com escalas computacionais desfavoráveis (por exemplo, $O(N^2)$ ou $O(N^4)$), tornando-os proibitivamente caros para aplicações em grande escala.

Metodologia
Os autores propõem Exphormer-XC, uma aproximação XC não local com escala linear baseada em um ansatz de transformador de grafo expansor. A metodologia envolve os seguintes componentes-chave:

1. Construção de Grafos em Grades Computacionais: Em vez de usar grafos moleculares (onde os nós são núcleos), a abordagem constrói um grafo diretamente na grade eletrônica computacional (grade de Becke) usada na DFT. O grafo $G$ consiste em vértices $V_{grid}$ representando pontos da grade e um pequeno conjunto de nós globais fictícios $V_{global}$.
2. Definições de Arestas: As arestas do grafo são definidas em três categorias para garantir escala linear enquanto mantêm a conectividade:
   • Arestas Locais ($E_{local}$): Conectam vizinhos radiais mais próximos e vizinhos angulares dentro de cascas de Lebedev com base na distância de Haversine.
   • Arestas Expansoras ($E_{exp}$): Utilizam um esquema de Friedman simplificado para criar uma estrutura de grafo esparsa e altamente conectada. Isso permite que o grafo tenha uma contagem de arestas com escala linear enquanto mantém um grande gap espectral (critério de Ramanujan), facilitando a propagação eficiente de informações através da grade.
   • Arestas Globais ($E_{global}$): Conectam um número fixo e pequeno de nós de reservatório globais a todos os nós da grade.
3. Arquitetura Neural: Um transformador multicamadas e multi-cabeça processa o grafo. As características dos nós de entrada incluem densidade eletrônica ($n$) e polarização de spin ($\zeta$). As características das arestas incluem distância euclidiana e tipo de aresta (local, expansora ou global).
4. Formulação do Funcional XC: O transformador produz um fator de melhoria $F_{exp}$ aplicado a uma densidade de energia XC local base $\epsilon_{XC}$. O funcional final é $\tilde{\epsilon}_{XC} = \epsilon_{XC}(1 + \beta F_{exp})$, onde $\beta$ é um parâmetro aprendível inicializado em zero para garantir transições suaves a partir da DFA base.
5. Estrutura de Treinamento: O modelo é treinado de forma autoconsistente dentro de um solucionador KS diferenciável (estendendo o pacote DQC), usando dados de Interação de Configuração Completa (FCI) como verdade fundamental.

Principais Resultados
O artigo avalia o Exphormer-XC em dois sistemas fortemente correlacionados de referência:

• Curva de Dissociação do Hidrogênio: O modelo recupera com sucesso a curva de dissociação correta para a molécula de H2, um regime onde DFAs semilocais e híbridas falham. Ao treinar em uma faixa de geometrias (fator de escala $S=1$ a $4.5$), o modelo alcança erros absolutos médios (MAE) inferiores a 1 kcal/mol no regime interpolativo.
• Estudo de Ablação: Os autores demonstram que todos os componentes da arquitetura são críticos. Especificamente:
  • Modelos puramente locais (NN-LDA) e convoluções de grafos padrão falham em capturar a curva.
  • A remoção de arestas expansoras ou incorporações de distância degrada significativamente o desempenho.
  • Embora nós globais não sejam estritamente necessários para atingir o limiar de precisão, sua exclusão atrasa substancialmente a convergência do treinamento (em ~21%).
• Sistema H4 Planar: O modelo é aplicado ao H4 planar próximo a uma configuração quadrada, um sistema conhecido por forte correlação estática e quase-degeneração.
  • DFAs padrão (por exemplo, PBE) preveem incorretamente um cume de energia agudo, enquanto a FCI prevê uma barreira parabólica.
  • O Exphormer-XC (restrito) captura a forma parabólica correta e energias mais próximas da FCI do que outras DFAs.
  • Limitação: O modelo exibe problemas de convergência (saltos estocásticos entre estados singleto e tripleto) próximo à configuração quadrada devido à quase-degeneração. Os autores observam que, embora o modelo capture a energetica de ambos os estados, o solucionador diferenciável usado não pode impor explicitamente a quebra de simetria para estabilizar o cálculo, uma capacidade presente em códigos padrão de FCI, mas ainda não no atual framework diferenciável.

Significado e Alegações
O artigo afirma apresentar a primeira DFA-ML com escala linear capaz de capturar com precisão a curva de dissociação do hidrogênio. A contribuição principal é a arquitetura Exphormer-XC, que melhora a escala de funcionais ML anteriores de $O(N^2)$ ou pior para escala linear ($O(N)$), mantendo a não localidade necessária para sistemas fortemente correlacionados.

Os autores argumentam que essa abordagem traça um caminho para funcionais ML que são tanto precisos para sistemas correlacionados difíceis quanto computacionalmente baratos o suficiente para escala. Eles enfatizam que a construção de grafos expansores é essencial para alcançar esse equilíbrio, pois topologias de grafos mais simples falham em convergir ou carecem da expressividade necessária. Embora o trabalho atual seja limitado a sistemas de teste específicos (H2 e H4) e enfrente desafios de convergência em regimes degenerados sem quebra de simetria explícita, os resultados sugerem que funcionais ML não locais com escala linear são uma alternativa viável à escala pobre de métodos anteriores de alta precisão.