AtomBench: A Benchmarking Framework for Generative Crystal Reconstruction Models in Conventional Superconductors

Resumo Técnico: AtomBench

Problema
O artigo aborda uma lacuna crítica na avaliação de modelos generativos para a ciência dos materiais: a falta de um framework padronizado e justo para avaliar tarefas de reconstrução de cristais (design inverso). Embora muitos modelos aleguem sucesso em geração de novo (avaliando estabilidade, unicidade e novidade), a capacidade de reconstruir com precisão uma estrutura cristalina conhecida a partir de informações parciais é frequentemente negligenciada.

Um problema primário identificado é que as comparações existentes frequentemente confundem o desempenho arquitetural com a assimetria de informação. Diferentes modelos são frequentemente avaliados sob condições em que possuem quantidades desiguais de informação sobre o cristal alvo (por exemplo, alguns recebem apenas a composição, enquanto outros recebem a composição mais um embedding latente da estrutura completa). Isso torna as comparações arquiteturais diretas inválidas. Além disso, métricas padrão como "taxa de correspondência" (match rate) e "RMSD correspondido" (matched RMSD) dependem de limiares rígidos (ex: pymatgen StructureMatcher), criando uma avaliação descontínua onde os modelos são pontuados apenas no subconjunto de estruturas que por acaso conseguem corresponder, o que pode mascarar falhas sistemáticas na maioria das previsões.

Metodologia
Os autores introduzem o AtomBench, um framework extensível e agnóstico a modelos, projetado para avaliar modelos generativos em uma tarefa de reconstrução de cristais bem definida utilizando supercondutores convencionais.

1. Datasets: O estudo utiliza dois conjuntos de dados gerados por DFT:

   • JARVIS Supercon-3D: 1.058 pares cristal/$T_c$, dominado por O, Al e Ti, com diversos sistemas cristalinos.
   • Alexandria DS-A/B: 8.253 pares cristal/$T_c$, dominado por Ti, Rh e Al, primariamente cúbico e tetragonal.
   • Ambos os datasets utilizam Teoria do Funcional da Densidade de Perturbação (DFPT) com protocolos específicos (ex: valores de $\mu^*$, pseudopotenciais).

2. Modelos Avaliados: Quatro arquiteturas generativas distintas foram treinadas e avaliadas:

   • AtomGPT: Um grande modelo de linguagem (LLM) ajustado (fine-tuned), o Mistral-7B, tratando cristais como texto.
   • CDVAE: Um autoencoder variacional de difusão usando redes neurais de grafos equivariantes SE(3).
   • FlowMM: Uma rede de correspondência de fluxo Riemanniano (Riemannian flow-matching).
   • MatterGen: Um modelo de difusão conjunta com redes de score equivariantes.
   • Variantes de Condicionamento: Para modelos que suportam tal função (AtomGPT, MatterGen), variantes foram treinadas com e sem condicionamento na temperatura crítica de supercondutividade ($T_c$).

3. Protocolo de Avaliação:

   • Divisão de Dados: 10% de cada dataset foi reservado para teste.
   • Regimes de Informação: Os modelos foram agrupados pelo nível de informação disponível antes da previsão:
     • Embedding de estrutura completa: CDVAE (recebe um vetor latente do alvo completo).
     • Composição + $T_c$: AtomGPT/MatterGen ($T_c$-condicionado).
     • Apenas composição: FlowMM, AtomGPT, MatterGen (não condicionado).
   • Métricas:
     • Parâmetros de Rede (Lattice Parameters): Erro Médio Absoluto (MAE) e Divergência de Kullback-Leibler (KLD) para distribuições.
     • Coordenadas Atômicas:
       • RMSD Correspondido (Matched RMSD): Erro médio sobre as estruturas que foram correspondidas com sucesso pelo StructureMatcher.
       • RMSD Corrigido Contínuo (ccRMSD): Uma métrica inovadora introduzida neste trabalho. Utiliza o descritor de Distância Mínima Média (AMD) e a distância de Chebyshev para medir a fidelidade geométrica local continuamente em todas as estruturas de teste, independentemente de uma correspondência discreta ser encontrada. Isso elimina o viés de descartar estruturas "difíceis de corresponder".

