Technische Samenvatting: Supervised Guidance Training voor oneindig-dimensionale Diffusiemodellen

Probleemstelling

Inverse problemen, met name die voortvloeiend uit partiële differentiaalvergelijkingen (PDE's), vereisen vaak het afleiden van onbekende parameters die bestaan in oneindig-dimensionale functieruimten. In het Bayesiaanse raamwerk is het doel om te steekproeven uit een posteriorverdeling over deze functies, geconditioneerd op ruisbeïnvloede waarnemingen. Hoewel scores gebaseerde diffusiemodellen (SDM's) recent zijn uitgebreid naar oneindig-dimensionale functieruimten om als expressieve priors te dienen, blijft de theorie voor het conditioneren van deze modellen om uit de posterior te steekproeven een open probleem.

Bestaande conditioneringstechnieken voor eindige dimensies, zoals de $h$-transformatie van Doob, vertrouwen op dichtheidsdecomposities en overgangsoperatoren die niet direct vertaalbaar zijn naar oneindige dimensies. Uitdagingen omvatten het niet-bestaan van Lebesgue-dichtheden, de potentiële singulariteit van maatstaven (Feldman-Hájek-stelling) bij het verschuiven van Gaussische maatstaven buiten hun Cameron-Martin-ruimte, en de moeilijkheid om differentieerbaarheid van overgangsoperatoren in Hilbertruimten te garanderen. Bovendien introduceert het direct discretiseren van data om methoden voor eindige dimensies toe te passen, resolutie-afhankelijke fouten en beperkt het de flexibiliteit van het raamwerk.

Methodologie

Theoretische Grondslag: Oneindig-dimensionale $h$-transformatie

De auteurs stellen een principieel raamwerk voor voor het conditioneren van oneindig-dimensionale diffusiemodellen met behulp van een uitbreiding van de $h$-transformatie van Doob. Zij stellen het bestaan vast van een conditioneel diffusieproces onder twee scenario's:

1. Scenario 1: De priorverdeling wordt ondersteund binnen een bol van de Cameron-Martin-ruimte van het Wiener-proces dat wordt gebruikt in de stochastische differentiaalvergelijking (SDE) voor diffusie.
2. Scenario 2: De prior is absoluut continu met betrekking tot een Gaussische maat, met specifieke grenzen voor de potentiefunctie en haar gradiënt.

Onder deze scenario's bewijzen de auteurs dat het tijdsomgekeerde conditionele proces voldoet aan een stochastische differentiaalvergelijking (SDE) waarbij de drift wordt gewijzigd door een geleidingsterm. Specifiek demonstreren zij dat de conditionele scorefunctie, $s_y(t, x)$, decomposeert in de onvoorwaardelijke score, $s(t, x)$, en een geleidingsterm, $C\nabla \log h_y(t, x)$:
$$s_y(t, x) = s(t, x) + C\nabla \log h_y(t, x)$$
Hierbij is $h_y$ de oneindig-dimensionale $h$-transformatie gedefinieerd via de conditionele verwachting van de likelihood-potentiaal.

Supervised Guidance Training (SGT)

Aangezien de geleidingsterm $C\nabla \log h_y$ onberekenbaar is vanwege de afhankelijkheid van de volledige overgangsregel van het oneindig-dimensionale proces, introduceren de auteurs Supervised Guidance Training (SGT). Dit is een trainingsdoelstelling zonder simulatie die is ontworpen om de geleidingsterm direct te leren.

• Doelstelling: De methode minimaliseert een denoising score matching-verlies. In plaats van te vertrouwen op op simulatie gebaseerde stochastische optimale controle (die lijdt aan hoge variantie en het steekproeven van trajecten vereist), maakt SGT gebruik van de bekende onvoorwaardelijke score $s(t, x)$ en de gezamenlijke verdeling van data en waarnemingen $(X_0, Y)$.
• Verliesfunctie: Het verlies is geformuleerd om het verschil te minimaliseren tussen de geleerde geleidingsterm $u(t, x, y)$ en de ware geleidingsterm. De doelstelling is:
  $$\mathbb{E} \left[ \left\| \frac{X_t - e^{-t/2}X_0}{1 - e^{-t}} + s(t, X_t) + u(t, X_t, Y) \right\|^2_K \right]$$
  waarbij $K$ een gekozen Hilbertruimte is (bijvoorbeeld de Cameron-Martin-ruimte).
• Parametrisatie: De geleidingsterm wordt geparametriseerd met behulp van een neurale netwerk. De auteurs onderzoeken het preconditioneren van de netwerkboutput met de covariantieoperator $C$ tot verschillende graden ($k \in \{0, 1/2, 1\}$) om regulariteit af te dwingen en af te stemmen op de ruisstructuur van de SDE.

Benaderingen en Vergelijkingen

Het artikel analyseert ook bestaande heuristische benaderingen:

• Stochastische Optimale Controle (SOC): Hoewel theoretisch sound, vereist SOC op simulatie gebaseerde training die computergewijs duur is en instabiel in de vroege trainingsfasen.
• Tweedie-benadering (FunDPS): De auteurs breiden de FunDPS-benadering uit (die de formule van Tweedie gebruikt om de schone data te schatten) naar Scenario 2 (absolute continuïteit), waardoor de restrictieve aannames van eerder werk worden versoepeld. Zij merken echter op dat deze benadering nog steeds rekenkundige overhead met zich meebrengt vanwege de noodzaak van backpropagatie door het scoremodel bij elke discretisatiestap.

Belangrijkste Bijdragen

1. Theoretische Uitbreiding: Het artikel biedt het eerste rigoureuze bewijs dat oneindig-dimensionale diffusiemodellen kunnen worden geconditioneerd via de $h$-transformatie van Doob onder algemene scenario's (Cameron-Martin-ondersteuning of absolute continuïteit), waarbij de decompositie van de conditionele score in een onvoorwaardelijke score en een geleidingsterm wordt vastgesteld.
2. Supervised Guidance Training (SGT): De auteurs stellen een nieuw, op simulatie gebaseerd trainingsalgoritme voor om de onberekenbare geleidingsterm te leren. Deze aanpak vermijdt de hoge variantie van op simulatie gebaseerde methoden en de rekenkosten van het differentiëren door het scoremodel bij elke stap (zoals vereist door op Tweedie gebaseerde benaderingen).
3. Resolutie-invariantie: Door het probleem in oneindige dimensies te behandelen en pas in de laatste stap te discretiseren, biedt het raamwerk een resolutie-invariante aanpak voor conditionering, waardoor de prestatiedegradering wordt vermeden die vaak wordt waargenomen bij het toepassen van architecturen voor eindige dimensies op hogere resoluties.
4. Generalisatie van FunDPS: Het werk bewijst dat op Tweedie gebaseerde benaderingen (FunDPS) gelden onder het minder restrictieve Scenario 2 (absolute continuïteit), waardoor hun toepasbaarheid wordt verbreed.

Experimentele Resultaten

De auteurs evalueren SGT op drie Bayesiaanse inverse problemen in functieruimten:

1. Sparce Waarnemingen: Het herstellen van een 1D-functie uit ruisbeïnvloede puntwaarnemingen.
2. Warmtevergelijking: Het herstellen van de initiële temperatuurverdeling van een warmtevergelijking uit ruisbeïnvloede waarnemingen op een later tijdstip.
3. Vorm Inpainting: Het reconstrueren van vormen (cijfer "3" uit MNIST) uit partiële Fourier-coëfficiënten (elliptische Fourier-descriptoren).

Vondsten:

• Prestaties: SGT presteert consistent beter dan de heuristische FunDPS-baseline wat betreft Root Mean Square Error (RMSE) en Energy Score (ES).
• Vergelijking met Conditionele Diffusie: SGT bereikt een nauwkeurigheid die vergelijkbaar is met een "conditioneel diffusie"-model dat vanaf nul is getraind (dat dient als bovengrens), ondanks dat het conditionele model aanzienlijk groter is en voor meer stappen is getraind.
• Robuustheid: In experimenten met imperfecte (willekeurig geïnitialiseerde of gedeeltelijk getrainde) onvoorwaardelijke basismodellen, toont SGT superieure robuustheid en herstelt het realistische steekproeven waar FunDPS volledig faalt.
• Efficiëntie: SGT vereist minder parameters en trainingsstappen dan het trainen van een volledig conditioneel model vanaf nul, aangezien alleen het geleidingsnetwerk wordt getraind terwijl een vooraf getraind onvoorwaardelijk model wordt hergebruikt.

Betekenis en Beweringen

Het artikel beweert de eerste methode in functieruimte voor het fine-tunen van getrainde diffusiemodellen aan te bieden om nauwkeurig uit een posteriorverdeling te steekproeven. De betekenis ligt in het overbruggen van de kloof tussen de theoretische eisen van oneindig-dimensionale Bayesiaanse inferentie en de praktische implementatie van diffusiemodellen.

De auteurs benadrukken dat hun raamwerk toelaat voor:

• Resolutie-invariantie: Het vermogen om steekproeven te genereren bij resoluties die verschillen van de trainingsdata zonder prestatiedegradering.
• Stabiliteit: Een trainingsregime dat stabiel is, zelfs wanneer het onderliggende onvoorwaardelijke model imperfect is.
• Flexibiliteit: Alternatieve opties voor discretisatie, training en verliesfuncties die niet worden beperkt door aannames over dichtheden voor eindige dimensies.

Het werk concludeert dat het expliciet leren van de oneindig-dimensionale geleidingsterm superieure nauwkeurigheid oplevert in vergelijking met heuristische benaderingen, en een principieel pad voorwaarts biedt voor het toepassen van diffusiemodellen op complexe inverse problemen in wetenschap en techniek.