Resumo Técnico: Benefícios Sinérgicos da Geração Conjunta de Moléculas e Predição de Propriedades

Formulação do Problema

O artigo aborda um desafio de longa data no aprendizado de máquina: desenvolver modelos que se destaquem simultaneamente em tarefas generativas e preditivas. No domínio do desenho de fármacos, o sucesso depende de equilibrar a geração de moléculas novas a partir de espaços químicos inexplorados com uma predição de propriedades robusta capaz de extrapolar para essas novas moléculas. Atualmente, a geração de moléculas e a predição de propriedades são abordadas predominantemente de forma separada. Embora modelos baseados em transformers sejam o estado da arte para ambas as tarefas individualmente, o treinamento conjunto impõe desafios significativos de arquitetura e otimização. Combinar partes generativas e preditivas em um único modelo frequentemente leva a super-regularização ou interferência de gradientes entre os objetivos generativos e preditivos, fazendo com que os modelos abdicem dos benefícios sinérgicos potenciais da aprendizagem conjunta.

Metodologia: Hyformer

Para abordar esses desafios, os autores propõem o Hyformer, um modelo conjunto baseado em transformers que unifica funcionalidades generativas e preditivas dentro de um único conjunto de parâmetros compartilhados.

Arquitetura

O Hyformer combina um decoder transformer autoregressivo com um encoder transformer bidirecional. A inovação central reside em um mecanismo de atenção alternada controlado por um token de tarefa $[TASK] \in \{[LM], [PRED], [MLM]\}$:

• Modo Generativo ($[LM]$): O modelo atua como um decoder usando uma máscara de atenção causal ($M_\rightarrow$), restringindo cada token a atender apenas ao seu contexto à esquerda.
• Modo Preditivo/Pré-treinamento ($[PRED], [MLM]$): O modelo atua como um encoder usando uma máscara de atenção bidirecional ($M_\leftrightarrow$), permitindo acesso total ao contexto.

Esquema alternado alivia a interferência de gradientes típica em modelos conjuntos. O Jacobiano da função softmax sob uma máscara causal altera os gradientes do módulo de atenção de forma diferente daquela sob uma máscara bidirecional, impedindo que os objetivos entrem em conflito durante as atualizações.

Estratégia de Treinamento

O Hyformer emprega um processo de treinamento em duas etapas:

1. Pré-treinamento Conjunto: O modelo é pré-treinado em dados não rotulados usando uma função de perda combinada:
   $$\ell_{Hyformer} = \ell_{LM} + \mu\ell_{MLM} + \eta\ell_{PRED}$$
   Aqui, $\ell_{LM}$ é a perda de modelagem de linguagem, $\ell_{MLM}$ é a perda de modelagem de linguagem mascarada e $\ell_{PRED}$ é uma perda preditiva baseada em descritores moleculares computáveis analiticamente (por exemplo, peso molecular, hidrofobicidade). O equilíbrio entre as tarefas é controlado por um vetor de probabilidades de tarefa $p_{[TASK]}$.
2. Ajuste Fino (Fine-tuning): O modelo é ajustado finamente em tarefas downstream rotuladas usando uma perda conjunta:
   $$\ell_{Hyformer} = \ell_{LM} + \lambda\ell_{pred}$$
   O modelo alterna entre tarefas generativas e preditivas durante o ajuste fino com base nas probabilidades de tarefa, omitindo a perda de modelagem de linguagem mascarada para focar na tarefa downstream, mantendo ao mesmo tempo as capacidades generativas.

Amostragem

O Hyformer suporta amostragem incondicional e condicional:

• Incondicional: Amostra moléculas $x \sim p_\theta(x)$ usando a parte autoregressiva.
• Condicional: Para gerar moléculas que satisfaçam uma condição de propriedade $Y$, o modelo amostra $K$ candidatos de $p_\theta(x)$ e aceita aqueles onde o preditor treinado conjuntamente $p_\theta(y|x)$ os classifica como possuindo a propriedade $Y$. Isso é enquadrado como uma variante da amostragem best-of-$K$, aproveitando o espaço de representação unificado para um alinhamento mais estreito entre geração e controle.

Principais Contribuições

