UltraDiff: Transferring High-Fidelity Priors to Compressed Latent Spaces for High-Resolution Image Generation¶

会议: NeurIPS 2026
arXiv: 2610.21202
OpenReview: https://openreview.net/forum?id=j164YLIssM
领域: 图像生成
关键词: 扩散Transformer, 高分辨率图像生成, 潜在空间压缩, 联合潜在扩散, 边际分数蒸馏

一句话总结¶

针对高分辨率生成中激进压缩潜空间（$f32$）易失真与歧义的问题，UltraDiff 提出联合建模 $f8$ 与 $f32$ 分布的双分支 MixTeacher，并通过原生 $f32$ 空间下的边际分数蒸馏，将高质量几何与细节先验无损迁移至极低开销的学生模型，实现单张 A100 上 6.2 秒生成 2K 图像。

研究背景与动机¶

高分辨率及超高清（UHD，如 2K、4K）图像生成已成为内容创作、VR 与游戏产业的核心需求。基于扩散 Transformer（DiT）的生成模型在保真度上展现出显著优势，但其自注意力机制的计算与内存复杂度随空间 token 数量呈二次方增长（$\mathcal{O}(N^2)$）。主流潜在扩散模型普遍采用标准 $f8$ 压缩（下采样因子 $f=8$），在面对 2K 乃至 4K 图像时，序列长度爆炸导致计算开销急剧飙升甚至引发显存溢出（OOM）。为此，学术界与工业界开始探索激进空间下采样（如 $f32$），将二维潜在 token 数量减少 16 倍，从而获得显著的推理提速。

然而，极高压缩比的效率提升伴随着扩散建模难度的骤增。在 $f32$ 潜空间中，模型被迫从极少的空间 token 重建复杂的自然场景与高频纹理，潜在目标分布表现出严重的欠定性与信息歧义。这导致现存的高压缩模型在微小结构、细节质感以及复杂文本对齐上普遍逊色于成熟的 $f8$ 系统。虽然直接采用预训练 $f8$ 模型参数对 $f32$ 模型进行静态初始化能够带来初步增益，但其性能很快陷入饱和，与高性能 $f8$ 教师之间依然存在巨大鸿沟；而传统的特征对齐蒸馏方法在跨分辨率网格对齐时，必须依赖启发式特征插值，不仅会引入棋盘效应与量化伪影，更缺乏针对去噪向量场的直接显式监督。

本文的切入角度是打破独立潜空间建模的思维定式，将高压缩生成转化为“联合潜在扩散学习”与“原生边际分数蒸馏”的两阶段范式。核心 idea：通过构建共享自注意力的双分支 MixTeacher 联合学习 $f8$ 与 $f32$ 联合分布，并在原生 $f32$ 空间对学生模型执行边际速度场蒸馏，利用高保真轨迹的条件期望天然滤除空间歧义与伪影。

方法详解¶

整体框架¶

UltraDiff 的整体流程由两阶段构成：联合潜在教师模型（MixTeacher）的构建与训练，以及向紧凑单分支学生模型的边际分数蒸馏（Marginal Score Distillation）。首先，对同一张输入图像分别提取 $f8$ 潜在表示（$C_8=16$ 通道）与 $f32$ 潜在表示（$C_{32}=64$ 通道），并在加噪后送入具备独立参数但共享自注意力的 MixTeacher 双分支架构中，配合随机路由策略学习联合流匹配速度场。随后，冻结 MixTeacher，将学生模型置于原生 $f32$ 空间，以教师模型在联合引导下投影至 $f32$ 分支的速度场为目标执行显式蒸馏，最后通过多分辨率流匹配适配至超高清生成。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像 I + 文本提示词"] --> B["双路径 VAE 编码<br/>f8 潜空间 (C=16) & f32 潜空间 (C=64)"]
    B --> C["MixTeacher 双分支联合潜在扩散<br/>独立参数 + 层级共享自注意力"]
    C --> D["随机路由训练策略<br/>联合模式 / f8-Solo / f32-Solo 混合采样"]
    D --> E["原生 f32 空间边际分数蒸馏<br/>条件期望投影过滤歧义伪影"]
    E --> F["单分支 f32 学生模型超分辨率适配<br/>1K 蒸馏 + 2K-4K 纯流匹配微调"]
    F --> G["超高清生成输出<br/>极速 2K / 4K 高保真图像"]

