Learning the Neighborhood: Contrast-Free Multimodal Self-Supervised Molecular Graph Pretraining

Riepilogo Tecnico: Learning the Neighborhood (C-FREE)

Problematica
Le rappresentazioni molecolari di alta qualità sono critiche per la predizione delle proprietà e la scoperta di farmaci, eppure i grandi dataset etichettati rimangono scarsi. Gli approcci esistenti di apprendimento auto-supervisionato (SSL) per i grafi molecolari soffrono spesso di specifiche limitazioni:

1. Limitazioni di Modalità: Molti metodi si affidano esclusivamente alla topologia 2D o alle sequenze 1D (SMILES), sottoutilizzando preziose informazioni strutturali 3D.
2. Complessità e Costo: I metodi contrastivi dipendono da campioni negativi accuratamente elaborati e grandi batch, che sono difficili da definire per strutture grafiche irregolari (ad esempio, distinguere isomeri chirali). I metodi generativi richiedono spesso la ricostruzione discreta nello spazio di input o una complessa tokenizzazione del grafo. I metodi predittivi latenti si affidano frequentemente ad aumenti costosi o procedure di clustering (ad esempio, METIS) e possono soffere di collasso della rappresentazione.
3. Scarsità di Dati: I dataset curati sono spesso non disponibili, particolarmente nei regimi a bassi dati, rendendo necessaria una strategia di pre-training robusta che si trasferisca efficacemente a nuovi domini chimici.

Metodologia: C-FREE
Gli autori introducono C-FREE (Contrast-Free Representation learning on Ego-nets), un framework auto-supervisionato multimodale che integra grafi 2D con ensemble di conformeri 3D. Il metodo adotta una strategia di apprendimento predittivo nello spazio latente, evitando campioni negativi, codifiche posizionali e ricostruzione dello spazio di input.

• Meccanismo Core (Ego-nets): Il framework segmenta i grafi molecolari in sottografi disgiunti utilizzando k-EgoNets a raggio fisso (dove $k \in \{3, 4\}$). Per un dato nodo ancora, il vicinato $k$-hop forma il sottografo di contesto, mentre i restanti archi e nodi formano il sottografo target. I nodi di confine vengono duplicati per garantire la disgiunzione degli archi.
• Codifica Multimodale:
  • Encoder 2D: Un Graph Isomorphism Network con caratteristiche di bordo (GINE) elabora la topologia del grafo.
  • Encoder 3D: Una rete equivariante (PaiNN o SchNet) elabora le coordinate 3D da molteplici conformeri.
  • Fusione: Le embedding a livello di nodo di entrambe le modalità vengono concatenate con token di classificazione apprendibili (hCLS) e di separazione (hSEP), formando una sequenza multimodale elaborata da un Transformer per catturare le dipendenze globali.
• Obiettivo Predittivo: Il modello allinea l'embedding del sottografo di contesto con quello del sottografo target.
  • Un Context Encoder elabora la vista di input.
  • Una Predictor Network (un Transformer leggero seguito da un MLP) mappa l'embedding del contesto nello spazio target.
  • Un Target Encoder (una Media Mobile Esponenziale, EMA, del context encoder) genera l'embedding target.
  • L'obiettivo minimizza la distanza media L2 tra l'embedding del contesto predetto e l'embedding target.
• Fine-Tuning: Durante i compiti a valle, il target encoder funge da backbone pre-addestrato. Il framework supporta due strategie di valutazione:
  1. C-FREEMOL: Linear probing su embedding dell'intera molecola.
  2. C-FREESUB: Aggregazione di embedding di sottografi usando DeepSets, simulando le Equivariant Subgraph Aggregation Networks (ESAN).

Contributi Chiave

