DIAMOND: Directed Inference for Artifact Mitigation in Flow Matching Models

Technische Samenvatting: DIAMOND – Directed Inference voor Artefact-mitigatie in Flow Matching Modellen

Probleemstelling
Ondanks de hoge getrouwheid van recente tekst-naar-afbeelding modellen, met name die gebaseerd op Rectified Flow-architecturen (bijv. FLUX.2) en Diffusiemodellen (bijv. SDXL), vertonen gegenereerde afbeeldingen frequent structurele en anatomische artefacten. Deze omvatten vervormde handen, extra ledematen, misvormde gezichten en incorrecte tekst. Bestaande mitigatiestrategieën lijden aan significante beperkingen:

• Post-hoc verfijning: Methoden die afbeeldingen na de generatie corrigeren, falen vaak om tijdens het kernvormingsproces in te grijpen.
• Invasieve modificaties: Technieken die model fine-tuning of LoRA-gewichten vereisen, zijn computationeel duur, moeilijk schaalbaar en missen vaak publiek beschikbare gewichten voor moderne architecturen.
• Heuristische beperkingen: Real-time detectiemethoden vertrouwen vaak op heuristische drempelwaarden of willekeurige latent-vervangingen, wat leidt tot onnauwkeurige lokalisatie of suboptimale correcties.

Methodologie
De auteurs stellen DIAMOND voor (Directed Inference for Artifact Mitigation in Flow Matching Models), een training-vrij, zero-shot framework dat artefacten corrigeert tijdens de inferentie-traject zonder de modelgewichten te wijzigen. De methode werkt volgens het principe van het reconstrueren van een schatting van de schone steekproef bij elke tijdstap om de generatie weg te leiden van artefactgevoelige latent-toestanden.

1. Schatting van de Schone Latent:
   In tegenstelling tot standaard benaderingen die direct op ruisige latents ($x_t$) opereren, schat DIAMOND de schone data-latent ($\hat{x}_{0,t}$) bij elke stap.

   • Voor Rectified Flow modellen (die snelheid $v_\theta$ voorspellen): $\hat{x}_{0,t} = x_t - t \cdot v_\theta(x_t, t)$.
   • Voor Diffusie modellen (die ruis $\epsilon_\theta$ voorspellen): $\hat{x}_{0,t} = x_t - \sigma_t \cdot \epsilon_\theta(x_t, t)$.
     Deze geschatte latent wordt gedecodeerd naar de pixelruimte ($\hat{I}_t = D(\hat{x}_{0,t})$) om ervoor te zorgen dat de artefact-detector opereert in zijn optimale domein (schone data) in plaats van op tussenliggende ruisige toestanden.

2. Artefact Detectie en Loss:
   Een differentieerbare, vooraf getrainde Artefact Detector (AD) (bijv. uit DiffDoctor) analyseert $\hat{I}_t$ om een artefact-waarschijnlijkheidsmasker $M_t$ te generen. Een artefact-loss ($L_a$) wordt berekend als de gemiddelde kwadratische fout tussen dit masker en een nul-masker (dat een artefact-vrije afbeelding representeert):
   $$L_a = \frac{1}{HW} \sum_{i,j} (M_{i,j}^t)^2$$

3. Gradiënt-gebaseerde Trajectcorrectie:
   De methode berekent de gradiënt van de loss met betrekking tot de huidige latent $x_t$ ($\nabla_{x_t} L_a$). Om stabiliteit te garanderen en te voorkomen dat de correctie-amplitude wordt gedomineerd door de grootte van het artefactgebied, wordt de gradiënt genormaliseerd:
   $$\delta_t = \lambda_t \frac{\nabla_{x_t} L_a}{\|\nabla_{x_t} L_a\|^2 + \epsilon}$$
   Deze verplaatsingsvector $\delta_t$ wordt afgetrokken van de standaard solver update-stap, wat het traject effectief wegstuurt van artefacten.

