通过显式建模数据流形几何的扩散图像生成
研究人员提出了MIND（数据流形感知图像扩散模型），通过将离散补丁标记化集成到连续扩散模型的得分函数中，显式建模流形几何。该方法在ImageNet 256×256上实现了FID 2.06（MIND-B，1.3亿参数），优于LlamaGen-3B（31亿参数）的FID。
图像生成模型的核心目标是能够从底层数据流形中采样高质量的数据点。然而，这一任务要求模型能够学习并解码一个密集、低维且紧凑的参数化空间，传统方法往往难以兼顾结构表达与生成效率。近日，来自研究团队的一篇新论文提出了数据流形感知图像扩散模型（MIND），为这一问题提供了全新的解决思路。
MIND是一个新颖的框架，其核心创新在于将离散补丁标记化技术无缝集成到连续扩散模型的得分函数中。通过这种方式，MIND能够显式地对数据流形的几何结构进行建模，充分利用离散标记在结构量化方面的优势，同时保持连续扩散模型并行生成的灵活性。这种结合使得模型在生成过程中能够更准确地捕捉数据的底层结构。
为了实现端到端的可微训练，研究团队引入了一种新颖的软top-k聚合机制。该机制可以高效地处理离散标记的选择问题，同时保证整个训练过程的梯度可以顺利回传。此外，为了应对Transformer主干网络在低维输入上表现出的频谱偏差，团队设计了双分支高频特征嵌入层。这一设计通过增强模型对高频信号的感知能力，有效缓解了频谱偏差带来的性能损失。
在推理阶段，MIND采用了一种多阶段过渡采样方案。该方案能够根据当前的时间步动态调整采样策略，从而在生成质量和效率之间取得更好的平衡。
实验部分，研究团队在ImageNet 256×256数据集上进行了广泛的评估。经过80个周期的训练，基础模型在无引导条件下实现了22.73的FID，这一成绩几乎是原始DiT-B/2基线（43.47 FID）的一半。与DiT和SiT基线相比，MIND在平均FID上分别降低了15.95和9.06，提升幅度显著。在有引导的图像生成任务中，MIND展现出了更强的竞争力。仅含1.3亿参数的MIND-B模型实现了2.06的FID，这一结果甚至超越了拥有31亿参数的LlamaGen-3B模型。进一步地，拥有7.15亿参数的MIND-XL模型将FID降低至1.95，达到了新的高度。
MIND为基于扩散的图像生成领域带来了全新的视角，其明确的几何建模思路有望推动该领域的进一步发展。研究团队表示，相关代码将公开提供，以促进学术界的交流与创新。