Uncertainty Estimation for Molecular Diffusion Models

Resumo Técnico: Estimativa de Incerteza para Modelos de Difusão Molecular

Definição do Problema
Os modelos de difusão tornaram-se o paradigma dominante para a geração de moléculas 3D. No entanto, existe um gargalo prático crítico: esses modelos carecem de um sinal fundamentado para identificar quando uma molécula gerada é quimicamente inválida ou instável. Enquanto modelos preditivos tipicamente utilizam a estimativa de incerteza para sinalizar saídas não confiáveis, os métodos existentes para modelos de difusão têm se concentrado primariamente no domínio de imagens. No domínio molecular, onde a avaliação downstream (ex: docking ou validação em laboratório úmido) é dispendiosa, confiar em filtros grosseiros baseados em regras ou na própria verossimilhança do modelo — que frequentemente é mal calibrada — permanece como o padrão. Há uma falta de trabalhos análogos que forneçam sinais de qualidade de amostra para modelos de difusão molecular.

Metodologia
Os autores propõem um método post-hoc para estimar a incerteza por amostra para modelos de difusão molecular pré-treinados sem retreinar o gerador. A abordagem baseia-se na intuição de que as amostras sobre as quais o modelo possui incerteza interna exibirão trajetórias de denoising mais variáveis.

1. Aproximação de Laplace: Para manter a compatibilidade com qualquer modelo pré-treinado, o método ajusta uma aproximação de Laplace aos parâmetros da rede de denoising. Especificamente, uma distribuição posterior $q(\theta) \approx \mathcal{N}(\theta | \hat{\theta}, \Sigma)$ é construída em torno dos parâmetros pré-treinados $\hat{\theta}$. A covariância $\Sigma$ é aproximada usando a matriz de informação de Fisher empírica (aplicada apenas às camadas finais para garantir escalabilidade) somada a um termo de precisão de prior.
2. Variabilidade da Predição de Ruído: Durante a trajetória de geração, para um subconjunto selecionado de timesteps $\mathcal{T}$, o método extrai $M$ amostras de pesos $\theta_m$ da posterior $q(\theta)$. Para cada amostra, o denoiser prediz o ruído $\epsilon^m_t = f_{\theta_m}(\hat{x}_t, t)$.
3. Agregação do Score de Incerteza: A incerteza intermediária $u_t$ em cada timestep é definida como a variância amostral elemento a elemento dessas predições de ruído através das $M$ amostras. Esta incerteza por passo é então agregada através dos timesteps selecionados, todos os átomos ($N_a$) e todas as dimensões de características ($D$) para produzir um único score de incerteza escalar $u(\hat{x})$ para a molécula gerada:
   $$u(\hat{x}) := \frac{1}{|\mathcal{T}| \cdot N_a \cdot D} \sum_{t \in \mathcal{T}} \sum_{n=1}^{N_a} \sum_{d=1}^{D} [u_t]_{n,d}$$
   Empiricamente, os autores descobrem que medir a incerteza apenas nos 10% finais do processo de geração (onde $t$ é pequeno) é suficiente para o desempenho, limitando o overhead computacional.

Principais Contribuições

• Primeira Aplicação à Difusão Molecular: Este trabalho estende as abordagens de incerteza de difusão do domínio de imagem para a geração molecular, abordando a necessidade específica de sinais de qualidade de amostra em domínios científicos com avaliação downstream dispendiosa.
• Framework Post-hoc: O método não requer retreinamento do modelo de difusão subjacente, tornando-o aplicável a arquiteturas de estado da arte existentes como EDM e GeoLDM.
• Escalonamento em Tempo de Teste: O artigo demonstra como o score de incerteza pode ser usado como um critério de filtragem em um cenário de escalonamento de tempo de teste (oversampling e filtragem) para melhorar a qualidade das moléculas geradas.

Resultados
Os autores avaliaram sua abordagem nos datasets QM9 e GEOM-Drugs utilizando modelos EDM e GeoLDM.

• Correlação com a Qualidade: No dataset QM9, o score de incerteza proposto exibe uma correlação negativa estatisticamente significativa com métricas de qualidade estabelecidas (estabilidade molecular, estabilidade atômica e validade). Isso indica que scores de incerteza mais altos correspondem a amostras de menor qualidade.
• Comparação com Verossimilhança: O score de incerteza supera consistentemente a verossimilhança negativa (NLL) do modelo de difusão como um preditor da qualidade da amostra. Por exemplo, no EDM/QM9, o score de incerteza alcançou uma correlação de Spearman de $-0,284$ com a estabilidade molecular, comparado a $-0,150$ para NLL.
• Desempenho de Escalonamento em Tempo de Teste: Ao gerar 20 mil moléculas e filtrar para manter as 10 mil com menor incerteza, os autores observaram melhorias significativas na estabilidade molecular ($\sim10\%$), estabilidade atômica ($\sim1\%$) e validade ($\sim5\%$) para o modelo EDM. Esses ganhos foram comparáveis à melhoria alcançada ao trocar a arquitetura do modelo subjacente de EDM para GeoLDM.
• Estudos de Ablação:
  • Aproximação de Fisher: Remover a aproximação de Fisher (substituindo a covariância por uma Gaussiana isotrópica) resultou em um desempenho quase idêntico, sugerindo que o score atua mais como uma medida de sensibilidade a pequenas perturbações do que como uma incerteza epistêmica Bayesiana estrita.
  • Seleção de Timestep: O sinal é mais forte na extremidade "limpa" da trajetória de geração (pequeno $t$), validando a escolha de focar a medição de incerteza nos passos finais.
• Limitações: O método não gerou melhorias significativas no dataset GEOM-Drugs, que é maior e mais complexo, indicando que a abordagem não transfere atualmente para datasets com maior complexidade conformacional.

Significância e Alegações
O artigo afirma dar um "primeiro passo" para estender as abordagens de incerteza de difusão para o domínio molecular. A principal significância reside em fornecer uma ferramenta post-hoc prática para filtrar moléculas geradas, melhorando assim a eficiência dos pipelines de avaliação downstream. Os autores posicionam seu score de incerteza como um substituto natural para verificadores baseados em NLL em métodos de escalonamento adaptativo em tempo de teste.

Os autores mantêm-se modestos quanto às suas alegações, observando explicitamente que a abordagem atualmente falha em datasets maiores e mais complexos (GEOM-Drugs) e que o mecanismo subjacente parece ser uma medida de sensibilidade em vez de verdadeira incerteza epistêmica Bayesiana. Eles sugerem que trabalhos futuros devem explorar o cálculo da incerteza em um espaço de características "semânticas" usando encoders moleculares pré-treinados, analogamente aos sucessos recentes em geração de imagens de alta dimensão.