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

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.

• The paper demonstrates that optimizing the initial noise vector significantly improves output diversity and mitigates mode collapse in diffusion models.
• It employs a composite loss combining set-level diversity and instance-level quality objectives, with a key focus on low-frequency noise modifications.
• Experimental results reveal marked improvements across benchmarks, maintaining image fidelity while enhancing semantic variety of outputs.

Noise Optimization for Mode Collapse Recovery in Pre-trained Diffusion Models

Introduction and Motivation

Modern diffusion models exhibit notable mode collapse even under varying random initializations: repeated sampling from fixed text prompts often leads to minimal output diversity. In practice, this undermines the utility of generative models in creative and exploratory tasks. Previous solutions, such as stronger guidance techniques and post-hoc candidate selection, provide limited remedy and often degrade image quality or impose significant computational overhead. The work "It's Never Too Late: Noise Optimization for Collapse Recovery in Trained Diffusion Models" (2601.00090) proposes an inference-time optimization framework directly targeting the initial noise vector for each generation batch, explicitly steering outputs towards maximally diverse yet high-fidelity sets, and systematically analyzes the properties and importance of latent noise initialization schemes.

Methodology: End-to-End Noise Optimization

The core of the proposed method is differentiable optimization of a batch of initial noise vectors $\{\mathbf{x}_0^{(i)}\}_{i=1}^B$, instead of independent random sampling. Each batch is optimized via a composite loss combining set-level diversity statistics (e.g., pairwise DINOv2 cosine distance, DPP, Vendi, DreamSim, LPIPS, color histograms, or low-res L2) and instance-level quality objectives (CLIPScore, HPSv2), along with a prior regularization term ensuring $\mathbf{x}_0$ remains in a high-density region of the prior ("norm matching"). The optimized objective formally couples batch-wide and sample-wise rewards:

Figure 1: The pipeline uses a diversity objective to iteratively update the initial noise, significantly increasing the variation of outputs for a given text prompt and diffusion model.

The optimization proceeds by backpropagating through the frozen generative sampler, progressively increasing the diversity of generated outputs while maintaining or exceeding pre-specified quality thresholds. Stopping criteria are based on diversity and quality targets or computational budgets.

Analysis of Noise Initialization: Frequency-domain Insights

A key contribution is the empirical and spectral analysis of how optimized noise latents evolve through iterations. Detailed frequency-decomposition reveals the majority of the changes induced by diversity-targeted optimization occur in the lowest frequency bands of the latent noise (i.e., large-scale structures), and that high-frequency content remains comparatively stable.

Figure 2: Iterative optimization primarily modifies the low-frequency (long spatial scale) content of the noise vector.

Motivated by the 1/f power-law statistics of natural images, the work proposes Pink Noise initialization (via frequency-domain filtering with $\alpha\in[0,1]$), resulting in larger coverage of the low-freq subspace. This produces consistently higher diversity in the resulting samples both under i.i.d. sampling and under subsequent optimization, outperforming standard white noise initializations.

Figure 3: Output variation (DINO diversity) grows with higher pink noise exponents, surpassing group-inference baselines with fewer optimization steps.

Experimental Results

Quantitative and Qualitative Diversity

Empirical evaluation on GenEval and T2I-CompBench benchmarks utilizing SDXL-Turbo, SANA-Sprint-1.6B, PixArt-$\alpha$, and Flux.1 [schnell] models demonstrates consistent, large improvements in output diversity (DINOv2, DreamSim, LPIPS) over both i.i.d sampling and prior search-based approaches such as Group Inference (GI) [gi]. The method maintains or even slightly improves CLIP-based prompt consistency and HPSv2 human-preference scores across all model backbones.

Sample generations highlight clear qualitative improvements: generated image sets exhibit wider semantic and structural variance (in colors, background, pose, and object instance), as visualized below.

