Geometry-Aware Ornstein–Uhlenbeck Diffusion via Structured Drift Design

Diffusion-based generative models have achieved remarkable success in high-dimensional data generation; however, they fundamentally rely on isotropic diffusion processes that destroy meaningful geometric structures in the forward process. For complex, multimodal, and highly correlated distributions—such as biologically constrained genetic data—isotropic noise merges distinct modes and distorts intrinsic dependencies. This forces the reverse process to recover structure from heavily degraded signals, leading to slow convergence, mode averaging, and biologically implausible samples. To address this, we introduce the Geometry-aware Ornstein–Uhlenbeck (GOU) process, a structured drift design that embeds data geometry into the forward and backward dynamics. By employing a variance-aware anisotropic drift, GOU contracts low-variance directions rapidly while preserving high-variance directions longer, maintaining key multimodal structures as stable channels over time. We provide rigorous theoretical analysis of the proposed dynamics, including the evolution of moments in the forward process and bounds on the initialization error of the backward process when using a Gaussian mixture model with explicitly derived means and covariance matrix. Furthermore, we theoretically analyze the optimal finite stopping time that prevents mode merging during the forward process. Experiments on synthetic and real-world biological data demonstrate that GOU achieves superior structural preservation and sample quality across both statistical and biological metrics.