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

Technical Summary: Interpreting GFlowNets for Drug Discovery

Problem Statement
Generative Flow Networks (GFlowNets), particularly those trained for molecular design like SynFlowNet, offer a powerful framework for exploring chemical space while adhering to synthetic constraints. However, their internal decision policies remain opaque, limiting their adoption in drug discovery where medicinal chemists require interpretable rationales for proposed structures. Existing interpretability methods are often designed for continuous latent spaces or likelihood-based models, not for the discrete, sequential, graph-based policies of GFlowNets. Furthermore, there is a critical need to distinguish between chemical knowledge genuinely learned by the policy and information that is trivially available through the graph architecture and atom featurization.

Methodology
The authors introduce a multi-scale interpretability framework applied to SynFlowNet (trained with a QED reward function) using frozen embeddings from a specific checkpoint. The framework integrates three complementary approaches:

1. Gradient-Based Saliency and Counterfactuals: Integrated Gradients (IG) are applied to the log-probability of the "Stop" action to generate atom-level saliency maps. High-saliency atoms are grouped into motifs (connected components and ring systems). Counterfactual perturbations are then applied to these motifs using RDKit transformation rules (e.g., halogen exchange, amide modification) to measure changes in QED ($\Delta$QED).
2. Sparse Autoencoders (SAEs):
   • Undercomplete (MLP-SAE): A compressive autoencoder (256 $\to$ 128) is used to identify dominant, sparse latent factors and test their correlation with physicochemical properties (e.g., polarity, lipophilicity).
   • Overcomplete (BatchTopK SAE): An overcomplete autoencoder (256 $\to$ 1024) with a BatchTopK sparsity mechanism ($k=64$) is employed to recover fine-grained feature detectors, including those in superposition, without $\ell_1$ regularization.
3. Probing and Controls: Shallow feedforward classifiers and strictly linear probes (Ridge/logistic regression) are trained on embeddings to predict the presence of chemical motifs (functional groups, halogens) and physicochemical properties. Crucially, the study employs rigorous controls:
   • Descriptor Baselines: Comparing probe performance against RDKit descriptors and ECFP4 fingerprints.
   • Random-Label Controls: Ensuring probes do not memorize noise.
   • Scaffold-Disjoint Splits: Preventing scaffold leakage in train/test splits.
   • Untrained-Network Control: Running identical probes on a randomly initialized network of the exact same architecture to isolate the contribution of training versus architecture.
   • Feature Ablation: Zeroing single features to test specificity in downstream decoding.

Key Results

• Saliency and Counterfactuals: Integrated Gradients highlight chemically meaningful regions (polar substituents, aromatic fragments). Counterfactual edits to these motifs yield small but measurable shifts in QED (e.g., +0.015 to +0.051), though the effect sizes are small relative to the QED range, and a controlled comparison against random edits remains future work.
• Latent Structure: The undercomplete SAE identifies factors strongly correlated with size and polarity, but an MLP predictor on these factors performs poorly on composite QED ($R^2=0.25$) compared to individual properties.
• Reconstruction and Sparsity: The overcomplete BatchTopK SAE reconstructs embeddings with low NMSE (0.0015) and maintains a strict batch-mean sparsity of 64 active features. It outperforms a matched linear baseline (NMSE 0.0026) at similar activity levels.
• Linear Decodability: Embeddings decode physicochemical properties well (e.g., QED $R^2 \approx 0.59$, size $R^2 \approx 0.996$). However, these scores are not significantly better than those achieved by a "leakage-free" RDKit descriptor baseline under random splits.
• The Untrained-Network Control (Decisive Finding): A randomly initialized network of the identical architecture decodes every probed property essentially as well as the trained policy. For the headline property (drug-likeness/QED), training provides a negligible gain of only +0.005 (0.581 untrained vs. 0.586 trained), which falls within the variance of random seeds. On molecular size, the untrained network is marginally better.
• Feature Specificity: In the overcomplete dictionary, individual runs reveal chemically enriched detectors (e.g., separate features for iodine, chlorine, bromine, fluorine, and boron). Single-feature ablation shows property-specific effects (specificity ratios of 8–23$\times$), but feature indices are not reproducible across seeds; only a small core of directions is stable.

Significance and Claims
The paper's most consequential result is a negative control: the linear decodability of physicochemical properties from GFlowNet embeddings reflects the graph architecture and atom featurization, not the chemical knowledge learned during training.

• What is NOT supported: High probe scores (even near-perfect AUROCs) are not evidence of acquired chemical understanding or learned structure-activity reasoning. The ability to decode properties is largely an artifact of the input representation and network architecture, as demonstrated by the untrained network matching the trained policy's performance.
• What IS supported:
  1. The overcomplete dictionary provides a valid, exact-sparsity decomposition of the embeddings.
  2. Individual runs contain chemically enriched substructure detectors (e.g., per-halogen features) that are stable across a small subset of seeds.
  3. These features have property-specific interventional effects on probe decoding.
  4. Drug-likeness (QED) becomes more linearly accessible in the sparse latent space than in the dense embedding.

Conclusion
The authors conclude that interpretability methods must be rigorously controlled against architectural baselines. For SynFlowNet, the framework successfully identifies fine-grained feature detectors and validates the robustness of the sparse representation. However, the study explicitly cautions against interpreting high decodability scores as proof of learned chemistry. The paper restricts its claims to the existence of a reproducible core of stable, chemically enriched features within the dictionary, while asserting that the model's ability to predict physicochemical properties is primarily a function of its architecture rather than its training. Future work should explore whether these detectors persist when SAEs are trained on untrained-network embeddings and extend the framework to multi-objective settings.