Discovering Crystal Structure Prediction Algorithms with an AI Co-Scientist

Resumo Técnico: Descoberta de Algoritmos de Predição de Estrutura Cristalina com um Co-Cientista de IA

Definição do Problema
A Predição de Estrutura Cristalina (CSP - Crystal Structure Prediction) é um desafio central na descoberta de materiais computacionais, exigindo a geração de estruturas cristalinas estáveis a partir de composições químicas. O problema é limitado pela periodicidade (cristais se repetem no espaço), coordenadas fracionárias que vivem em um toro, simetria de grupo espacial e a existência de múltiplos polimorfos para uma única composição. Embora avanços recentes em modelagem generativa (ex: modelos de difusão, flow matching) tenham melhorado a CSP, eles frequentemente dependem de formulações contínuas que requerem amostragem estocástica ou iterativa, carecendo de um espaço de decisão discreto finito para uma decodificação gananciosa (greedy decoding) confiável. Além disso, agentes de pesquisa de IA existentes tipicamente operam dentro de especificações de problemas fixas, otimizando código ou hiperparâmetros para uma determinada família de modelos, em vez de descobrir novos princípios de modelagem através de diferentes domínios.

Metodologia: O Sistema HACO
Os autores introduzem o HACO (Human–AI Co-discovery system), um framework projetado para descobrir algoritmos científicos através de busca trans-domínio e direcionamento humano esparso. O HACO opera sobre dois princípios fundamentos:

1. Transferência Trans-Domínio: Em vez de perturbar modelos existentes específicos de um domínio, o agente busca através de frameworks de modelagem generativa de diversos campos (visão, linguagem, etc.) para identificar princípios transferíveis.
2. Interação Humana Esparsa: Humanos fornecem objetivos de alto nível ou mecanismos (ex: "melhorar a cobertura de polimorfos" ou "incorporar simetria"), enquanto o agente lida com a implementação, treinamento, avaliação e refinamento.

O HACO emprega uma busca estruturada em árvore orquestrada por um LLM (Claude Opus 4.7). A busca procede em três estágios:

• Estágio de Descoberta: Busca pela metodologia generativa (orçamento de 2 horas por candidato).
• Estágio de Escalonamento: Otimiza a metodologia selecionada com orçamentos maiores (orçamento de 12 horas).
• Otimização de Amostragem: Ajusta o algoritmo de amostragem.

Os operadores principais incluem idea (propõe frameworks trans-domínio), draft (instancia o código), improve (realiza mudanças coerentes) e debug. O operador idea identificou especificamente o MaskGIT, um modelo generativo mascarado da visão computacional, como um candidato promissor para CSP.

O Algoritmo Resultante: MaskGXT
Através do processo HACO, o sistema desenvolveu o MaskGXT (Masked Generative Crystal Transformer). Ao contrário dos modelos de difusão contínuos, o MaskGXT formula a CSP como uma tarefa de geração de tokens discretos.

• Tokenização: Um cristal é representado como uma sequência de $1 + 6 + 5N$ tokens:
  • 1 Token de Grupo Espacial.
  • 6 Tokens de rede (lattice) (comprimentos e ângulos quantizados).
  • $5N$ tokens de sítio (3 tokens de coordenadas quantizadas, 1 token de posição de Wyckoff e 1 token de tipo de átomo).
  • Tipos de átomos são fornecidos como informação lateral (conditioning).
