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

Technisches Resümee: Interpretation von GFlowNets für die Wirkstoffforschung

Problemstellung
Generative Flow Networks (GFlowNets), insbesondere solche, die für das molekulare Design wie SynFlowNet trainiert wurden, bieten einen leistungsstarken Rahmen zur Exploration des chemischen Raums unter Einhaltung synthetischer Constraints. Die internen Entscheidungsstrategien (Policies) dieser Modelle bleiben jedoch opak, was ihre Anwendung in der Wirkstoffforschung einschränkt, wo Medizinalchemiker interpretierbare Begründungen für vorgeschlagene Strukturen benötigen. Bestehende Interpretationsmethoden sind oft für kontinuierliche latente Räume oder Likelihood-basierte Modelle konzipiert, nicht für die diskreten, sequentiellen, graphbasierten Policies von GFlowNets. Zudem besteht ein kritischer Bedarf daran, zwischen echtem chemischem Wissen, das von der Policy gelernt wurde, und Informationen zu unterscheiden, die trivial durch die Grapharchitektur und die Atom-Featurisierung verfügbar sind.

Methodik
Die Autoren führen ein Multi-Skalen-Interpretations-Framework ein, das auf SynFlowNet (trainiert mit einer QED-Belohnungsfunktion) unter Verwendung gefrorener Embeddings eines spezifischen Checkpoints angewendet wird. Das Framework integriert drei komplementäre Ansätze:

1. Gradientenbasierte Salienz und Kontrafaktika: Integrated Gradients (IG) werden auf die Log-Wahrscheinlichkeit der „Stop“-Aktion angewendet, um Atom-basierte Salienz-Maps zu generieren. Hoch-saliante Atome werden zu Motiven (verbundene Komponenten und Ringsysteme) gruppiert. Kontrafaktische Perturbationen werden anschließend auf diese Motive unter Verwendung von RDKit-Transformationsregeln (z. B. Halogenaustausch, Amid-Modifikation) angewendet, um Änderungen in QED ($\Delta$QED) zu messen.
2. Sparse Autoencoder (SAEs):
   • Undercomplete (MLP-SAE): Ein komprimierender Autoencoder (256 $\to$ 128) wird verwendet, um dominante, spärliche latente Faktoren zu identifizieren und deren Korrelation mit physikochemischen Eigenschaften (z. B. Polarität, Lipophilie) zu testen.
   • Overcomplete (BatchTopK SAE): Ein überkompletter Autoencoder (256 $\to$ 1024) mit einem BatchTopK-Sparsitätsmechanismus ($k=64$) wird eingesetzt, um feingranulare Feature-Detektoren, einschließlich solcher in Superposition, ohne $\ell_1$-Regularisierung zu rekonstruieren.
3. Probing und Kontrollen: Flache Feedforward-Klassifikatoren und strikt lineare Probes (Ridge-/Logistische Regression) werden auf den Embeddings trainiert, um das Vorhandensein chemischer Motive (funktionelle Gruppen, Halogene) und physikochemischer Eigenschaften vorherzusagen. Entscheidend ist, dass die Studie rigorose Kontrollen anwendet:
   • Deskriptor-Baselines: Vergleich der Probe-Performance gegen RDKit-Deskriptoren und ECFP4-Fingerprints.
   • Random-Label-Kontrollen: Sicherstellung, dass Probes kein Rauschen memorieren.
   • Scaffold-Disjoint Splits: Verhindern von Scaffold-Leakage in den Train/Test-Splits.
   • Untrained-Network Control: Ausführung identischer Probes auf einem zufällig initialisierten Netzwerk mit exakt derselben Architektur, um den Beitrag von Training versus Architektur zu isolieren.
   • Feature-Ablation: Nullsetzen einzelner Features, um die Spezifität in der nachgeschalteten Dekodierung zu testen.

Wichtigste Ergebnisse

