Diffusion Image Generation with Explicitly Modeling of Data Manifold Geometry¶

会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/xddgit/MIND
领域: 图像生成
关键词: 图像生成、扩散模型、数据流形、图像离散化、超球面几何

一句话总结¶

针对连续扩散模型在无界欧氏空间中扩散而忽略数据流形几何先验的问题，本文提出了流形感知图像扩散模型（MIND），通过将离散词元映射到紧凑超球面流形，结合可微 Soft Top-$k$ 聚合桥接与动态多阶段迁移采样，在仅 130M 参数量下取得了媲美数十亿参数离散及连续生成模型的生成质量。

研究背景与动机¶

现代深度图像生成模型以连续扩散模型（如 DDPM、LDM）和扩散 Transformer（DiT）为代表，取得了令人瞩目的图像合成质量。然而，现有连续扩散模型普遍假设潜空间处于无界欧氏空间中，数据点拥有无限连续状态。尽管现代生成模型的理论分析广泛依赖于“低维流形假设”——即高维自然图像本质上聚集在低维紧凑的拓扑流形上，但现有的连续生成方法除极少数复杂的黎曼流形扩散探索外，几乎未曾显式建模这一内在的几何先验，导致模型在无界空间中采样效率受限且难以充分利用流形几何约束。

若直接将连续自编码器（如 VAE）提取的连续潜表示强行约束或投影到严格的低维拓扑流形（例如低维超球面表面），会导致灾难性的表征崩溃（Representation Collapse）。原因在于连续自编码器高度依赖特征向量的模长（Magnitude）来编码图像的结构与纹理信息，超球面硬性要求的 $\ell_2$ 归一化彻底破坏了模长自由度，造成严重的感知失真与重建失败。与之相反，离散视觉词元化方法（如 VQ-GAN、IBQ、GigaTok）将图像压缩为离散词元索引，以类别属性编码信息，其表征完全依赖于角向分离度（Angular Separation）而非模长。当离散词元被锚定到低维紧凑超球面上时，天然对模长破坏具有免疫力。然而，基于离散词元的自回归与掩码生成范式强制引入一维或随机因果解码顺序，破坏了视觉数据固有的全局双向依赖与全图并行渐进去噪优势。

本文的切入角度是融合离散词元在低维几何流形参数化上的抗崩溃鲁棒性，以及连续扩散模型全局并行去噪的生成灵活性。核心 idea：提出流形感知图像扩散模型（MIND），显式将离散视觉词元投影至连续低维超球面流形上进行前向加噪与反向扩散，设计可微 Soft Top-$k$ 聚合算子解决离散-连续边界的可微训练问题，并辅以缓解低维谱偏差的高频投影层与多阶段动态迁移采样策略。

方法详解¶

整体框架¶

MIND 的整体流水线涵盖预训练分词器离散量化、超球面流形几何参数化、可微软聚合前向/反向扩散过程，以及推理阶段的动态多阶段迁移采样。给定自然图像 $I \in \mathbb{R}^{C \times H \times W}$，首先由冻结的预训练离散分词器（如 IBQ 或 GigaTok）将其量化为离散词元序列 $k \in \mathbb{R}^N$（词表大小为 $V$）。随后，流形投影算子 $P_\theta$ 将离散词元映射为半径为 $R$ 的低维超球面表面上的连续特征 $x_0 \in \mathbb{R}^{N \times L}$，并沿子空间进行正则化切分。

在前向加噪阶段，超球面上的连续表征 $x_0$ 遵循加噪调度注入高斯噪声生成含噪潜变量 $x_t$；去噪主干网络（基于 DiT 架构）接收含噪特征并经过双分支高频特征增强层以克服低维谱偏差，输出关于离散词表的全维度未归一化 logits $\tilde{x}_{0t} \in \mathbb{R}^{N \times V}$。在训练阶段，通过 Soft Top-$k$ 聚合与重投影算子将 logits 映射回超球面潜空间生成去噪特征 $\dot{x}_{0t}$，构建端到端可微的交叉熵与均方误差混合监督；在推理阶段，设计三阶段动态迁移采样算子在去噪进程中自适应切换软采样、熵驱动混合采样和贪婪确定性采样，最终经 Detokenizer 重建出高保真图像。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像与离散分词<br/>Image -> VQ Tokenizer -> 离散词元 k"] --> B["超球面流形参数化与投影<br/>P_theta 投影至超球面表面 x_0"]
    B --> C["双分支高频特征映射<br/>RFF 基础分支 + 零初始化残差正弦分支"]
    C --> D["DiT 扩散去噪预测<br/>预测未归一化分类 Logits x_tilde_0t"]
    D --> E["可微 Soft Top-k 聚合桥接<br/>Top-k Softmax 加权与 L2 球面投影得到 x_dot_0t"]
    E --> F["动态多阶段迁移采样<br/>软采样 -> 熵驱动混合采样 -> 贪婪采样"]
    F --> G["Detokenizer 解码输出<br/>高质量生成图像"]