• Discretização e Suavização: Coordenadas fracionárias contínuas são mapeadas para bins discretos em um círculo (módulo 1). Para preservar a estrutura métrica, o modelo utiliza suavização de rótulo ordinal e circular para coordenadas e suavização ordinal padrão para parâmetros de rede, em vez de tratar os bins como categorias não ordenadas.
• Tratamento de Simetria: O modelo aprende a atribuir tokens de grupo espacial e Wyckoff consistentes através de diferentes origens de célula unitária e ordenações atômicas via aumentos on-line (deslocamentos de origem normalizadores euclidianos e permutações intra/inter-órbita).
• Arquitetura: Um tronco Transformer bidirecional com atenção auto-referencial QK-normalizada e redes feedforward SwigLU. Inclui uma cabeça de regressão leve para prever offsets de sub-bin, refinando a precisão geométrica além da resolução do bin discreto.
• Amostragem:
  • Amostragem Estratificada por Grupo Espacial: O modelo prevê uma distribuição posterior sobre grupos espaciais. Para cobertura de polimorfos, a geração é ramificada em fluxos separados para os grupos espaciais de maior posterioridade.
  • Decodificação Gananciosa Baseada em Confiança: Em vez de amostragem estocástica, o modelo desmascara deterministicamente os tokens com maior confiança (argmax), permitindo a geração eficiente de estruturas de alta probabilidade.

Principais Contribuições

1. Framework HACO: Demonstra que um co-cientista de IA, guiado por objetivos humanos esparsos, pode identificar um princípio de modelagem transferível (MaskGIT da visão) e adaptá-lo a um domínio científico (CSP) para produzir um algoritmo de estado da arte.
2. MaskGXT: Introduz um transformer generativo mascarado discreto para cristais que unifica rede, coordenadas, grupos espaciais e posições de Wyckoff em uma única sequência de tokens, permitindo decodificação gananciosa confiável e tratamento explícito de simetria.
3. Ganhos de Desempenho: Mostra que um único modelo MaskGXT atinge taxas de correspondência (match rates) de estado da arte em benchmarks padrão e melhora significativamente a cobertura de polimorfos.

Resultados Experimentais
O modelo foi avaliado nos benchmarks MP-20 e MPTS-52, bem como no MP-20 polymorph split (que testa a cobertura de múltiplas estruturas por composição).

• CSP Padrão (Um-para-Um): O MaskGXT alcançou a maior taxa de correspondência (MR) no MP-20 (73,79% sem filtro) e no MPTS-52 (36,75% sem filtro), superando baselines fortes como Crystalite, MCFlow e FlowMM. Também alcançou o menor RMSE de par correspondente.
• Cobertura de Polimorfos (METRe): No MP-20 polymorph split, o MaskGXT atingiu 79,06% de precisão de Match-Everyone-to-Reference (METRe), comparado aos 70,87% do baseline mais forte (Crystalite). Este aumento é atribuído à amostragem estratificada por grupo espacial, que aloca o orçamento de geração para distintos grupos espaciais de alta probabilidade.
• Estudos de Ablação:
  • Remover os tokens de simetria causou uma queda significativa no METRe no split de polimorfos, confirmando a necessidade de modelagem explícita de simetria.
  • Remover o offset de regressão de sub-bin teve efeito mínimo na taxa de correspondência, mas aumentou substancialmente o RMSE, provando seu papel na precisão geométrica.
  • A decodificação gananciosa melhorou a precisão de estrutura única, enquanto a amostragem estratificada foi crítica para a cobertura de polimorfos.

Significância e Alegações
O artigo afirma que em domínios caracterizados por "validação barata, rápida e bem alinhada" (como a CSP com correspondência de estrutura automatizada), co-cientistas de IA interativos guiados por transferência podem contribuir para a descoberta de algoritmos científicos. Especificamente, o sistema obteceu sucesso ao:

• Identificar um princípio de modelagem de um campo não relacionado (modelagem generativa mascarada na visão) e adaptá-lo às restrições da cristalografia.
• Combinar essa transferência com orientação humana direcionada (ex: solicitar geração consciente de simetria) para refinar o algoritmo.
• Produzir um algoritmo competitivo que avança os benchmarks padrão sem exigir que o humano especifique detalhes de implementação de baixo nível.

Os autores observam que o sucesso deste ciclo depende da estrutura específica do problema de CSP (datasets fixos, ciclos de treinamento rápidos e métricas de validação confiáveis). Eles sugerem que estender esta abordagem para domínios com experimentos físicos ou objetivos de proxy fracos exigiria uma avaliação mais estruturada e um direcionamento humano mais forte.