关键设计¶

1. MixTeacher 双分支联合潜在建模：以跨空间注意力桥接异构分辨率 为了让模型在扩散学习阶段自发捕捉 $f8$ 高频细节与 $f32$ 紧凑语义之间的统计相关性，作者构建了统一扩散模型 MixTeacher，参数化为 $\Phi$。给定输入图像 $\mathbf{I} \in \mathbb{R}^{3 \times H \times W}$，分别提取其无噪潜变量 $\mathbf{x}_0^{\text{f8}} \in \mathbb{R}^{16 \times \frac{H}{8} \times \frac{W}{8}}$ 与 $\mathbf{x}_0^{\text{f32}} \in \mathbb{R}^{64 \times \frac{H}{32} \times \frac{W}{32}}$。MixTeacher 包含平行的 $f8$ 分支与 $f32$ 分支，两分支拥有各自独立的 Patch Embedding、QKV 投影矩阵、前馈网络（FFN）、层归一化与输出投影头，以此兼容截然不同的通道维度与特征分布。然而，在骨干网络的每一层 Transformer 内部，两分支的 token 序列被拼接在一起并通过统一的自注意力机制进行全局交互。这种机制允许 $f32$ 分支的稀疏 token 密集查询来自 $f8$ 分支的高保真几何线索，从而在无须外部网格重采样的前提下动态消解压缩带来的结构歧义。

2. 随机路由训练策略：消除单模态协同依赖与表示崩溃 若在训练中始终保持双分支联合激活，由于 $f8$ 空间包含更加完备且容易拟合的纹理信息，$f32$ 分支极易产生“协同适应”（co-adaptation）依赖捷径，即退化为完全依赖 $f8$ 上下文的被动解码器，而丧失自身的单边际生成能力。为此，MixTeacher 引入随机路由机制，通过分支激活掩码 $(m_{\text{f8}}, m_{\text{f32}}) \in \{(1, 1), (1, 0), (0, 1)\}$ 控制前向与反向传播。优化目标定义为三种模式的期望： $$\mathcal{L}_{\text{Mix}} = \mathbb{E}_{t, \mathbf{x}, (m_{\text{f8}}, m_{\text{f32}})} \left[ \sum_{k \in \{\text{f8}, \text{f32}\}} m_k \cdot \|\mathbf{v}_\Phi^k(\mathbf{x}_t^k) - \mathbf{u}_t^k\|^2 \right]$$ 其中三种拓扑模式分别为：（1）联合模式（Joint Mode, $m_{\text{f8}}=1, m_{\text{f32}}=1$），激活跨分支自注意力以学满联合分数；（2）f8-Solo 模式（$m_{\text{f8}}=1, m_{\text{f32}}=0$），保证原始高质量先验不受污染；（3）f32-Solo 模式（$m_{\text{f8}}=0, m_{\text{f32}}=1$），强制 $f32$ 分支在无任何 $f8$ 输入辅助下独立去噪。作者探索了最佳采样概率组合 $(p_{\text{joint}}=0.4, p_{\text{f32}}=0.4, p_{\text{f8}}=0.2)$，使 $f32$ 分支既能充分吸收跨空间交互带来的高保真监督，又保留坚实的独立边际生成能力。

3. 原生 f32 空间边际分数蒸馏：基于条件期望滤除压缩伪影 在获得训练完备的 MixTeacher 后，下游推理需要摆脱对昂贵 $f8$ 分支的依赖。学生模型 $\theta$ 仅在紧凑的 $f32$ 潜空间中运行。不同于传统在中间特征图做 MSE 拟合的间接正则化，UltraDiff 直接在原生 $f32$ 速度场上构建显式蒸馏目标： $$\mathcal{L}_{\text{distill}}(\theta) = \mathbb{E}_{t, (\mathbf{x}_t^{\text{f32}}, \mathbf{x}_t^{\text{f8}}) \sim q_t} \Big[ \|\mathbf{v}_\theta(\mathbf{x}_t^{\text{f32}}, t) - \mathbf{v}_\Phi^{\text{f32}}(\mathbf{x}_t^{\text{f32}}, \mathbf{x}_t^{\text{f8}}, t)\|^2 \Big]$$ 理论上，在最小化该 $L_2$ 目标时，学生模型学到的是教师模型在给定 $f32$ 观测时关于真实高维 $f8$ 分布的条件期望： $$\mathbf{v}_\theta^*(\mathbf{x}_t^{\text{f32}}, t) = \mathbb{E}_{q(\mathbf{x}^{\text{f8}} | \mathbf{x}^{\text{f32}})} \Big[ \mathbf{v}_\Phi^{\text{f32}}(\mathbf{x}_t^{\text{f32}}, \mathbf{x}_t^{\text{f8}}, t) \Big]$$ 这表明学生模型在 $f32$ 空间学到了联合流场投影出的边际向量场。由于 MixTeacher 的跨分支注意力已经在联合去噪中将条件分布收敛为一个极为明确、无模糊的确定性信号，这一数学投影如同在概率流上施加了“结构滤波器”，将由于下采样 32 倍造成的随机混叠与欠定歧义“积分消除”（integrate out），引导学生模型朝向具有高保真全局一致性的流形收敛。

