Abstract

Current flow matching models learn to transport the source i.i.d. Gaussian noise into the target distribution of natural images, yet this source distribution carries no notion of spatial structure. Images however are fundamentally local since nearby pixels are strongly correlated. By sampling the noise independently, we hypothesize that models are implicitly encouraged to exploit less noisy neighbors as context during training, partially bypassing the need to properly learn the true local structure of images. The source distribution, in other words, works against the inductive bias of the image domain. To ameliorate this design discrepancy, we propose StructFlow which encodes spatial locality directly into the source by having the pixels within a small region share a common noise component. This structured source produces transport paths that are geometrically aligned with image regions — enabling properties that generic flow matching struggles to provide: fine-grained local editing that naturally respects boundaries, robust structure preservation, and smooth semantic interpolation between images. We show that these benefits also extend to large pre-trained models, demonstrating that StructFlow can even be incorporated through a lightweight post-training phase. Comprehensive experiments on multiple datasets, in unconditional, class and text-conditioned regimes, using different diffusion transformer architectures confirm that StructFlow not only offers competitive image generation quality, but also significantly improves localized controllable re-synthesis.

Method

Standard flow matching samples its source noise i.i.d. per pixel, ignoring the fact that nearby pixels in natural images are strongly correlated — forcing the model to reconstruct this local structure entirely on its own. StructFlow instead tries to enforce this locality into the initial noise distribution. During training, each image is first partitioned into local regions using SLIC superpixels; every region is assigned a shared "anchor" noise vector, and each pixel then samples its own noise around its region's anchor, so nearby pixels are correlated while different regions stay independent. No other part of the flow matching objective or inference procedure is modified. Because strongly correlated noise can destabilize training, we introduce three techniques to make it practical: cosine-induced spatial anchoring, which replaces random per-segment anchors with a smooth, deterministic cosine field to reduce stochastic variance; progressive coherence annealing, which starts training near a standard Gaussian source and gradually tightens the locality constraint — this also lets a pretrained flow matching model be turned into a StructFlow model via lightweight fine-tuning; and mixed-coherence training, which varies the coherence strength during training so a single model spans the full quality-editability trade-off at inference. Algorithms 1 and 2 below summarize the noise sampling and training procedure.


Results

Qualitative Results on Different DiT Models. StructFlow produces high-quality and diverse samples across unconditional, class-conditional, and text-to-image generation settings, on both a DiT-XL model trained from scratch and a SANA model post-trained with StructFlow.

Structure Preservance Generation and Part Resampling. StructFlow enables fine-grained local edits by resampling noise only in selected regions while keeping the rest fixed (left), and generates diverse samples that preserve the overall structure when conditioned on the same superpixel mask (right).

BibTeX

@article{zarei2026spatially,
  title={Spatially-Grounded Flow Matching: Structured Source Distributions for Image Generation},
  author={Zarei, Arman and Kalayeh, Mahdi M},
  journal={arXiv preprint arXiv:2608.15452},
  year={2026}
}