DIAMOND: Directed Inference for Artifact Mitigation in Flow Matching Models

Alicja Polowczyk1*, Agnieszka Polowczyk1*, Piotr Borycki2, Joanna Waczyńska2, Jacek Tabor2, Przemysław Spurek2,3
1Silesian University of Technology 2Jagiellonian University 3IDEAS Research Institute
*Indicates Equal Contribution
Paper Code
DIAMOND teaser visualization
We propose DIAMOND, an inference-time trajectory correction mechanism to mitigate artifacts in Rectified Flow and Diffusion Models. First, we generate an image D(x1) from the initial probability distribution sample, which then undergoes the generative process. The Base Trajectory (red) leads to an image containing artifacts, whereas our corrected DIAMOND Trajectory (green) results in the artifact-free image. At timestep t, we estimate the final image x̂0,t (gray dashed line) and use an Artifact Detector to apply a gradient-based trajectory correction (purple arrows), shifting it away from artifact region.

Artifact Detector Activation Across Trajectory Steps

Detection is performed only on the estimated x̂0,t; it is not applied to xt.

x̂0,t overlay Drag down to reveal the overlay
t0 x0 step t0 x0 overlay step t0
t2 x0 step t2 x0 overlay step t2
t4 x0 step t4 x0 overlay step t4
t6 x0 step t6 x0 overlay step t6
t8 x0 step t8 x0 overlay step t8
xt overlay Drag down to reveal the overlay
t0 xt step t0 xt overlay step t0
t2 xt step t2 xt overlay step t2
t4 xt step t4 xt overlay step t4
t6 xt step t6 xt overlay step t6
t8 xt step t8 xt overlay step t8

Qualitative Results

Baseline DIAMOND
FLUX.2 [dev]
Baseline result 14
DIAMOND result 15
FLUX1.[dev]
Baseline result 16
DIAMOND result 17
FLUX1. [schnell]
Baseline result 18
DIAMOND result 19
SDXL
Baseline result 20
DIAMOND result 21

Additional Results Across Three Datasets

Each example smoothly transitions between Baseline and DIAMOND.

People example 1 baseline Baseline DIAMOND
People example 2 baseline Baseline DIAMOND
People example 3 baseline Baseline DIAMOND
People example 4 baseline Baseline DIAMOND

Overview of DIAMOND

DIAMOND pipeline overview
Inference-Time Pipeline

Our technique employs an inference-time pipeline designed to mitigate artifacts without additional training.

Artifact suppression is achieved by correcting the trajectory using gradients derived from a pixel-wise segmentation loss. During inference, the latent representation is iteratively updated using controlled shifts along the trajectory, enabling progressive correction of artifact-prone regions.

Abstract

Despite impressive results from recent text-to-image models like FLUX, visual and anatomical artifacts remain a significant hurdle for practical and professional use. Existing methods for artifact reduction, typically work in a post-hoc manner, consequently failing to intervene effectively during the core image formation process. Notably, current techniques require problematic and invasive modifications to the model weights, or depend on a computationally expensive and time-consuming process of regional refinement. To address these limitations, we propose DIAMOND, a training-free method that applies trajectory correction to mitigate artifacts during inference. By reconstructing an estimate of the clean sample at every step of the generative trajectory, DIAMOND actively steers the generation process away from latent states that lead to artifacts. Furthermore, we extend the proposed method to standard Diffusion Models, demonstrating that DIAMOND provides a robust, zero-shot path to high-fidelity, artifact-free image synthesis without the need for additional training or weight modifications in modern generative architectures.

BibTeX

@misc{polowczyk2026diamonddirectedinferenceartifact,
  title={DIAMOND: Directed Inference for Artifact Mitigation in Flow Matching Models},
  author={Alicja Polowczyk and Agnieszka Polowczyk and Piotr Borycki and Joanna Waczyńska and Jacek Tabor and Przemysław Spurek},
  year={2026},
  eprint={2602.00883},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2602.00883}
}