TriForces: Augmenting Atomistic GNNs for Transferable Representations

TriForces é um framework de três fluxos agnóstico a modelos que combina aprendizado auto-supervisionado com representações separadas de composição e estrutura para aprimorar significativamente a transferibilidade e a eficiência de dados de redes neurais de grafos atômicos para potenciais interatômicos de aprendizado de máquina.

O essencial

Imagine que você está tentando ensinar um chef robô a cozinhar.

O Problema: O Chef "Tamanho Único"
Atualmente, os cientistas utilizam modelos de IA poderosos (chamados MLIPs) para prever como os átomos se comportam, como a quantidade de energia que um material possui ou o quão difícil é mover os átomos. Esses modelos são treinados com grandes quantidades de dados provenientes de supercomputadores (DFT).

No entanto, esses modelos têm uma falha: são como um chef que memorizou o sabor exato de um prato específico, mas esqueceu por que ele tinha aquele sabor. Se você pedir que cozinhe um prato ligeiramente diferente (um novo tipo de material), eles têm dificuldade. Eles confundem os ingredientes (composição) com a forma da panela (estrutura). Se você mudar os ingredientes, eles ficam confusos quanto à forma, e vice-versa. Isso os torna ruins em aprender novas tarefas rapidamente, especialmente quando não há muitos dados para ensiná-los.

A Solução: TriForces (A Cozinha de Três Fluxos)
Os autores apresentam o TriForces, uma nova maneira de construir esses chefs de IA. Em vez de um único cérebro gigante tentando lembrar de tudo ao mesmo tempo, eles dividem o cérebro em três "fluxos" ou departamentos especializados:

1. O Fluxo de Ingredientes (Composição): Este departamento olha apenas para o que está na panela (por exemplo, "Temos 2 Hidrogênios e 1 Oxigênio"). Ele ignora completamente a forma. Ele aprende a química.
2. O Fluxo de Forma (Estrutura): Este departamento olha apenas para como os átomos estão dispostos no espaço (por exemplo, "Eles estão em um triângulo"). Ele ignora o que os átomos realmente são. Ele aprende a geometria.
3. O Fluxo de Interação: Este é o chef principal que pega as anotações dos departamentos de Ingredientes e Forma e as combina para prever o resultado final (energia ou força).

O Segredo: Aprendizado Auto-Supervisionado
Antes que o modelo seja solicitado a prever uma propriedade específica, os autores o treinam usando um jogo chamado "Aprendizado Auto-Supervisionado". Pense nisso como uma sessão de prática onde a IA tem que:

• Dessensibilizar: Olhar para uma imagem levemente quebrada ou ruidosa de uma molécula e corrigi-la.
• Mascaramento: Cobrir um ingrediente e adivinhar o que era com base nos vizinhos.
• Correspondência: Olhar para duas versões ligeiramente diferentes da mesma molécula e perceber que são a mesma coisa.

Esse treinamento força a IA a organizar seu conhecimento de forma ordenada. Ela aprende que "ingredientes" pertencem a uma pasta e "formas" a outra, em vez de misturá-los todos juntos.

Por Que Isso Importa (Os Resultados)
O artigo mostra que essa nova cozinha de "Três Fluxos" funciona muito melhor do que as antigas cozinhas de "Um Cérebro":

• Aprendizado Mais Rápido: Quando recebe uma pequena quantidade de novos dados (como 20.000 exemplos em vez de milhões), o TriForces aprende muito mais rápido e comete menos erros. É como um chef que pode aprender uma nova receita após prová-la uma vez, em vez de precisar cozinhá-la mil vezes.
• Melhor Memória: A IA não esquece o que aprendeu. Ela pode transferir seu conhecimento de um tipo de material para outro sem ficar confusa.
• Conhecimento Pesquisável: Como a IA mantém "ingredientes" e "formas" separados, você pode pedir que ela encontre materiais que tenham a mesma aparência, mas ingredientes diferentes, ou materiais com os mesmos ingredientes, mas formas diferentes. Os modelos antigos não conseguiam fazer isso porque seu conhecimento estava muito misturado.

