Interpreting GFlowNets for Drug Discovery: What probes can and cannot show

Resumo Técnico: Interpretando GFlowNets para Descoberta de Fármacos

Definição do Problema
Redes de Fluxo Generativo (GFlowNets), particularmente aquelas treinadas para o design molecular como a SynFlowNet, oferecem um arcabouço poderoso para explorar o espaço químico enquanto aderem a restrições sintéticas. No entanto, suas políticas de decisão internas permanecem opacas, limitando sua adoção na descoberta de fármacos, onde químicos medicinais exigem justificativas interpretáveis para as estruturas propostas. Métodos de interpretabilidade existentes são frequentemente projetados para espaços latentes contínuos ou modelos baseados em verossimilhança, não para as políticas discretas, sequenciais e baseadas em grafos das GFlowNets. Além disso, há uma necessidade crítica de distinguir entre o conhecimento químico genuinamente aprendido pela política e a informação que está trivialmente disponível através da arquitetura do grafo e da featurização dos átomos.

Metodologia
Os autores introduzem um arcabouço de interpretabilidade multiescala aplicado à SynFlowNet (treinada com uma função de recompensa QED) usando embeddings congelados de um checkpoint específico. O arcabouço integra três abordagens complementares:

1. Saliência Baseada em Gradiente e Contrafatuais: Gradientes Integrados (IG) são aplicados ao log-probabilidade da ação "Stop" para gerar mapas de saliência ao nível do átomo. Átomos de alta saliência são agrupados em motivos (componentes conectados e sistemas de anéis). Perturbações contrafatuais são então aplicadas a esses motivos usando regras de transformação do RDKit (ex: troca de halogênio, modificação de amida) para medir mudanças no QED ($\Delta$QED).
2. Autoencoders Esparsos (SAEs):
   • Subcompleto (MLP-SAE): Um autoencoder compressivo (256 $\to$ 128) é usado para identificar fatores latentes esparsos e dominantes e testar sua correlação com propriedades físico-químicas (ex: polaridade, lipofilicidade).
   • Sobrecompleto (BatchTopK SAE): Um autoencoder sobrecompleto (256 $\to$ 1024) com um mecanismo de esparsidade BatchTopK ($k=64$) é empregado para recuperar detectores de características finas, incluindo aqueles em superposição, sem regularização $\ell_1$.
3. Probing e Controles: Classificadores feedforward rasos e sondas estritamente lineares (probes de regressão Ridge/logística) são treinados nos embeddings para prever a presença de motivos químicos (grupos funcionais, halogênios) e propriedades físico-químicas. Crucialmente, o estudo emprega controles rigorosos:
   • Baselines de Descritores: Comparação do desempenho das sondas contra descritores RDKit e impressões digitais (fingerprints) ECFP4.
     implementação de controles de rótulos aleatórios para garantir que as sondas não memorizem ruído.
   • Divisões Disjuntas de Scaffold: Prevenção de vazamento de scaffold nas divisões de treino/teste.
   • Controle de Rede Não Treinada: Execução de sondas idênticas em uma rede aleatoriamente inicializada com a exata mesma arquitetura para isolar a contribuição do treinamento versus a arquitetura.
   • Ablação de Características: Zeramento de características únicas para testar a especificidade na decodificação subsequente.

Principais Resultados

• Saliência e Contrafatuais: Os Gradientes Integrados destacam regiões quimicamente significativas (substituintes polares, fragmentos aromáticos). Edições contrafatuais a esses motivos produzem mudanças pequenas, mas mensuráveis no QED (ex: +0,015 a +0,051), embora os tamanhos de efeito sejam pequenos em relação ao intervalo do QED, e uma comparação controlada contra edições aleatórias permanece como trabalho futuro.
• Estrutura Latente: O SAE subcompleto identifica fatores fortemente correlacionados com tamanho e polaridade, mas um preditor MLP sobre esses fatores apresenta desempenho pobre em QED composto ($R^2=0,25$) comparado a propriedades individuais.
• Reconstrução e Esparsidade: O SAE sobrecompleto BatchTopK reconstrói os embeddings com baixo NMSE (0,0015) e mantém uma esparsidade de média de lote estrita de 64 características ativas. Ele supera um baseline linear correspondente (NMSE 0,0026) em níveis de atividade similares.
• Decodificabilidade Linear: Os embeddings decodificam propriedades físico-químicas bem (ex: QED $R^2 \approx 0,59$, tamanho $R^2 \approx 0,996$). No entanto, esses scores não são significativamente melhores do que os alcançados por um baseline de descritor RDKit "livre de vazamento" sob divisões aleatórias.
• O Controle de Rede Não Treinada (Achado Decisivo): Uma rede aleatoriamente inicializada da mesma arquitetidade decodifica cada propriedade sondada tão bem quanto a política treinada. Para a propriedade principal (semelhança com fármacos/QED), o treinamento fornece um ganho negligenciável de apenas +0,005 (0,581 não treinado vs. 0,586 treinado), o que entra na variância de sementes aleatórias. Em termos de tamanho molecular, a rede não treinada é marginalmente melhor.
• Especificidade de Característica: No dicionário sobrecompleto, execuções individuais revelam detectores quimicamente enriquecidos (ex: detectores separados para iodo, cloro, bromo, flúor e boro). A ablação de característica única mostra efeitos específicos de propriedade (razões de especificidade de 8–23$\times$), mas os índices das características não são reproduzíveis entre sementes; apenas um pequeno núcleo de direções é estável.

Significância e Alegações
O resultado mais consequente do artigo é um controle negativo: a decodificabilidade linear de propriedades físico-químicas a partir dos embeddings da GFlowNet reflete a arquitetura do grafo e a featurização dos átomos, não o conhecimento químico aprendido durante o treinamento.

• O que NÃO é sustentado: Altos scores de sondagem (mesmo com AUROCs próximos de 1) não são evidência de compreensão química adquirida ou raciocínio de relação estrutura-atividade. A capacidade de decodificar propriedades é, em grande parte, um artefato da representação de entrada e da arquitetura da rede, conforme demonstrado pela rede não treinada que iguala o desempenho da política treinada.
• O que É sustentado:
  1. O dicionário sobrecompleto fornece uma decomposição de esparsidade exata e válida dos embeddings.
  2. Execuções individuais contêm detectores de subestruturas quimicamente enriquecidos (ex: características por halogênio) que são estáveis através de um pequeno subconjunto de sementes.
  3. Essas características possuem efeitos intervencionistas específicos de propriedade na decodificação da sonda.
  4. A semelhança com fármacos (QED) torna-se mais linearmente acessível no espaço latente esparso do que no embedding denso.

Conclusão
Os autores concluem que métodos de interpretabilidade devem ser rigorosamente controlados contra baselines arquiteturais. Para a SynFlowNet, o arcabouço identifica com sucesso detectores de características finas e valida a robustez da representação esparsa. No entanto, o estudo explicitamente alerta contra interpretar altos scores de decodificabilidade como prova de química aprendida. O artigo restringe suas alegações à existência de um núcleo reproduzível de características estáveis e quimicamente enriquecidas dentro do dicionário, enquanto afirma que a capacidade do modelo de prever propriedades físico-químicas é primariamente uma função de sua arquitetura, e não de seu treinamento. Trabalhos futuros devem explorar se esses detectores persistem quando os SAEs são treinados em embeddings de redes não treinadas e estender o arcabouço para configurações de múltiplos objetivos.