Bias mitigation in graph diffusion models

1. Il Problema: Due Bias Critici nei Modelli di Diffusione su Grafi

Gli autori identificano che i modelli di diffusione su grafi esistenti (come GDSS, MOOD, GSDM, HGDM) soffrono di due problemi fondamentali che degradano la qualità della generazione:

• Bias di Inizio Inverso (Reverse-Starting Bias):
  • Teoria vs. Realtà: Idealmente, il processo di diffusione in avanti (forward) dovrebbe perturbare i dati fino a raggiungere una distribuzione Gaussiana standard, e il processo inverso (reverse) dovrebbe iniziare da tale distribuzione.
  • Il Problema: A causa delle limitazioni nella scala dei dati e nella capacità di apprendimento delle reti, i modelli su grafi adottano strategie conservative. Di conseguenza, la distribuzione di massima perturbazione nel processo forward ($x_T$) non raggiunge la Gaussiana standard, ma rimane in uno stato a "basso rumore" (low-noise state).
  • La Discrepanza: Tuttavia, il campionamento inverso inizia sempre da una Gaussiana standard. Questo disallineamento tra il punto di partenza reale del forward e quello ipotetico del reverse crea un bias significativo che compromette la generazione.
• Bias di Esposizione (Exposure Bias):
  • Durante il processo inverso, il modello utilizza le sue stesse previsioni precedenti (rumorose) come input per il passo successivo, invece dei dati reali puliti. Gli errori di previsione della rete di score si accumulano e propagano, peggiorando la qualità del campione finale.

2. Metodologia Proposta: S++

Il paper propone un approccio completo chiamato S++ che mitiga entrambi i bias senza modificare l'architettura della rete o introdurre nuovi componenti (come generatori o discriminatori aggiuntivi). La soluzione si basa su due pilastri:

A. Allineamento tramite Campionamento di Langevin (Mitigazione del Reverse-Starting Bias)

Invece di forzare la distribuzione forward a diventare una Gaussiana standard (cosa difficile da apprendere), il metodo allinea il punto di partenza del campionamento inverso alla distribuzione di massima perturbazione effettiva del forward.

• Meccanismo: Utilizza un algoritmo di Langevin sampling guidato dalla rete di score pre-addestrata $s_{\theta, T}(\cdot)$ al tempo $T$.
• Funzionamento: Invece di iniziare da $N(0, I)$, il campionamento inizia con iterazioni di Langevin che spingono il punto di partenza verso le regioni ad alta densità della distribuzione reale $q(x_T)$. Questo allinea il punto di inizio inverso con la realtà del processo forward, eliminando il bias di inizio.

B. Correzione dello Score basata sulla Differenza (Mitigazione dell'Exposure Bias)

Per correggere gli errori di previsione della rete di score senza riaddestrarla, gli autori introducono un meccanismo di correzione basato sulla differenza tra due score.

• Creazione di una "Pseudo-Rete":
  1. Si usa la rete pre-addestrata $s_{\theta, t}$ per generare un batch di campioni.
  2. Si addestra una rete di score "pseudo" ($s_{\psi, t}$) su questi campioni generati. Poiché i campioni generati contengono bias, la pseudo-rete impara a prevedere con una deviazione maggiore rispetto alla rete originale.
• Correzione dello Score:
  • Si calcola la differenza tra lo score della rete originale e quello della pseudo-rete: $\Delta s = s_{\theta, t} - s_{\psi, t}$.
  • Questa differenza contiene informazioni utili sul dato originale $x_0$ (come illustrato geometricamente nel paper).
  • Lo score viene corretto con la formula:
    $$s'_{\theta, t} = \frac{s_{\theta, t} + \lambda (s_{\theta, t} - s_{\psi, t})}{\omega}$$
    dove $\lambda$ controlla la direzione (spingendo verso lo score vero) e $\omega$ aggiusta la magnitudine.
• Vantaggio: Questo approccio sfrutta la differenza tra le due reti per correggere il bias direzionale e di magnitudine, migliorando l'accuratezza dello score durante tutto il processo di campionamento.

3. Contributi Chiave

1. Analisi Sistematica: È il primo lavoro a identificare e affrontare sistematicamente il reverse-starting bias nei modelli di diffusione su grafi, dimostrando che è una causa primaria della scarsa qualità generativa.
2. Soluzione Plug-and-Play: Il metodo S++ non richiede modifiche alla rete neurale, né l'aggiunta di nuovi moduli (come GAN o discriminatori). Può essere integrato in qualsiasi modello di diffusione esistente.
3. Efficienza e Robustezza: La correzione dello score è particolarmente efficace all'inizio del processo (tempo $T$), permettendo di ridurre la lunghezza della catena di campionamento e accelerare la generazione senza perdere qualità.
4. Validazione Estesa: Il metodo è stato testato su diversi modelli (GDSS, GSDM, HGDM, MOOD), dataset (grafi generici e molecolari) e task, ottenendo risultati State-of-the-Art (SOTA).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su dataset generici (Community-small, Enzymes, Grid) e dataset molecolari (QM9, ZINC250k).

• Qualità Generativa: S++ supera significativamente i modelli baseline. Ad esempio, su QM9, la metrica FCD (Frechet ChemNet Distance) scende da 4.584 (GDSS-OC) a 1.661 (GDSS-OC-S++), indicando una distribuzione dei dati generati molto più simile a quella reale.
• Validità Molecolare: La percentuale di molecole valide senza correzione post-processing aumenta drasticamente (es. dal 73.5% al 94.0% su QM9).
• Efficienza: Il metodo permette di ottenere risultati ottimali con meno passi di diffusione inversa. Su QM9, anche con soli 100 passi (invece di 1000), S++ mantiene prestazioni vicine all'ottimale, mentre i baseline crollano.
• Generazione di Farmaci: Nel task di generazione di molecole con proprietà specifiche (MOOD), S++ mostra un miglior "hit ratio" e punteggi di docking migliori rispetto al baseline MOOD.

5. Significato e Impatto

Questo lavoro è significativo perché:

• Ridefinisce la comprensione dei bias: Sposta l'attenzione dai classici errori di previsione (exposure bias) alla discrepanza fondamentale tra la distribuzione di fine processo forward e l'inizio del reverse nei grafi.
• Semplicità ed Efficacia: Dimostra che problemi complessi di generazione su grafi possono essere risolti con strategie di campionamento intelligenti e correzioni matematiche dello score, senza la necessità di architetture più complesse o costose.
• Generalizzabilità: Essendo indipendente dall'architettura della rete, S++ offre una soluzione universale per migliorare l'intero ecosistema dei modelli di diffusione su grafi, rendendoli più affidabili per applicazioni critiche come la scoperta di farmaci e la generazione di materiali.

In sintesi, il paper presenta S++ come una soluzione elegante e potente che risolve le incongruenze fondamentali nei modelli di diffusione su grafi, portando a una generazione di dati di qualità superiore, più veloce e robusta.