DriftingMol: Decoder-Coupled Drift for One-Pass Property-Conditional Molecular Generation

Resumo Técnico: DriftingMol

Declaração do Problema
O artigo aborda o desafio da geração molecular condicional a propriedades, que exige produzir moléculas válidas e diversas que respondam a valores contínuos de propriedades-alvo (por exemplo, QED, LogP) com baixos custos de amostragem. Métodos existentes enfrentam um trade-off: abordagens de amostragem iterativa (como difusão ou aprendizado por reforço) são computacionalmente caras na inferência, enquanto decodificadores no espaço latente ou baseados em grafos frequentemente perdem sinais contínuos de propriedades quando a decodificação discreta colapsa estados latentes próximos. Além disso, garantir a validade molecular (correção sintática e de valência) permanece uma restrição discreta que complica o controle contínuo de propriedades.

Metodologia
DriftingMol introduz um framework de duas etapas que adapta "modelos de deriva" a um espaço latente molecular SELFIES.

1. Pipeline de Duas Etapas:

   • Etapa 1: Um $\beta$-VAE SELFIES é treinado no conjunto de dados ZINC250K. O codificador mapeia sequências de tokens para um espaço latente de 256 dimensões ($z$), e um decodificador Transformer de um único shot prevê posições SELFIES em paralelo. Esta etapa estabelece um espaço latente contínuo regularizado onde a validade é herdada da representação SELFIES, e não aprendida pelo gerador.
   • Etapa 2: O VAE é congelado. Um gerador no estilo DiT ($f_\theta$) é treinado para mapear ruído gaussiano, condições de propriedade opcionais ($c$) e uma escala de orientação embutida ($\alpha$) para um vetor latente $\hat{z}$ em uma única passagem direta. A molécula final é obtida decodificando $\hat{z}$ através do decodificador SELFIES congelado.

2. Deriva Acoplada ao Decodificador:
   A inovação central é a definição do campo de deriva. Em vez de usar um extrator de características externo, o DriftingMol utiliza a representação oculta do decodificador VAE congelado ($\phi$) como o mapa de características de deriva.

   • Campo de Deriva: Um campo de deriva baseado em kernel é construído entre latentes gerados ($\hat{z}_i$) e latentes de referência ($z_{ref, j}$) com base em suas características de decodificador: $V_{\tau, i} = \sum_j W^\tau_{ij} (\phi(z_{ref, j}) - \phi(\hat{z}_i))$.
   • Acoplamento de Gradiente: Crucialmente, embora os pesos do decodificador estejam congelados, o caminho do gradiente da perda de deriva através do mapa de características do decodificador $\phi(\hat{z})$ de volta aos parâmetros do gerador $\theta$ é mantido aberto. Isso cria uma "métrica de pullback" onde as atualizações do gerador são ponderadas pela sensibilidade estrutural local do decodificador.
   • Diversidade: Um termo de repulsão latente $k$-NN ($L_{zdiv}$) é adicionado à função de perda para evitar o colapso de modos e garantir unicidade.

3. Objetivo de Treinamento:
   O gerador minimiza uma perda de deriva ($L_{drift}$) que puxa as características de decodificador do latente gerado em direção a uma distribuição-alvo definida por referências compatíveis com propriedades, combinada com a regularização de diversidade.

Principais Contribuições

• Pipeline de Deriva Latente SELFIES: Um framework de baixo custo para geração condicional a propriedades que separa a validade (tratada por SELFIES/VAE) da controlabilidade.
• Deriva Acoplada ao Decodificador: Um mecanismo que utiliza o decodificador VAE congelado como espaço de características de deriva, preservando o caminho de gradiente através do decodificador até o gerador. Isso alinha o gradiente de deriva com o processo real de decodificação molecular.
• Evidências de Ablação: Um estudo abrangente de ablação com 15 condições demonstrando que o acoplamento ao decodificador é o principal motor do controle de propriedades, superando variantes de deriva no espaço latente, com características aleatórias e com características externas.
• Eficiência: O método alcança viés de propriedades com apenas uma avaliação do gerador mais uma passagem pelo decodificador congelado, evitando denoising iterativo ou integração de ODE.

Resultados
Avaliado no ZINC250K com condicionamento de QED e multi-propriedade:

• Controle de QED: A configuração padrão do DriftingMol alcança uma correlação de Spearman de QED ($\rho$) de 0,493 com 94,7% de unicidade. A variante mais forte acoplada ao decodificador (sem regularização de diversidade) atinge $\rho = 0,510$, mas com unicidade reduzida (74,1%).
• Condicionamento Multi-Propriedade: Sob um protocolo de condicionamento de quatro propriedades compatível (QED, SA, LogP, MolWt), a deriva acoplada ao decodificador alcança uma correlação média de Spearman de 0,598.
• Descobertas de Ablação:
  • Necessidade de Acoplamento: Modelos que desvinculam o decodificador ou usam stop-gradients produzem correlação de QED próxima de zero ($\rho \approx 0,01$) e unicidade muito baixa.
  • Qualidade das Características: Características treinadas externamente (por exemplo, LatentMAE) ou características aleatórias resultam em correlações significativamente menores ($\rho \approx 0,28\text{--}0,33$), indicando que características preditivas sozinhas são insuficientes sem o Jacobiano do decodificador.
  • Diversidade: Remover o termo de diversidade $z$ aumenta ligeiramente a correlação, mas reduz drasticamente a unicidade, confirmando o papel do termo na prevenção do colapso de modos.
• Baselines: Baselines com mesma espinha dorsal (VAE, WGAN-GP, DDPM, Flow-Matching) treinadas no mesmo cache latente alcançaram correlações significativamente menores ($\rho < 0,16$).

Significado e Alegações
O artigo alega que a deriva acoplada ao decodificador é um mecanismo eficiente para geração molecular enviesada por propriedades. O significado reside no insight matemático de que o decodificador congelado fornece não apenas um espaço de características, mas uma métrica de pullback específica ($J_\phi^\top J_\phi$) que define a geometria na qual o gerador é otimizado. Essa métrica alinha as atualizações do gerador com a sensibilidade estrutural do decodificador molecular.

Os autores enfatizam que:

1. Validade é de Nível de Representação: A validade é garantida pela representação SELFIES e pelo decodificador congelado, não aprendida pelo gerador condicional.
2. Eficiência de Uma Passagem: O método evita os altos custos de inferência da amostragem iterativa.
3. Viés de Propriedade vs. Exatidão: O sistema é projetado para geração enviesada por propriedades (enriquecendo bibliotecas em direção a regiões-alvo) e não para otimização exata de propriedades. Não garante sintetizabilidade ou atividade biológica, e as moléculas geradas requerem filtragem a jusante.
4. Especificidade do Mecanismo: Os ganhos de desempenho são atribuídos especificamente ao acoplamento de gradiente através do decodificador, e não meramente ao uso de um VAE pré-treinado ou extratores de características externos.