损失函数 / 训练策略¶

UltraDiff 的整体训练管线细分为三个阶段： 1. 阶段 1（$f32$ 迁移预训练）：使用预训练 $f8$ 模型的内部权重初始化 $f32$ DiT（随机初始化输入 Patch Embedding 与输出 Head），在 $0.37\text{B}$ 过滤图文对上训练 120K 步，得到基准权重 f32-phase1。 2. 阶段 2（MixTeacher 联合训练）：$f8$ 分支加载成熟 $f8$ 教师权重，$f32$ 分支加载 f32-phase1，在 512 与 1024 分辨率下按照混合路由概率训练 5K 步。 3. 阶段 3（学生模型蒸馏与分辨率缩放）：冻结 MixTeacher，学生模型加载 f32-phase1。在 512 与 1024 分辨率（50K 步）采用蒸馏与流匹配复合目标： $$\mathcal{L}_{\text{total}} = \lambda \mathcal{L}_{\text{distill}} + \mathbb{E}_{t, \mathbf{x}} \left[ \|\mathbf{v}_\theta(\mathbf{x}_t^{\text{f32}}) - \mathbf{u}_t^{\text{f32}}\|^2 \right]$$ 在后续 10K 步的超分辨率缩放（1K、2K、3K、4K 混合分辨率）中，为规避高分辨率下教师推理的巨大开销与潜在不稳定性，完全停用蒸馏损失（$\lambda=0$），仅采用标准流匹配目标进行纯数据驱动微调。全阶段使用 AdamW 优化器（$\text{LR}=1\times 10^{-4}$，500 步预热后余弦衰减）及 FSDP 分布式并行。

实验关键数据¶

主实验¶

在 2048×2048 与 4096×4096 超高分辨率生成基准上，UltraDiff 与主流高分辨率 T2I 模型进行了全面对比。下表汇总了生成质量与推理延迟表现（统一单张 NVIDIA A100 GPU 测试）：

分辨率	方法	参数量 (B)	潜空间尺度	延迟 / 采样步数	FID ↓	pFID ↓	Aesthetic ↑	CLIP ↑
$2048^2$	Ultrapixel	5.1	$16 \times \frac{H}{32} \times \frac{W}{32}$	12s / 20	46.60	43.03	6.24	30.21
$2048^2$	Diffusion-4k	12.0	$16 \times \frac{H}{42} \times \frac{W}{42}$	23s / 50	37.79	42.43	6.28	29.21
$2048^2$	Sana-2K	1.6	$32 \times \frac{H}{32} \times \frac{W}{32}$	5.8s / 20	45.99	42.99	6.18	30.13
$2048^2$	HunyuanImage-2.1	17.0	$64 \times \frac{H}{32} \times \frac{W}{32}$	82s / 50	43.45	48.72	6.44	30.17
$2048^2$	UltraDiff (Ours)	2.0	$64 \times \frac{H}{32} \times \frac{W}{32}$	6.2s / 20	39.52	42.17	6.40	30.53
$4096^2$	Ultrapixel	5.1	$16 \times \frac{H}{42} \times \frac{W}{42}$	73s / 20	155.17	57.14	5.91	27.39
$4096^2$	Diffusion-4k	12.0	$16 \times \frac{H}{32} \times \frac{W}{32}$	122s / 50	162.94	75.70	5.68	27.20
$4096^2$	Sana-4K	1.6	$32 \times \frac{H}{32} \times \frac{W}{32}$	21.3s / 20	149.20	56.25	6.12	29.97
$4096^2$	UltraDiff (Ours)	2.0	$64 \times \frac{H}{32} \times \frac{W}{32}$	28.3s / 20	146.60	54.27	6.34	30.14