Principais Contribuições

1. Framework AtomBench: Um pacote Python de código aberto (atombench) que padroniza a avaliação de modelos de design inverso, suportando submissão direta ao JARVIS-Leaderboard.
2. Métrica ccRMSD: A introdução de uma métrica contínua, independente de taxa de correspondência, que fornece uma avaliação mais robusta da fidelidade geométrica local do que o RMSD correspondido tradicional.
3. Avaliação de Informação-Teórica: Um agrupamento rigoroso de modelos por regimes de "entropia condicional" (quantidade de informação conhecida) para garantir comparações arquiteturais justas.
4. Análise Sistemática: A primeira comparação em larga escala dessas arquiteturas específicas na reconstrução de supercondutores, analisando o impacto do condicionamento de $T_c$ e da simetria do sistema cristalino.

Resultos

• Desempenho por Métrica:
  • Coordenadas Atômicas: O MatterGen alcançou a melhor fidelidade de reconstrução (menor ccRMSD e menor RMSD correspondido), seguido pelo AtomGPT. CDVAE e FlowMM apresentaram erros significativamente maiores.
  • Parâmetros de Rede: O CDVAE foi o mais preciso na reconstrução de parâmetros de rede (menor MAE e KLD), seguido pelo AtomGPT e MatterGen. O FlowMM foi o menos preciso.
  • Velocidade vs. Precisão: O FlowMM foi o mais rápido na inferência, mas o menos preciso. O MatterGen ofereceu a maior precisão de coordenadas, porém com um custo computacional mais elevado.
• Impacto do Condicionamento de $T_c$: O condicionamento na temperatura crítica ($T_c$) não melhorou consistentemente a fidelidade de reconstrução para o AtomGPT ou MatterGen, sugerindo que um escalar único de propriedade adiciona pouca informação estrutural além da composição neste contexto.
• Tendências de Sistemas Cristalinos: O erro de reconstrução variou conforme o sistema cristalino. Sistemas cúbicos (alta simetria, menos graus de liberdade) exibiram os menores erros, enquanto sistemas monoclínicos e hexagonais mostraram erros maiores, particularmente no parâmetro de rede $c$.
• Discrepância de Métricas: O estudo destacou que o RMSD correspondido pode ser enganoso; por exemplo, o FlowMM teve uma taxa de correspondência muito baixa (<10%), tornando seu RMSD correspondido estatisticamente não confiável, enquanto o ccRMSD forneceu uma estimativa de erro significativa para todo o dataset.

Significância e Alegações
O artigo afirma que o AtomBench fornece uma correção necessária ao campo, expondo como a assimetria de informação pode inflar artificialmente as métricas de desempenho. Ao introduzir o ccRMSD, os autores defendem uma mudança do paradigma de "corresponder ou descartar" para medidas de fidelidade contínas que considerem toda a distribuição das previsões.

Os autores afirmam modestamente que seus achados não classificam universalmente as arquiteturas, mas sim ilustram como a fidelidade de reconstrução varia sob diferentes regimes de informação de inferência. Eles enfatizam que a escolha do modelo deve depender da aplicação específica:

• FlowMM para amostragem de alto rendimento (high-throughput) onde a velocidade é primordial.
• CDVAE para tarefas que exigem alta precisão nos parâmetros de rede.
• MatterGen para tarefas que exigem alta fidelidade no posicionamento de coordenadas atômicas.

O trabalho conclui que, embora os modelos atuais (CDVAE, AtomGPT, MatterGen) alcancem níveis de precisão adequados para inicializar cálculos de DFT, o campo requer métricas contínuas como o ccRMSD para avaliar e melhorar adequadamente os modelos generativos para a descoberta de materiais.