Bias mitigation in graph diffusion models

Título: Mitigação de Viés em Modelos de Difusão em Grafos

Autores: Meng Yu e Kun Zhan (Lanzhou University)

1. O Problema: Viés em Modelos de Difusão em Grafos

O artigo identifica que os modelos de difusão em grafos existentes (como GDSS, MOOD, GSDM, HGDM) sofrem de dois problemas fundamentais que degradam a qualidade da geração:

1. Viés de Início Reverso (Reverse-Starting Bias):

   • Causa: Devido às limitações de escala dos dados de grafos e à capacidade de aprendizado das redes, os modelos truncam o processo de difusão direta (forward) para evitar instabilidades em estados de alto ruído. Consequentemente, a distribuição de perturbação máxima no final do processo direto ($x_T$) não atinge a distribuição Gaussiana Padrão ($N(0, I)$), mas permanece em um estado de "baixo ruído" (low-noise state).
   • O Erro: No entanto, durante a amostragem reversa (geração), os modelos começam arbitrariamente a partir de uma distribuição Gaussiana Padrão.
   • Consequência: Existe uma incompatibilidade fundamental entre o ponto de partida da amostragem e a distribuição real esperada pelo modelo treinado, introduzindo um viés severo que prejudica a qualidade final.

2. Viés de Exposição (Exposure Bias):

   • Causa: Durante o processo de difusão direta, o modelo é treinado com amostras ruidosas reais ($x_t$). Na amostragem reversa, o modelo gera amostras preditas ($\hat{x}_t$) que contêm erros de previsão da rede de pontuação (score network).
   • Consequência: Esses erros se acumulam e propagam ao longo das etapas de amostragem. O modelo é exposto a dados que não seguem a distribuição de treinamento, levando a uma degradação progressiva da qualidade.

2. Metodologia Proposta (S++)

Os autores propõem uma abordagem abrangente chamada S++ que mitiga ambos os vieses sem modificar a arquitetura da rede ou introduzir novos componentes (como geradores ou discriminadores adicionais). A metodologia consiste em duas etapas principais:

A. Alinhamento de Início Reverso via Amostragem de Langevin

Para resolver o Reverse-Starting Bias, o método não força a distribuição direta a atingir o Gaussiano (o que exigiria re-treinamento complexo), mas sim ajusta o ponto de partida da amostragem reversa.

• Mecanismo: Utiliza a rede de pontuação treinada $s_{\theta, T}(\cdot)$ para guiar uma amostragem de Langevin a partir do Gaussiano Padrão.
• Objetivo: Essa amostragem itera para encontrar uma distribuição estável que se alinhe com a distribuição de perturbação máxima real do processo direto ($q(x_T|x_0)$).
• Benefício: Isso estabelece um novo ponto de início reverso que está no estado de "baixo ruído", onde o modelo é mais robusto a erros de previsão, reduzindo a sensibilidade inicial e mitigando o viés de início.

B. Correção de Pontuação baseada em Diferença (Score Correction)

Para resolver o Exposure Bias e refinar a distribuição estável da amostragem de Langevin, propõe-se um mecanismo de correção de pontuação.

• Conceito:
  1. Usa-se a rede de pontuação pré-treinada ($s_{\theta}$) para gerar um lote de amostras pseudo-dados ($\hat{x}_0$).
  2. Treina-se uma rede de pontuação pseudo ($s_{\psi}$) usando esses dados gerados. Como os dados gerados contêm viés, a rede $s_{\psi}$ aprende esse viés.
  3. Calcula-se a diferença de pontuação entre a rede original e a rede pseudo: $\Delta s = s_{\theta} - s_{\psi}$.
• Correção: A diferença de pontuação contém informações sobre o dado original verdadeiro ($x_0$). O método ajusta a pontuação predita adicionando essa diferença ponderada:
  $$s'_{\theta,t} = s_{\theta,t} + \lambda (s_{\theta,t} - s_{\psi,t})$$
  Onde $\lambda$ controla a direção da correção e um fator $\omega$ ajusta a magnitude.
• Resultado: Isso "puxa" a direção viésada da pontuação em direção à direção não viésada (verdadeira), melhorando a precisão da previsão e reduzindo o acúmulo de erro durante a amostragem.

3. Contribuições Principais

1. Primeira Abordagem Sistemática: É a primeira vez que os vieses de início reverso e de exposição em modelos de difusão em grafos são analisados e mitigados conjuntamente de forma sistemática.
2. Método Plug-and-Play: A solução não requer alterações na arquitetura da rede neural, nem introduz novos componentes aprendíveis (como GANs ou discriminadores). Pode ser integrada a qualquer modelo de difusão em grafos existente.
3. Alta Eficiência e Qualidade: O método permite a geração de amostras de alta qualidade com menos passos de difusão reversa, acelerando o processo de inferência.
4. Validação Abrangente: O método foi validado em múltiplos modelos (GDSS, GSDM, HGDM, MOOD), diferentes tipos de SDEs (VPSDE, VESDE) e diversos conjuntos de dados.

4. Resultados Experimentais

Os experimentos foram realizados em conjuntos de dados genéricos (Community-small, Enzymes, Grid) e moleculares (QM9, ZINC250k).

• Geração de Grafos Genéricos: O S++ superou todos os baselines (incluindo versões com corretores tradicionais) em métricas de distribuição (MMD de Grau, Coeficiente de Agrupamento e Órbitas). Em muitos casos, o S++ sem corretores (apenas com alinhamento) superou os baselines com corretores complexos.
• Geração Molecular:
  • Qualidade: Redução significativa na distância Fréchet ChemNet (FCD) e no MMD NSPDK, indicando maior similaridade com a distribuição real de moléculas.
  • Validade: Aumento na taxa de validade das moléculas geradas (sem correção química).
  • Descoberta de Fármacos: No modelo MOOD-S++, houve melhoria no "hit ratio" (proporção de moléculas únicas e promissoras) e na pontuação de acoplamento (docking score) para alvos proteicos específicos (PARP1, BRAF, etc.).
• Eficiência: O método conseguiu manter desempenho próximo ao ótimo mesmo com uma redução drástica no número de passos de amostragem (ex: de 1000 para 100 passos no QM9), enquanto os modelos baselines degradavam severamente.

5. Significado e Conclusão

O artigo demonstra que o viés de início reverso é um problema crítico e único em modelos de difusão em grafos, distinto dos modelos de imagem, devido à natureza esparsa e interconectada dos dados de grafos e às limitações de treinamento em estados de alto ruído.

A proposta S++ oferece uma solução elegante e eficiente:

• Elimina a necessidade de re-treinamento de modelos complexos.
• Resolve a raiz do problema de exposição ao alinhar o início da geração com a realidade do modelo treinado.
• Estabelece um novo estado da arte (SOTA) em tarefas de geração de grafos e moléculas, tornando os modelos de difusão mais robustos e práticos para aplicações reais em descoberta de fármacos e modelagem de redes.