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

1. Problem Statement

Contemporary text-to-image diffusion models (e.g., SDXL-Turbo, Flux.1) suffer from mode collapse. When sampling multiple images from the same text prompt using different random seeds (i.i.d. noise), the models often produce nearly identical outputs with minimal visual variation.

• Current Limitations: Existing solutions rely on:
  • Guidance mechanisms: Steering the model via modified conditioning (e.g., Particle Guidance, DPP), which often trades off quality for diversity.
  • Search-based methods: Generating a large pool of candidates and pruning them (e.g., "Best-of-N"), which is computationally expensive and memory-intensive.
  • Prompt engineering: Modifying the text input, which cannot easily capture all variations in the output space.
• Core Question: Can we directly optimize the initial noise ($x_0$) to recover from mode collapse and generate diverse, high-quality images without altering the pre-trained model or the user prompt?

2. Methodology

The authors propose an end-to-end noise optimization framework that treats the initial noise vector as a learnable parameter. The process involves:

A. Noise Optimization Objective

Given a fixed prompt $c$ and a frozen diffusion model $g_\theta$, the method optimizes a batch of initial noise samples $\{x_0^{(i)}\}_{i=1}^B$. The objective function minimizes a loss $\mathcal{L}$ composed of three terms:

1. Quality Reward ($r_s$): Ensures individual images adhere to the prompt and maintain high fidelity (e.g., using CLIPScore or HPSv2).
2. Diversity Metric ($v_\mathcal{B}$): Maximizes the dissimilarity between images in the batch. The paper explores various metrics:
   • Pairwise: DINOv2 patch features, DreamSim, LPIPS, Color Histograms.
   • Set-level: Determinantal Point Processes (DPP) and Vendi Score, which measure the effective number of unique items in a set.
3. Regularization ($\text{reg}$): Prevents the optimized noise from drifting too far from the prior distribution $N(0, I)$. This is achieved by penalizing deviations from the $\chi_d$ distribution of the noise norm.

The optimization is performed via backpropagation through the frozen sampler $g_\theta$.

B. Noise Initialization Strategy (Pink Noise)

A key insight from the paper is the analysis of noise evolution. During optimization, the authors observed that the most significant changes occur in the low-frequency components of the noise spectrum, while high frequencies remain relatively stable.

• Hypothesis: Natural images follow a $1/f$ power spectrum (more power in low frequencies). Standard white noise is spectrally flat.
• Solution: The authors introduce Pink Noise initialization. They transform white noise into the frequency domain via FFT, apply a power-scaling factor of $f^{-\alpha/2}$ (where $\alpha \approx 0.2$), and transform back.
• Benefit: This initialization aligns the starting noise with the statistical properties of natural images and the optimization dynamics, leading to faster convergence and higher diversity.

C. Optimization Modes

• Batched Optimization: Jointly optimizes a set of images (e.g., 4) to maximize set-level diversity.
• Sequential Optimization: Generates images one by one, optimizing each to differ from previously generated outputs. This reduces memory overhead and allows for generating large diverse sets.

3. Key Contributions

1. Noise Optimization for Diversity: Demonstrates that directly optimizing initial noise is a superior strategy for recovering from mode collapse compared to guidance-based or search-based methods.
2. Set-Level Diversity Objectives: Identifies that DPP and Vendi Score are more effective than simple pairwise metrics for optimizing diverse sets, as they prevent the "average pairwise" metric from being gamed by making one image very different while keeping others similar.
3. Pink Noise Initialization: Proposes a simple, parameter-free strategy (using pink noise) that consistently improves diversity across different models (SDXL-Turbo, Flux.1) and baselines, even without further optimization.
4. Quality-Diversity Trade-off: Shows that their method can increase diversity without sacrificing image quality or prompt adherence, overcoming the typical trade-off seen in guidance-based methods.

4. Experimental Results

The method was evaluated on GenEval, DPG-Bench (complex prompts), and T2I-CompBench using models like SDXL-Turbo and Flux.1 [schnell].

• Quantitative Performance:
  • Diversity: Outperforms i.i.d. sampling and the state-of-the-art search method [Parmar et al., 2025] across all metrics (DINO, DreamSim, LPIPS, Vendi Score).
  • Quality: Maintains or slightly improves image quality scores (HPSv2, CLIPScore) compared to baselines.
  • Efficiency: With pink noise initialization, the method reaches higher diversity scores than the baseline search method in significantly fewer iterations (e.g., 6–8 iterations vs. 9–12 for white noise).
• Qualitative Results:
  • Generates diverse variations in object pose, color, style, and background while strictly adhering to complex prompts.
  • Successfully handles complex prompts (e.g., "A picturesque autumn scene...") where i.i.d. sampling fails to produce variation.
• User Study: Human evaluators preferred the outputs of the proposed method (especially those optimized with DPP/Vendi) over i.i.d. sampling and the search baseline in 85–91% of cases.

5. Significance

This work shifts the paradigm of diversity generation in diffusion models from sampling strategies (waiting for a good seed) or guidance (modifying the generation path) to input optimization.

• General Applicability: The approach is model-agnostic and works on both step-distilled (fast) and standard diffusion models.
• Efficiency: By optimizing the noise directly, it avoids the massive memory overhead of generating and storing hundreds of candidates for pruning.
• Insight into Latent Space: The discovery that low-frequency noise components drive diversity provides a new understanding of how diffusion models interpret initial noise, suggesting that "structure" is established early in the frequency domain.

In conclusion, the paper presents a robust, efficient, and high-quality solution to the mode collapse problem, proving that "it's never too late" to optimize the input noise to unlock the full diversity potential of trained diffusion models.