• Salienz und Kontrafaktika: Integrated Gradients heben chemisch bedeutsame Regionen hervor (polare Substituenten, aromatische Fragmente). Kontrafaktische Edits an diesen Motiven führen zu kleinen, aber messbaren Verschiebungen im QED (z. B. +0,015 bis +0,051), wenngleich die Effektstärken im Verhältnis zum QED-Bereich gering sind; ein kontrollierter Vergleich gegenüber zufälligen Edits bleibt eine zukünftige Aufgabe.
• Latente Struktur: Der Undercomplete-SAE identifiziert Faktoren, die stark mit Größe und Polarität korrelieren, aber ein MLP-Prädiktor auf diesen Faktoren schneidet bei zusammengesetztem QED ($R^2=0,25$) im Vergleich zu Einzel-Eigenschaften schlecht ab.
• Rekonstruktion und Sparsität: Der Overcomplete BatchTopK SAE rekonstruiert Embeddings mit niedrigem NMSE (0,0015) und behält eine strikte Batch-Mean-Sparsität von 64 aktiven Features bei. Er übertrifft eine gematchte lineare Baseline (NMSE 0,0026) bei ähnlichem Aktivitätsniveau.
• Lineare Dekodierbarkeit: Die Embeddings dekodieren physikochemische Eigenschaften gut (z. B. QED $R^2 \approx 0,59$, Größe $R^2 \approx 0,996$). Diese Scores sind jedoch nicht signifikant besser als die, die durch eine „leakage-freie“ RDKit-Deskriptor-Baseline unter Random-Splits erreicht wurden.
• Die Untrained-Network Control (Entscheidender Befund): Ein zufällig initialisiertes Netzwerk der identischen Architektur dekodiert jede geprüfte Eigenschaft im Wesentlichen so gut wie die trainierte Policy. Für die Haupt-Eigenschaft (Drug-likeness/QED) liefert das Training einen vernachlässigbaren Gewinn von nur +0,005 (0,581 ungetrainiert vs. 0,586 trainiert), was innerhalb der Varianz der Random Seeds liegt. Bei der molekularen Größe ist das ungetrainierte Netzwerk geringfügig besser.
• Feature-Spezifität: Im Overcomplete-Dictionary offenbaren einzelne Durchläufe chemisch angereicherte Detektoren (z. B. separate Features für Iod, Chlor, Brom, Fluor und Bor). Die Ablation einzelner Features zeigt eigenschaftsspezifische Effekte (Spezifitätsverhältnisse von 8–23$\times$), aber die Feature-Indizes sind über verschiedene Seeds hinweg nicht reproduzierbar; nur ein kleiner Kern von Richtungen ist stabil.

Bedeutung und Ansprüche
Das bedeutendste Ergebnis des Papers ist eine Negative-Control: Die lineare Dekodierbarkeit physikochemischer Eigenschaften aus den GFlowNet-Embeddings spiegelt die Grapharchitektur und die Atom-Featurisierung wider, nicht das während des Trainings gelernte chemische Wissen.

• Was NICHT gestützt wird: Hohe Probe-Scores (selbst nahe perfekter AUROCs) sind kein Beleg für erworbenes chemisches Verständnis oder gelerntes Struktur-Aktivitäts-Schlussfolgern. Die Fähigkeit, Eigenschaften zu dekodieren, ist weitgehend ein Artefakt der Repräsentation der Eingangsdaten und der Netzwerkarchitektur, wie die Übereinstimmung der Leistung des ungetrainierten Netzwerks mit der trainierten Policy zeigt.
• Was WAS IST gestützt wird:
  1. Das Overcomplete-Dictionary bietet eine gültige, exakt-spärliche Zerlegung der Embeddings.
  2. Einzelne Durchläufe enthalten chemisch angereicherte Substruktur-Detektoren (z. B. pro-Halogen-Features), die über einen kleinen Teilbereich der Seeds hinweg stabil sind.
  3. Diese Features haben eigenschaftsspezifische interventionelle Effekte auf die Probe-Dekodierung.
  4. Drug-likeness (QED) wird im spärlichen latenten Raum linear zugänglicher als im dichten Embedding.

Faz Schlussfolgerung
Die Autoren kommen zu dem Schluss, dass Interpretationsmethoden streng gegen architektonische Baselines kontrolliert werden müssen. Für SynFlowNet identifiziert das Framework erfolgreich feingranulare Feature-Detektoren und validiert die Robustheit der spärlichen Repräsentation. Die Studie warnt jedoch explizit davor, hohe Dekodierbarkeits-Scores als Beweis für gelerntes chemisches Wissen zu interpretieren. Das Paper beschränkt seine Ansprüche auf die Existenz eines reproduzierbaren Kerns stabiler, chemisch angereicherter Features innerhalb des Dictionarys, während es gleichzeitig behauptet, dass die Fähigkeit des Modells, physikochemische Eigenschaften vorherzusagen, primär eine Funktion seiner Architektur und nicht seines Trainings ist. Zukünftige Arbeiten sollten untersuchen, ob diese Detektoren bestehen bleiben, wenn SAEs auf ungetrainierten Netzwerk-Embeddings trainiert werden, und das Framework auf Multi-Objective-Settings erweitern.