Technische Samenvatting: FlowADMM: Plug-and-play ADMM met Flow-gebaseerde Renoise-Denoise Priors

Probleemstelling

Inverse problemen in beeldverwerking, zoals ruisreductie, deblurring, super-resolutie en inpainting, zijn doorgaans slecht gesteld en worden geformuleerd als het herstellen van een beeld $x$ uit gedegradeerde metingen $y = Ax + \nu$. Hoewel deep learning krachtige directe mappingen biedt, missen deze vaak interpreteerbaarheid en theoretische garanties. Klassieke variatiebenaderingen formuleren dit als het minimaliseren van een term voor data-trouw plus een regularisatieterm, vaak opgelost via proximaal algoritmes zoals Forward-Backward Splitting (FBS).

Plug-and-Play (PnP) methoden verbeteren hierop door de proximaal operator van een handgemaakte regularisator te vervangen door een geleerde denoiser. Onlangs hebben generatieve flow- en diffusiemodellen tijdsafhankelijke families van denoisers $\{D_t\}$ geleverd die zijn aangepast aan specifieke ruisniveaus. Een belangrijke uitdaging doet zich echter voor: tijdens iteratieve reconstructie wijken de tussentijdse iteraten vaak af van de ruisverdeling $p_t$ waarvoor de denoiser is getraind. Bestaande flow-gebaseerde PnP-methoden (zoals PnP-Flow, Flower) lossen dit op door expliciete renoise-denoise-stappen op te nemen, waarbij de huidige iteraat wordt verstoord naar het juiste ruisniveau voordat gedenoist wordt. Deze methoden middelen vaak over meerdere ruismonsters om het proces te stabiliseren.

Ondanks hun empirische succes worden bestaande renoise-denoise PnP-methoden geformuleerd als stochastische algoritmes, waardoor hun convergentiegedrag theoretisch moeilijk te analyseren is. Specifiek is de deterministische operator die impliciet wordt gegenereerd door het middelen van de denoiser-output over de latente ruisverdeling niet gekarakteriseerd, waardoor de convergentie-eigenschappen van deze methoden slecht begrepen blijven.

Methodologie

De Gemiddelde Renoise-Denoise Operator

De auteurs identificeren en formaliseren de deterministische operator die ten grondslag ligt aan stochastische flow-gebaseerde PnP-methoden. Gegeven een voorgetraind flow-snelheidsnetwerk $v_\theta$ dat een tijdsafhankelijke MMSE-denoiser definieert $D_t = I + (1-t)v_\theta^t$, omvat de standaard stochastische stap het renoisen van een iteraat $x$ met ruis $\epsilon \sim \mathcal{N}(0, I)$ om $tx + (1-t)\epsilon$ te verkrijgen, gevolgd door denoising.

Het artikel definieert de gemiddelde renoise-denoise operator $\bar{S}_t(x)$ als de verwachting van dit proces over de ruisverdeling:
$$\bar{S}_t(x) = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I)} [D_t(tx + (1-t)\epsilon)]$$
Deze operator vertegenwoordigt de deterministische limiet van de stochastische updates die worden gebruikt door methoden zoals PnP-Flow. De auteurs analyseren hun residu-vorm $R_t(x) = \bar{S}_t(x) - x$, en interpreteren deze als een vectorveld dat lokale correcties biedt richting het data-maand. Onder ideale MMSE-aannames zijn punten op het data-maand vaste punten van deze operator.

FlowADMM Algoritme

Gebaseerd op dit inzicht stellen de auteurs FlowADMM voor, dat de gemiddelde renoise-denoise operator integreert in het Alternating Direction Method of Multipliers (ADMM)-kader. In tegenstelling tot FBS, staat ADMM meerdere prior-operatoren toe en behandelt het niet-differentieerbare data-termen zonder operator-afhankelijke stapgrootte-beperkingen.

Het algoritme lost het beperkte probleem $\min_{x,z} F_y(x) + R(z)$ onder de voorwaarde $x=z$ op door de uitgebreide Lagrangiaan te minimaliseren. De updates zijn:

1. Data-consistentie: $x^{k+1} = \text{prox}_{\tau F_y}(z^k - u^k)$
2. Prior-update: $z^{k+1} = \bar{S}_{t_k}(x^{k+1} + u^k)$
3. Dual-update: $u^{k+1} = u^k + x^{k+1} - z^{k+1}$

Omdat de verwachting in $\bar{S}_t$ onberekenbaar is, wordt deze benaderd via Monte Carlo-steekproeven. De auteurs stellen een tijdstap-afhankelijke steekproefstrategie voor: het gebruik van minder monsters in vroege iteraties (waar grove correcties worden gemaakt en de operator minder gevoelig is voor ruis) en het verhogen van het aantal monsters in latere iteraties (waar fijne details worden verfijnd en precisie cruciaal is).

Convergentie-analyse

Het artikel vestigt convergentie-garanties voor FlowADMM onder zwakke Lipschitz-voorwaarden op het onderliggende flow-netwerk:

