Bias mitigation in graph diffusion models

Probleemstelling: Twee vormen van Bias in Graf Diffusiemodellen

De auteurs identificeren dat bestaande graf-diffusiemodellen (zoals GDSS, GSDM, HGDM) lijden aan twee fundamentele bias-problemen die de generatiekwaliteit aanzienlijk verminderen:

1. Reverse-Starting Bias (Terugwaartse Start-Bias):

   • Oorzaak: In theorie zou de voorwaartse diffusie (het toevoegen van ruis) de data moeten transformeren naar een standaard Gaussische verdeling. Echter, vanwege beperkingen in de schaal van grafdata en de leercapaciteit van netwerken, trunceren bestaande modellen het voorwaartse proces. De maximale verstoring (bij tijdstip $T$) bereikt de standaard Gaussische verdeling niet; het blijft vaak in een "low-noise" staat.
   • Het Conflict: Ondanks dat de voorwaartse verdeling afwijkt, starten deze modellen het terugwaartse proces (sampling) altijd vanuit een perfecte standaard Gaussische verdeling ($N(0, I)$).
   • Gevolg: Er ontstaat een fundamenteel misverstand tussen de startpositie van de sampling en de werkelijke verdeling van de data na het voorwaartse proces, wat leidt tot een slechte generatiekwaliteit.

2. Exposure Bias (Blootstellingsbias):

   • Oorzaak: Tijdens het voorwaartse proces ziet het model de "ware" verstoide data $x_t$. Tijdens het terugwaartse proces moet het model echter denoisen op basis van zijn eigen voorspellingen $\hat{x}_t$. Door voorspellingsfouten in het score-netwerk accumuleren deze fouten tijdens het iteratieve proces.
   • Gevolg: De fouten stapelen zich op naarmate het samplingproces vordert, wat de uiteindelijke kwaliteit van de gegenereerde grafen of moleculen verslechtert.

Methodologie: De S++ Aanpak

Het paper introduceert S++, een methode die beide biases aanpakt zonder de onderliggende netwerkarchitectuur te wijzigen of extra componenten (zoals discriminators of generators) toe te voegen. De aanpak bestaat uit twee kernonderdelen:

1. Langevin Sampling voor Reverse-Starting Alignment

Om de reverse-starting bias op te lossen, starten de auteurs het samplingproces niet vanuit een standaard Gaussische verdeling, maar vanuit de werkelijke maximale verstoorde verdeling van het voorwaartse proces.

• Techniek: Ze gebruiken een nieuw ontworpen Langevin sampling-algoritme.
• Werking: Het algoritme gebruikt de score $s_{\theta, T}(\cdot)$ (geschat bij het eindtijdstip $T$) om te leiden naar de hoge-dichtheidsgebieden van de werkelijke maximale verstoorde verdeling $q(x_T|x_0)$.
• Resultaat: Dit creëert een nieuwe startpunt voor het terugwaartse proces dat perfect aansluit bij de werkelijke uitgang van het voorwaartse proces, waardoor de start-bias wordt geëlimineerd. Bovendien bevindt deze verdeling zich vaak in een "low-noise" staat, waar het model minder gevoelig is voor voorspellingsfouten, wat indirect ook de exposure bias vermindert.

2. Score Correctie op Basis van Score-verschil (S++ Bias Correction)

Om de exposure bias verder te reduceren en de stabiliteit van de Langevin sampling te verbeteren, introduceren ze een score-correctiemechanisme.

• Concept: Ze trainen een "pseudo" score-netwerk ($s_{\psi, t}$) op een dataset van gegenereerde samples (die per definitie een bias hebben ten opzichte van de ware data).
• Score-verschil: Ze berekenen het verschil tussen het oorspronkelijke score-netwerk ($s_{\theta, t}$) en het pseudo-netwerk ($s_{\psi, t}$). Dit verschil bevat informatie over de ware data $x_0$.
• Correctieformule: De geschatte score wordt aangepast door dit verschil te gebruiken als een richtingsvector:
  $$\hat{s}_{\theta, t} = s_{\theta, t} + \lambda (s_{\theta, t} - s_{\psi, t})$$
  Waarbij $\lambda$ de stapgrootte voor de correctie bepaalt.
• Schaal aanpassing: Vervolgens wordt de geschaalde score vermenigvuldigd met een factor $\omega$ om de grootte van de correctie te optimaliseren.
• Voordeel: Deze methode trekt de voorspelde score dichter naar de optimale (ware) score, zonder het netwerk opnieuw te hoeven trainen.

Belangrijkste Bijdragen

1. Systematische Analyse: Het is de eerste studie die systematisch zowel de reverse-starting bias als de exposure bias in graf-diffusiemodellen analyseert en aanpakt.
2. Innovatieve Oplossing: De combinatie van Langevin sampling (voor start-alignment) en score-correctie op basis van score-verschil (voor exposure bias) biedt een krachtige, plug-and-play oplossing.
3. Architectuur-onafhankelijk: De methode vereist geen wijzigingen in het bestaande netwerkontwerp (spatial, spectral, of hyperbolisch) en introduceert geen extra leercomponenten.
4. State-of-the-Art Prestaties: De methode is gevalideerd op meerdere modellen (GDSS, GSDM, HGDM, MOOD), datasets (algemene grafen, enzymen, QM9, ZINC250k) en taken, waarbij overtuigend betere resultaten werden behaald dan de bestaande state-of-the-art.

Resultaten

De experimenten tonen aan dat S++ aanzienlijke verbeteringen levert:

• Generieke Grafen: Op datasets zoals Community-small, Enzymes en Grid verbetert S++ de statistische distributiemetingen (MMD voor graad, clustering en orbits) aanzienlijk. Het presteert vaak beter dan bestaande methoden die correctors gebruiken, en doet dit met een lagere rekenkosten.
• Moleculaire Generatie: Op QM9 en ZINC250k datasets verbetert S++ de validiteit van de gegenereerde moleculen en verlaagt het de Frechet ChemNet Distance (FCD) en NSPDK MMD aanzienlijk (lagere waarden zijn beter).
• Efficiëntie: S++ kan hoge kwaliteit samples genereren met minder stappen in het terugwaartse proces. Bijvoorbeeld, op QM9 behaalde S++ met slechts 100 stappen bijna dezelfde prestaties als de baseline met 1000 stappen.
• Drug Discovery: In de MOOD-taken (ontdekking van nieuwe moleculen met specifieke eigenschappen) behaalde S++ een hogere "novel hit ratio" en betere docking scores dan de baselines.

Significantie

Dit paper is significant omdat het een fundamentele beperking in graf-diffusiemodels oplost die eerder over het hoofd werd gezien. Bestaande oplossingen voor beeld-diffusiemodellen (zoals het aanpassen van het ruis-schema) werken niet voor grafen vanwege de complexe structuur en schaalbeperkingen van grafdata.

De voorgestelde S++-methode biedt een universeel, lichtgewicht en effectief kader om de kwaliteit van graf-generatie te verbeteren. Het bewijst dat het aanpassen van het samplingproces (via Langevin alignment en score-correctie) vaak effectiever en efficiënter is dan het opnieuw trainen van complexe netwerken. Dit opent de deur voor robuustere en snellere generatieve modellen voor complexe grafstructuren, met directe toepassingen in chemie, farmacologie en netwerkanalyse.