Timestep Rescheduling in Diffusion Inversion

Technische Samenvatting: Timestep Rescheduling in Diffusion Inversion

Probleemstelling
Diffusion inversion is het proces waarbij een reële afbeelding $z_0$ wordt teruggekaatst naar de Gaussische latente ruimte ($z_T$) van een diffusiemodel, wat taken zoals beeldreconstructie en bewerking mogelijk maakt. Hoewel Denoising Diffusion Implicit Models (DDIM) een deterministische en efficiënte formulering bieden voor dit proces, is de inversie inherent benaderend. De standaardaanpak behelst het analytisch omkeren van de denoising-stap, wat afwijkingen introduceert die over de tijd accumuleren, wat leidt tot merkbare reconstructiefouten en een verslechterde bewerkingsgetrouwheid (fidelity).

Bestaande methoden proberen deze fouten primair te mitigeren door vastpuntproblemen (fixed-point problems) op te lossen om de inversievergelijking te verfijnen. Deze benaderingen negeren echter grotendeels een cruciale factor: de invloed van timestep scheduling (de selectie en tussenruimte van diffusie-timesteps) op de inversiegetrouwheid. De meeste pipelines maken gebruik van uniform verdeelde timesteps, wat de auteurs als suboptimaal beschouwen. Het artikel stelt dat de distributie van timesteps een significante impact heeft op de accumulatie van globale inversiefouten, terwijl deze relatie minder onderzocht is vergeleken met de rol ervan bij diffusion sampling of training.

Methodologie: TRDI (Timestep Rescheduling in Diffusion Inversion)
De auteurs stellen TRDI voor, een niet-uniforme timestep scheduler die ontworpen is om de totale inversiefout te minimaliseren zonder dat er extra parameters of computationele overhead nodig zijn tijdens de inferentiefase. De methodologie is gebaseerd op een theoretische analyse van de inversiefout en geïmplementeerd via een twee-fasen "coarse-to-fine" strategie.

1. Theoretisch Inzicht:
   De auteurs leiden af dat de door een timestep jump $\Delta t$ geïnduceerde inversiefout kan worden geherformuleerd als een geschaald vastpuntprobleem. De foutgrootte $\delta$ is evenredig aan een schalingscoëfficiënt $c_\alpha(t, \Delta t)$ vermenigvuldigd met de discrepantie tussen opeenvolgende modelvoorspellingen ($\Delta \epsilon_\theta$).

   • Kernobservatie: De schalingscoëfficiënt vertoont een parabolische trend ten opzichte van de timestep-index. Fouten concentreren zich zowel bij kleine timesteps (vroege diffusiestappen) als bij grote timesteps (late stappen), met een dalpunt in het midden. Bovendien induceren grotere stapgroottes ($\Delta t$) inherent grotere fouten.

2. Global Rescaling (Coarse Adjustment):
   Om de parabolische foutdistributie aan te pakken, introduceren de auteurs een mechanisme voor globale herschaling. Ze modificeren de oorspronkelijke uniforme timestep-sequentie $t_k$ met behulp van een machtswet-transformatie die wordt gecontroleerd door een hyperparameter $\gamma$:
   $$t_k = t_1 + (t_K - t_1) \left( \frac{k-1}{K-1} \right)^\gamma$$

   • Wanneer $\gamma > 1$, worden de timesteps uitgerekt naar de eerdere stappen (hierdoor wordt meer resolutie toegewezen aan de gebied met hoge fouten).
   • Wanneer $\gamma < 1$, worden de timesteps dichter bij elkaar geplaatst richting het einde.
     Deze stap bereidt de scheduler voor op fijnmazige optimalisatie door de allocatie van stappen globaal aan te passen.

