Machine-learned, finite temperature Fermi-operator expansions suitable for GPUs and AI-hardware

Resumo Técnico: Expansões de Operador de Fermi de Temperatura Finita Aprendidas por Máquina

Declaração do Problema
Os cálculos de estrutura eletrônica, particularmente dentro da Teoria do Funcional da Densidade de Kohn-Sham (KS-DFT), são limitados computacionalmente pelo custo de escalonamento cúbico da diagonalização da matriz Hamiltoniana para resolver o problema de autovalores. Embora esquemas recursivos de expansão de operador de Fermi, como o método de Projeção Espectral de Segunda Ordem (SP2), ofereçam uma maneira de calcular a matriz densidade diretamente sem diagonalização, implementações eficientes existentes estão restritas à temperatura eletrônica zero. A temperatura zero, a matriz densidade é idempotente (ocupações são estritamente 0 ou 1). No entanto, muitos sistemas físicos — como metais ou sistemas em temperaturas eletrônicas elevadas — requerem ocupações orbitais fracionárias para modelar com precisão estados próprios degenerados ou alisamento térmico.

Tentativas anteriores de generalizar o SP2 para temperaturas finitas envolveram o truncamento da recursão para introduzir alisamento térmico. No entanto, essas expansões truncadas são inerentemente aproximadas, falhando em reproduzir a função de Fermi exata, particularmente próximo ao potencial químico, onde a precisão é crítica. Métodos alternativos, como expansões de Chebyshev ou aproximantes de Padé, exigem ordens polinomiais proibitivamente altas para suprimir oscilações de Gibbs ou incorrem em sobrecarga computacional significativa devido à resolução repetida de sistemas lineares.

Metodologia
Os autores propõem uma estrutura que generaliza o método recursivo SP2 para temperaturas finitas, mapeando sua estrutura algébrica sobre arquiteturas de redes neurais profundas (DNN). A ideia central é que as atualizações recursivas do SP2 assemelham-se às camadas de uma rede neural. Ao tratar os coeficientes de expansão como pesos e vieses treináveis, os autores constroem modelos de aprendizado de máquina capazes de aproximar a função de distribuição de Fermi com ocupações fracionárias em temperaturas arbitrárias.

Componentes metodológicos-chave incluem:

1. Arquiteturas de Rede Neural:

   • MLSP2 (SP2 Aprendido por Máquina): Uma generalização do SP2 onde as regras de atualização quadrática ($X^2$ ou $2X-X^2$) são substituídas por polinômios quadráticos aprendíveis ($ax^2 + bx + c$) com um termo acumulador. Isso permite que o modelo aproxime o alisamento térmico exato da função de Fermi em vez de uma função degrau truncada.
   • Max-SP2: Uma arquitetura mais expressiva incorporando "conexões de salto" (skip connections), onde cada camada é o quadrado de uma combinação linear de todas as camadas anteriores.
   • Skip-SP2: Uma versão comprimida do Max-SP2 usando uma memória finita de camadas recentes e acumuladores para equilibrar expressividade e uso de memória.

2. Aproximação de Entropia:
   Os autores também desenvolvem um esquema recursivo para aproximar a função de entropia eletrônica, $s(x)$, necessária para calcular a energia livre eletrônica. Eles utilizam um produto escalado da função de Fermi e seu complemento, $f(x)(1-f(x))$, como uma estimativa inicial, que é então refinada via uma expansão quadrática recursiva treinada para corresponder à segunda derivada da entropia verdadeira no potencial químico.

3. Treinamento e Otimização:
   Os modelos são treinados em entradas escalares dentro do intervalo unitário $[0, 1]$ em vez de matrizes completas, usando o algoritmo de Levenberg–Marquardt com aceleração geodésica. Os dados de treinamento são amostrados com um peso proporcional à derivada da função de Fermi para minimizar o erro máximo próximo ao potencial químico.

4. Redimensionamento Afim e Transferibilidade:
   Uma inovação crítica é o uso de redimensionamento afim para eliminar a necessidade de retreinamento quando os parâmetros de simulação mudam. Ao normalizar o Hamiltoniano ($H'$), o potencial químico ($\mu'$) e a temperatura inversa ($\beta'$), um único modelo treinado em parâmetros específicos $(\beta_0, \mu_0)$ pode ser aplicado a uma ampla "região de validade" de outros parâmetros. Isso é alcançado redimensionando o Hamiltoniano de entrada para corresponder às condições de treinamento, permitindo que o mesmo conjunto de pesos seja usado em temperaturas e potenciais químicos variados durante uma simulação.

5. Implementação de Hardware:
   Os algoritmos são otimizados para GPUs modernas e hardware de IA (especificamente Tensor Cores da NVIDIA). Os autores aproveitam a aritmética de precisão mista (FP16/FP32) para realizar operações de quadrado de matriz eficientemente, explorando a simetria do Hamiltoniano para reduzir o número de multiplicações e transferências de dados necessárias.

Resultados Principais

• Precisão: Os modelos MLSP2 alcançam erros na ordem de $10^{-7}$ para a aproximação da função de Fermi, superando significativamente os métodos SP2 truncados (que têm erros em torno de $10^{-2}$) e igualando a precisão da diagonalização de dupla precisão em muitos regimes.
• Desempenho: Em uma GPU NVIDIA RTX 6000 Ada, a abordagem MLSP2 demonstra um aumento de velocidade de 16 vezes sobre a diagonalização de dupla precisão (usando cuSOLVER) para tamanhos de matriz intermediários e um aumento de 9 vezes para matrizes maiores. Mesmo comparado à diagonalização de precisão simples, o MLSP2 oferece um aumento de velocidade de 2x a 5x, mantendo estabilidade e precisão superiores.
• Escalabilidade: O método depende exclusivamente de kernels altamente otimizados de multiplicação de matriz-matriz, evitando a diagonalização explícita. O número de camadas necessário para alcançar uma precisão-alvo escala logaritmicamente com a temperatura inversa ($\beta$), permitindo computação eficiente mesmo em temperaturas baixas.

Significado e Alegações
O artigo afirma que esta abordagem fornece uma solução robusta e generalizável para cálculos de estrutura eletrônica em temperatura finita que evita os gargalos computacionais da diagonalização. Ao generalizar o SP2 por meio de aprendizado de máquina, os autores permitem o cálculo de matrizes densidade para sistemas com ocupações fracionárias a uma fração do custo dos métodos tradicionais.

O significado reside na capacidade de realizar simulações dinâmicas de temperatura finita (como dinâmica molecular quântica) onde o potencial químico e a temperatura eletrônica flutuam entre os passos de tempo. A estratégia de redimensionamento afim garante que um único modelo pré-treinado possa ser reutilizado ao longo de uma simulação sem retreinamento, tornando o método prático para aplicações em grande escala. Além disso, a abordagem é especificamente adaptada para explorar as características de desempenho do hardware de IA moderno (Tensor Cores), oferecendo aumentos de velocidade substanciais sobre rotinas de diagonalização otimizadas por fornecedores, mantendo alta precisão numérica.