1. Nuovo Task di Pre-training: Un obiettivo predittivo ampiamente applicabile basato su sottografi k-EgoNet che utilizza sia viste 2D che 3D senza costosi aumenti artigianali o campionamento negativo.
2. Robustezza: Il framework performa fortemente sia in contesti multimodali (2D + 3D) che in contesti solo 2D, dimostrando un trasferimento efficace anche in assenza di conformeri.
3. Semplicità ed Espressività: Lo schema di addestramento evita mismatch tra pre-training e fine-tuning e campioni negativi. Teoricamente, quando utilizza una testa DeepSets, C-FREE è provabilmente più espressivo del test 1-Weisfeiler-Leman (1-WL), eguagliando l'espressività di ESAN senza richiedere obiettivi gerarchici o clustering.
4. Prestazioni allo Stato dell'Arte: L'approccio raggiunge risultati di alto livello su MoleculeNet e altri benchmark, superando i metodi SSL contrastivi, generativi e altri metodi multimodali.

Risultati Sperimentali
Il modello è stato pre-addestrato sul dataset GEOM (circa 330k molecole con conformazioni 3D) ed è stato valutato su molteplici benchmark:

• MoleculeNet (Frozen Backbone): C-FREE ha ottenuto la migliore performance media su 8 dataset, superando i baseline in 6 dei 8 compiti. La variante multimodale (C-FREEMM) ha generalmente superato le varianti solo 2D, sebbene il C-FREE solo 2D abbia comunque superato tutti gli altri baseline in media.
• MoleculeNet (Full Fine-Tuning): Utilizzando un encoder PaiNN pre-addestrato su GEOM, C-FREE ha eguagliato o superato modelli foundation più grandi (es. UniMol, GEM) addestrati su dataset significativamente più grandi (fino a 19M di molecole), raggiungendo la migliore ROC-AUC media.
• QM9 e ZINC: Su QM9, C-FREE ha superato l'enorme UniMol2 (1.1B di parametri) su 4 dei 6 target nonostante avesse un numero di parametri molto inferiore (9.1M). Su ZINC, ha superato GraphJEPA.
• Efficienza dei Label: Sul dataset Drugs-75K, il pre-training ha fornito chiari vantaggi nei regimi a bassi label (1%, 10% di label), superando costantemente i modelli addestrati da zero.
• Ablazioni:
  • Modalità: L'informazione 3D ha fornito un segnale più forte rispetto al solo 2D, ma la combinazione di entrambi ha prodotto i risultati migliori.
  • Predictor: Rimuovere il predictor ha causato il collasso dell'addestramento; un Transformer predictor ha superato un MLP.
  • EMA: Tassi di decadimento iniziali più bassi per il target encoder EMA hanno migliorato la qualità della rappresentazione.
  • Raggio EgoNet: Un raggio $k=5$ ha prodotto la migliore performance, suggerendo che sottografi di contesto e target di dimensioni comparabili siano ottimali, mentre $k=1$ non è riuscito a catturare un segnale strutturale sufficiente.

Significatività e Rivendicazioni
Il paper sostiene che C-FREE rappresenti il primo framework predittivo multimodale non contrastivo e non generativo per l'apprendimento di rappresentazioni molecolari. La sua significatività risiede nel:

• Eliminare la Complessità: Rimuove la necessità di campioni negativi, grandi batch, clustering costosi e codifiche posizionali, che sono comuni colli di bottiglia nell'SSL per grafi.
• Sfruttare la Geometria 3D: Integra efficacemente la diversità conformazionale 3D, che è cruciale per le proprietà dipendenti dalla forma molecolare, senza richiedere complessi obiettivi generativi.
• Generalizzazione: Il metodo dimostra che il pre-addestramento auto-supervisionato può fornire inizializzazioni forti che si trasferiscono efficacemente a diversi domini chimici e contesti con pochi dati, evidenziando l'importanza delle rappresentazioni informate dal 3D anche quando la topologia 2D è disponibile.

Gli autori osservano che, sebbene il loro modello sia più piccolo di alcuni modelli foundation, le sue prestazioni suggeriscono che le scelte architettoniche e la pipeline di pre-training (specificamente il task predittivo basato su EgoNet) sono altamente efficaci. Suggeriscono che il lavoro futuro potrebbe includere lo scaling su dataset più grandi, l'incorporazione di altre modalità come SMILES e l'esplorazione di encoder sensibili alla chiralità.