It's Never Too Late: Noise Optimization for Collapse Recovery in Trained Diffusion Models

TL;DR

This work shows that a simple noise optimization objective can mitigate mode collapse while preserving the fidelity of the base model and demonstrates that noise optimization yields superior results in terms of generation quality and variety.

Abstract

Contemporary text-to-image models exhibit a surprising degree of mode collapse, as can be seen when sampling several images given the same text prompt. While previous work has attempted to address this issue by steering the model using guidance mechanisms, or by generating a large pool of candidates and refining them, in this work we take a different direction and aim for diversity in generations via noise optimization. Specifically, we show that a simple noise optimization objective can mitigate mode collapse while preserving the fidelity of the base model. We also analyze the frequency characteristics of the noise and show that alternative noise initializations with different frequency profiles can improve both optimization and search. Our experiments demonstrate that noise optimization yields superior results in terms of generation quality and variety.

Figures (16)

• Figure 1: Repeatedly sampling from text-to-image models using a fixed text prompt produces surprisingly little visual variation (top row) in both Stable Diffusion SDXL-Turbo sdxlturbo (left) and Flux.1 [schnell] flux (right). Our approach (bottom row) directly optimizes the initial noise to recover from mode collapse, producing diverse outputs.
• Figure 2: We optimize the noise initialization to increase visual diversity given a fixed text prompt and diffusion model. Starting from i.i.d. noise samples, we generate a set of images. Using a diversity objective (e.g. DINO dissimilarity) and optionally a quality reward (e.g. HPSv2), we update the noise to produce output images that capture more diversity per text prompt. Our method supports optimizing over a variety of objective ensembles.
• Figure 3: Example images generated with SDXL-Turbo using different optimization objectives for the prompt "a photo of a teddy bear" (top row: $\text{i.i.d.}$ samples). Additional examples are included in the Appendix (\ref{['fig:supp_objectives_1', 'fig:supp_objectives_2']}).
• Figure 4: Image generations using our noise optimization approach for SDXL-Turbo yields improved diversity within generated image sets compared to $\text{i.i.d}$ sampling and gi. Pink noise initializations (b) give more diverse generations than standard white noise (a). Ours uses the DINO diversity objective (similar to \ref{['tab:main_table']} and \ref{['tab:pink_table']}).
• Figure 5: Sequential image generations using our noise optimization approach for Flux.1 [schnell] yields improved diversity of generated image sets compared to $\text{i.i.d}$ sampling. Our approach scales to large image sets by sequentially generating diverse images.