Score-based diffusion models for accurate crystal-structure inpainting and reconstruction of hydrogen positions

Resumo Técnico: Modelos de difusão baseados em score para inpainting e reconstrução precisa de estruturas cristalinas de hidrogênio

Definição do Problema
A determinação precisa das posições atômicas para átomos de hidrogênio em materiais cristalinos permanece um desafio significativo na ciência dos materiais computacional. Embora o hidrogênio seja abundante e crítico para aplicações que variam desde o armazenamento de energia até a indústria de semicondutores, seu baixo poder de espalhamento de raios X e elétrons torna difícil localizá-lo via difração de raios X padrão. Consequentemente, os bancos de dados cristalográficos históricos (como o MC3D, derivado do COD, ICSD e MPDS) frequentemente contêm estruturas onde as posições do hidrogênio são omitidas, inseridas via heurísticas ou determinadas com incerteza significativa. Embora avanços experimentais (por exemplo, cristalografia quântica) e espalhamento de nêutrons caros possam resolver essas posições, os bancos de dados existentes ainda dependem fortemente de dados incompletos ou aproximados. O problema central abordado é a reconstrução dessas posições de hidrogênio ausentes dado uma estrutura cristalina "hospedeira" conhecida (a estrutura de hidrogênio não-hidrogênio) para gerar estruturas cristalinas precisas, estáveis e completas.

Metodologia
Os autores propõem uma abordagem generativa combinando modelos de difusão baseados em score com técnicas adaptadas da inpainting de imagens de visão computacional.

1. Arquitetura do Modelo: O estudo utiliza o MatterGen da Microsoft, um modelo de difusão baseado em score pré-treinado, originalmente projetado para gerar novas estruturas cristalinas. Os autores retreinam e adaptam este modelo para a tarefa específica de inpainting (preenchimento de dados ausentes) em vez de geração total.
2. Estratégias de Inpainting: Diversas abordagens são testadas:
   • MatterGen Baseline: O modelo original aplicado à tarefa com rede e tipos atômicos fixos.
   • pos-only: Um modelo retreinado para realizar o denoising apenas das posições atômicas, mantendo a rede e os tipos fixos.
   • pos-only-RePaint: O modelo retreinado combinado com o algoritmo RePaint, que reamostra iterativamente regiões conhecidas para harmonizá-las com as regiões preenchidas (inpainted).
   • pos-only-TD (Time-Dependent): Uma nova adaptação do conceito TD-Paint da visão computacional. Este modelo é treinado para lidar com níveis de ruído variáveis por sítio dentro de uma única estrutura. Ele condiciona o processo de denoising em posições "limpas" conhecidas (ruído zero) enquanto prevê iterativamente os sítios de hidrogênio ausentes (ruído alto). Isso permite que o modelo utilize diretamente a estrutura hospedeira conhecida sem adicionar ruído a ela.
3. Fluxo de Trabalho:
   • Partir de uma estrutura hospedeira (sítios de hidrogênio removidos).
   • Inicializar aleatoriamente $N_H$ sítios de hidrogênio na célula unitária.
   • Aplicar o modelo de difusão pos-only-TD por $N_{steps}$ iterações para prever as posições do hidrogênio.
   • Realizar uma relaxação restrita dos sítios de hidrogênio usando um Potencial Interatômico de Aprendizado de Máquina (MLIP), especificamente o NequIP, mantendo os átomos hospedeiros e a célula fixos.
   • Realizar uma relaxação total de todas as posições (incluindo átomos hospedeiros) usando NequIP.
4. Métricas de Avaliação: O sucesso é definido de duas maneiras:
   • Correspondência Estrutural: A estrutura prevista coincide com a referência (da base de dados) dentro das tolerâncias do StructureMatcher do pymatgen.
   • Correspondência LES (Lower Energy or Structural): A previsão ou coincide com a estrutura ou encontra uma configuração que é energeticamente mais estável que a referência (de acordo com DFT ou MLIP).

