Resumo Técnico: Sobre o Espaço de Design da Adaptação Online de Difusão Discreta para Otimização Molecular

Declaração do Problema

A otimização molecular geralmente começa com um modelo generativo pré-treinado (por exemplo, um modelo de difusão discreta) que captura um prior amplo sobre estruturas químicas válidas. No entanto, o objetivo prático no momento do teste não é amostrar deste prior, mas identificar moléculas que maximizem uma função de recompensa específica da tarefa $r(x)$ (por exemplo, afinidade de ligação ou fitness proteico) usando um orçamento limitado de oráculo. O oráculo, que avalia a recompensa, é frequentemente custoso (por exemplo, ensaios de laboratório úmido ou pipelines de docking caros).

O desafio central é a adaptação de modelo online: direcionar um gerador pré-treinado $p_{\theta_0}$ em direção a regiões de alta recompensa do espaço químico, mantendo a validade e a diversidade. Isso envolve um ciclo fortemente acoplado onde o sistema deve decidir:

1. Aquisição: Quais candidatos avaliar dado o orçamento do oráculo.
2. Adaptação: Como converter recompensas observadas em atualizações de modelo.
3. Reuso de Feedback: Quais dados históricos reter para treinamento.
4. Debiasing de Prior: O quanto deslocar o gerador de seus modos pré-treinados sem colapsar a validade.

Embora esses componentes tenham sido estudados isoladamente, permanece incerto se eles se complementam, interferem ou tornam-se redundantes quando integrados em um ciclo completo de adaptação online.

Metodologia

Os autores propõem um estudo controlado de um "harness" modular para a adaptação online de modelos de difusão discreta. O ciclo (Figura 1) alterna entre amostrar candidatos, selecionar um lote para avaliação do oráculo e realizar o ajuste fino (fine-tuning) do gerador. O estudo avalia cinco escolhas específicas de design de "plug-in" (Tabela 1) através de nove tarefas: seis tarefas de afinidade de ligação de pequenas moléculas e três tarefas de fitness proteico.

Componentes Principais de Design

1. Direcionamento (Aquisição):

   • Amostragem de Thompson: Em vez de uma seleção gananciosa baseada na previsão da média de um modelo substituto, os autores utilizam a amostragem de Thompson para selecionar candidatos. Isso equilibra exploração e explotação ao amostrar estimativas de recompensa $\hat{r}_i = \mu_i + \epsilon_i \sigma_i$, onde $\sigma_i$ representa o desacordo do ensemble (incerteza). Isso incentiva a avaliação de candidatos com alta incerteza, mesmo que sua previsão média seja menor.

2. Adaptação (Modelagem de Recompensa & Debiasing):

   • Modelagem de Recompensa por CVaR: Para focar as atualizações na cauda de alta recompensa em vez da média, a recompensa é moldada usando o Valor em Risco Condicional (CVaR). A recompensa é transformada para $\tilde{r}(x) = (r(x) - Q_\tau)_+ / (1-\tau)$, onde $Q_\tau$ é o $\tau$-quantil (definido como 0,8). Isso concentra as atualizações de gradiente nas amostras de melhor desempenho.
   • Regularização de Entropia de Densidade (Debiasing do Modelo): Para evitar que o modelo colapse nos modos de alta densidade já existentes no prior, um regularizador é adicionado: $\log \tilde{r}_{debias}(x) = \log r(x) - \gamma \log p_{\theta_t}(x)$. Como computar a verossimilhança marginal exata $\log p_{\theta_t}(x)$ para difusão discreta é intratável, os autores utilizam um estimador de Monte Carlo computacionalmente eficiente baseado no ELBO da difusão (uma aproximação de passagem única) para aproximar o log-verossimilhança negativa.

3. Correções de Descompasso de Objetivo:

   • Buffer de Replay: Para evitar que o modelo esqueça descobertas de alta recompensa feitas em rodadas iniciais conforme a distribuição de amostragem muda, um buffer FIFO de capacidade fixa com expulsão por prioridade (removendo primeiro as amostras de menor recompensa) é utilizado. Mini-lotes são extraídos via amostragem estratificada para garantir que outliers de alta recompensa não sejam diluídos.
   • Penalidade de Saída Inválida: Para manter a validade (por exemplo, strings SMILES válidas), candidatos que falham na análise (por exemplo, pelo RDKit) recebem uma penalidade negativa fixa $r_{inv}$. Isso puxa o gerador de volta para o manufoldo molecular válido quando a exploração produz estruturas inválidas.

Configuração Experimental

• Tarefas: Otimização de pequenas moléculas para ligação a seis alvos proteicos (usando GenMol como o gerador e FlashAffinity/Boltz-2 como oráculos) e otimização de fitness proteico para três famílias (GB1, CreiLOV, TrpB).
• Baselines: O ciclo online completo é comparado contra:
  • Ajuste fino offline (treinamento em um conjunto de dados estático).
  • Busca em tempo de inferência (Beam Search, MCTS, Best-of-N, etc.) com um gerador fixo.
  • Versões com ablação do ciclo online (removendo um componente por vez).
