Resumo Técnico: Otimização Molecular Amortizada via Otimização de Política Relativa por Grupos (AMORTIX)

1. Declaração do Problema

Na descoberta de fármacos, o desenho molecular frequentemente exige a modificação de uma estrutura inicial fixa (por exemplo, decoração de esqueleto, desenho de conectores) em vez de gerar moléculas do zero. Essas tarefas são tipicamente avaliadas usando "oráculos" caros (funções de pontuação), que variam de proxies heurísticos a simulações de alta fidelidade, como perturbação de energia livre.

Os métodos atuais de última linha seguem predominantemente um paradigma de otimização por instância. Algoritmos como Algoritmos Genéticos (Mol GA) e modelos de difusão discreta guiada (GenMol) realizam uma otimização iterativa baseada em busca para cada nova estrutura de entrada. Essa abordagem requer milhares de chamadas ao oráculo por tarefa de desenho, resultando em altos custos computacionais que escalam linearmente com o número de estruturas iniciais.

Embora a otimização amortizada — aprender uma política transferível para gerar moléculas otimizadas em uma única passagem forward sem chamadas ao oráculo no tempo de inferência — ofereça uma solução teórica, os métodos existentes lutam na prática. Eles falham em generalizar para diversas restrições estruturais ou sofrem com instabilidade no treinamento. Os autores identificam a causa raiz como a heterogeneidade na dificuldade de otimização: algumas estruturas iniciais produzem derivados de alta pontuação facilmente, enquanto outras são quimicamente restritas e produzem pontuações modestas independentemente da modificação. Métodos padrão de aprendizado por reforço (RL) que utilizam baselines globais (por exemplo, REINFORCE) ou críticos de valor aprendidos (por exemplo, PPO) falham em estabilizar o treinamento sob essa variância, pois os sinais de recompensa tornam-se incomparáveis entre diferentes estruturas iniciais.

2. Metodologia: AMORTIX

Os autores propõem o AMORTIX, um modelo Transformer de Grafo amortizado projetado para otimizar estruturas moleculares sob restrições estruturais estritas com zero chamadas ao oráculo no tempo de inferência.

2.1 Arquitetura da Política

O AMORTIX emprega um Transformer de Grafo apenas com decodificador que opera diretamente sobre representações de grafos moleculares. O processo de geração é modelado como um Processo de Decisão de Markov (MDP) sequencial com um espaço de ação hierárquico:

1. Seleção de Operação: Escolher adicionar um novo átomo, modificar um átomo existente ou terminar.
2. Seleção de Alvo: Selecionar o átomo específico no grafo atual para conectar.
3. Especificação de Ligação: Determinar a ordem da ligação (simples, dupla, tripla, etc.).

Para garantir validade química, uma máscara de validade é aplicada a cada passo, zerando ações que violariam regras de valência atômica. Isso garante que todas as estruturas geradas sejam quimicamente válidas por construção. O modelo é pré-treinado no banco de dados ChEMBL via forçamento supervisionado de professor para aprender plausibilidade química antes do ajuste fino por RL.

2.2 Treinamento via Otimização de Política Relativa por Grupos (GRPO)

A inovação central reside na adaptação da Otimização de Política Relativa por Grupos (GRPO) ao domínio molecular para abordar a heterogeneidade das estruturas iniciais.

• Formação de Grupos: Para cada passo de treinamento, um lote de estruturas iniciais é amostrado. Para cada estrutura inicial $S_i$, a política gera um grupo de $G$ trajetórias completadas (moléculas).
• Vantagem Relativa: Em vez de usar uma baseline global ou um crítico de valor, a vantagem $A_{i,j}$ para uma trajetória $j$ da estrutura $i$ é calculada centralizando sua recompensa $R(O_{i,j})$ em relação à recompensa média $\mu_i$ de seu grupo específico:
  $$A_{i,j} = R(O_{i,j}) - \mu_i, \quad \text{onde } \mu_i = \frac{1}{G} \sum_{k=1}^{G} R(O_{i,k})$$
• Normalização: Os autores omitiram a divisão pelo desvio padrão do grupo e a normalização pelo comprimento da trajetória. Isso evita que estruturas onde todas as completuras pontuam de forma similar (seja trivialmente fácil ou quase impossível) recebam peso de gradiente desproporcional e evita penalizar completuras concisas.
• Atualização: A política é atualizada usando um único passo de gradiente por lote, agregando contribuições através de todos os passos de geração. Essa abordagem desacopla o sinal de aprendizado das variações de dificuldade entre estruturas.

3. Contribuições Principais

1. Identificação de Modos de Falha: O artigo identifica que a alta variância nos sinais de recompensa causada pela dificuldade heterogênea de estruturas moleculares é um modo de falha primário para RL amortizado em otimização com restrições estruturais. Estudos de ablação demonstram que nem o REINFORCE com baseline global nem o PPO com crítico de valor estabilizam efetivamente o treinamento nesse cenário.
2. Framework AMORTIX: Uma política baseada em Transformer de Grafo treinada com RL relativa por grupos que generaliza para esqueletos fora da distribuição com zero chamadas ao oráculo no tempo de inferência.
3. Superioridade de Desempenho: O AMORTIX supera tanto as linhas de base de otimização amortizada quanto as de otimização por instância em tarefas de decoração de esqueleto orientadas a objetivos, exigindo significativamente menos chamadas totais ao oráculo ao escalar para múltiplas estruturas iniciais.