Principais Resultados
O método foi avaliado em dois conjuntos de dados derivados do banco de dados MC3D:

• Conjunto de Dados DFT: Estruturas com posições de hidrogênio conhecidas que foram relaxadas usando Teoria do Funcional da Densidade (DFT).

• Conjunto de Dados EXP: As estruturas experimentais originais antes da curadoria por DFT.

• Comparação de Desempenho: O modelo pos-only-TD superou significativamente o baseline MatterGen, o modelo pos-only e a abordagem pos-only-RePaint.

  • O RePaint melhorou a variância, mas aumentou o custo computacional (exigindo ~1000 passos contra 300) sem aumentar significativamente a taxa de correspondência.
  • O modelo pos-only-TD alcançou uma taxa de correspondência estrutural de 88% em um subconjunto onde um método "pinball" baseado em DFT atingiu ~87%, mas com uma abordagem de otimização global mais eficiente.

• Taxas de Sucesso:

  • Ao gerar uma única amostra por estrutura, o método alcançou altas taxas de correspondência.
  • Ao gerar múltipl de amostras (ex: $k=10$ ou $k=30$) e selecionar a configuração mais estável, a taxa de correspondência LES excedeu 97% para o conjunto de dados DFT+MLIP e 97,2% para o conjunto de dados EXP+MLIP.
  • Ao excluir estruturas marcadas como "teóricas" nos bancos de dados de origem, a taxa de sucesso no conjunto de dados EXP excedeu 99%.

• Estabilidade e Generalização:

  • Muitas previsões "não correspondentes" foram encontradas como sendo energeticamente mais estáveis do que as estruturas de referência originais, sugerindo que as referências experimentais nos bancos de dados eram subótimas ou incorretas.
  • O modelo demonstrou generalização zero-shot para estruturas com até 40 átomos por célula unitária, apesar de ter sido treinado apenas em estruturas com até 20 átomos. Embora o desempenho tenha caído ligeiramente para sistemas maiores, a taxa de correspondência LES permaneceu acima de 90%.

• Eficiência Computacional: A abordagem de difusão é significativamente mais rápida que o método pinball baseado em DFT, que sofre com explosão combinatória e altos custos computacionais.

Significância e Alegações
O artigo afirma que a transferência de técnicas de inpainting de imagem (especificamente TD-Paint) para a ciência dos materiais permite uma conclusão de estruturas hospedeiras mais rápida e precisa do que modelos de difusão não condicionados ou abordagens tradicionais baseadas em DFT.

• Curadoria de Bancos de Dados: O método fornece uma ferramenta robusta para curar e estender bancos de dados de materiais (como o MC3D) ao reconstruir confiavelmente as posições de hidrogênio ausentes, permitindo assim a triagem de alto rendimento de materiais contendo hidrogênio.
• Descoberta de Configurações Estáveis: A abordagem não apenas reconstrói estruturas conhecidas; ela frequentemente identifica configurações que são mais estáveis do que as relatadas em bancos de dados experimentais, potencialmente corrigindo erros históricos ou inserções heurísticas.
• Transferibilidade Agnóstica ao Hidrogênio: Os modelos são treinados de maneira "agnóstica ao hidrogênio". Os autores afirmam que esta arquitetura pode ser facilmente transferida para outras tarefas de inpainting, como a intercalação de íons (ex: Li, Na) em materiais de cátodo ou a previsão de sítios de parada de múons, sem a necessidade de retreinamento.
• Base para Trabalhos Futuros: O artigo posiciona este método como uma base que pode ser ajustada (fine-tuned) para estruturas maiores ou aplicações específicas, em vez de exigir um retreinamento completo.

Os autores enfatizam que sua definição de taxa de sucesso inclui tanto a recuperação estrutural quanto a descoberta de configurações mais estáveis, enquadrando esta última não como uma falha de reconstrução, mas como a previsão bem-sucedida de um estado fisicamente mais preciso.