DIAMOND: Directed Inference for Artifact Mitigation in Flow Matching Models

Technical Summary: DIAMOND – Directed Inference for Artifact Mitigation in Flow Matching Models

Problem Statement
Despite the high fidelity of recent text-to-image models, particularly those based on Rectified Flow architectures (e.g., FLUX.2) and Diffusion Models (e.g., SDXL), generated images frequently exhibit structural and anatomical artifacts. These include distorted hands, extra limbs, malformed faces, and incorrect text. Existing mitigation strategies suffer from significant limitations:

• Post-hoc refinement: Methods that correct images after generation often fail to intervene during the core formation process.
• Invasive modifications: Techniques requiring model fine-tuning or LoRA weights are computationally expensive, difficult to scale, and often lack publicly available weights for modern architectures.
• Heuristic limitations: Real-time detection methods often rely on heuristic thresholds or random latent replacements, leading to inaccurate localization or sub-optimal corrections.

Methodology
The authors propose DIAMOND (Directed Inference for Artifact Mitigation in Flow Matching Models), a training-free, zero-shot framework that corrects artifacts during the inference trajectory without modifying model weights. The method operates on the principle of reconstructing an estimate of the clean sample at every timestep to guide the generation away from artifact-prone latent states.

1. Clean Latent Estimation:
   Unlike standard approaches that operate directly on noisy latents ($x_t$), DIAMOND estimates the clean data latent ($\hat{x}_{0,t}$) at each step.

   • For Rectified Flow models (predicting velocity $v_\theta$): $\hat{x}_{0,t} = x_t - t \cdot v_\theta(x_t, t)$.
   • For Diffusion models (predicting noise $\epsilon_\theta$): $\hat{x}_{0,t} = x_t - \sigma_t \cdot \epsilon_\theta(x_t, t)$.
     This estimated latent is decoded to pixel space ($\hat{I}_t = D(\hat{x}_{0,t})$) to ensure the artifact detector operates in its optimal domain (clean data) rather than on intermediate noisy states.

2. Artifact Detection and Loss:
   A differentiable, pre-trained Artifact Detector (AD) (e.g., from DiffDoctor) analyzes $\hat{I}_t$ to generate an artifact probability mask $M_t$. An artifact loss ($L_a$) is computed as the mean squared error between this mask and a zero-mask (representing an artifact-free image):
   $$L_a = \frac{1}{HW} \sum_{i,j} (M_{i,j}^t)^2$$

3. Gradient-Based Trajectory Correction:
   The method computes the gradient of the loss with respect to the current latent $x_t$ ($\nabla_{x_t} L_a$). To ensure stability and prevent the correction magnitude from being dominated by the size of the artifact region, the gradient is normalized:
   $$\delta_t = \lambda_t \frac{\nabla_{x_t} L_a}{\|\nabla_{x_t} L_a\|^2 + \epsilon}$$
   This displacement vector $\delta_t$ is subtracted from the standard solver update step, effectively steering the trajectory away from artifacts.

4. Dynamic Scheduling ($\lambda_t$):
   A power-law schedule controls the intensity of the correction ($\lambda_t$). Based on the observation that structural artifacts appear in early, noise-dominated stages while high-frequency details are refined later, the schedule applies strong correction early in the generation process and decays it towards the end.

5. Identity Preservation (Optional):
   To prevent the correction from altering the global image structure or prompt alignment, an optional regularization term ($L_{rec}$) can be applied. This term minimizes the Mean Absolute Error (MAE) between the corrected image and the baseline image in non-artifact regions.

Key Contributions

• Training-Free Framework: DIAMOND is the first zero-shot method to integrate geometric trajectory guidance directly into the inference process of Rectified Flow and Diffusion models without requiring additional training or weight modifications.
• On-the-Fly Estimation: The method introduces a novel mechanism to estimate clean latents ($\hat{x}_{0,t}$) during inference, enabling reliable artifact detection on intermediate steps where standard detectors typically fail.
• Generalization: The approach is demonstrated to work effectively across both Rectified Flow models (FLUX.1, FLUX.2) and standard Diffusion models (SDXL).
• Gradient Normalization: The paper identifies and implements gradient normalization to decouple the update strength from the artifact's spatial size, ensuring consistent correction regardless of the artifact's extent.

Experimental Results
Quantitative evaluations were conducted on three datasets (animals, words, people) using models like FLUX.1 [dev], FLUX.1 [schnell], FLUX.2 [dev], and SDXL.

• Artifact Reduction: On the "words" dataset with FLUX.1 [dev], DIAMOND reduced the Mean Artifact Frequency from a 100% baseline to 9.85%, significantly outperforming baselines like DiffDoctor (46.97%) and HPSv2 (56.82%).
• Pixel Ratio: The Artifact Pixel Ratio (APR) was reduced to near-zero levels (e.g., 0.009% on the words dataset) compared to baselines.
• Quality Preservation: The method maintained high CLIP-T scores (text-image alignment) and ImageReward scores (human preference), indicating that artifact removal did not degrade semantic consistency or visual quality.
• Identity Preservation: Ablation studies showed that while unregularized DIAMOND effectively removes artifacts, adding the $L_{rec}$ term (with $\alpha=0.1$) optimally balances artifact suppression with the preservation of the base model's identity.

Significance and Claims
The authors claim that DIAMOND provides a robust, zero-shot path to high-fidelity, artifact-free image synthesis. By intervening directly during the sampling process through trajectory correction, the method overcomes the computational costs of fine-tuning and the inaccuracies of post-hoc heuristics. The paper emphasizes that this approach allows for fine-grained spatial supervision and precise correction of localized artifacts (such as hands or text) while maintaining the global structure and prompt adherence of the generated image. The method is presented as a practical solution for professional applications where anatomical correctness is critical, such as e-commerce and marketing visuals.