关键设计¶

1. 超球面流形参数化与投影算子：基于离散表征彻底消除几何瓶颈崩溃

针对连续潜变量在映射至低维紧凑拓扑空间时因模长压缩引发的表征崩溃缺陷，MIND 利用离散分词的角向区分特性，构建了专用的超球面投影算子 $P_\theta$。投影层首先将词表大小为 $V$ 的离散索引映射为 $L$ 维连续向量，并将其分解为若干 $d_{\text{sub}}$ 维超球面子空间进行子向量 $\ell_2$ 正则化，使得每个空间位置的连续表征 $x_0(n, :)$ 严格落在满足 $\sum_{l=0}^{L-1} x_0^2(n, l) = R^2$ 的半径为 $R$ 的超球面表面上。前向扩散过程以球面几何为基底对连续表征进行连续加噪： $$x_t = c_1 \sqrt{1 - t} \cdot w + c_2 \sqrt{t} \cdot x_0$$ 其中 $w \sim \mathcal{N}(0, I)$ 为高斯白噪声，$t \in [0, 1]$。该设计保证了数据点紧密、稠密且具有明确几何拓扑约束地分布于低维紧致空间中，从根本上兼顾了离散词元的抗畸变特性与连续扩散空间的可渐进去噪性。

2. 可微 Soft Top-$k$ 聚合桥接：打通离散 Logits 与连续超球面流形的反向传播

去噪骨干网络 $s_\phi$ 在输出端预测的是针对离散词表的 unnormalized logits $\tilde{x}_{0t} \in \mathbb{R}^{N \times V}$，而扩散状态的递推需要位于超球面表面的连续向量 $\dot{x}_{0t}$。直接使用离散 argmax 或 Multinomial 采样会导致梯度截断无法端到端反向传播，而全局 Softmax 加权又会引入大量长尾噪声类别的负面干扰。为此，本文设计了 Soft Top-$k$ 聚合与球面投影算子：对每个位置的 logits 提取 Top-$k$ 候选集合 $I_k$，在候选子集内计算归一化注意力权重 $\alpha_i = \exp(x_i) / \sum_{j \in I_k} \exp(x_j)$。由于超球面上点的凸组合必然落入球体内部而非球面上，算子进一步执行显式 $\ell_2$ 投影重定标： $$\dot{x}_{0t}(n) = \frac{\sum_{i \in I_k} \alpha_i P_\theta(i)}{\left\| \sum_{i \in I_k} \alpha_i P_\theta(i) \right\|_2}$$ 这一机制构成了一条光滑、可导的桥梁，使去噪网络既能接受离散层面的分类交叉熵（CE）监督，又能计算连续潜变量空间中的均方误差（MSE）对齐损失。

3. 双分支高频特征映射：克服 Transformer 对低维紧致输入的谱偏差

标准的扩散 Transformer（DiT）在接收低维空间连续坐标或低维流形输入时存在严重的“谱偏差”（Spectral Bias），倾向于仅学习低频平滑信号而忽略细粒度高频几何细节。为解决此问题，输入特征在送入主干前经由双分支协同映射：基础分支采用随机傅里叶特征（RFF），利用大方差高斯矩阵 $B \sim \mathcal{N}(0, \sigma^2)$ 进行投影并拼接正弦/余弦激活 $f^{\text{base}} = [\sin(2\pi x_t B), \cos(2\pi x_t B)]$，经浅层 MLP 变换为匹配 DiT 维度的 $h_{\text{base}}$；残差分支则通过确定性正弦位置编码得到 $f^{\text{freq}}$，输入带有 LayerNorm 和 SiLU 激活的深层 MLP 中。极其关键的是，该残差 MLP 的最后一层线性权重被严格初始化为全零，在训练初期保证梯度平稳流动并杜绝激进数值发散，随着训练展开平滑注入多尺度高频细节。

