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

1. Probleemstelling

Moderne tekst-naar-beeld diffusiemodellen (bijv. SDXL-Turbo, Flux.1) lijden onder mode collapse. Bij het genereren van meerdere afbeeldingen vanuit dezelfde tekstprompt met verschillende willekeurige zaden (i.i.d. ruis) produceren de modellen vaak bijna identieke outputs met minimale visuele variatie.

• Huidige beperkingen: Bestaande oplossingen vertrouwen op:
  • Stuurmechanismen: Het sturen van het model via gewijzigde conditionering (bijv. Particle Guidance, DPP), wat vaak kwaliteit inwisselt voor diversiteit.
  • Zoekmethoden: Het genereren van een groot aantal kandidaten en deze te cullen (bijv. "Best-of-N"), wat rekenkundig duur en geheugenintensief is.
  • Prompt engineering: Het aanpassen van de tekstinput, wat niet eenvoudig alle variaties in de outputruimte kan vastleggen.
• Kernvraag: Kunnen we de initiële ruis ($x_0$) direct optimaliseren om te herstellen van mode collapse en diverse, hoogwaardige afbeeldingen te genereren zonder het vooraf getrainde model of de gebruikersprompt te wijzigen?

2. Methodologie

De auteurs stellen een end-to-end ruisoptimalisatiekader voor dat de initiële ruisvector behandelt als een leerbare parameter. Het proces omvat:

A. Doelfunctie voor ruisoptimalisatie

Gegeven een vaste prompt $c$ en een bevroren diffusiemodel $g_\theta$, optimaliseert de methode een batch van initiële ruisstalen $\{x_0^{(i)}\}_{i=1}^B$. De doelfunctie minimaliseert een verlies $\mathcal{L}$ dat bestaat uit drie termen:

1. Kwaliteitsbeloning ($r_s$): Zorgt ervoor dat individuele afbeeldingen voldoen aan de prompt en een hoge fideliteit behouden (bijv. met CLIPScore of HPSv2).
2. Diversiteitsmeting ($v_\mathcal{B}$): maximaliseert de ongelijkheid tussen afbeeldingen in de batch. Het paper onderzoekt diverse metingen:
   • Paarsgewijs: DINOv2 patch-features, DreamSim, LPIPS, Kleurenhistogrammen.
   • Op set-niveau: Determinantal Point Processes (DPP) en Vendi Score, die het effectieve aantal unieke items in een set meten.
3. Regularisatie ($\text{reg}$): Voorkomt dat de geoptimaliseerde ruis te ver afwijkt van de priorverdeling $N(0, I)$. Dit wordt bereikt door afwijkingen van de $\chi_d$-verdeling van de ruisnorm te straffen.

De optimalisatie wordt uitgevoerd via backpropagation door de bevroren sampler $g_\theta$.

B. Strategie voor ruisinitialisatie (Pink Noise)

Een belangrijke inzicht uit het paper is de analyse van ruis-evolutie. Tijdens de optimalisatie observeerden de auteurs dat de meest significante veranderingen optreden in de laagfrequente componenten van het ruisspectrum, terwijl hoge frequenties relatief stabiel blijven.

• Hypothese: Natuurlijke afbeeldingen volgen een $1/f$-powerspectrum (meer vermogen in lage frequenties). Standaard witte ruis is spectraal vlak.
• Oplossing: De auteurs introduceren Pink Noise-initialisatie. Ze transformeren witte ruis via FFT naar het frequentiedomein, passen een power-scaling factor toe van $f^{-\alpha/2}$ (waarbij $\alpha \approx 0.2$), en transformeren terug.
• Voordeel: Deze initialisatie brengt de startende ruis in overeenstemming met de statistische eigenschappen van natuurlijke afbeeldingen en de optimalisatiedynamiek, wat leidt tot snellere convergentie en hogere diversiteit.

C. Optimalisatiemodi

