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

Sintesi Tecnica: Interpretare i GFlowNet per la Drug Discovery

Enunciato del Problema
Le Generative Flow Networks (GFlowNet), in particolare quelle addestrate per il design molecolare come SynFlowNet, offrono un framework potente per esplorare lo spazio chimico rispettando i vincoli sintetici. Tuttavia, le loro politiche decisionali interne rimangono opache, limitandone l'adozione nella drug discovery dove i chimici medicinali richiedono razionali interpretabili per le strutture proposte. I metodi di interpretabilità esistenti sono spesso progettati per spazi latenti continui o modelli basati sulla verosimiglianza (likelihood), non per le politiche discrete, sequenziali e basate su grafi delle GFlowNet. Inoltre, esiste una necessità critica di distinguere tra la conoscenza chimica genuinamente appresa dalla policy e le informazioni che sono banalmente disponibili attraverso l'architettura del grafo e la featurizzazione degli atomi.

Metodologia
Gli autori introducono un framework di interpretabilità multi-scala applicato a SynFlowNet (addestrata con una funzione di reward QED) utilizzando embedding congelati da un checkpoint specifico. Il framework integra tre approcci complementari:

1. Saliency e Controfatuali basati su Gradiente: Integrated Gradients (IG) viene applicato alla log-probabilità dell'azione "Stop" per generare mappe di saliency a livello di atomo. Gli atomi ad alta saliency vengono raggruppati in motivi (componenti connesse e sistemi ad anello). Vengono poi applicate perturbazioni controfatuali a questi motivi utilizzando regole di trasformazione RDKit (ad es., scambio di alogeni, modifica di ammidi) per misurare i cambiamenti in QED ($\Delta$QED).
2. Sparse Autoencoders (SAE):
   • Undercomplete (MLP-SAE): Un autoencoder compressivo (256 $\to$ 128) viene utilizzato per identificare fattori latenti dominanti e sparsi, e per testare la loro correlazione con proprietà fisico-chimiche (ad es., polarità, lipofilia).
   • Overcomplete (BatchTopK SAE): Un autoencoder overcomplete (256 $\to$ 1024) con un meccanismo di sparsità BatchTopK ($k=64$) è impiegato per recuperare detector di feature fini, inclusi quelli in sovrapposizione (superposition), senza regolarizzazione $\ell_1$.
3. Probing e Controlli: Classificatori feedforward superficiali e probe strettamente lineari (regressione Ridge/logistica) sono addestrati sugli embedding per predire la presenza di motivi chimici (gruppi funzionali, alogeni) e proprietà fisico-chimiche. Crucialmente, lo studio impiega controlli rigorosi:
   • Baseline di Descrittori: Confronto delle prestazioni dei probe rispetto ai descrittori RDKit ed ECFP4.
   • Controlli con Etichette Casuali: Per garantire che i probe non memorizzino il rumore.
   • Split Disgiunti per Scaffold: Per prevenire la fuga di scaffold (scaffold leakage) negli split train/test.
   • Controllo della Rete Non Addestrata: Esecuzione di probe identici su una rete inizializzata casualmente con la stessa esatta architettura per isolare il contributo dell'addestramento rispetto all'architettura.
   • Ablazione delle Feature: Azzeramento di singole feature per testare la specificità nel decoding a valle.

Risultati Chiave

• Saliency e Controfatuali: Integrated Gradients evidenzia regioni chimicamente significative (sostituenti polari, frammenti aromatici). Le edizioni controfatuali a questi motivi producono spostamenti piccoli ma misurabili nel QED (ad es., +0.015 a +0.051), sebbene le dimensioni dell'effetto siano piccole rispetto al range del QED, e un confronto controllato contro edizioni casuali rimane un lavoro futuro.
• Struttura Latente: L'SAE undercomplete identifica fattori fortemente correlati con dimensione e polarità, ma un predittore MLP su questi fattori performa scarsamente su QED composito ($R^2=0.25$) rispetto alle singole proprietà.
• Ricostruzione e Sparsità: L'SAE overcomplete BatchTopK ricostruisce gli embedding con un basso NMSE (0.0015) e mantiene una stretta sparsità di media di batch di 64 feature attive. Supera una baseline lineare corrispondente (NMSE 0.0026) a livelli di attività simili.
• Decodificabilità Lineare: Gli embedding decodificano bene le proprietà fisico-chimiche (ad es., QED $R^2 \approx 0.59$, dimensione $R^2 \approx 0.996$). Tuttavia, questi punteggi non sono significativamente migliori di quelli ottenuti da una baseline di descrittori RDKit "priva di leakage" sotto split casuali.
• Il Controllo della Rete Non Addestrata (Risultato Decisivo): Una rete inizializzata casualmente con la stessa identica architettura decodifica ogni proprietà sondata quasi altrettanto bene quanto la policy addestrata. Per la proprietà principale (drug-likeness/QED), l'addestramento fornisce un guadagno trascurabile di soli +0.005 (0.581 non addestrato vs 0.586 addestrato), che rientra nella varianza dei seed casuali. Sulla dimensione molecolare, la rete non addestrata è marginalmente migliore.
• Specificità delle Feature: Nel dizionario overcomplete, singole esecuzioni rivelano detector chimicamente arricchiti (ad es., feature separate per iodio, cloro, bromo, fluoro e boro). L'ablazione di singola feature mostra effetti specifici per proprietà (rapporti di specificità di 8–23$\times$), ma gli indici delle feature non sono riproducibili tra i seed; solo un piccolo nucleo di direzioni è stabile.

Significatività e Claim
Il risultato più consequenziale del paper è un controllo negativo: la decodificabilità lineare delle proprietà fisico-chimiche dagli embedding delle GFlowNet riflette l'architettura del grafo e la featurizzazione degli atomi, non la conoscenza chimica appresa durante l'addestramento.

• Cosa NON è supportato: Alti punteggi di probe (anche AUROC quasi perfetti) non sono prova di un acquisito apprendimento chimico o di ragionamento struttura-attività. La capacità di decodificare le proprietà è in gran parte un artefatto della rappresentazione di input e dell'architettura della rete, come dimostrato dal fatto che la rete non addestrata eguaglia le prestazioni della policy addestrata.
• Cosa È supportato:
  1. Il dizionario overcomplete fornisce una decomposizione valida e a sparsità esatta degli embedding.
  2. Singole esecuzioni contengono detector di sottostrutture chimicamente arricchite (ad es., feature per ogni alogeno) che sono stabili attraverso un piccolo sottoinsieme di seed.
  3. Queste feature hanno effetti interventistici sulle proprietà del decoding dei probe.
  4. La drug-likeness (QED) diventa più linearmente accessibile nello spazio latente sparso rispetto all'embedding denso.

Conclusione
Gli autori concludono che i metodi di interpretabilità devono essere rigorosamente controllati rispetto alle baseline architettoniche. Per SynFlowNet, il framework identifica con successo detector di feature fini e valida la robustezza della rappresentazione sparsa. Tuttavia, lo studio mette esplicitamente in guardia contro l'interpretazione di alti punteggi di decodificabilità come prova di chimica appresa. La capacità del modello di predire le proprietà fisico-chimiche è principalmente una funzione della sua architettura piuttosto che del suo addestramento. Il lavoro futuro dovrebbe esplorare se questi detector persistono quando gli SAE sono addestrati su embedding di reti non addestrate e estendere il framework a setting multi-obiettivo.