4. Dynamische Scheduling ($\lambda_t$):
   Een power-law schema controleert de intensiteit van de correctie ($\lambda_t$). Gebaseerd op de observatie dat structurele artefacten verschijnen in de vroege, door ruis gedomineerde fasen terwijl hoogfrequente details later worden verfijnd, past het schema een sterke correctie toe in het vroege stadium van het generatieproces en neemt deze af richting het einde.

5. Identiteitsbehoud (Optioneel):
   Om te voorkomen dat de correctie de globale beeldstructuur of prompt-alignement verandert, kan een optionele regularisatieterm ($L_{rec}$) worden toegepast. Deze term minimaliseert de Mean Absolute Error (MAE) tussen de gecorrigeerde afbeelding en de baseline-afbeelding in niet-artefact regio's.

Belangrijkste Bijdragen

• Training-vrij Framework: DIAMOND is de eerste zero-shot methode die geometrische traject-sturing direct integreert in het inferentieproces van Rectified Flow en Diffusiemodellen zonder aanvullende training of gewichtsmodificaties te vereisen.
• On-the-Fly Schatting: De methode introduceert een nieuw mechanisme om schone latents ($\hat{x}_{0,t}$) te schatten tijdens inferentie, wat betrouwbare artefact-detectie mogelijk maakt op tussenliggende stappen waar standaard detectoren doorgaans falen.
• Generalisatie: Het is aangetoond dat de aanpak effectief werkt over zowel Rectified Flow modellen (FLUX.1, FLUX.2) als standaard Diffusiemodellen (SDXL).
• Gradiënt Normalisatie: Het paper identificeert en implementeert gradiënt-normalisatie om de update-sterkte te ontkoppelen van de ruimtelijke grootte van het artefact, wat een consistente correctie garandeert ongeacht de omvang van het artefact.

Experimentele Resultaten
Kwantitatieve evaluaties werden uitgevoerd op drie datasets (dieren, woorden, mensen) met modellen zoals FLUX.1 [dev], FLUX.1 [schnell], FLUX.2 [dev] en SDXL.

• Reductie van Artefacten: Op de "woorden" dataset met FLUX.1 [dev] verminderde DIAMOND de Mean Artifact Frequency van een 100% baseline naar 9,85%, waarmee het significant beter presteert dan baselines zoals DiffDoctor (46,97%) en HPSv2 (56,82%).
• Pixel Ratio: De Artifact Pixel Ratio (APR) werd teruggebracht tot bijna-nul niveaus (bijv. 0,009% op de woorden dataset) vergeleken met baselines.
• Kwaliteitsbehoud: De methode behield hoge CLIP-T scores (tekst-beeld alignement) en ImageReward scores (menselijke voorkeur), wat aangeeft dat de verwijdering van artefacten de semantische consistentie of visuele kwaliteit niet heeft verslechterd.
• Identiteitsbehoud: Ablatie-studies toonden aan dat terwijl onge-regulariseerde DIAMOND effectief artefacten verwijdert, het toevoegen van de $L_{rec}$ term (met $\alpha=0,1$) optimaal balanceert tussen artefact-onderdrukking en het behoud van de identiteit van het basismodel.

Betekenis en Claims
De auteurs claimen dat DIAMOND een robuuste, zero-shot weg biedt naar hoogwaardige, artefact-vrije beeld-synthese. Door direct tijdens het samplingsproces in te grijpen via trajectcorrectie, overwint de methode de computationele kosten van fine-tuning en de onnauwkeurigheden van post-hoc heuristieken. Het paper benadrukt dat deze aanpak fijnmazige ruimtelijke supervisie en precieze correctie van gelokaliseerde artefacten (zoals handen of tekst) mogelijk maakt, terwijl de globale structuur en de prompt-getrouwheid van de gegenereerde afbeelding behouden blijven. De methode wordt gepresenteerd als een praktische oplossing voor professionele toepassingen waar anatomische correctheid cruciaal is, zoals in e-commerce en marketing visuals.