MAD: Manifold Attracted Diffusion

Technische Samenvatting: Manifold Attracted Diffusion (MAD)

Probleemstelling
Score-gebaseerde diffusiemodellen hebben zich gevestigd als state-of-the-art methoden voor beeldsynthese en presteren vaak beter dan GANs en VAEs op het gebied van steekproefkwaliteit en diversiteit. Er ontstaat echter een aanzienlijke beperking wanneer trainingsdata corrupt is door ruis uit meetfouten, compressie-artefacten of datacollectieprocessen. Standaard diffusiemodellen die direct op dergelijke ruisige datasets worden getraind, leren de corruptie te reproduceren, waardoor gegenereerde monsters de ruiskenmerken van de trainingsdata behouden. Hoewel bestaande benaderingen voor ruisige training vaak het trainingsproces zelf vereisen om te wijzigen (bijv. het aanpassen van verliesfuncties of het vereisen van een mix van schone en ruisige data), is er een gebrek aan efficiënte methoden die schone monsters kunnen genereren uit een ruisige distributie met behulp van gevestigde trainingsframeworks en vooraf getrainde netwerken zonder hertraining.

Methodologie
De voorgestelde methode, Manifold Attracted Diffusion (MAD), pakt dit aan door de inferentieprocedure van score-gebaseerde modellen te wijzigen in plaats van het trainingsproces. De benadering is gebaseerd op de manifold hypothesis (variëteetshypothese), die stelt dat betekenisvolle data zich bevindt op een laagdimensionale manifold binnen een hoogdimensionale ambiance-ruimte. De auteurs interpreteren ruis als variaties met een lage amplitude in richtingen loodrecht (off-manifold) op deze structuur, terwijl betekenisvolle datavariaties optreden langs de manifold (on-manifold).

De kern van de technische innovatie is de introductie van de extended score (uitgebreide score).

• Definitie: Voor een waarschijnlijkheidsmaat $p$ en een parameter $\gamma > 0$, wordt de extended score $H_\gamma p$ gedefinieerd als $(1 + \gamma)S(p * g_\gamma) + \gamma \frac{d}{d\gamma}S(p * g_\gamma)$, waarbij $S$ de standaard score-operator ($\nabla \log p$) is en $g_\gamma$ een Gaussische kernel.
• Eigenschappen: In tegenstelling tot de standaard score, die ongedefinieerd is voor Dirac-delta-distributies (die data op een lagerdimensionale manifold vertegenwoordigen met nul variantie in bepaalde richtingen), is de extended score wel gedefinieerd. Cruciaal is dat voor een Dirac-delta de extended score zich gedraagt als de score van een Gaussische verdeling met een eenheidsvariantie. Dit maakt het mogelijk dat het inferentieproces de richtingen met een lage variantie (off-manifold) behandelt alsof ze een niet-nul variantie hebben, wat effectief een "soft thresholding"-effect toepast.
• Inferentie-algoritme: De auteurs stellen een iteratief inferentie-algoritme voor dat een benadering van de extended score gebruikt. Gegeven een vooraf getraind denoiser-netwerk $D_\theta$ (dat de standaard score benadert), wordt de extended score benaderd door het netwerk te evalueren bij een licht geperturbeerd ruisniveau en een eindige differentie te berekenen om de afgeleide met betrekking tot $\gamma$ te benaderen.
• Implementatie: Het algoritme introduceert hyperparameters ($a, b, p$) die de sterkte van de "manifold attraction" controleren. Door deze aan te passen, onderdrukt de inferentieprocedure kleine variaties in off-manifold richtingen (ruis) terwijl grotere variaties in on-manifold richtingen (signaal) behouden blijven. De methode is compatibel met gevestigde frameworks, zoals de probability flow ODE-formulering door Karras et al. (2022), en kan gebruikmaken van vooraf getrainde netwerken.

Belangrijkste Bijdragen