3. Local Rescheduling via Dynamic Programming (Fine Adjustment):
   Om de geaccumuleerde fout verder te minimaliseren, gebruiken de auteurs een windowed dynamic programming (DP) benadering.

   • Cost Map: Een cost map $E[k, t]$ wordt geïnitialiseerd om de minimale geaccumuleerde fout bij het bereiken van stap-index $t$ bij de $k$-de inversiestap bij te houden.
   • Recursie: Binnen een sliding window van breedte $2d+1$ rond de huidige stap, zoekt het algoritme recursief naar de lokaal optimale transitie tussen aangrenzende timesteps die de som van de vorige geaccumuleerde fout en de huidige stap-coëfficiënt minimaliseert.
   • Backtracking: De optimale sequentie van timesteps wordt verkregen door vanaf de laatste stap naar de eerste stap terug te werken (backtracking).

Het gehele rescheduling-proces wordt offline uitgevoerd (of als een pre-processing stap) op basis van de bekende ruisschema (noise schedule), wat betekent dat het verwaarloosbare runtime-overhead toevoegt aan het daadwerkelijke beeldinversie- of bewerkingsproces.

Belangrijkste Bijdragen

• Theoretische Formulering: Het artikel leidt een theoretische verbinding af tussen timestep jumps en inversiefout, waarbij de fout wordt geframed als een geschaald vastpuntprobleem dat afhankelijk is van het ruisschema.
• Adaptieve Scheduler: Het stelt een coarse-to-fine rescheduling framework voor (TRDI) dat adaptief grotere stappen toewijst aan gebieden met lage fouten en fijnere stappen aan gebieden met hoge fouten, waardoor efficiëntie en nauwkeurigheid effectief worden gebalanceerd.
• Plug-and-Play Verbetering: De methode is ontworpen als een "off-the-shelf" verbetering die geen extra parameters, training of significante computationele overhead vereist, waardoor het compatibel is met bestaande inversie-pipelines.

Experimentele Resultaten
De auteurs valideren TRDI via uitgebreide experimenten op beeldreconstructie (MSCOCO dataset) en beeldbewerking (PIE-Bench) met diverse diffusiemodellen (Stable Diffusion v1.5, SDXL, en SDXL Turbo) en baseline inversiemethoden (DDIM, ReNoise, NPI, GNRI).

• Reconstructie: Bij de integratie van TRDI op SD v1.5 verbeterden de metrieken consistent over alle baselines heen. Bijvoorbeeld, met de DDIM-baseline steeg de PSNR van 20.07 naar 20.21, en de LPIPS nam af (wat de perceptuele kwaliteit verbetert) van 193.97 naar 187.85. Kwalitatieve resultaten toonden correcties in structurele details zoals treinramen en de vorm van fruit, die baselines niet accuraat konden reconstrueren.
• Bewerking: Op de PIE-Bench verbeterde TRDI aanzienlijk de structurele getrouwheid en achtergrondbehoud, terwijl de CLIP-gelijkenis hoog bleef. Met name voor few-step modellen zoals SDXL Turbo (4 stappen), waar achtergronddegradatie doorgaans ernstig is, bereikte TRDI substantiële verbeteringen in achtergrondbehoud (bijvoorbeeld een reductie van 17,43% in Structure Distance voor DDIM).
• Ablatie-studies: De studie bevestigde dat een schalingfactor $\gamma \approx 1.05$ en een window breedte $d=8$ de beste balans boden tussen structurele getrouwheid en semantische bewerkbaarheid.

Betekenis en Claims
Het artikel claimt dat timestep scheduling een kritische, maar voorheen over het hoofd geziene factor is in de getrouwheid van diffusion inversion. Door de parabolische relatie tussen de grootte van de timestep/index en de foutaccumulatie te onthullen, demonstreren de auteurs dat een strategische, niet-uniforme allocatie van computationele stappen de globale inversiefouten aanzienlijk kan verminderen.

De auteurs positioneren TRDI niet als een vervanging voor sterkere inversieformuleringen (zoals fixed-point iteraties of stochastische inversie), maar als een complementaire, lichtgewicht scheduler-level retrofit. Het stelt bestaande deterministische inversie-pipelines in staat om superieure reconstructie- en bewerkingsresultaten te behalen zonder de computationele kosten te verhogen of modeltraining te vereisen. Het werk benadrukt dat het optimaliseren van wanneer men stappen zet even belangrijk is als hoe men ze zet.