Diffusion Image Generation with Explicit Modeling of Data Manifold Geometry

2026-05-25

AI Summary

Existing image generation models struggle to accurately capture the low-dimensional, compact structure of data manifolds. To address this limitation, this work proposes MIND—a novel framework that uniquely integrates discrete image patch tokenization with continuous diffusion modeling to explicitly represent manifold geometry. MIND preserves a Transformer backbone while enhancing generation quality through a differentiable soft top-k aggregation mechanism, a dual-branch high-frequency feature embedding module, and a multi-stage dynamic sampling strategy. Experimental results demonstrate that MIND-B achieves an unprecedented unguided FID of 22.73 on ImageNet 256×256, nearly halving the score of DiT-B/2; under classifier guidance, it attains an FID of 2.06, outperforming LlamaGen-3B despite using 24 times fewer parameters. Moreover, MIND-XL further reduces the guided FID to 1.95.

Abstract

Image generative models aim to sample data points from the underlying data manifold, a task that requires learning and decoding a dense, low-dimensional, and compact parameterization space. To achieve this, we propose the Data Manifold-aware Image diffusioN moDel (MIND), a novel framework that explicitly models manifold geometry by integrating discrete patch tokenization into the score function of a continuous diffusion model. This approach successfully leverages both the structural quantification capabilities of discrete tokens and the parallel generation flexibility of continuous diffusion. Moreover, we enable end-to-end differentiable training via a novel soft top-$k$ aggregation mechanism and introduce dual-branch high-frequency feature embedding layers to alleviate the spectral bias of transformer backbones on low-dimensional inputs. Furthermore, for inference, we design a multi-stage transition sampling scheme that dynamically adjusts the sampling scheme based on timestep. Extensive experiments on ImageNet 256$\times$256 demonstrate the effectiveness of MIND. After 80-epoch training, our base model achieves an FID of 22.73 without guidance, nearly halving the 43.47 FID of the vanilla DiT-B/2 baseline. The proposed method reduces FID by 15.95 and 9.06 on average compared with the baselines DiT and SiT, respectively. For image generation on ImageNet-256$\times$256 with guidance, the proposed MIND-B with only 130M parameters achieves an FID of 2.06, superpassing the LlamaGen-3B with 3.1B parameters. The proposed MIND-XL with 715M parameters further reduces the FID to 1.95. Our MIND introduces a fresh perspective on diffusion-based image generation, paving the way for future research and innovation in this community. The code will be publicly available.

Problem

• data manifold
• diffusion image generation
• manifold geometry
• score function
• discrete-continuous modeling

Innovation

• manifold geometry
• diffusion model
• discrete-continuous hybrid
• soft top-k aggregation
• multi-stage sampling