4. 动态多阶段迁移采样策略：在生成全程自适应权衡探索多样性与局部确定性

推理阶段需要将起始白噪声 $x_1 \sim \mathcal{N}(0, c_1^2 I)$ 逐步引导至离散码本的高概率流形上。若自始至终采用单一采样策略，会导致早期模式坍塌或晚期生成伪影。MIND 提出了以时间步阈值 $\rho_1, \rho_2$ 划分的动态多阶段迁移采样机制 $S(\cdot)$： - 初级阶段（$t < \rho_1$）：直接沿用训练期的 Soft 连续聚合机制，允许连续状态在超球面邻域内自由探索扩散方向； - 中间过渡阶段（$\rho_1 \le t < \rho_2$）：执行熵驱动自适应混合采样。首先计算离散分布的香农熵 $H = -\sum p \log p$，据此动态调控采样温度 $\tau_{\text{adj}} = \tau \cdot (2.5 e^{-H/3} + 0.6)$。当模型预测置信度低（熵值高）时自适应降低温度压制不确定性，并在 Top-$k$ 与 Nucleus (Top-$p$) 截断的交集空间中进行多项式采样； - 终末阶段（$t \ge \rho_2$）：完全收缩为贪婪确定性采样 $\hat{k}_t = \arg\max(\tilde{x}_{0t})$，彻底剥离随机性扰动，锁定清晰锐利的局部高频结构。

损失函数 / 训练策略¶

MIND 采用端到端联合优化策略，总损失由作用于未归一化 logits 的离散交叉熵损失与作用于重投影连续去噪潜变量的均方误差损失线性加权构成： $$\mathcal{L}(s_\phi, P_\theta) = \mathbb{E}_{k \sim p_{\text{data}}, t \in [0, 1]} \left[ \text{CE}(\tilde{x}_{0t}, k) + \lambda \cdot \text{MSE}(\dot{x}_{0t}, x_0) \right]$$ 其中权重超参数 $\lambda$ 调节连续几何对齐与离散语义分类的相对强度。训练时引入概率为 $p_{\text{drop}} = 0.1$ 的类别标签随机丢弃以支持推理期的无分类器引导（CFG）；模型采用 AdamW 优化器，并在 DiT 主干上结合分块梯度检查点（Chunked Gradient Checkpointing）显著降低显存开销。

实验关键数据¶

主实验¶

在 ImageNet $256 \times 256$ 官方基准上，作者在严格受限的 80-epoch 预算下进行了系统对比，并在收敛训练下与各类代表性架构开展了系统级比较。

模型	范式/架构	参数量	训练周期 (Epochs)	CFG 引导系数	FID ↓	IS ↑	Precision ↑	Recall ↑
DiT-S/2	连续扩散	33M	80	w/o CFG	68.40	N/A	N/A	N/A
SiT-S/2	插值/流模型	33M	80	w/o CFG	57.64	24.78	0.41	0.60
MIND-S (Ours)	几何流形扩散	~35M	80	w/o CFG	40.72	31.51	0.48	0.61
DiT-B/2	连续扩散	130M	80	w/o CFG	43.47	N/A	N/A	N/A
SiT-B/2	插值/流模型	130M	80	w/o CFG	33.02	43.71	0.53	0.63
DiT-L/2	连续扩散	458M	80	w/o CFG	23.33	N/A	N/A	N/A
MIND-B (Ours)	几何流形扩散	~130M	80	w/o CFG	22.73	56.17	0.55	0.58
DiT-B/2	连续扩散	130M	80	cfg = 1.5	20.01	73.00	0.65	0.56
SiT-B/2	插值/流模型	130M	80	cfg = 1.5	16.88	84.26	0.66	0.56
MIND-B (Ours)	几何流形扩散	~130M	80	cfg = 1.5	12.15	100.16	0.71	0.51
DiT-B/2	连续扩散	130M	80	cfg = 2.0	10.95	119.12	0.76	0.47
SiT-B/2	插值/流模型	130M	80	cfg = 2.0	9.07	137.07	0.77	0.47
MIND-B (Ours)	几何流形扩散	~130M	80	cfg = 2.0	7.97	154.20	0.81	0.43

扩展与收敛系统级对比¶

在充分训练和更大模型规格下，MIND 展现出超越数十亿参数大型自回归与掩码模型的极佳效率。