• Gebatchte optimalisatie: Optimaliseert gezamenlijk een set afbeeldingen (bijv. 4) om set-niveau diversiteit te maximaliseren.
• Sequentiële optimalisatie: Genereert afbeeldingen één voor één, waarbij elke afbeelding wordt geoptimaliseerd om te verschillen van eerder gegenereerde outputs. Dit vermindert het geheugentekort en maakt het mogelijk om grote, diverse sets te genereren.

3. Belangrijkste bijdragen

1. Ruisoptimalisatie voor diversiteit: Toont aan dat het direct optimaliseren van initiële ruis een superieure strategie is voor herstel van mode collapse in vergelijking met op geleiding gebaseerde of op zoek gebaseerde methoden.
2. Diversiteitsdoelen op set-niveau: Identificeert dat DPP en Vendi Score effectiever zijn dan eenvoudige paarsgewijze metingen voor het optimaliseren van diverse sets, omdat ze voorkomen dat de "gemiddelde paarsgewijze" meting wordt omzeild door één afbeelding zeer verschillend te maken terwijl de anderen gelijk blijven.
3. Pink Noise-initialisatie: Stelt een eenvoudige, parameterloze strategie voor (met pink noise) die consistent de diversiteit verbetert over verschillende modellen (SDXL-Turbo, Flux.1) en baselines, zelfs zonder verdere optimalisatie.
4. Trade-off tussen kwaliteit en diversiteit: Toont aan dat hun methode de diversiteit kan verhogen zonder in te leveren op beeldkwaliteit of prompt-gevolgzaamheid, waarmee de typische trade-off die wordt gezien bij op geleiding gebaseerde methoden wordt overwonnen.

4. Experimentele resultaten

De methode werd geëvalueerd op GenEval, DPG-Bench (complexe prompts) en T2I-CompBench met modellen zoals SDXL-Turbo en Flux.1 [schnell].

• Kwantitatieve prestaties:
  • Diversiteit: Presteert beter dan i.i.d. sampling en de state-of-the-art zoekmethode [Parmar et al., 2025] over alle metingen (DINO, DreamSim, LPIPS, Vendi Score).
  • Kwaliteit: Behoudt of verbetert lichtjes de beeldkwaliteitsscores (HPSv2, CLIPScore) in vergelijking met baselines.
  • Efficiëntie: Met pink noise-initialisatie bereikt de methode hogere diversiteitsscores dan de baselines zoekmethode in aanzienlijk minder iteraties (bijv. 6–8 iteraties versus 9–12 voor witte ruis).
• Kwalitatieve resultaten:
  • Genereert diverse variaties in objectpose, kleur, stijl en achtergrond terwijl strikt wordt voldaan aan complexe prompts.
  • Hanteert succesvol complexe prompts (bijv. "Een schilderachtig herfsttafereel...") waarbij i.i.d. sampling faalt om variatie te produceren.
• Gebruikersstudie: Menselijke beoordelaars gaven de outputs van de voorgestelde methode (vooral die geoptimaliseerd met DPP/Vendi) de voorkeur boven i.i.d. sampling en de zoekbaseline in 85–91% van de gevallen.

5. Betekenis

Dit werk verschuift het paradigma van diversiteitsgeneratie in diffusiemodellen van samplingstrategieën (wachten op een goed zaad) of geleiding (aanpassen van het generatiepad) naar inputoptimalisatie.

• Algemene toepasbaarheid: De aanpak is modelonafhankelijk en werkt op zowel stap-gedistilleerde (snelle) als standaard diffusiemodellen.
• Efficiëntie: Door de ruis direct te optimaliseren, wordt de enorme geheugenoverhead van het genereren en opslaan van honderden kandidaten voor culling vermeden.
• Inzicht in de latente ruimte: De ontdekking dat laagfrequente ruiscomponenten diversiteit aandrijven, biedt een nieuw begrip van hoe diffusiemodellen initiële ruis interpreteren, wat suggereert dat "structuur" vroeg in het frequentiedomein wordt vastgesteld.

Kortom, het paper presenteert een robuuste, efficiënte en hoogwaardige oplossing voor het probleem van mode collapse, en bewijst dat het "nooit te laat" is om de inputruis te optimaliseren om het volledige diversiteitspotentieel van getrainde diffusiemodellen te ontsluiten.