1. Novo Modelo Conjunto: Introdução do Hyformer, que unifica tarefas generativas e preditivas em um único conjunto de parâmetros usando um mecanismo de atenção alternada.
2. Demonstração de Benefícios Sinérgicos: Evidências empíricas mostrando que o Hyformer supera as linhas de base em:
   • Geração condicional de moléculas.
   • Predição de propriedades fora da distribuição (OOD).
   • Aprendizado de representação molecular (via sondagem).
3. Paridade de Desempenho: Mostrar que o Hyformer rivaliza com o desempenho de modelos puramente generativos e puramente preditivos de última geração, provando que o treinamento conjunto não exige um sacrifício no desempenho de tarefas individuais.
4. Aplicação no Mundo Real: Um caso de uso na descoberta de novos peptídeos antimicrobianos (AMPs), demonstrando a utilidade da modelagem conjunta no desenho de fármacos.

Resultados Experimentais

Geração Condicional de Moléculas

Avaliado no conjunto de dados GuacaMol, o Hyformer ajustado finamente conjuntamente alcançou o menor Desvio Absoluto Médio (MAD) e Desvio Padrão (DP) em propriedades alvo (QED, SA, logP), mantendo alta validade. Superou significativamente variantes não conjuntas e linhas de base como MolGPT e GraphGPT, confirmando que o ajuste fino conjunto melhora a controlabilidade.

Predição de Propriedades Fora da Distribuição (OOD)

Na tarefa de Identificação de Hits do benchmark Lo-Hi, o Hyformer alcançou a maior média de AUPRC em todos os conjuntos de dados (DRD2, HIV, KDR, Sol), superando linhas de base baseadas em fingerprints (k-NN, GB, SVM, MLP) e modelos de aprendizado profundo (Chemformer, Graphformer). Isso demonstra a capacidade do modelo de generalizar para andaimes moleculares não vistos durante o treinamento.

Aprendizado de Representação

Usando um novo protocolo de sondagem no benchmark MoleculeNet, representações extraídas do Hyformer pré-treinado conjuntamente foram as mais preditivas tanto na sondagem linear quanto na KNN. O Hyformer alcançou o melhor desempenho em 4 de 10 conjuntos de dados para sondagem linear e 5 de 10 para sondagem KNN, superando modelos preditivos de última geração como Uni-Mol e outros modelos conjuntos.

Geração Incondicional e Desenho de AMP

• Geração Incondicional: Nos benchmarks GuacaMol e MOSES, o Hyformer consistentemente permaneceu na fronteira de Pareto, equilibrando fidelidade distribucional (FCD, divergência KL), validade e unicidade, rivalizando com modelos baseados em SMILES e baseados em grafos de última geração.
• Desenho de Peptídeos Antimicrobianos: Em uma tarefa real de desenho de fármacos, o Hyformer gerou peptídeos com alta aptidão e probabilidade de AMP, superando significativamente as linhas de base (PepCVAE, AMPGAN, HydrAMP, AMP-Diffusion) enquanto mantinha perplexidade e diversidade competitivas. A análise de atenção revelou que o modelo priorizou aminoácidos biologicamente relevantes (Arginina, Lisina, Triptofano, Prolina) associados à atividade antimicrobiana.

Significado e Alegações

O artigo afirma que o Hyformer unifica com sucesso a geração molecular e a predição de propriedades para representações baseadas em SMILES, abrindo caminhos para integração em pipelines reais de descoberta de fármacos. Os autores enfatizam que os benefícios sinérgicos da aprendizagem conjunta — especificamente o controle aprimorado sobre o processo generativo, a robustez a dados OOD e a aprendizagem de representações preditivas de alto nível — são alcançáveis sem sacrificar o desempenho de tarefas individuais.

Os autores reconhecem compensações inerentes, observando que parâmetros compartilhados podem limitar a especialização específica de tarefas. Eles sugerem que trabalhos futuros poderiam explorar arquiteturas de atenção dinâmicas ou modulares para alocar capacidade de forma mais flexível. Além disso, observam que a extensão na qual esses benefícios sinérgicos se aplicam a outras modalidades, como estruturas 3D ou transcriptômica, permanece uma questão em aberto. O trabalho é apresentado como um passo em direção a ferramentas que facilitam o desenho de novos medicamentos potenciais, em vez de uma solução para aplicações maliciosas.