在 1024×1024 标准基准上，UltraDiff（2.0B）取得 85.4 DPG、0.68 GenEval 和 9.39 HPSv3，综合对齐度与画质超越 PixArt-$\Sigma$（80.5 / 0.54）、SD3-medium（84.1 / 0.62）与 Sana（84.5 / 0.66），并逼近 12B 参数的 FLUX-dev（83.8 / 0.67）。

消融实验¶

为严格解耦训练机制的贡献，作者在相同 DiT 架构、相同 $f32$ VAE 及数据预算下进行了同框架消融，并与不同蒸馏基线进行对比：

实验组 / 蒸馏策略	训练范式 / 监督来源	CLIP ↑	ImageReward ↑	GenEval ↑	潜空间 Token 规格
$f32$ 从头训练	纯随机初始化 + 流匹配	27.21	0.70	0.54	$H/32 \times W/32$
$f32$ 经 $f8$ 参数初始化	静态权重迁移 + 流匹配	28.91	0.91	0.61	$H/32 \times W/32$
外部 DINOv2 特征蒸馏	REPA 特征网格对齐	27.53	0.06	0.33	$H/32 \times W/32$
显式 $f8$ 特征蒸馏	跨网格特征插值拟合	27.87	0.05	0.31	$H/32 \times W/32$
UltraDiff 学生模型	MixTeacher 边际分数蒸馏	29.27	0.98	0.72	$H/32 \times W/32$
$f8$ 教师锚点 (Teacher Anchor)	原始高保真密集潜空间	29.25	1.08	0.73	$H/8 \times W/8$

同时，在控制推理效率实验中，20 步采样下 UltraDiff 在 $1024^2$ 耗时 3.2s（比 $f8$ 锚点的 27.3s 提速 8.5 倍）；在 $2048^2$ 耗时 6.2s（比 $f8$ 锚点的 249.8s 提速 40.3 倍）；在 $4096^2$ 下 $f8$ 锚点直接显存溢出（OOM），而 UltraDiff 稳定在 28.3s 完成生成。

关键发现¶

边际分数蒸馏打破了初始化增益的饱和瓶颈：从头训练的 $f32$ 模型 GenEval 仅为 0.54，单纯引入 $f8$ 静态权重初始化可提升至 0.61，但无法进一步提升；引入 MixTeacher 边际分数蒸馏后跃升至 0.72，几乎完全抹平了与 $f8$ 锚点（0.73）之间的性能差距。
特征对齐在跨潜空间分辨率下存在负迁移：基于 DINOv2 或 $f8$ 隐藏状态的特征对齐均遭遇崩溃（GenEval 骤降至 0.33 与 0.31），根源在于 $f32$ 与 $f8$（或 DINOv2）存在 4 倍的网格错配，通过卷积或插值对齐时引入严重的网格伪影；而 UltraDiff 统一在原生 $f32$ 向量场监督，规避了插值弊端。
随机路由的配比至关重要：若完全只跑联合模式（$p_{\text{joint}}=1.0$），$f32$ 分支退化严重（GenEval 仅 0.28）；只有保持 $p_{\text{f32}}=0.4$ 的独立边缘分支训练，才能使 $f32$ 分支在混合模式下逼出 0.74 的极高分数。

亮点与洞察¶

将潜空间知识传递建模为联合分数投影：避开了传统对 VAE 潜在表示做硬性对齐的死胡同，创新地利用双分支联合扩散生成器充当桥梁，通过条件期望在概率测度空间化解了分辨率失真。
高效的解耦训练与渐进式超分微调：仅在 1K 分辨率阶段进行双分支教师蒸馏，到了 2K–4K 阶段平滑切换为纯数据驱动的流匹配，既规避了构建超高清双分支教师的显存爆炸，又让学生模型充分泛化到大尺度画幅。
极低的适配成本：整个适配管线仅耗费 463 GPU-days，仅占直接从头训练低质量 $f32$ 基线计算量（1792 GPU-days）的 26%，展示了极高工程实用性。

