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

Résumé Technique : Interprétation des GFlowNets pour la découverte de médicaments

Énoncé du Problème
Les réseaux de flux génératifs (GFlowNets), particulièrement ceux entraînés pour la conception moléculaire comme SynFlowNet, offrent un cadre puissant pour explorer l'espace chimique tout en respectant les contraintes de synthèse. Cependant, leurs politiques de décision internes restent opaques, ce qui limite leur adoption dans la découverte de médicaments où les chimistes médicinaux exigent des raisonnements interprétables pour les structures proposées. Les méthodes d'interprétabilité existantes sont souvent conçues pour des espaces latents continus ou des modèles basés sur la vraisemblance, et non pour les politiques discrètes, séquentielles et basées sur des graphes des GFlowNets. De plus, il existe un besoin critique de distinguer la connaissance chimique véritablement apprise par la politique des informations qui sont trivialement disponibles via l'architecture du graphe et la featurisation des atomes.

Méthodologie
Les auteurs introduisent un cadre d'interprétabilité multi-échelle appliqué à SynFlowNet (entraîné avec une fonction de récompense QED) en utilisant des plongements (embeddings) gelés provenant d'un point de contrôle spécifique. Le cadre intègre trois approches complémentaires :

1. Salience basée sur le gradient et contre-factuels : Les gradients intégrés (Integrated Gradients - IG) sont appliqués au log-probabilité de l'action « Stop » pour générer des cartes de saillance au niveau de l'atome. Les atomes à haute saillance sont regroupés en motifs (composantes connexes et systèmes cycliques). Des perturbations contre-factuelles sont ensuite appliquées à ces motifs en utilisant les règles de transformation de RDKit (ex: échange d'halogène, modification d'amide) pour mesurer les changements de QED ($\Delta$QED).
2. Auto-encodeurs creux (Sparse Autoencoders - SAE) :
   • Sous-complet (MLP-SAE) : Un auto-encodeur compressif (256 $\to$ 128) est utilisé pour identifier des facteurs latents dominants et creux, et pour tester leur corrélation avec les propriétés physico-chimiques (ex: polarité, lipophilie).
   • Surcomplet (BatchTopK SAE) : Un auto-encodeur surcomplet (256 $\to$ 1024) avec un mécanisme de parcimonie BatchTopK ($k=64$) est employé pour récupérer des détecteurs de caractéristiques fins, y compris ceux en superposition, sans régularisation $\ell_1$.
3. Sondage et Contrôles : Des classifieurs de type feedforward peu profonds et des sondes strictement linéaires (régression Ridge/logistique) sont entraînés sur les plongements pour prédire la présence de motifs chimiques (groupes fonctionnels, halogènes) et de propriétés physico-chimiques. Crucialement, l'étude emploie des contrôles rigoureux :
   • Baselines de descripteurs : Comparaison des performances des sondes par rapport aux descripteurs RDKit et aux empreintes ECFP4.
   • Contrôles d'étiquettes aléatoires : S'assurer que les sondes ne mémorisent pas le bruit.
   • Séparations disjointes de squelettes (Scaffold-Disjoint Splits) : Prévenir la fuite de squelettes (scaffold leakage) dans les séparations train/test.
   • Contrôle de réseau non entraîné : Exécution de sondes identiques sur un réseau initialisé aléatoirement avec la même architecture exacte pour isoler la contribution de l'entraînement par rapport à l'architecture.
   • Ablation de caractéristiques : Mise à zéro de caractéristiques uniques pour tester la spécificité lors du décodage en aval.

Résultats Clés