Em Resumo
O TriForces é uma estrutura que desmonta o trabalho complexo de entender átomos em três tarefas mais simples: conhecer os ingredientes, conhecer a forma e saber como eles funcionam juntos. Ao treinar a IA para manter essas tarefas separadas e praticar com "jogos de adivinhação" (aprendizado auto-supervisionado), o modelo torna-se uma ferramenta muito mais flexível, eficiente e precisa para descobrir novos materiais.

Os autores disponibilizaram seu código e modelos pré-treinados para que outros cientistas possam usar essa "cozinha de três fluxos" para construir melhores IAs para a ciência dos materiais.

Resumo técnico

Resumo Técnico: TriForces

Declaração do Problema

Potenciais Interatômicos de Aprendizado de Máquina (MLIPs) baseados em Redes Neurais em Grafos Geométricas (GNNs) alcançaram alta precisão na previsão de propriedades atômicas quando treinados em conjuntos de dados de Teoria do Funcional da Densidade (DFT) em grande escala. No entanto, sua utilidade prática é prejudicada por duas limitações primárias:

1. Transferibilidade Inconsistente: MLIPs frequentemente falham em transferir efetivamente para novas químicas ou tarefas, particularmente quando o ajuste fino é realizado em conjuntos de dados pequenos, caros e específicos de tarefas. Modelos pré-treinados frequentemente lutam com tarefas diagnósticas simples (por exemplo, identificar sistemas cristalinos ou elementos majoritários) e exibem "esquecimento catastrófico" ou desempenho de transferência instável entre domínios.
2. Baixa Reutilização de Representações: MLIPs atuais otimizam representações para alvos de previsão específicos (energia e forças) em vez de para reutilização geral. Consequentemente, essas representações frequentemente entrelaçam composição (química) e estrutura (geometria), tornando-as inadequadas para análise exploratória, recuperação de vizinhos mais próximos ou busca de similaridade decomposta. Objetivos supervisionados padrão encorajam representações suficientes para regressão, mas não organizadas para preservar informações de composição e estrutura acessíveis.

Metodologia: A Framework TriForces

Os autores propõem TriForces, uma framework agnóstica a modelos que aprimora GNNs geométricas existentes com uma arquitetura de três fluxos e aprendizado auto-supervisionado (SSL) multi-objetivo. A inovação central é a fatoração explícita das representações atômicas em três componentes distintos:

1. Arquitetura de Três Fluxos

Em vez de um único vetor latente, o TriForces decompõe a representação em nível de nó $h_i$ em três fluxos concatenados:

• Fluxo de Composição ($h^{comp}$): Codifica informações químicas sem coordenadas. Processa o conjunto de elementos atômicos únicos e suas contagens estequiométricas usando um Transformer com atenção ponderada por contagem. Este fluxo preserva contagens absolutas de elementos para codificar o tamanho do sistema e a escala de energia, independentemente da geometria.
• Fluxo de Estrutura ($h^{struct}$): Codifica geometria sem identidade de elemento (agnóstico a tipos). Constrói descritores locais invariantes à rotação inspirados em SOAP (Sobreposição Suave de Posições Atômicas), utilizando funções de base radial, harmônicos esféricos e cortes em múltiplas escalas. Este fluxo captura motivos geométricos reutilizáveis e padrões topológicos, aprimorados por passagem de mensagens invariante.
• Fluxo de Interação ($h^{int}$): Uma GNN geométrica base padrão (por exemplo, MACE, eSEN, Orb-v3) que captura o acoplamento entre composição e geometria, preservando a expressividade da arquitetura original.

2. Pré-treinamento Auto-Supervisionado

Para organizar o espaço latente e melhorar a transferibilidade, o TriForces emprega uma estratégia de pré-treinamento SSL multi-objetivo usando aumentações estocásticas (ruído de posição, mascaramento de elementos, variação de grafos e rotações). A framework combina três objetivos complementares:

• Não-Reconstrução (LeJEPA): Alinha embeddings de duas visões aumentadas da mesma estrutura em níveis de nó e de grafo. Isso impõe invariância às aumentações e organiza o espaço latente global sem exigir gradientes de parada ou codificadores de momento.
• Denoising (Remoção de Ruído): Treina o modelo para recuperar posições atômicas limpas a partir de entradas ruidosas. Isso estabiliza representações geométricas e fornece implicitamente aumento de rotação.
• Mascaramento: Prevê tipos atômicos mascarados com base na geometria e composição circundantes. Isso encoraja o modelo a aprender padrões composicionais e contexto.

