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

Technische Samenvatting: Interpretatie van GFlowNets voor Drug Discovery

Probleemstelling
Generative Flow Networks (GFlowNets), met name die getraind voor moleculaire vormgeving zoals SynFlowNet, bieden een krachtig kader voor het verkennen van de chemische ruimte terwijl zij zich houdt aan synthetische beperkingen. De interne beslissingsbeleid (policies) van deze netwerken blijft echter opaak, wat de adoptie ervan in drug discovery beperkt, waar medicinale chemici interpreteerbare rationale nodig hebben voor voorgestelde structuren. Bestaande interpretatiemethoden zijn vaak ontworpen voor continue latente ruimtes of waarschijnlijkheidsgebaseerde modellen, en niet voor de discrete, sequentiële, graafgebaseerde beleidsregels van GFlowNets. Bovendien is er een kritieke behoefte om onderscheid te maken tussen chemische kennis die daadwerkelijk door het beleid is geleerd en informatie die triviaal beschikbaar is via de graafarchitectuur en atoom-featurisatie.

Methodologie
De auteurs introduceren een multi-schaal interpretatiekader toegepast op SynFlowNet (getraind met een QED-beloningsfunctie) met behulp van bevroren embeddings van een specifiek checkpoint. Het kader integreert drie complementaire benaderingen:

1. Gradiënt-gebaseerde Saliency en Counterfactuals: Integrated Gradients (IG) worden toegepast op de log-waarschijnlijkheid van de "Stop"-actie om atoom-niveau saliency maps te generen. Atomen met een hoge saliency worden gegroepeerd in motieven (verbonden componenten en ringsystemen). Counterfactual perturbaties worden vervolgens toegepast op deze motieven met behulp van RDKit-transformatie-regels (bijv. halogeenuitwisseling, amide-modificatie) om veranderingen in QED ($\Delta$QED) te meten.
2. Sparse Autoencoders (SAEs):
   • Undercomplete (MLP-SAE): Een compressieve autoencoder (256 $\to$ 128) wordt gebruikt om dominante, ijlere (sparse) latente factoren te identificeren en hun correlatie met fysisch-chemische eigenschappen (bijv. polariteit, lipofiliteit) te testen.
   • Overcomplete (BatchTopK SAE): Een overcomplete autoencoder (256 $\to$ 1024) met een BatchTopK sparsity-mechanisme ($k=64$) wordt ingezet om fijnmazige feature-detectoren te herstellen, inclusief die in superpositie, zonder $\ell_1$-regularisatie.
3. Probing en Controles: Ondiepe feedforward classificaties en strikt lineaire probes (Ridge/logistische regressie) worden getraind op embeddings om de aanwezigheid van chemische motieven (functionele groepen, halogenen) en fysisch-chemische eigenschappen te voorspellen. Cruciaal is dat de studie rigoureuze controles hanteert:
   • Descriptor Baselines: Vergelijking van probe-prestaties tegenover RDKit-descriptors en ECFP4-fingerprints.
   • Random-Label Controls: Zorgen dat probes geen ruis memoriseren.
   • Scaffold-Disjoint Splits: Voorkomen van scaffold-lekkage in train/test-splits.
   • Untrained-Network Control: Het draaien van identieke probes op een willekeurig geïnitialiseerd netwerk met exact dezelfde architectuur om de bijdrage van training versus architectuur te isoleren.
   • Feature Ablatie: Het op nul zetten van individuele features om specificiteit in downstream decoding te testen.

Belangrijkste Resultaten