模型	模型类型	参数量	训练周期 (Epochs)	FID ↓	IS ↑	Precision ↑	Recall ↑
LlamaGen-3B	离散自回归 (AR)	3.1B	300	2.18	263.33	0.81	0.58
RQTran.-re	离散自回归 (AR)	3.8B	N/A	3.80	323.70	N/A	N/A
VAR-d24	离散多尺度 (VAR)	1.0B	200-350	2.09	312.90	0.82	0.59
DiT-XL/2	连续扩散 (Diff.)	675M	1400	2.27	278.20	0.83	0.57
SiT-XL/2	连续流模型 (Flow)	675M	1400	2.06	270.30	0.82	0.59
SimDiff	连续扩散 (Diff.)	2.0B	800	2.77	211.80	N/A	N/A
MIND-B (Ours)	几何流形扩散	~130M	1000	2.18	258.46	0.80	0.58
MIND-B-G (Ours)	几何流形扩散 + GigaTok	~130M	1600	2.06	268.03	0.78	0.62
MIND-XL (Ours)	几何流形扩散	~715M	1000	1.97	297.93	0.78	0.61
MIND-XL-G (Ours)	几何流形扩散 + GigaTok	~715M	1600	1.95	293.78	0.75	0.67

关键发现¶

80-Epoch 快速收敛表现突出：在无引导生成中，MIND-B 仅经 80 个 epoch 训练即获得 22.73 的 FID，较 DiT-B/2（43.47）下降了近一半（降低 20.74），甚至直接超越了参数量高出数倍的 DiT-L/2（458M，FID 23.33）。
极佳的参数效率：在有引导生成的收敛设定下，仅 130M 参数的 MIND-B 配合 GigaTok 分词器取得了 2.06 的 FID，全面击败参数量达 3.1B 的 LlamaGen-3B（FID 2.18）以及 2.0B 参数的 SimDiff（FID 2.77）。规格扩大至 MIND-XL（715M）后，FID 进一步刷新至 1.95。
高保真与精确语义对齐：在 cfg=2.0 下，MIND-B 的 Precision 达 0.81，显著高于 SiT-B/2（0.77），生成样本完全避免了基线中常见的物体结构断裂和拓扑扭曲。

亮点与洞察¶

离散分类抗崩溃与连续扩散全局性的完美统一：深刻揭示了低维紧凑流形上连续潜在特征“模长依赖脆弱性”导致坍塌的数学物理本质，创造性地用具有角向分离不变性的离散词元作为几何流形载体，彻底破除了传统扩散模型与离散生成范式二元对立的壁垒。
端到端可微 Soft Top-$k$ 球面投影设计：避免了传统直通估计（STE）带来的高梯度方差，又避免了全空间 Softmax 导致球内严重塌缩的问题，通过局域注意力聚合与 $\ell_2$ 球面投影保持了流形切空间的梯度一致性。
低维输入谱偏差的精准矫正：针对低维流形输入导致 Transformer 高频衰减的理论痛点，采用随机傅里叶特征加零初始化深层正弦残差分支，为紧凑表示的生成式主干设计提供了极有价值的工程示范。

局限与展望¶

采样计算开销依然偏高：推理阶段依赖 250 步 SDE 迭代求解，生成耗时仍显著高于单步或少步一致性模型与 Rectified Flow 模型。
矢量量化带来的固有信息损耗：尽管离散词元提供了强健的流形几何特性，但离散分词器（Tokenizer）重建的固有下采样信息损失依然存在上限，未来可通过在流形空间中引入高频残差调制项加以弥补。
高级流形动力学扩展空间：当前轨迹生成基于标准超球面扩散设定，未来可进一步结合测地线流动匹配（Geodesic Flow Matching）、均值流（Mean Flow）或自监督表征对齐（如 REPA 范式）进一步提升收敛速度与生成表现。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [首次将离散视觉分词显式锚定到低维超球面几何并结合连续扩散，理论洞察深刻，软投影与动态采样机制极具独创性]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖 80-epoch 短期对比与充分收敛系统级对比，横跨小/基准/超大规格与多种分词器，对比 baseline 丰富扎实]
写作质量: ⭐⭐⭐⭐⭐ [数学推导严谨，崩溃对比实验与理论动机阐述极具说服力，图表清晰整洁]
价值: ⭐⭐⭐⭐⭐ [打破了离散自回归与连续扩散的割裂界限，证明了显式流形几何先验对提升扩散模型参数效率与生成质量的关键作用]