• Métricas: O desempenho é medido pela melhor recompensa encontrada, médias do top-k e recompensa normalizada, avaliadas contra chamadas cumulativas de oráculo e horas de GPU. Métricas secundárias incluem validade, diversidade (Tanimoto), QED e acessibilidade sintética (SA).

Principais Resultados

1. Complementaridade dos Componentes

O estudo descobre que os componentes propostos fornecem benefícios complementares, particularmente para tarefas de pequenas moléculas onde os candidatos de alta recompensa estão longe do prior pré-treinado.

• Pequenas Moléculas: A aquisição (amostragem de Thompson), a modelagem de CVaR e a Regularização de Entropia de Densidade produzem os ganhos de recompensa mais claros. A amostragem de Thompson melhora a alocação de chamadas do oráculo; o CVaR foca as atualizações na cauda de recompensa; e o debiasing empurra a exploração para longe dos modos de alta densidade do prior.
• Fitness Proteico: Os efeitos de ablação são geralmente mais fracos. Os priors específicos da família para proteínas já cobrem regiões de alto fitness, tornando os grandes deslocamentos de distribuição menos críticos. No entanto, a amostragem de Thompson ainda fornece ganhos consistentes ao melhorar a alocação do oráculo.

2. Mecanismos de Estabilidade

• Buffers de Replay e Penalidades de Validade atuam como estabilizadores em vez de impulsionadores primários de recompensa. O buffer de replay alinha os gradientes com a distribuição média de tempo das amostras avaliadas, evitando a perda de descobertas de alta recompensa iniciais. A penalidade de validade evita o "reward hacking", onde o modelo gera strings quimicamente inválidas, um modo de falha que ocorre rapidamente sem a penalidade.

3. Adaptação Online vs. Busca em Tempo de Inferência

O ciclo online completo supera significativamente as baselines de busca em tempo de inferência (ex: Beam Search, MCTS) sob orçamentos equivalentes de chamadas de oráculo e horas de GPU.

• Eficiência: Em tarefas de pequenas moléculas, o ciclo online atinge o escore final da melhor baseline de busca em aproximadamente 58× menos chamadas de oráculo. Em tarefas de proteínas, a aceleração é de 16×.
• Raciocínio: Métodos de busca gastam o orçamento gerando e pontuando candidatos de um gerador fixo. Em contraste, a adaptação online reutiliza o feedback do oráculo para atualizar o próprio gerador, multiplicando o valor de cada chamada de oráculo. Métodos de busca também consomem muito computação para geração, deixando pouco tempo para passos de gradiente de ajuste fino dentro de um orçamento fixo de tempo de execução (wall-clock).

4. Análise de Deslocamento de Distribuição

Os ganhos de desempenho são mais pronunciados quando a solução ótima requer um grande deslocamento do prior pré-treinado.

• Regularização de Entropia de Densidade é o principal motor do deslocamento de distribuição, empurrando a massa do gerador para longe das regiões de alta densidade do prior em direção a áreas de baixa densidade e alta recompensa.
• CVaR e Amostragem de Thompson facilitam esse deslocamento condicionalmente, dependendo de onde os candidatos de alta recompensa estão localizados em relação ao prior.
• As tarefas de proteína mostram deslocamentos de distribuição mais fracos porque os priors específicos da família já estão concentrados perto de regiões de alto fitness.

Significância e Alegações

O artigo afirma fornecer uma receita prática para otimização molecular eficiente em termos de feedback: ajuste fino online com aquisição, modelagem de recompensa, debiasing, replay e controle de validade.

• Novidade: Este trabalho integra e avalia sistematicamente esses componentes dentro de um único framework de difusão discreta, demonstrando que eles não são redundantes, mas sim mutuamente reforçadores.
• Eficiência: O método proposto supera tanto o ajuste fino offline quanto as baselines de busca em tempo de inferência, particularmente em cenários onde o oráculo é caro e as moléculas alvo estão fora da distribuição (out-of-distribution) em relação ao modelo pré-treinado.
• Robustez: A receita mostra-se robusta através de diferentes objetivos de ajuste fino (DDPP-LB e VIDD) e diferentes domínios (pequenas moléculas vs. proteínas), embora a magnitude dos ganhos varie com base na distância entre o prior e a distribuição alvo.

Os autores concluem que, embora os componentes sejam mais críticos para o design de pequenas moléculas (onde o prior é amplo e o alvo é distante), o framework fornece uma abordagem generalizável para direcionar modelos generativos discretos em tarefas de descoberta científica com feedback limitado.