DIAMOND: Directed Inference for Artifact Mitigation in Flow Matching Models

Sintesi Tecnica: DIAMOND – Inferenza Diretta per la Mitigazione degli Artefatti nei Modelli di Flow Matching

Problematica
Nonostante l'alta fedeltà dei recenti modelli di text-to-image, in particolare quelli basati su architetture Rectified Flow (es. FLUX.2) e Diffusion Models (es. SDXL), le immagini generate presentano frequentemente artefatti strutturali e anatomici. Questi includono mani distorte, arti extra, volti malformati e testo errato. Le attuali strategie di mitigazione soffrono di limitazioni significative:

• Raffinamento post-hoc: I metodi che correggono le immagini dopo la generazione spesso non riescono a intervenire durante il processo di formazione principale.
• Modifiche invasive: Le tecniche che richiedono il fine-tuning del modello o pesi LoRA sono computazionalmente costose, difficili da scalare e spesso mancano di pesi pubblicamente disponibili per le moderne architetture.
• Limitazioni euristiche: I metodi di rilevamento in tempo reale spesso si affidano a soglie euristiche o sostituzioni casuali dei latent, portando a una localizzazione imprecisa o a correzioni sub-ottimali.

Metodologia
Gli autori propongono DIAMOND (Directed Inference for Artifact Mitigation in Flow Matching Models), un framework zero-shot, privo di addestramento, che corregge gli artefatti durante la traiettoria di inferenza senza modificare i pesi del modello. Il metodo opera sul principio della ricostruzione di una stima del campione pulito ad ogni timestep per guidare la generazione lontano dagli stati latenti inclini agli artefatti.

1. Stima del Latent Pulito:
   A differenza degli approcci standard che operano direttamente sui latent rumorosi ($x_t$), DIAMOND stima il latent dei dati puliti ($\hat{x}_{0,t}$) ad ogni step.

   • Per i modelli Rectified Flow (che predicono la velocità $v_\theta$): $\hat{x}_{0,t} = x_t - t \cdot v_\theta(x_t, t)$.
   • Per i modelli Diffusion (che predicono il rumore $\epsilon_\theta$): $\hat{x}_{0,t} = x_t - \sigma_t \cdot \epsilon_\theta(x_t, t)$.
     Questo latent stimato viene decodificato nello spazio dei pixel ($\hat{I}_t = D(\hat{x}_{0,t})$) per garantire che il rilevatore di artefatti operi nel suo dominio ottimale (dati puliti) piuttosto che su stati intermedi rumorosi.

2. Rilevamento degli Artefatti e Loss:
   Un Artifact Detector (AD) differenziabile e pre-addestrato (es. da DiffDoctor) analizza $\hat{I}_t$ per generare una maschera di probabilità dell'artefatto $M_t$. Una loss di artefatto ($L_a$) viene calcolata come l'errore quadratico medio tra questa maschera e una maschera nulla (che rappresenta un'immagine priva di artefatti):
   $$L_a = \frac{1}{HW} \sum_{i,j} (M_{i,j}^t)^2$$

3. Correzione della Traiettoria Basata su Gradiente:
   Il metodo calcola il gradiente della loss rispetto al latent corrente $x_t$ ($\nabla_{x_t} L_a$). Per garantire la stabilità e impedire che l'entità della correzione sia dominata dalla dimensione della regione dell'artefatto, il gradiente viene normalizzato:
   $$\delta_t = \lambda_t \frac{\nabla_{x_t} L_a}{\|\nabla_{x_t} L_a\|^2 + \epsilon}$$
   Questo vettore di spostamento $\delta_t$ viene sottratto dallo step di aggiornamento del solver standard, deviando efficacemente la traiettoria lontano dagli artefatti.

4. Scheduling Dinamico ($\lambda_t$):
   Uno schedule a legge di potenza controlla l'intensità della correzione ($\lambda_t$). Sulla base dell'osservazione che gli artefatti strutturali appaiono nelle fasi iniziali, dominate dal rumore, mentre i dettagli ad alta frequenza vengono raffinati successivamente, lo schedule applica una correzione forte all'inizio del processo di generazione e la fa decadere verso la fine.

5. Preservazione dell'Identità (Opzionale):
   Per evitare che la correzione alteri la struttura globale dell'immagine o l'allineamento con il prompt, può essere applicato un termine di regolarizzazione opzionale ($L_{rec}$). Questo termine minimizza l'Errore Assoluto Medio (MAE) tra l'immagine corretta e l'immagine di base nelle regioni non affette da artefatti.

Contributi Chiave

• Framework Training-Free: DIAMOND è il primo metodo zero-shot che integra la guida geometrica della traiettoria direttamente nel processo di inferenza di modelli Rectified Flow e Diffusion senza richiedere addestramento aggiuntivo o modifiche ai pesi.
• Stima On-the-Fly: Il metodo introduce un nuovo meccanismo per stimare i latent puliti ($\hat{x}_{0,t}$) durante l'inferenza, consentendo un rilevamento affidabile degli artefatti sugli step intermedi dove i rilevatori standard tipicamente falliscono.
• Generalizzazione: L'approccio dimostra di funzionare efficacemente sia su modelli Rectified Flow (FLUX.1, FLUX.2) che su modelli Diffusion standard (SDXL).
• Normalizzazione del Gradiente: Il paper identifica e implementa la normalizzazione del gradiente per disaccoppiare la forza dell'aggiornamento dalla dimensione spaziale dell'artefatto, garantendo una correzione coerente indipendentemente dall'estensione dell'artefatto.

Risultati Sperimentali
Valutazioni quantitative sono state condotte su tre dataset (animali, parole, persone) utilizzando modelli come FLUX.1 [dev], FLUX.1 [schnell], FLUX.2 [dev] e SDXL.

• Riduzione degli Artefatti: Sul dataset "words" con FLUX.1 [dev], DIAMOND ha ridotto la Frequenza Media degli Artefatti da un baseline del 100% al 9,85%, superando significativamente i baseline come DiffDoctor (46,97%) e HPSv2 (56,82%).
• Rapporto di Pixel: L'Artifact Pixel Ratio (APR) è stato ridotto a livelli quasi nulli (es. 0,009% sul dataset delle parole) rispetto ai baseline.
• Preservazione della Qualità: Il metodo ha mantenuto alti punteggi CLIP-T (allineamento testo-immagine) e ImageReward (preferenza umana), indicando che la rimozione degli artefatti non ha degradato la coerenza semantica o la qualità visiva.
• Preservazione dell'Identità: Gli studi di ablazione hanno mostrato che, mentre l'unregulated DIAMOND rimuove efficacemente gli artefatti, l'aggiunta del termine $L_{rec}$ (con $\alpha=0,1$) bilancia in modo ottimale la soppressione degli artefatti con la preservazione dell'identità del modello base.

Significatività e Claim
Gli autori affermano che DIAMOD fornisce una via robusta e zero-shot verso una sintesi di immagini ad alta fedeltà e priva di artefatti. Intervenendo direttamente durante il processo di campionamento attraverso la correzione della traiettoria, il metodo supera i costi computazionali del fine-tuning e le imprecisioni delle euristiche post-hoc. Il paper sottolinea che questo approccio consente una supervisione spaziale fine e una correzione precisa di artefatti localizzati (come mani o testo) mantenendo la struttura globale e l'aderenza al prompt dell'immagine generata. Il metodo è presentato come una soluzione pratica per applicazioni professionali dove la correttezza anatomica è critica, come l'e-commerce e la grafica pubblicitaria.