WaveDiT：分布感知的小波流匹配方法用于高效 3D 脑部 MRI 合成 WaveDiT: Distribution-Aware Wavelet Flow Matching for Efficient 3D Brain MRI Synthesis
在小波系数空间做流匹配，配合状态感知不确定性建模，单卡 10 步合成全分辨率 3D 脑 MRI
前置知识
3D 离散小波变换 (3D DWT)
把三维体数据 $\mathbb{R}^{D\times H\times W}$ 分解为 1 个低频近似子带 (LLL) 和 7 个方向高频细节子带的可逆线性变换，每一维空间分辨率减半，因此每个子带大小为 $D/2 \times H/2 \times W/2$，且能通过逆 DWT 完美重建。
WaveDiT 整个方法的载体就是 3D Haar 小波系数空间，不理解小波就无法理解为何低频与高频子带需要分开建模、以及为何体积缩小 8 倍反而能保留解剖结构。
流匹配 (Flow Matching)
用一条从噪声 $x_0\sim\mathcal{N}(0,I)$ 到数据 $x_1$ 的连续轨迹 $x_t=(1-t)x_0+tx_1$ 训练速度场 $v_\theta$，使其逼近目标速度；推理时用 ODE 求解器沿轨迹把噪声推到数据。
WaveDiT 的核心生成范式就是 CFM 条件流匹配，它决定了训练损失形式、推理 ODE 求解以及为何能比扩散少几个数量级的步数得到结果。
异方差不确定性 (Heteroscedastic Uncertainty)
预测误差方差随输入而变化，不再假设观测噪声固定。Kendall & Gal 把这种与数据相关的方差建模为可学习变量，把 MSE 替换为 $e^{-s}\|v-\hat v\|^2/2 + s/2$ 的贝叶斯损失，$s$ 由网络预测。
Morpheus 模块本质上就是在频域复现这套异方差回归思想，但用在小波高频子带的优化上，直接关系到本文最关键的工程贡献。
扩散 Transformer (DiT) 与 Hourglass 注意力
DiT 用 patch token + 全局自注意力替代 UNet，Hourglass DiT (HDiT) 在浅层用局部窗口注意力、深层用全局注意力，兼顾局部结构与长程依赖，并配合 AdaRMSNorm 注入条件。
WaveDiT 把 HDiT 扩展到 3D 切片，并提出分时空因子化注意力，是其骨干架构的来源与对比基线。
研究动机
全分辨率 3D 脑 MRI 合成对数据增强与衰老/疾病建模意义重大，但现有路径在算力或保真度上两头吃亏。论文明确指出，像素空间扩散 (WDM、MD、3DMD) 在 $\sim 10^6$ 体素量级上需要 1000 步迭代，单卡训练需要数天；潜在扩散 (MLDM、BS) 通过有损压缩把体积压进潜空间，节省算力却可能丢失皮层细微解剖结构或在重建阶段引入伪影。另一类小波方法 (FlowLet) 虽然避免了潜空间压缩，但把所有子带一视同仁地施加均匀 MSE 损失，没有考虑不同子带统计量差异极大这一事实：作者在表 1 中实测到 LLL 子带峰度仅约 5.06、HHH 子带峰度却高达 269.44，相差近 90 倍，且高频方差随空间位置变化跨越约 8 个数量级。结果是均匀损失在组织边界过度惩罚、在均匀区域又欠惩罚，训练和采样过程都对这一统计不匹配敏感，导致合成 MRI 在下游脑龄预测和 95 区 ROI 分析中难以保持解剖可信度。
本文的目标是本文的具体目标是提出一个名为 WaveDiT 的条件流匹配框架，在 3D Haar 小波系数空间直接进行全分辨率 3D 脑 MRI 合成，要求 (1) 在单张 H100 GPU 上、仅 10 步 ODE 求解即可在 $\sim 1$ 秒内生成一个完整体积；(2) 显式建模小波子带的重尾、异方差结构，使优化和条件化都依据子带当前状态的可靠性进行自适应调节；(3) 在 FID/MMD 等全局分布指标、脑龄预测 MAE 以及 95 区 ROI 的 iMAE/KL/Dice 多层级评估上同时优于现有扩散、潜在扩散和小波基线，并可作为数据增强手段帮助下游 3D DenseNet 把 MAE 降到比纯真实数据训练更低的水平。
与已有工作不同的是，本文的独特切入角度是把"状态感知不确定性"从回归任务系统地引入到基于小波的高分辨率 3D 生成模型中。已有工作要么在像素空间统一加权 (WDM/FlowLet)，要么用潜空间压缩规避问题 (MLDM/BS)，都回避了小波子带本身的统计异质性。WaveDiT 通过一个轻量辅助网络 Morpheus 实时从当前 $x_t$ 的 6 阶统计量 (均值、标准差、最大绝对值、L2、偏度、峰度) 预测子带级对数方差 $s_\psi(x_t,t)$，再把这一信号同时注入贝叶斯异方差损失和骨干的频率条件通路。这一"分布感知"的设计在小波域生成中属于首次系统化尝试，使模型能在 10 步推理预算下兼顾高频细节与全局一致性。
核心方法
整体思路可以一句话概括为：先把 3D MRI 体数据用单层 3D Haar DWT 分解为 8 个子带，再在小波空间训练一个条件流匹配模型，模型的预测精度和条件信号都通过一个统计特征驱动的 Morpheus 模块动态调度。技术上，输入体 $x_1$ 与高斯噪声 $x_0$ 在 CFM 轨迹 $x_t=(1-t)x_0+tx_1$ 上插值；Morpheus 用 6 个统计量 + 时间嵌入经两层 MLP 输出 8 个子带的对数方差 $s_\psi(x_t,t)\in\mathbb{R}^8$；骨干是基于 HDiT 的 WaveDiT，把小波张量重塑为 $B\cdot D'$ 个 2D 切片做 patch token 化，浅层用 $K\times K$ 邻域注意力处理局部结构，深层做分时空因子化注意力 (intra-slice + inter-slice) 恢复体积一致性；训练损失是异方差贝叶斯目标 $\mathcal{L}=\mathbb{E}[\tfrac12 e^{-s}\|v_\theta-v_{\text{target}}\|^2+\tfrac12 s]$，同时 $s$ 投影为频率提示与时间/年龄/切片嵌入一起经 AdaRMSNorm 调制骨干；推理时用二阶 Heun 求解器 10 步推到 $t=1$，再经逆 DWT 重建到体素空间。
WaveDiT 相对已有方法有两点本质区别。其一是"在小波域而不是潜空间"做流匹配：相比 WDM/MOTFM 等像素级扩散保留了 8 倍空间压缩下的可逆性，避免了 VAE 类有损编码；相比 FlowLet 等已有小波方法，又把均匀损失换成子带级异方差贝叶斯目标。其二是把不确定性建模从被动式 (训练后处理) 变成主动式状态机：Morpheus 不是事后修正误差，而是直接参与 (a) 损失加权 (高频降权、低频保权)，(b) 骨干条件化 (频率提示 $s$ 与时间/年龄嵌入共同走 AdaRMSNorm)，相当于让网络在每一步都"知道自己对哪个频段有多确定"，并据此调整生成方向。
方法步骤详情
方法流程可拆为五个环节。① 3D 分解：单层 3D Haar DWT 把 $x_1\in\mathbb{R}^{B\times 1\times D\times H\times W}$ 映射到 $w\in\mathbb{R}^{B\times 8\times D'\times H'\times W'}$，$D'=D/2$ 等；作者实测 224³ 输入得到 $D'=H'=W'=112$，$N\approx 1.4\times 10^6$ token。② 切片+token 化：$w$ 沿深度维重塑为 $B\cdot D'$ 张 2D 切片，$8\times 8$ patch 嵌入后叠加固定随机 Fourier 切片位置编码、正弦时间编码、年龄元数据嵌入和 Morpheus 频率提示 $s$，构成全局条件向量。③ Morpheus 调度：对每个子带 $c$ 提取 $\mu_c,\sigma_c,\max|x_c|,\|x_c\|_2/\sqrt{N},\gamma_1,\gamma_2$ 共 6 维统计量，与时间嵌入一起送入两层 width=128 的 SiLU MLP，得到 $s_\psi(x_t,t)\in\mathbb{R}^8$。④ WaveDiT 骨干：浅层 2D 滑动窗口注意力 (window $K=7$, head dim 64) 捕获局部边缘，深层做因子化分时空注意力——先在每个 2D 切片内做全局自注意力建模左右对称/脑室-皮层关系，再沿深度轴对同一 $(h,w)$ 位置做注意力以恢复 Z 轴连续性；两级深度各 2、宽度 1024、FFN 4096，总参数 $\sim 142$M，AdaRMSNorm 注入 $s$。⑤ 训练与推理：用 AdamW (lr $10^{-4}$、wd $10^{-2}$、bs=4、dropout 0.1) 联合优化 $\theta$ 和 $\psi$ 共 200 轮；推理阶段 Heun 二阶求解器从 $t=0$ 推到 $t=1$ 仅需 10 步、约 1 秒，再经逆 3D Haar DWT 重建为体素空间 3D MRI。
技术新颖性
技术新颖性体现在三个层次的耦合。层次一是小波域的 CFM 而非扩散：把 FlowLet 使用的扩散思路换成条件流匹配，并配合二阶 Heun 求解器，把 1000 步扩散压缩到 10 步而不损失质量，作者在表 2b 中通过 CFM/RFM/OTFM 三种流匹配目标的消融证明了 CFM 路径最匹配小波系数的统计结构。层次二是 Morpheus 的状态感知异方差损失：把 $e^{-s}\|\cdot\|^2/2+s/2$ 这种原本用于回归的不确定性加权迁移到生成任务的子带级别，并用峰度等高阶统计量动态估计 $s$，是首个在小波生成模型中显式做"分布感知"的工作。层次三是 3D HDiT 的因子化分时空注意力：把 HDiT 从 2D 扩展到 3D 时不直接做全 3D 自注意力，而是先用 2D 切片批处理再做深度轴 cross-slice 注意力，复杂度从 $\mathcal{O}(N^2)\approx 2\times 10^{12}$ 降到 $\mathcal{O}(D'(H'W')^2+H'W'D'^2)\approx 1.8\times 10^{10}$，约 110× 缩减，使 224³ 体积分辨率在单 H100 上以 bs=1、~12GB VRAM 训练成为可能。
实验结果
实验在 5989 名认知正常受试者 (5.9–95.5 岁) 的多中心数据上进行，每个模型合成 3000 个体积用于评估，结论可以分成五点。第一，全局分布指标上，WaveDiT-CFM 10 步以 FID 0.0039、MMD 0.00010、MS-SSIM 0.830 同时拿下最佳三项，显著优于 1000 步的 WDM (FID 0.0045) 和 10 步的 FlowLet (FID 0.0117)，FID 越低越好、MS-SSIM 越低代表多样性越高。第二，下游脑龄预测 (BAP) 上，WaveDiT-CFM 用作数据增强后，3D DenseNet 在 >44 岁保留集上的 MAE 降到 $2.44\pm 3.19$ 年，比纯真实数据训练的 2.92 还要低 0.48 年，击败了所有条件基线 (FlowLet 2.72、MLDM 2.68、BS 3.08、3DMD 3.10)。第三，95 区 ROI 分析中，WaveDiT-CFM 同样拿到最低 iMAE 42.49、最低 KLD 0.96 和最高 Dice 0.46，这说明全局指标优秀的 WDM (Dice 0.36) 在解剖一致性上反而有差距，作者借此论证了多层级评估协议的必要性。第四，Morpheus 消融显示，去掉 Morpheus 后 FID 从 0.0039 退化到 0.0295、BAP MAE 从 2.44 升至 2.98、Dice 从 0.46 跌到 0.41，验证了状态感知不确定性加权的关键作用。第五，目标函数消融中，CFM 在 FID 和 BAP 上均优于 RFM (FID 0.1811) 和 OTFM (FID 0.0157)，而体素空间变体 (Voxel-space ablation) 虽然 FID 0.0045 与小波版本接近、BAP MAE 3.04 明显恶化、Dice 0.43 也低于 0.46，证实小波域本身提供了体素空间无法保留的解剖结构。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 全局分布 (全脑 MRI 体积) | FID ↓ (10 bootstrap 平均) | 0.0039 (WaveDiT-CFM, 10 步) | 0.0045 (WDM, 1000 步) / 0.0117 (FlowLet, 10 步) | FID 比 WDM 再降 13%，比 FlowLet 降 67%，且只需 1/100 的步数 |
| 全局分布 (全脑 MRI 体积) | MS-SSIM ↓ (多样性) | 0.830 (WaveDiT-CFM, 10 步) | 0.918 (MOTFM) / 0.908 (WDM) / 0.869 (FlowLet) | MS-SSIM 越低多样性越好，WaveDiT 比 WDM 低 8.6%，比 MOTFM 低 9.6%，无多样性坍缩 |
| 脑龄预测 (3D DenseNet 数据增强) | 测试 MAE (年) ↓ | 2.44 ± 3.19 | 2.92 (纯真实数据) / 2.72 (FlowLet) / 2.68 (MLDM) / 2.76 (MOTFMa) | 比纯真实数据基线再降 0.48 年，相对 FlowLet 降 10.3%、相对 MOTFMa 降 11.6% |
| 95 区 ROI 解剖一致性 | Dice ↑ (95 区平均) | 0.46 ± 0.16 | 0.44 (FlowLet) / 0.42 (MOTFMa/MOTFM) / 0.36 (WDM) / 0.35 (3DMD) | Dice 比最佳非 Ours 基线 (FlowLet) 提升 0.02，比 WDM 提升 0.10 |
| 95 区 ROI 解剖一致性 | iMAE ↓ (强度平均绝对误差) | 42.49 ± 12.96 | 45.18 (FlowLet) / 48.57 (WDMa) / 62.91 (3DMD) | 比 FlowLet 低 6.0%、比 WDMa 低 12.5%、比 3DMD 低 32.4% |
| Morpheus 消融 | FID / BAP MAE | 0.0039 / 2.44 | 0.0295 / 2.98 (w/o Morpheus) | Morpheus 使 FID 降 7.6×、BAP MAE 降 0.54 年，验证子带级异方差加权的必要性 |
局限与改进
作者在结论中坦承两个核心限制。第一，实验仅在 T1 加权 MRI 上验证，没有覆盖 T2、FLAIR、CT 等其他 3D 成像模态，而不同模态的解剖对比度差异很大，子带统计量也会变化，Morpheus 是否需要重新校准并未给出答案。第二，整套评估完全基于 FID/MMD/MS-SSIM、BAP MAE、ROI iMAE/KL/Dice 等定量代理指标，没有引入放射科医生的盲法读片评估，因此"临床可用性"只能间接推断。我自己再补充两点：(1) 论文依赖单一 backbone (HDiT)、单一 backbone 宽度 (1024) 和单一 patch 尺寸 (8×8) 的配置，没有给出 HDiT 深度/宽度/patch 的消融，难以判断是结构还是 Morpheus 起主导；(2) 推理时全部使用 10 步 Heun，质量与步数的 Pareto 曲线没有呈现，"高效"这个卖点缺一张 1/5/10/20 步的对比图；(3) 数据集虽然多中心 (OpenBHB+ADNI+OASIS3) 但都是公共欧美队列，跨种族跨扫描仪的鲁棒性没有验证，而小波系数对 MRI 偏置场与配准误差非常敏感。
独立分析的弱点
从独立视角看，论文至少有以下三点可以补强。其一，Morpheus 只输出 8 维子带级 $s$，没有空间分辨率；表 1 显示 HF 子带局部方差在空间上跨越 8 个数量级，全子带共享一个 $s$ 显然过粗，可以换成 (子带, 切片, patch) 级别的张量 log-variance，并用 chunked 稀疏估计避免 VRAM 爆炸，对应改进方向是引入空间门控机制或局部 3D 卷积头。其二，训练只用单一 MSE-based CFM 目标，没有引入对抗感知损失 (例如 LPIPS-Med 或医学预训练特征匹配) 来强化皮层细节，而 ROI Dice 在所有方法中都只有 0.4 出头，提示皮层/皮层下结构的一致性仍有明显缺口；可在不破坏异方差加权的前提下，给高频子带附加感知正则。其三，推理采用固定 10 步 Heun，没有自适应步数或早停；若能在 FID 上升超过阈值时停止迭代，可在 5–7 步内维持 95% 质量，进一步压缩到 0.5 秒以内；改进方向是参考 Consistency Model 的自适应离散化或 Karras 信号的 $\kappa$ 估计。其四，全文没有给出 95 区 ROI 的细分结果 (例如海马、杏仁核等小结构)，无法判断小解剖结构是否依然失败，可补充 per-ROI 表 + 小结构热力图。
未来方向
作者明确列出的两个方向是 (1) 把 WaveDiT 拓展到 CT、其他 MRI 模态及其他 3D 影像域，(2) 引入比年龄更丰富的条件信号 (性别、疾病标签、扫描协议、站点 ID)。基于成果可延伸的方向还包括：(a) 把小波域 + 异方差生成范式迁移到 4D 动态 MRI 或胎儿脑发育时间序列上，把时间作为额外坐标的子带；(b) 把 Morpheus 的统计特征与下游 BAP/分割模型的不确定性联合训练，实现"生成-评估-校准"端到端优化；(c) 在扩散-流匹配混合框架下把 WaveDiT 作为高分辨率精修模块，与低分辨率潜扩散级联，进一步降低算力门槛；(d) 与隐私保护学习结合，把生成 MRI 替换或补充真实数据用于联邦训练，绕开医疗数据共享壁垒；(e) 探索小波基选择 (Haar→Daubechies→Bior) 对子带统计与生成质量的影响，作者目前固定为单层 Haar，可能不是最优。
复现评估
复现友好度处于中上水平。论文在 GitHub (https://github.com/sisinflab/WaveDiT) 公开了代码、训练脚本与预训练权重；数据来自三个公开数据集 (OpenBHB、ADNI、OASIS-3)，其中 ADNI 需要申请，整体获取成本可控；预处理链 (偏置场校正 N4、仿射配准到 MNI152、颅骨剥离、[0,1] 归一化) 全部使用标准工具 (ANTs、FreeSurfer 风格)，复现门槛低。算力要求相对友好：单卡 H100 即可在 ~12 GB VRAM、bs=1、200 epoch、~26 小时内完成训练；推理只需 10 步 Heun、约 1 秒/体积。复现难度集中在三点：① ADNI 数据申请有审批周期；② Morpheus 6 维统计量的具体实现细节 (例如 $\max|x_c|$ 是否去均值) 需要从代码确认；③ 评价管线依赖 Medical ResNet-50 特征提取器与 3D DenseNet 脑龄回归器，两者的训练超参未在正文完全披露，需要查代码。整体而言，相比同级别 3D 医学生成工作，WaveDiT 的算力门槛和工程细节可访问性都明显更友好。
论文图表