4. Resultados Experimentais

Os autores avaliam o AMORTIX em três tarefas distintas:

4.1 Benchmark PMO (Eficiência de Amostragem)

No benchmark PMO (22 tarefas com um orçamento de 10.000 chamadas ao oráculo), o AMORTIX foi executado como um otimizador por instância (usando REINFORCE padrão com baseline global, já que a heterogeneidade da estrutura inicial está ausente na geração de novo para tarefa única).

• Resultado: O AMORTIX alcançou uma pontuação AUC acumulada top-10 de 16,433, classificando-se em primeiro lugar entre políticas amortizadas e em segundo lugar no geral entre todos os paradigmas, ficando atrás apenas do otimizador por instância de difusão discreta GenMol.

4.2 Decoração de Esqueleto Orientada a Objetivos

Esta avaliação testa a generalização para esqueletos não vistos (fora da distribuição) para otimização de alvo único (GSK3$\beta$, JNK3) e múltiplos alvos (Kinase MPO).

• Configuração: Os esqueletos foram divididos por decomposição de Murcko e agrupamento para garantir que os esqueletos de teste fossem topologicamente distintos dos dados de treinamento.
• Resultados: O AMORTIX alcançou as maiores pontuações em todas as três tarefas, superando linhas de base amortizadas (LibINVENT, GraphXForm, DrugEx v3) e otimizadores por instância (Mol GA, GenMol).
  • Kinase MPO: O AMORTIX pontuou 0,619, superando significativamente o LibINVENT (0,455) e os otimizadores por instância Mol GA (0,461) e GenMol (0,479).
  • Estabilidade: Estudos de ablação confirmaram que o AMORTIX-REINFORCE e o AMORTIX-PPO exibiram instabilidade severa (alto desvio padrão nas estimativas de vantagem) em comparação com o estável AMORTIX-GRPO.
• Eficiência de Custo: O orçamento fixo de treinamento do AMORTIX (~50.000 chamadas ao oráculo) é recuperado após processar aproximadamente 100 esqueletos de teste contra o Mol GA e 150 contra o GenMol. Além desse ponto de equilíbrio, o AMORTIX oferece custos cumulativos ao oráculo ordens de magnitude menores.

4.3 Estudo de Caso: Transferência de Pró-fármaco

O modelo foi ajustado finamente em apenas quatro fármacos exemplares para aprender uma regra de modificação para melhorar a penetração no sistema nervoso central via desenho de pró-fármaco.

• Resultado: Aplicado a 52 fármacos pais não vistos, o AMORTIX gerou com sucesso pró-fármacos estruturalmente limpos. Ele recuperou três motivos cliváveis distintos (éster alquílico, carbamato alquílico, éster benzílico) sem instrução explícita.
• Desempenho: 31% das gerações foram pró-fármacos "puros" (apenas junções cliváveis) e 63% continham pelo menos uma alça clivável, com um aumento médio previsto na penetração da barreira hematoencefálica (BBB) de +0,26.

5. Significado e Alegações

O artigo alega que o AMORTIX demonstra que o ônus computacional da otimização molecular pode ser efetivamente deslocado da busca no tempo de inferência para a aprendizagem no tempo de treinamento. Ao adaptar o GRPO para normalizar recompensas dentro de grupos de completuras que compartilham a mesma estrutura inicial, o método estabiliza a aprendizagem na presença de dificuldade heterogênea.

Os autores posicionam o AMORTIX como uma alternativa escalável aos métodos de busca iterativa, permitindo desenho de alto rendimento sem o custo proibitivo de avaliações repetidas do oráculo. Eles observam que, embora o processo de treinamento permaneça intensivo em oráculos, a natureza amortizada da política torna-o altamente eficiente para cenários envolvendo muitas estruturas iniciais.

6. Limitações

Os autores reconhecem explicitamente várias limitações:

• Topologia 2D: O framework opera exclusivamente em topologias moleculares 2D, não capturando explicitamente dinâmicas conformacionais 3D ou estereoquímica.
• Conectividade: A máscara de ação atual assume um único esqueleto inicial conectado. A extensão para subgrafos desconectados (por exemplo, ligação de fragmentos) requer modificações arquitetônicas.
• Ações Aditivas: O processo de geração é restrito a ações aditivas (adição de átomos/ligações). Não suporta deleção ou substituição de átomos, que são necessárias para tarefas como salto de esqueleto.
• Custo de Treinamento: Embora a inferência seja barata, o processo de treinamento requer dezenas de milhares de chamadas ao oráculo, necessitando de validação adicional para quantificar precisamente o trade-off de amortização para casos de uso específicos.