Resumo Técnico: NEAT – Transformer de Conjunto Autoregressivo, Eficiente e Guiado por Vizinhança para Geração Molecular 3D

Declaração do Problema

Os fluxos de trabalho computacionais atuais para a descoberta de moléculas bioativas dependem fortemente de triagem virtual ou geração de novo. Embora a triagem virtual seja limitada pelo tamanho finito de bibliotecas químicas e pelo aumento dos custos computacionais, a modelagem generativa oferece uma alternativa promissora. No entanto, a geração de estruturas moleculares 3D apresenta desafios únicos. As moléculas são inerentemente tridimensionais, e suas propriedades são frequentemente determinadas pela conformação.

Os modelos de geração 3D mais avançados existentes utilizam predominantemente técnicas de difusão e correspondência de fluxo. Esses métodos frequentemente sofrem com escalamento quadrático nas atualizações de passagem de mensagens, carecem de adaptabilidade inerente a tamanhos variáveis de moléculas (muitas vezes exigindo contagens fixas de átomos determinadas na inicialização) e lutam para gerenciar inserções dinâmicas de átomos durante a remoção de ruído.

Modelos autoregressivos oferecem uma alternativa flexível que suporta naturalmente contagens variáveis de átomos e requer apenas uma chamada de codificador por token. No entanto, os transformers autoregressivos padrão exigem uma ordenação sequencial de tokens. Em grafos moleculares, os átomos não possuem uma ordem sequencial inerente. Trabalhos anteriores abordaram isso usando ordenações canônicas de átomos, mas essas abordagens não são invariantes à permutação em relação aos átomos. Elas enviesam a previsão do próximo token para convenções de ordenação específicas, limitando a capacidade do modelo de generalizar ou realizar tarefas como conclusão de prefixo, onde a ordem inicial dos átomos pode não alinhar-se com a ordem canônica da molécula final. Por outro lado, abordagens de Redes Neurais em Grafos (GNN) que garantem invariância à permutação frequentemente produzem moléculas com baixa validade química.

Metodologia

Os autores propõem o NEAT (Transformer de Conjunto Autoregressivo, Eficiente e Guiado por Vizinhança), o primeiro modelo transformer autoregressivo para geração molecular 3D que respeita a invariância à permutação de átomos sem depender de ordenações canônicas de átomos.

Arquitetura Central

O NEAT trata grafos moleculares como conjuntos de átomos, e não como sequências. A arquitetura consiste em três componentes principais:

1. Módulo de Aumento de Dados (Orientação por Vizinhança): Em vez de codificações posicionais sequenciais, o modelo usa uma estratégia de amostragem para imitar o crescimento molecular. Dado um grafo molecular, um subgrafo conectado (conjunto fonte) é amostrado. O conjunto alvo é definido como os "nós de fronteira" (vizinhos de 1-hop) do conjunto fonte. Isso fornece um conjunto de continuições de tokens válidas para o próximo passo, independentemente da ordem em que os átomos foram adicionados.
2. Tronco Transformer de Conjunto: O conjunto fonte é codificado usando um transformer de conjunto (baseado em NanoGPT/GPT-2 com atenção bidirecional e sem incorporações posicionais). Os tipos de átomos são incorporados linearmente, e as posições são codificadas usando recursos de Fourier. O agrupamento aditivo agrega essas representações de nós em uma representação de nível de conjunto ($z_n$). Este design garante invariância à permutação em relação à ordem dos átomos no conjunto fonte.
3. Cabeça de Correspondência de Fluxo: A representação de nível de conjunto $z_n$ é usada para prever o próximo token, fatorizada em tipo de átomo e posição: $p(a_{n+1}, x_{n+1} | z_n) = p(a_{n+1} | z_n) \cdot p(x_{n+1} | a_{n+1}, z_n)$.
   • Tipo de Átomo: Modelado como uma distribuição categórica via uma camada linear e softmax.
   • Posição do Átomo: Modelada via Correspondência de Fluxo Condicional (CFM). O modelo aprende um campo vetorial para transportar uma distribuição inicial de posição aleatória para as verdadeiras posições atômicas.

Estratégia de Treinamento

• Orientação por Vizinhança: Durante o treinamento, o modelo é supervisionado não por um único próximo token, mas pelo conjunto de todos os nós de fronteira válidos. Isso elimina a necessidade de canonicização.
• Forçamento do Professor: A previsão do tipo de átomo e a recuperação da posição são desacopladas. O modelo prevê o tipo, e a cabeça de correspondência de fluxo aprende a posição condicionada a esse tipo.
• Aumento: Para melhorar a robustez contra o acúmulo de erros, ruído gaussiano é injetado em uma fração dos átomos do conjunto fonte durante o treinamento, e as posições são recentralizadas em relação ao conjunto fonte ruidoso.
• Função de Perda: A perda total é a soma da perda de entropia cruzada para tipos de átomos e da perda de correspondência de fluxo para posições. Um problema de atribuição linear é resolvido para acoplar distribuições fonte e alvo, minimizando custos de transporte e evitando caminhos cruzados.

