Diffusion Image Generation with Explicitly Modeling of Data Manifold Geometry

Image generative models aim to sample data points fromthe underlying data manifold, a task that requires learning and decod-ing a dense, low-dimensional, and compact parameterization space. Toachieve this, we propose the Data Manifold-aware Image diffusioN moDel(MIND), a novel framework that explicitly models manifold geometry byintegrating discrete patch tokenization into the score function of a con-tinuous diffusion model. This approach successfully leverages both thestructural quantification capabilities of discrete tokens and the parallelgeneration flexibility of continuous diffusion. Moreover, we enable end-to-end differentiable training via a novel soft top-k aggregation mechanismand introduce dual-branch high-frequency feature embedding layers toalleviate the spectral bias of transformer backbones on low-dimensionalinputs. Furthermore, for inference, we design a multi-stage transitionsampling scheme that dynamically adjusts the sampling scheme based ontimestep. Extensive experiments on ImageNet 256×256 demonstrate theeffectiveness of MIND. After 80-epoch training, our base model achievesan FID of 22.73 without guidance, nearly halving the 43.47 FID of thevanilla DiT-B/2 baseline. The proposed method reduces FID by 15.95and 9.06 on average compared with the baselines DiT and SiT, respec-tively. For image generation on ImageNet 256×256 with guidance, theproposed MIND-B with only 130M parameters achieves an FID of 2.06,surpassing LlamaGen-3B with 3.1B parameters. Our MIND-XL with715M parameters further reduces the FID to 1.95. Our MIND intro-duces a fresh perspective on diffusion-based image generation, pavingthe way for future research and innovation in this community. The codeis available here: https://github.com/xddgit/MIND.