Graph Flow Matching: Enhancing Image Generation with Neighbor-Aware Flow Fields

TL;DR

Graph Flow Matching (GFM) improves image generation by incorporating neighbor-aware information into flow models using graph neural modules, enhancing quality and diversity without significant computational overhead.

Contribution

GFM introduces a reaction-diffusion approach that integrates local neighbor information into flow matching networks, enhancing image generation quality and diversity.

Findings

Consistently improves FID and recall across five benchmarks.

Operates efficiently within the latent space of pretrained VAEs.

Enhances existing flow matching models with minimal additional cost.

Abstract

Flow matching casts sample generation as learning a continuous-time velocity field that transports noise to data. Existing flow matching networks typically predict each point's velocity independently, considering only its location and time along its flow trajectory, and ignoring neighboring points. However, this pointwise approach may overlook correlations between points along the generation trajectory that could enhance velocity predictions, thereby improving downstream generation quality. To address this, we propose Graph Flow Matching (GFM), a lightweight enhancement that decomposes the learned velocity into a reaction term -- any standard flow matching network -- and a diffusion term that aggregates neighbor information via a graph neural module. This reaction-diffusion formulation retains the scalability of deep flow models while enriching velocity predictions with local context, all at minimal additional computational cost. Operating in the latent space of a pretrained variational autoencoder, GFM consistently improves Fr\'echet Inception Distance (FID) and recall across five image generation benchmarks (LSUN Church, LSUN Bedroom, FFHQ, AFHQ-Cat, and CelebA-HQ at $256\times256$), demonstrating its effectiveness as a modular enhancement to existing flow matching architectures.