• Salience et Contre-factuels : Les gradients intégrés mettent en évidence des régions chimiquement significatives (substituants polaires, fragments aromatiques). Les éditions contre-factuelles de ces motifs produisent des décalages faibles mais mesurables de QED (ex: +0,015 à +0,051), bien que les tailles d'effet soient faibles par rapport à la plage de QED, et une comparaison contrôlée contre des éditions aléatoires reste un travail futur.
• Structure Latente : Le SAE sous-complet identifie des facteurs fortement corrélés à la taille et à la polarité, mais un prédicteur MLP sur ces facteurs performe mal sur le QED composite ($R^2=0,25$) comparé aux propriétés individuelles.
• Reconstruction et Parcimonie : Le SAE surcomplet BatchTopK reconstruit les plongements avec un NMSE faible (0,0015) et maintient une parcimonie stricte de moyenne de lot de 64 caractéristiques actives. Il surpasse une baseline linéaire appariée (NMSE 0,0026) à des niveaux d'activité similaires.
• Décodabilité Linéaire : Les plongements décodent bien les propriétés physico-chimiques (ex: QED $R^2 \approx 0,59$, taille $R^2 \approx 0,996$). Cependant, ces scores ne sont pas significativement meilleurs que ceux obtenus par une baseline de descripteurs RDKit « sans fuite » sous des séparations aléatoires.
• Le Contrôle du Réseau Non Entraîné (Résultat Décisif) : Un réseau initialisé aléatoirement de l'architecture identique décode chaque propriété sondée aussi bien que la politique entraînée. Pour la propriété principale (similitude avec un médicament/QED), l'entraînement n'apporte qu'un gain négligeable de seulement +0,005 (0,581 non entraîné vs 0,586 entraîné), ce qui entre dans la variance des graines aléatoires. Sur la taille moléculaire, le réseau non entraîné est légèrement meilleur.
• Spécificité des Caractéristiques : Dans le dictionnaire surcomplet, des exécutions individuelles révèlent des détecteurs enrichis chimiquement (ex: caractéristiques distinctes pour l'iode, le chlore, le brome et le fluor). L'ablation d'une caractéristique unique montre des effets spécifiques aux propriétés (ratios de spécificité de 8 à 23×), mais les indices de caractéristiques ne sont pas reproductibles entre les graines ; seul un petit noyau de directions est stable.

Signification et Revendications
Le résultat le plus important de l'article est un contrôle négatif : la décodabilité linéaire des propriétés physico-chimiques à partir des plongements de GFlowNet reflète l'architecture du graphe et la featurisation des atomes, et non la connaissance chimique apprise pendant l'entraînement.

• Ce qui n'est PAS soutenu : Des scores de sondage élevés (même des AUROC proches de la perfection) ne sont pas une preuve de compréhension chimique acquise ou de raisonnement de structure-activité appris. La capacité de décoder des propriétés est largement un artefact de la représentation d'entrée et de l'architecture du réseau, comme le démontre le fait que le réseau non entraîné égale la performance de la politique entraînée.
• Ce qui EST soutenu :
  1. Le dictionnaire surcomplet fournit une décomposition exacte-par-parcimonie des plongements.
  2. Les exécutions individuelles contiennent des détecteurs de sous-structures enrichis chimiquement (ex: caractéristiques par halogène) qui sont stables à travers un petit sous-ensemble de graines.
  3. Ces caractéristiques ont des effets interventionnels spécifiques aux propriétés sur le décodage des sondes.
  4. La similitude avec un médicament (QED) devient plus linéairement accessible dans l'espace latent creux que dans le plongement dense.

Conclusion
Les auteurs concluent que les méthodes d'interprétabilité doivent être rigoureusement contrôlées par rapport aux baselines architecturales. Pour SynFlowNet, le cadre identifie avec succès des détecteurs de caractéristiques fins et valide la robustesse de la représentation creuse. Cependant, l'étude met explicitement en garde contre l'interprétation de scores de décodabilité élevés comme preuve de chimie apprise. L'article limite ses revendications à l'existence d'un noyau reproductible de caractéristiques stables et enrichies chimiquement au sein du dictionnaire, tout en affirmant que la capacité du modèle à prédire les propriétés physico-chimiques est principalement une fonction de son architecture plutôt que de son entraînement. Les travaux futurs devraient explorer si ces détecteurs persistent lorsque les SAE sont entraînés sur des plongements de réseaux non entraînés et étendre le cadre aux contextes multi-objectifs.