1. Formele Definitie en Analyse: Het artikel biedt een rigoureuze wiskundige definitie van de extended score, analyseert het gedrag ervan op zowel standaard distributies als Dirac-delta-maten, en demonstreert het vermogen om de score-operator te generaliseren naar gedegenereerde distributies.
2. Efficiënt Inferentie-algoritme: De auteurs presenteren een inferentie-algoritme dat ruis in gegenereerde monsters vermindert zonder dat een speciale trainingsprocedure nodig is. Dit maakt de methode compatibel met bestaande trainingspipelines en vooraf getrainde modellen.
3. Empirische Validatie: De effectiviteit van de benadering wordt aangetoond via toy-problemen, synthetische data en real-world datasets, inclusief FFHQ, AFHQ, ImageNet en cryo-elektronenmicroscopie (Cryo-EM) data.

Resultaten

• Toy en Synthetische Data: In 2D illustratieve voorbeelden slaagt de extended score erin om trajecten van Gaussische verdelingen met een lage variantie te laten instorten naar hun gemiddelden (wat Dirac-delta's nabootst), terwijl de spreiding van componenten met een hoge variantie behouden blijft. Op synthetische datasets met toegevogeerde Gaussische ruis en pixelverwijdering presteert MAD aanzienlijk beter dan standaard score-inferentie, waarbij monsters worden gegenereerd die de oorspronkelijke schone vormen herstellen.
• Real-World Beelddata: Experimenten op FFHQ, AFHQv2 en ImageNet laten zien dat het verhogen van de manifold attraction parameter ($b$) resulteert in schonere achtergronden en scherpere primaire kenmerken, wat effectief hoogfrequente ruis en secundaire variaties wegfiltert.
• Kwantitatieve Metrieken: Op CIFAR-10 en FFHQ-datasets die corrupt zijn met Gaussische, shot- en pixel-verwijderingsruis, bereikt MAD aanzienlijk lagere Fréchet Inception Distance (FID) scores vergeleken met standaard inferentie. Opvallend genoeg komen de standaard inferentie FID's nauw overeen met de ruisige trainingsdata FID's, wat bevestigt dat standaard modellen corruptie reproduceren, terwijl MAD succesvol denoiset.
• Cryo-EM Toepassing: Toegepast op de EMPIAR-11618 dataset (ruisige single-particle Cryo-EM beelden), genereerde MAD monsters die sterk leken op de verfijnde vormen verkregen met traditionele methoden, ondanks dat het netwerk uitsluitend op ruisige data was getraind.

Betekenis en Claims
Het artikel beweert dat MAD een praktische oplossing biedt voor "blind generative denoising", waarbij het specifieke degradatiemodel onbekend is en alleen ruisige data beschikbaar is. De betekenis van de methode ligt in het vermogen om de geometrische structuur van datamanifolds te benutten om ruis te filteren tijdens de generatiefase, in plaats van te vertrouwen op complexe aanpassingen tijdens de training.

De auteurs zijn bescheiden over de reikwijdte van hun claims:

• Ze beweren niet een inverse probleem op te lossen (zoals het denoisen van een enkele gegeven afbeelding), maar eerder te leren van een ruisige dataset om nieuwe, schone monsters uit de onderliggende distributie te produceren.
• Ze erkennen dat het bereiken van state-of-the-art prestaties in specifieke domeinen (zoals Cryo-EM) het integreren van meer domeinspecifieke voorkennis vereist, wat buiten de scope van dit werk valt.
• Ze merken op dat hoewel de methode effectief is, de computationele kosten maximaal het dubbele zijn van standaard inferentie (vanwege de extra netwerk-evaluatie), hoewel empirische observaties suggereren dat het gebruik van de extended score alleen in de tweede helft van de inferentiestappen volstaat.

Het werk positioneert de extended score als een geometrisch gefundeerd mechanisme om de generatiekwaliteit te verbeteren, potentieel toepasbaar zelfs op modellen getraind op schone data, en opent wegen voor toekomstig onderzoek naar adaptieve parametervariatie en integratie met conditionele diffusiemodellen.