Generative Molecular Morphing for Flexible-Size Design via Unbalanced Optimal Transport

Resumo Técnico: Morfismo Molecular Generativo para Design de Tamanho Flexível via Transporte Ótimo Desbalanceado

1. Definição do Problema

Os modelos generativos atuais para o design molecular 3D, particularmente as abordagens de difusão e baseadas em fluxo, enfrentam uma limitação algorítmica crítica: eles operam em um número fixo de átomos ($n$). Essa restrição impede que os modelos aprendam a distribuição conjunta das propriedades moleculares e do tamanho molecular. Como muitas propriedades químicas estão intrinsecamente ligadas ao número de átomos, fixar o tamanho restringe artificialmente a expressividade e a controlabilidade do modelo. Consequentemente, esses modelos têm dificuldade em navegar em relações complexas entre estrutura e propriedades, falham em integrar efetivamente priors estruturais (como scaffolds) e não conseguem gerar moléculas válidas fora da distribuição de treinamento (OOD) em relação ao tamanho.

2. Metodologia

Os autores introduzem o Morph, um modelo generativo de tamanho flexível para grafos geométricos que ajusta dinamicamente o número de nós (átomos) durante o processo de geração. A metodologia integra dinâmica contínua com saltos discretos, formulada como um processo de salto-fluxo (jump-flow process).

Componentes Principais:

• Espaço de Estados: O modelo opera em um espaço de estados $\mathcal{G}$ definido como a união disjunta de espaços quocientes para grafos geométricos de vários tamanhos, contabilizando simetrias $SE(3)$ (translação/rotação) e $S_n$ (permutação).
• Dinâmica Generativa: O processo combina:
  1. Fluxo Contínuo: Governado por um campo de velocidade $v_t$ que move as posições dos átomos.
  2. Saltos Discretos: Governados por uma Cadeia de Markov de Tempo Contínuo (CTMC) com densidade de taxa $u_t$, permitindo inserção, deleção e substituição de tipos de nós.
• Treinamento via Correspondência de Gerador Condicional (CGM): Para contornar a intratabilidade do gerador marginal, o modelo utiliza CGM. Ele aprende a regredir a velocidade condicional e as taxas de salto ao longo de caminhos entre endpoints pareados $(g_0, g_1)$, onde $g_0$ é um grafo prior e $g_1$ é uma molécula alvo.
• Transporte Ótimo Desbalanceado (UOT) para Alinhamento: Uma inovação fundamental é o uso de UOT para alinhar grafos geométricos de diferentes tamanhos ($n_0 \neq n_1$) antes de definir os caminhos de interpolação.
  • O objetivo do UOT minimiza uma função de custo envolvendo custos de transporte, penalidades de deleção de massa ($w_{del}$) e penalidades de inserção de massa ($w_{ins}$).
  • Os custos de edição de arestas são definidos como zero, justificados pela capacidade de inferir a topologia de ligações a partir da estrutura 3D e valência.
  • O plano de transporte resultante é discretizado usando o algoritmo de Hungarian em um espaço aumentado para determinar os conjuntos exatos de edições (correspondências, inserções, deleções).
  • As posições correspondentes são alinhadas rotacionalmente via algoritmo de Kabsch para garantir invariância $SE(3)$.
• Caminhos de Interpolação: Com base na correspondência UOT, o modelo define caminhos de probabilidade condicional para três modalidades:
  • Nós Correspondentes: As posições seguem um caminho Gaussiano; os tipos seguem um scheduler.
  • Nós de Inserção: Nós são amostrados em tempos de evento específicos com posições extraídas de uma Gaussiana centrada no alvo.
  • Nós de Deleção: Nós persistem até um tempo de evento específico, após o qual eles e suas arestas são removidos.
• Função de Perda: O objetivo de treinamento minimiza uma divergência de Bregman entre os geradores condicionais previstos e alvo, compreendendo termos para movimento (MSE), substituição (entropia cruzada), inserção (NLL de Poisson e GMM), deleção (entropia cruzada binária) e previsão de carga.

3. Principais Contribuições

1. Extensão para Grafos Geométricos 3D: Este trabalho estende técnicas de geração de sequências de tamanho flexível para grafos geométricos 3D, permitindo a integração fluida de priors estruturais (ex: scaffolds) no processo generativo.
2. Algoritmo de Correspondência Baseado em UOT: Os autores propõem um algoritmo de correspondência baseado em Transporte Ótimo Desbalanceado que alinha grafos geométricos de diferentes tamanhos para definir caminhos de probabilidade alvo para o treinamento, lidando com a natureza trans-dimensional do problema.
3. Melhoria na Controlabilidade e Geração OOD: O artigo demonstra que o Morph iguala modelos de tamanho fixo de estado da arte (SOTA) em geração incondicional, enquanto oferece uma controlabilidade de propriedades superior. Crucialmente, ele permite a geração bem-sucedida Out-of-Distribution (OOD) em regimes (ex: tamanhos de moléculas maiores) onde modelos anteriores falham.

4. Resultados

• Design De-novo: No conjunto de dados QM9, o Morph alcança métricas competitivas (Estabilidade de Átomos, Estabilidade Molecular, Validade, Unicidade) em comparação com modelos de tamanho fixo SOTA (FlowMol, MiDi, EQGAT-diff, SemlaFlow), frequentemente utilizando menos ou o mesmo Número de Avaliações de Função (NFE). Também apresenta bom desempenho no mais complexo GEOM-Drugs.
• Controle de Propriedades (Steering):
  • Condicionamento de Tamanho: Quando condicionado a contagens de átomos significativamente maiores que a distribuição de treinamento (OOD), o Morph mantém alta validade, gerando moléculas com até 40% mais átomos do que o maior exemplo de treinamento. Modelos de tamanho fixo falham em extrapolar neste regime.
  • Condicionamento de LogP: O modelo consegue direcionar o tamanho das moléculas geradas com base no coeficiente de partição logarítmico (logP), recuperando a marginal de tamanho da distribuição de treinamento e gerando moléculas válidas em regiões de baixo suporte ou logP OOD.
• Decoração de Scaffold: Em uma tarefa de otimização de lead, o Morph consegue estender pequenos scaffolds químicos em moléculas completas enquanto preserva a substrutura central, alcançando alta validade e unicidade.

5. Significância e Alegações

O artigo afirma que o Morph aborda uma limitação fundamental nos designs moleculares generativos atuais: a incapacidade de adaptar o tamanho molecular durante a geração. Ao aprender a distribuição conjunta de tamanho e propriedades, o Morph desbloqueia:

• Flexibilidade de Amostragem Inigualável: A capacidade de gerar moléculas de vários tamanhos a partir de um único modelo.
• Potencial de Descoberta Melhorado: A capacidade de explorar espaços químicos fora da distribuição de treinamento (OOD), o que é essencial para descobrir novas moléculas com propriedades específicas e inéditas.
• Integração de Prior Estrutural: A capacidade de usar scaffolds como pontos de partida para expansão flexível, um requisito comum na descoberta de fármacos.

Os autores reconhecem limitações, incluindo o aumento dos custos de treinamento devido à complexidade combinatória, a necessidade de mais épocas para atingir o desempenho SOTA e a dependência atual de altos passos de integração devido à natureza CTMC do modelo. Eles posicionam o Morph como um bloco de construção promissor para acelerar a descoberta nas ciências químicas, especificamente para tarefas que exigem geração de tamanho flexível.