One More Step: A Versatile Plug-and-Play Module for Rectifying Diffusion Schedule Flaws and Enhancing Low-Frequency Controls

Abstract

It is well known that many open-released foundational diffusion models have difficulty in generating images that substantially depart from average brightness, despite such images being present in the training data. This is due to an inconsistency: while denoising starts from pure Gaus-sian noise during inference, the training noise schedule retains residual data even in the final timestep distribution, due to difficulties in numerical conditioning in main-stream formulation, leading to unintended bias during in-ference. To mitigate this issue, certain ∊-prediction mod-els are combined with an ad-hoc offset-noise methodology. In parallel, some contemporary models have adopted zero-terminal SNR noise schedules together with v -prediction, which necessitate major alterations to pre-trained models. However, such changes risk destabilizing a large multitude of community-driven applications anchored on these pre-trained models. In light of this, our investigation revisits the fundamental causes, leading to our proposal of an inno-vative and principled remedy, called One More Step (OMS). By integrating a compact network and incorporating an ad-ditional simple yet effective step during inference, OMS ele-vates image fidelity and harmonizes the dichotomy between training and inference, while preserving original model pa-rameters. Once trained, various pre-trained diffusion mod-els with the same latent domain can share the same OMS module. Codes and models are released at here.