Inferência

A geração começa com um átomo inicial aleatório. Em cada passo, o conjunto atual de átomos é codificado, o próximo tipo de átomo é previsto e uma posição inicial aleatória é amostrada. O campo vetorial de correspondência de fluxo é então integrado (usando métodos de Euler ou Euler-Maruyama) para determinar as coordenadas 3D finais. Este processo se repete até que um token de parada seja previsto.

Principais Contribuições

1. Geração Autoregressiva Invariante à Permutação: O NEAT introduz uma estratégia de treinamento inovadora usando orientação por vizinhança e transformers de conjunto para alcançar geração molecular 3D que é invariante à ordenação de átomos, superando as limitações da ordenação canônica em modelos autoregressivos.
2. Eficiência: Ao tratar a entrada como um conjunto e usar uma única chamada de codificador por token, o NEAT evita o escalamento quadrático da passagem de mensagens baseada em difusão e a necessidade de múltiplos passos de integração por passo de geração de átomo encontrados em algumas bases de correspondência de fluxo.
3. Capacidade de Conclusão de Prefixo: A natureza agnóstica à ordem do NEAT permite conclusão de prefixo de alta qualidade (gerando moléculas ao redor de um andaime dado) sem exigir que o prefixo se alinhe a uma ordenação canônica específica, uma tarefa onde modelos baseados em sequências como o QUETZAL lutam.
4. Desempenho de Estado da Arte: O modelo alcança resultados competitivos ou superiores em benchmarks padrão (QM9 e GEOM-Drugs) em comparação com bases autoregressivas e de difusão/correspondência de fluxo.

Resultados Experimentais

O modelo foi avaliado nos conjuntos de dados QM9 (moléculas orgânicas pequenas) e GEOM-Drugs (moléculas semelhantes a fármacos).

• Qualidade de Geração: No QM9, o NEAT superou modelos autoregressivos anteriores (G-SchNet, G-SphereNet, Symphony, Geo2Seq) em termos de validade e unicidade xyz2mol. Alcançou uma validade xyz2mol de 95,2% e unicidade de 91,6%, ficando em segundo lugar em validade e unicidade de consulta. No GEOM, o NEAT alcançou a maior validade de consulta (99,7%) e a segunda melhor validade xyz2mol (93,9%) e unicidade (91,6%) entre modelos autoregressivos, igualando o desempenho de modelos de correspondência de fluxo de estado da arte como o FlowMol3.
• Conclusão de Prefixo: Em uma tarefa gerando 1.000 moléculas a partir de 100 prefixos diversos, o NEAT superou significativamente o QUETZAL. O NEAT alcançou 97,8% de validade de consulta e 92,7% de validade xyz2mol, enquanto o QUETZAL alcançou apenas 27,8% e 28,1%, respectivamente. Isso demonstra a capacidade do NEAT de generalizar características estruturais independentes das convenções de ordenação de átomos.
• Eficiência de Tempo de Execução: O NEAT demonstra vantagens significativas de velocidade. Ele roda aproximadamente 5 vezes mais rápido que o SemlaFlow e 16 vezes mais rápido que o FlowMol3. Em cenários de pior caso, a complexidade do NEAT é $O(BN^3)$ (onde $N$ é a contagem de átomos), enquanto modelos de difusão/fluxo escalonam como $O(BTN^2)$ (onde $T$ é o número de passos de integração). Como $N$ é tipicamente menor que $T$ para moléculas semelhantes a fármacos, o NEAT oferece uma vantagem prática de velocidade.

Significado e Alegações

O artigo afirma que o NEAT fornece uma base sólida para alavancar arquiteturas de transformers em tarefas complexas de design de fármacos, especificamente para geração molecular 3D condicional. Ao remover a dependência de ordenações canônicas de átomos, o modelo permite tarefas flexíveis como geração baseada em andaimes (conclusão de prefixo) que eram anteriormente difíceis para abordagens autoregressivas.

Os autores observam que, embora a arquitetura atual não seja explicitamente equivariante a SE(3), o uso de aumento de rotação aleatória durante o treinamento compensa efetivamente isso, resultando em apenas uma variância menor nas previsões. O trabalho sugere que integrar o NEAT com um transformer equivariante a SE(3) é uma direção promissora para pesquisas futuras.

Os autores enfatizam que, embora o modelo gere estruturas 3D de alta qualidade, a geração de moléculas é apenas um passo em um pipeline complexo que requer validação extensa, síntese e testes. Eles também reconhecem limitações, como a dependência de informações de aresta para orientação por vizinhança (o que pode limitar a aplicabilidade a cristais inorgânicos) e a falta de extrapolação além da contagem máxima de átomos vista durante o treinamento.