A perda final de pré-treinamento é uma soma ponderada desses três objetivos.

Contribuições Principais

1. Decomposição Arquitetural: Um design de três fluxos que separa explicitamente composição, estrutura e interação, garantindo que ambos os fatores sejam preservados por design em vez de entrelaçados.
2. Estratégia Híbrida de Pré-treinamento: Uma abordagem auto-supervisionada que combina objetivos baseados em reconstrução (denoising, mascaramento) com aprendizado de predição latente (LeJEPA) para estruturar o espaço de embeddings para melhor transferência a jusante.
3. Recuperação Interpretável: A capacidade de realizar busca de similaridade direcionada nos espaços de embeddings composicionais, estruturais ou conjuntos, permitindo comparação de materiais com base em critérios específicos (por exemplo, apenas química ou apenas estrutura).
4. Validação Empírica: Experimentos extensos em múltiplas arquiteturas (Orb-v3, eSEN, MACE) e benchmarks (OMat24, MatBench, QM9) demonstrando melhoria na eficiência de dados, desempenho de transferência e qualidade de representação.

Resultados

• Desempenho de Transferência (OMat24): Em regimes de dados limitados, o TriForces supera significativamente as linhas de base. Com 20K amostras, reduz o Erro Absoluto Médio (MAE) de energia em 57% em comparação com o modelo base. Melhora o MAE de força em todos os tamanhos de amostra e reduz erros de tensão.
• Eficiência de Dados: O TriForces alcança erros menores em cada tamanho de conjunto de dados (de 20K a 2M amostras), com os ganhos mais significativos observados em configurações de baixos dados.
• Desempenho em Benchmarks:
  • MatBench: Variantes do TriForces alcançam os melhores resultados gerais em 6 de 8 tarefas, superando tanto linhas de base de pré-treinamento auto-supervisionado quanto rotuladas por DFT. Por exemplo, o MAE de Fônons melhorou de 57,8 para 19,5 cm$^{-1}$.
  • MatBench Discovery: O TriForces eSEN-sm alcança um MAE de energia comparável a um modelo eSEN-30M-OAM muito maior, enquanto usa 60% menos parâmetros e treina até 5$\times$ mais rápido.
  • QM9: O pré-treinamento em entradas químicas diversas (bulk + moléculas) consistentemente reduz o MAE em comparação com linhas de base apenas de bulk ou sem SSL.
• Qualidade de Representação: Sondagem linear em embeddings congelados mostra que o TriForces preserva informações fundamentais (sistema cristalino, elemento majoritário, número de coordenação) que MLIPs padrão perdem. O TriForces alcança 96–100% de precisão na classificação de sistema cristalino e elemento majoritário, enquanto as linhas de base lutam (55–73%).
• Recuperação: A framework permite recuperação eficaz k-NN, onde o fluxo de composição se destaca na recuperação de conjuntos de elementos e o fluxo de estrutura se destaca na recuperação de grupo espacial, uma capacidade ausente em modelos de fluxo único.

Significado e Alegações

O artigo posiciona o TriForces não meramente como um método auto-supervisionado, mas como uma framework arquitetural cujas representações são ainda aprimoradas por SSL.

• Dependência de Regime: Os autores afirmam que a fatoração de fluxos fornece os ganhos dominantes em configurações supervisionadas em grande escala, enquanto o SSL é mais valioso para transferência de baixos dados, organização de representação e tarefas de recuperação.
• Desacoplamento: Ao separar composição e estrutura, o TriForces aborda os problemas de "transferência frágil" e "difícil reutilização" dos modelos fundamentais atuais. Permite que modelos aprendam representações organizadas para análise (por exemplo, sondagem, recuperação) além da previsão.
• Praticidade: A framework é agnóstica a modelos e plug-in, tornando-a imediatamente aplicável a arquiteturas atômicas existentes ou novas. Os autores lançam checkpoints pré-treinados e código para facilitar a reutilização em modelagem atômica a jusante.

O trabalho sugere que futuros modelos fundamentais atômicos devem ir além de arquiteturas de predição de fluxo único em direção a representações fatoradas que preservam explicitamente os fatores físicos distintos de sistemas químicos.