局限与展望¶

依赖私有强基座与多阶段训练：论文依赖未开源的高性能私有 $f8$ 模型作为锚点，且管线涉及预训练、教师联合训练、蒸馏与超分缩放四个子阶段，工程链路相对复杂。
VAE 仍保持解耦冻结状态：当前 $f8$ 与 $f32$ VAE 分别独立训练并冻结，未来若与可兼容 VAE（Compatible VAE）联合微调，可能进一步减轻扩散骨干网络的对齐负担。
文本编码与超大容量拓展：当前模型沿用 T5-XXL 编码器，参数量固定在 2.0B，尚未在 >10B 的超大容量架构或多模态 VLM 条件注入下进行极端极限评测。

评分¶

新颖性: ⭐⭐⭐⭐⭐（首创利用双分支联合潜在扩散与原生边际流场蒸馏解决高压缩潜空间失真）
实验充分度: ⭐⭐⭐⭐⭐（包含 1K–4K 分辨率评测、完整的内部解耦对照、严格的跨分辨率消融及显存延迟基准）
写作质量: ⭐⭐⭐⭐⭐（数学推导严谨清晰，系统架构与消融逻辑高度自洽）
价值: ⭐⭐⭐⭐⭐（大幅削减 2K/4K 超高清图像生成的计算门槛，具极高产业与工程落地价值）

分辨率	方法	参数量 (B)	潜空间尺度	延迟 / 采样步数	FID ↓	pFID ↓	Aesthetic ↑	CLIP ↑
\(2048^2\)	Ultrapixel	5.1	\(16 \times \frac{H}{32} \times \frac{W}{32}\)	12s / 20	46.60	43.03	6.24	30.21
\(2048^2\)	Diffusion-4k	12.0	\(16 \times \frac{H}{42} \times \frac{W}{42}\)	23s / 50	37.79	42.43	6.28	29.21
\(2048^2\)	Sana-2K	1.6	\(32 \times \frac{H}{32} \times \frac{W}{32}\)	5.8s / 20	45.99	42.99	6.18	30.13
\(2048^2\)	HunyuanImage-2.1	17.0	\(64 \times \frac{H}{32} \times \frac{W}{32}\)	82s / 50	43.45	48.72	6.44	30.17
\(2048^2\)	UltraDiff (Ours)	2.0	\(64 \times \frac{H}{32} \times \frac{W}{32}\)	6.2s / 20	39.52	42.17	6.40	30.53
\(4096^2\)	Ultrapixel	5.1	\(16 \times \frac{H}{42} \times \frac{W}{42}\)	73s / 20	155.17	57.14	5.91	27.39
\(4096^2\)	Diffusion-4k	12.0	\(16 \times \frac{H}{32} \times \frac{W}{32}\)	122s / 50	162.94	75.70	5.68	27.20
\(4096^2\)	Sana-4K	1.6	\(32 \times \frac{H}{32} \times \frac{W}{32}\)	21.3s / 20	149.20	56.25	6.12	29.97
\(4096^2\)	UltraDiff (Ours)	2.0	\(64 \times \frac{H}{32} \times \frac{W}{32}\)	28.3s / 20	146.60	54.27	6.34	30.14

实验组 / 蒸馏策略	训练范式 / 监督来源	CLIP ↑	ImageReward ↑	GenEval ↑	潜空间 Token 规格
\(f32\) 从头训练	纯随机初始化 + 流匹配	27.21	0.70	0.54	\(H/32 \times W/32\)
\(f32\) 经 \(f8\) 参数初始化	静态权重迁移 + 流匹配	28.91	0.91	0.61	\(H/32 \times W/32\)
外部 DINOv2 特征蒸馏	REPA 特征网格对齐	27.53	0.06	0.33	\(H/32 \times W/32\)
显式 \(f8\) 特征蒸馏	跨网格特征插值拟合	27.87	0.05	0.31	\(H/32 \times W/32\)
UltraDiff 学生模型	MixTeacher 边际分数蒸馏	29.27	0.98	0.72	\(H/32 \times W/32\)
\(f8\) 教师锚点 (Teacher Anchor)	原始高保真密集潜空间	29.25	1.08	0.73	\(H/8 \times W/8\)