• Vaste tijdschema: Voor een constante tijdstap $t$ is convergentie gegarandeerd als de Lipschitz-constante $L_v$ van het flow-netwerk voldoet aan $L_v < \frac{1}{1-t}$. Deze voorwaarde is aanzienlijk zwakker dan de strikte contractie ($L < 1$) die door veel andere PnP-methoden wordt vereist, vooral naarmate $t \to 1$.
• Niet-stationaire tijdschema's: De analyse wordt uitgebreid naar tijdschema's $(t_k)_k$ die convergeren naar een limiet $t_{max} < 1$. De auteurs bewijzen convergentie naar een vast punt van de limietoperator $T_{t_{max}}$, mits de rij operatoren de limiet voldoende snel benadert en het flow-netwerk Lipschitz-continu is met betrekking tot de tijd $t$.

Belangrijkste Bijdragen

1. Theoretische Karakterisering: De auteurs identificeren en analyseren de deterministische gemiddelde renoise-denoise operator die ten grondslag ligt aan stochastische flow-gebaseerde PnP-methoden, en karakteriseren hun residu-structuur en Lipschitz-gedrag.
2. Algoritmevoorstel: Zij stellen FlowADMM voor, een PnP-algoritme dat deze gemiddelde operator direct integreert in het ADMM-kader.
3. Convergentie-garanties: Zij leiden convergentiecriteria af voor FlowADMM onder zwakke Lipschitz-aannames, die zowel constante als niet-constante (tijdsvariabele) tijdstap-schema's omvatten.
4. Steekproefstrategie: Zij stellen een grof-naar-fijn Monte Carlo-steekproefstrategie voor die het aantal monsters in latere iteraties verhoogt, gemotiveerd door het gedrag van de gemiddelde operator.
5. Empirische Prestaties: Zij tonen state-of-the-art prestaties aan over meerdere inverse problemen, terwijl ze minder data-consistentie-evaluaties vereisen dan eerdere stochastische middelingbenaderingen.

Experimentele Resultaten

De auteurs evalueerden FlowADMM op de CelebA- en AFHQ-Cat-datasets over vijf taken: ruisreductie, deblurring, super-resolutie, willekeurige inpainting en box-inpainting.

• Prestaties: FlowADMM behaalde state-of-the-art resultaten in vervormingsmetrieken (PSNR en SSIM) op 7 van de 10 taken in vergelijking met recente flow- en diffusie-gebaseerde PnP-methoden (waaronder PnP-Flow, Flower, OT-ODE, D-Flow en DiffPIR). Opmerkelijke verbeteringen werden waargenomen in super-resolutie en inpainting-taken.
• Efficiëntie: Hoewel het het totale aantal flow-evaluaties (netwerkvoorwaartse passes) van de leidende concurrent Flower5-OT evenaarde, vereiste FlowADMM aanzienlijk minder evaluaties van de data-term proximaal operator. Dit komt omdat Flower5-OT volledige reconstructietrajecten middelt, terwijl FlowADMM alleen de prior-stap binnen elke iteratie middelt.
• Ablatiestudies: Experimenten bevestigden dat het "late-heavy" steekproefschema (het verhogen van monsters in latere iteraties) de prestaties significant verbetert ten opzichte van constante steekproeven. Bovendien verbeterde het toepassen van dit schema op PnP-Flow de resultaten, maar FlowADMM bleef het overtreffen, wat suggereert dat de ADMM-formulering zelf een voordeel biedt.
• Beperkingen: De convergentie-analyse is gebaseerd op Lipschitz-constanten van het voorgetrainde netwerk die niet expliciet worden bestraft tijdens training. Bovendien gelden de theoretische garanties voor de deterministische gemiddelde operator, terwijl de praktische implementatie stochastische Monte Carlo-benaderingen gebruikt.

Betekenis en Claims

Het artikel claimt dat door de deterministische operator die ten grondslag ligt aan stochastische renoise-denoise-methoden expliciet te formaliseren, de kloof tussen empirisch succes en theoretisch begrip in flow-gebaseerde PnP wordt overbrugd. Het werk biedt:

• Bewezen Convergentie: In tegenstelling tot eerdere stochastische PnP-methoden, wordt FlowADMM geleverd met convergentie-garanties onder realistische Lipschitz-voorwaarden op het flow-netwerk.
• Algoritmische Efficiëntie: De integratie in ADMM staat een efficiënter gebruik van rekenkracht toe (minder data-consistentiestappen) in vergelijking met traject-middelingsmethoden.
• Operator-theoretisch Inzicht: De analyse suggereert dat het "grof-naar-fijn" gedrag van de gemiddelde operator kan worden benut via adaptieve steekproefschema's, een principe dat zowel ADMM- als FBS-gebaseerde benaderingen ten goede komt.

De auteurs concluderen dat FlowADMM een robuust, theoretisch onderbouwd kader biedt voor het oplossen van inverse problemen met flow-gebaseerde priors, hoewel zij erkennen dat toekomstig werk zich kan richten op het afdwingen van Lipschitz-beperkingen tijdens training om stabiliteitsgaranties te versterken.