An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations

Probleemstelling
Diffusiemodellen (DM's) hebben een uitzonderlijk vermogen getoond bij het oplossen van computationele beeldvormingsinverse problemen door complexe beeldpriors te modelleren. De praktische toepassing ervan wordt echter ernstig beperkt door een afhankelijkheid van grootschalige, hoogwaardige, schone datasets voor pre-training. In veel real-world scenario's, met name in de wetenschappelijke en biomedische beeldvorming (bijv. röntgen CT, cryo-EM), is het verkrijgen van dergelijke schone data onhaalbaar of evenredig duur. Daarentegen zijn gecorrumpeerde observaties (ruis, vervaging of maskering) met differentieerbare voorwaartse modellen overvloedig aanwezig. Dit creëert een "kip-ei"-dilemma: het trainen van een nauwkeurig diffusiemodel vereist schone beelden, terwijl het reconstrueren van schone beelden uit gecorrumpeerde observaties een goed getraind diffusiemodel vereist. Bestaande methoden voor leren van gecorrumpeerde data, zoals AmbientDiffusion of SURE-Score, worstelen vaak met specifieke soorten corruptie (bijv. grootschalige inpainting of significante ruis) of slagen er niet in te generaliseren over diverse inverse problemen.

Methodologie: EMDiffusion
De auteurs stellen EMDiffusion voor, een nieuw Expectation-Maximization (EM) raamwerk dat is ontworpen om beeldvormingsinverse problemen gezamenlijk op te lossen en een diffusiemodel direct te trainen vanaf gecorrumpeerde observaties. De methode werkt via een iteratieve cyclus van twee stappen:

1. Initialisatie: Het proces begint met een "vage" diffusieprior die is getraind op een zeer kleine set schone beelden (slechts 10–50 beelden). Deze initialisatie is cruciaal om te voorkomen dat het model inklapt naar een distributie die uitsluitend wordt gedefinieerd door de gecorrumpeerde inputs. De auteurs merken op dat deze initiële schone beelden niet uit dezelfde distributie als de doeltaak hoeven te komen (out-of-distribution data kan een effectieve initializer dienen).

2. E-stap (Expectation): Gebruikmakend van het huidige diffusiemodel als prior, voert het algoritme Adaptive Diffusion Posterior Sampling uit om schone beelden te reconstrueren uit de gecorrumpeerde observaties.

   • Standaard posterior sampling faalt vaak in de vroege iteraties omdat de prior zwak is, wat leidt tot "mode collapse" waarbij samples de observatie negeren en terugkeren naar de trainingsdistributie.
   • Om dit te mitigeren, introduceert EMDiffusion een hyperparameter $\lambda$ die de data-likelihood term in de reverse-time stochastic differential equation (SDE) herschaalt.
   • $\lambda$ wordt in elke E-stap adaptief afgestemd om de dataverlies te minimaliseren ($\|y - A(\hat{x}_0)\|_2^2$), waardoor de invloed van de imperfecte prior wordt gebalanceerd tegenover de geobserveerde data. Dit maakt de generatie van schonere reconstructies mogelijk, zelfs wanneer de prior aanvankelijk slecht is.

3. M-stap (Maximization): Het algoritme werkt de gewichten van het diffusiemodel bij met behulp van de gereconstrueerde beelden (posterior samples) die in de E-stap zijn gegenereerd.

   • Deze stap lijkt op standaard denoising score matching, waarbij het model leert de scorefunctie van de gereconstrueerde "schone" data te benaderen.
   • Om de training te versnellen en overfitting op ruizige vroege reconstructies te voorkomen, hanteert de methode een twee-fasen trainingsstrategie: aanvankelijk verfijnt het de geërfde gewichten van de vorige iteratie; zodra de reconstructiekwaliteit verbetert (aangegeven door een optimale $\lambda$ die onder 1 zakt), wordt het model opnieuw geïnitialiseerd en vanaf nul getraind op de volledige set posterior samples.

Belangrijkste Bijdragen

• EM Raamwerk voor Gecorrumpeerde Data: Het artikel introduceert het eerste algemene raamwerk om score-gebaseerde diffusiemodellen te trainen vanaf willekeurig gecorrumpeerde observaties met behulp van een EM-benadering, waardoor de afhankelijkheid van grote schone datasets effectief wordt doorbroken.
• Adaptieve Posterior Sampling: Het voorstel van een adaptieve schalfactor ($\lambda$) voor de likelihood term in de reverse SDE adresseert het mode collapse probleem dat inherent is aan de vroege stadia van EM-iteraties, wat stabiele reconstructie mogelijk maakt, zelfs met zwakke priors.
• Iteratieve Verfijning: De methode demonstreert dat een kleine hoeveelheid schone data kan fungeren als een impliciete regularisator, waardoor het model progressief kan convergeren naar de ware schone datadistributie door middel van afwisselende reconstructie- en trainingsstappen.
• Generaliseerbaarheid: In tegen tegenstelling tot voorgaande werken die specifiek zijn afgestemd op bepaalde corrupties (zo.v. masking), biedt EMDiffusion een verenigde aanpak die toepasbaar is op random inpainting, denoising en deblurring.

Experimentele Resultaten
De auteurs valideren EMDiffusion op de CIFAR-10 en CelebA datasets over drie taken: random inpainting (60% pixel masking), Gaussian denoising en Gaussian deblurring.

• Inpainting: EMDiffusion presteert significant beter dan baselines zoals AmbientDiffusion en SURE-Score, en bereikt een reconstructiekwaliteit (PSNR/LPIPS) die vergelijkbaar is met Diffusion Posterior Sampling (DPS) dat een prior gebruikt die is getraind op de volledige schone dataset.
• Denoising & Deblurring: De methode overtreft SURE-Score en self-supervised baselines (Noise2Self), waarbij fijne details (bijv. auto-wielen, gelaatstrekken) worden hersteld die andere methoden gladstrijken.
• Generatieve Kwaliteit: De geleerde diffusiemodellen bereiken state-of-the-art FID-scores vergeleken met andere methoden getraind op gecorrumpeerde data. Opvallend genoeg presteert een model dat via EM is geleerd van 50.000 gecorrumpeerde beelden beter dan een model dat is getraind op ongeveer 15.000 schone beelden.
• Efficiëntie: De methode bereikt deze resultaten met een standaard DDPM-architectuur (35.7M parameters), wat aanzienlijk kleiner is dan de architecturen die worden gebruikt door concurrerende baselines (bijv. de 113M+ parameters van AmbientDiffusion).

Significantie en Claims
Het artikel claimt dat EMDiffusion aantoont dat het informatie-theoretisch mogelijk is om de onderliggende schone datadistributie primair te leren van gecorrumpeerde metingen. De significantie ligt in:

1. Overwinnen van Datatekort: Het biedt een praktische oplossing voor domeinen waar schone data niet beschikbaar is, door gebruik te maken van de overvloed aan gecorrumpeerde observaties.
2. Robuuste Initialisatie: Het stelt vast dat een minimale set schone beelden (zelfs out-of-distribution) voldoende is om de manifold te initialiseren, wat voorkomt dat het model inklapt naar een vertekende distributie van gecorrumpeerde inputs.
3. State-of-the-Art Prestaties: De methode bereikt nieuwe state-of-the-art prestaties in inverse beeldvormingstaken zonder dat er vooraf getrainde modellen op grote schone datasets nodig zijn, waarmee de kloof tussen theoretische priors en praktische dataconstraints wordt overbrugd.

De auteurs blijven bescheiden over de reikwijdte en erkennen dat de methode steunt op de aanname dat het voorwaartse model bekend en differentieerbaar is, en dat toekomstig werk zich kan uitbreiden naar het leren van onbekende voorwaartse modellen of ruisstatistieken.