Timestep Rescheduling in Diffusion Inversion

Technical Summary: Timestep Rescheduling in Diffusion Inversion

Problem Statement
Diffusion inversion is the process of mapping a real image $z_0$ back to the Gaussian latent space ($z_T$) of a diffusion model, enabling tasks like image reconstruction and editing. While Denoising Diffusion Implicit Models (DDIM) provide a deterministic and efficient formulation for this, the inversion process is inherently approximate. The standard approach involves analytically reversing the denoising step, which introduces deviations that accumulate over time, leading to noticeable reconstruction errors and degraded editing fidelity.

Existing methods primarily attempt to mitigate these errors by solving fixed-point problems to refine the inversion equation. However, these approaches largely overlook a critical factor: the influence of timestep scheduling (the selection and spacing of diffusion timesteps) on inversion fidelity. Most pipelines utilize uniformly spaced timesteps, which the authors argue is suboptimal. The paper posits that the distribution of timesteps significantly impacts the accumulation of global inversion errors, yet this relationship remains underexplored compared to its role in diffusion sampling or training.

Methodology: TRDI (Timestep Rescheduling in Diffusion Inversion)
The authors propose TRDI, a nonuniform timestep scheduler designed to minimize overall inversion error without requiring additional parameters or computational overhead during the inference stage. The methodology is grounded in a theoretical analysis of the inversion error and implemented through a two-stage "coarse-to-fine" strategy.

1. Theoretical Insight:
   The authors derive that the inversion error induced by a timestep jump $\Delta t$ can be reformulated as a scaled fixed-point problem. The error magnitude $\delta$ is proportional to a scaling coefficient $c_\alpha(t, \Delta t)$ multiplied by the discrepancy between consecutive model predictions ($\Delta \epsilon_\theta$).

   • Key Observation: The scaling coefficient exhibits a parabolic trend relative to the timestep index. Errors are concentrated at both small timesteps (early diffusion steps) and large timesteps (late steps), with a dip in the middle. Furthermore, larger step sizes ($\Delta t$) inherently induce larger errors.

2. Global Rescaling (Coarse Adjustment):
   To address the parabolic error distribution, the authors introduce a global rescaling mechanism. They modify the original uniform timestep sequence $t_k$ using a power-law transformation controlled by a hyperparameter $\gamma$:
   $$t_k = t_1 + (t_K - t_1) \left( \frac{k-1}{K-1} \right)^\gamma$$

   • When $\gamma > 1$, timesteps are stretched toward earlier steps (allocating more resolution to the high-error early region).
   • When $\gamma < 1$, timesteps become denser toward the end.
     This step prepares the schedule for fine-grained optimization by globally adjusting the allocation of steps.

3. Local Rescheduling via Dynamic Programming (Fine Adjustment):
   To further minimize accumulated error, the authors employ a windowed dynamic programming (DP) approach.

   • Cost Map: A cost map $E[k, t]$ is initialized to track the minimum accumulated error for reaching step index $t$ at the $k$-th inversion step.
   • Recursion: Within a sliding window of width $2d+1$ around the current step, the algorithm recursively searches for the locally optimal transition between adjacent timesteps that minimizes the sum of the previous accumulated error and the current step's error coefficient.
   • Backtracking: The optimal sequence of timesteps is retrieved by backtracking from the final step to the first.

The entire rescheduling process is performed offline (or as a pre-processing step) based on the known noise schedule, meaning it adds negligible runtime overhead to the actual image inversion or editing process.

Key Contributions

• Theoretical Formulation: The paper derives a theoretical connection between timestep jumps and inversion error, framing the error as a scaled fixed-point problem dependent on the noise schedule.
• Adaptive Scheduler: It proposes a coarse-to-fine rescheduling framework (TRDI) that adaptively allocates larger steps to low-error regions and finer steps to high-error regions, effectively balancing efficiency and accuracy.
• Plug-and-Play Enhancement: The method is designed as an "off-the-shelf" enhancement that requires no extra parameters, training, or significant computational overhead, making it compatible with existing inversion pipelines.

Experimental Results
The authors validate TRDI through extensive experiments on image reconstruction (MSCOCO dataset) and image editing (PIE-Bench) using various diffusion models (Stable Diffusion v1.5, SDXL, and SDXL Turbo) and baseline inversion methods (DDIM, ReNoise, NPI, GNRI).

• Reconstruction: On SD v1.5, integrating TRDI consistently improved metrics across all baselines. For instance, with the DDIM baseline, PSNR increased from 20.07 to 20.21, and LPIPS decreased (improving perceptual quality) from 193.97 to 187.85. Qualitative results showed corrections in structural details like train windows and fruit shapes that baselines failed to reconstruct accurately.
• Editing: On the PIE-Bench, TRDI significantly improved structure preservation and background preservation while maintaining high CLIP similarity. Notably, for few-step models like SDXL Turbo (4 steps), where background degradation is typically severe, TRDI achieved substantial improvements in background preservation (e.g., reducing Structure Distance by 17.43% for DDIM).
• Ablation Studies: The study confirmed that a scaling factor $\gamma \approx 1.05$ and a window width $d=8$ provided the best trade-off between structural fidelity and semantic editability.

Significance and Claims
The paper claims that timestep scheduling is a critical, yet previously overlooked, factor in diffusion inversion fidelity. By revealing the parabolic relationship between timestep size/index and error accumulation, the authors demonstrate that a strategic, nonuniform allocation of computational steps can significantly reduce global inversion errors.

The authors position TRDI not as a replacement for stronger inversion formulations (like fixed-point iterations or stochastic inversion) but as a complementary, lightweight scheduler-level retrofit. It allows existing deterministic inversion pipelines to achieve superior reconstruction and editing results without increasing computational costs or requiring model retraining. The work emphasizes that optimizing when to take steps is as important as how to take them.