• Saliency en Counterfactuals: Integrated Gradients benadrukken chemisch betekenisvolle regio's (polaire substituenten, aromatische fragmenten). Counterfactual edities aan deze motieven leveren kleine maar meetbare verschuivingen in QED op (bijv. +0,015 tot +0,051), hoewel de effectgroottes klein zijn ten opzichte van het QED-bereik; een gecontroleerde vergelijking tegenover willekeurige edities blijft toekomstig werk.
• Latente Structuur: De undercomplete SAE identificeert factoren die sterk correleren met grootte en polariteit, maar een MLP-voorspeller op deze factoren presteert slecht op samengestelde QED ($R^2=0,25$) vergeleken met individuele eigenschappen.
• Reconstructie en Sparsity: De overcomplete BatchTopK SAE reconstrueert embeddings met een lage NMSE (0,0015) en handhaaft een strikte batch-gemiddelde sparsity van 64 actieve features. Het presteert beter dan een gematchte lineaire baseline (NMSE 0,0026) bij vergelijkbare activiteitsniveaus.
• Lineaire Decodabiliteit: Embeddings decoderen fysisch-chemische eigenschappen goed (bijv. QED $R^2 \approx 0,59$, grootte $R^2 \approx 0,996$). Deze scores zijn echter niet significant beter dan de scores behaald door een "lekkage-vrije" RDKit-descriptor baseline onder willekeurige splits.
• De Untrained-Network Control (Beslissende Bevinding): Een willekeurig geïnitialiseerd netwerk van de identieke architectuur decodeert elke geteste eigenschap nagenoeg even goed als het getrainde beleid. Voor de belangrijkste eigenschap (drug-likeness/QED) levert training een verwaarloosbare winst op van slechts +0,005 (0,581 ongetraind versus 0,586 getraind), wat binnen de variantie van de random seeds valt. Wat moleculaire grootte betreft, is het ongetrainde netwerk marginaal beter.
• Feature Specificiteit: In de overcomplete dictionary onthullen individuele runs chemisch verrijkte detectoren (bijv. aparte features voor jodium, chloor, broom, fluor en boor). Single-feature ablatie toont eigenschap-specifieke effecten (specificiteitsratio's van 8–23$\times$), maar de feature-indices zijn niet reproduceerbaar over verschillende seeds; slechts een kleine kern van richtingen is stabiel.

Betekenis en Claims
Het meest consequente resultaat van het paper is een negatieve controle: de lineaire decodabiliteit van fysisch-chemische eigenschappen uit GFlowNet embeddings weerspiegelt de graafarchitectuur en atoom-featurisatie, en niet de chemische kennis die tijdens de training is geleerd.

• Wat NIET wordt ondersteund: Hoge probe-scores (zelfs nabij perfecte AUROCs) zijn geen bewijs van verworven chemisch begrip of geleerd structuur-activiteit redeneren. Het vermogen om eigenschappen te decoderen is grotendeels een artefact van de inputrepresentatie en de netwerkarchitectuur, zoals aangetoond door het feit dat het ongetrainde netwerk de prestaties van het getrainde beleid even goed evenaart.
• Wat WEL wordt ondersteund:
  1. De overcomplete dictionary biedt een geldige, exacte-sparsity decompositie van de embeddings.
  2. Individuele runs bevatten chemisch verrijkte substructuur-detectoren (bijv. per-halogeen features) die stabiel zijn over een kleine subset van seeds.
  3. Deze features hebben eigenschap-specifieke interventionele effecten op de probe-decoding.
  4. Drug-likeness (QED) wordt toegankelijker lineair in de sparse latente ruimte dan in de dense embedding.

Conclusie
De auteurs concluderen dat interpretatiemethoden rigoureus gecontroleerd moeten worden tegen architecturale baselines. Voor SynFlowNet slaagt het framework erin om fijnmazige feature-detectoren te identificeren en de robuustheid van de sparse representatie te valideren. Echter, de studie waarschuwt expliciet tegen het interpreteren van hoge decodabiliteitsscores als bewijs van geleerde chemie. Het paper beperkt de claims tot het bestaan van een reproduceerbare kern van stabiele, chemisch verrijkte features binnen de dictionary, terwijl het stelt dat het vermogen van het model om fysisch-chemische eigenschappen te voorspellen primair een functie is van de architectuur in plaats van de training. Toekomstig werk zou moeten onderzoeken of deze detectoren standhouden wanneer SAE's worden getraind op de embeddings van ongetrainde netwerken, en het framework uitbreiden naar multi-objective settings.