Figure 4: Image generations using the proposed noise optimization for SDXL-Turbo demonstrate greatly improved intra-batch diversity compared to i.i.d sampling.

Across multiple backbones and benchmarks, the method demonstrates higher diversity with negligible or positive effects on alignment and perceptual quality. Human preference studies corroborate quantitative metrics, showing strong subject preference for batchwise diversity under DPP and Vendi objective variants.

Objective and Hyperparameter Ablations

Optimization with different diversity objectives (DINOv2, DreamSim, LPIPS, DPP, Vendi, color histograms, L2) reveals each primarily maximizes its intended metric, but substantial cross-metric improvement is observed, indicating that setwise diversity in one perceptual space typically transfers to others. Set-level objectives such as DPP and Vendi are found preferable by annotators, attributed to their invariance to singleton outlier samples.

Illustrative results show that, for a given prompt, optimization over DINO, DreamSim, or LPIPS features meaningfully shifts both content and layout, and is robust against collapse patterns that are otherwise prevalent.

Figure 5: Example SDXL-Turbo generations for the prompt "a photo of a teddy bear", comparing i.i.d sampling and various objective-driven optimization outputs.

Scaling and Efficiency

The optimization pipeline is efficient: for a modest batch size ($B=4$), fewer than 15 steps suffice to outperform group-inference approaches operating with an initial pool size of 64–128. Scaling to larger models and larger sets is feasible, with sequential generation further improving scalability and memory efficiency.

Noise Trajectory and Pink Noise Advantage

Noise trajectory analysis, spatial heatmaps, and frequency decomposition across optimization steps confirm that larger modifications are localized in regions corresponding to semantic shifts in the outputs, with maximal effect at low frequencies.

Figure 6: Spatial heatmaps (L2 norm change) indicate alignment between predominant modifications in the noise and evolving semantic content of generations.

Notably, pink noise initialization both increases attainable diversity and decreases the number of optimization iterations required, and this holds for all tested architectures and baseline strategies.

Failure Analysis

The work documents several failure modes for specific objectives: for example, optimization for LPIPS or color histogram diversity may produce sets with technically high diversity but limited semantic value (e.g., plain backgrounds, blurred features, missing objects in some samples). These cases can be mitigated by incorporating stricter prompt-alignment or set-level objectives with higher semantic correspondence.

Figure 7: Characteristic failure cases. Top: Blurred image maximizes distance, but not semantic diversity; Middle: Plain backgrounds inflate color diversity; Bottom: LPIPS fails to instill missing semantic content.

Implications and Future Directions

The findings firmly establish that inference-time, batchwise noise optimization is a general, architecture-agnostic technique for recovering collapsed modes in generative diffusion models, yielding marked improvements in practical diversity without architectural modification or retraining. The results clarify the mechanisms behind mode collapse recovery, revealing the fundamental importance of low-frequency noise structure in spanning broader solution manifolds.

Practical implications include:

• Post-hoc diversity enhancement for user-facing content generation without model finetuning
• Efficient candidate set construction for tasks requiring exploration or inspiration (design, creative work)
• Scalable batch and sequential pipelines for high-throughput generation scenarios
• Insights transferable to related generative regimes (video, 3D, conditional sequence)

Theoretical implications include:

• The necessity to revisit initialization strategies for future generative models
• The potent role of spectral analysis in model interpretability and diversity control
• A new axis for controlling the quality/diversity trade-off via initialization frequency content

Future developments may involve joint optimization with reward models for compositionality and control, integration with RLHF/PRM techniques, and exploration of learned, adaptive noise priors beyond simple pink/white filters.

Conclusion

The work presents an effective, theoretically informed framework for overcoming mode collapse in trained diffusion models via end-to-end, batch-level noise optimization, leveraging frequency-domain analysis for further gains. The method provides strong numerical, qualitative, and user-validated improvements in output diversity, and establishes a compelling direction for inference-time enhancement of generative models (2601.00090).