Hi-DiT: Hybrid Latent-Pixel Diffusion Transformer for Image Generation¶
会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/HiDream-ai/Hi-DiT
领域: 图像生成
关键词: 扩散模型, Diffusion Transformer, 混合潜空间像素空间, 时间门控注入, 高频细节恢复
一句话总结¶
Hi-DiT 在单一参数共享的 Diffusion Transformer 内统一了潜空间与像素空间扩散,依托去噪时序异质性在低噪声阶段通过时间门控动态激活像素分支与分层子像素预测头,在 ImageNet 与 MS-COCO 上取得 SOTA 的生成质量与高频细节。
研究背景与动机¶
当前扩散模型在生成效率与细节保真度之间面临着难以调和的范式冲突。主流的潜空间扩散模型(LDMs)通过预训练的变分自编码器(VAE)将高维像素压缩到紧凑的低维空间,从而大幅降低计算开销并提升训练稳定性;然而,固定的空间下采样(如 16 倍压缩)与通道瓶颈不可避免地丢弃了高频纹理与边缘信息,导致 VAE 解码后的图像存在高频细节模糊与伪影,生成质量受制于自编码器的重建上限。与之相对,直接在原始像素空间建模的扩散模型虽然免除了 VAE 压缩带来的信息损耗,但必须在超高维、高熵的像素空间中同时学习低频全局结构与复杂高频变化,引发严重的容量竞争与优化难度,在同等算力下收敛缓慢且生成质量往往落后于潜空间模型。
这种困境的物理根源在于去噪过程的时序异质性(Temporal Heterogeneity):在扩散过程的早期大噪声阶段,网络主要负责建立粗粒度的全局布局与语义结构;而在晚期低噪声阶段,去噪重点则完全转向微观纹理与高频细节的精细合成。现有架构采用时间同质的单一去噪器,强行用同一组网络参数在全时序范围处理截然不同的生成目标,造成了明显的目标干扰与容量浪费。
本文的切入角度是:与其在纯潜空间或纯像素空间非此即彼,不如让紧凑的潜空间流在早期专注全局语义规划,并将像素级的建模能力精确留给晚期的细节合成。核心 idea:构建参数共享的双流 Diffusion Transformer(Hi-DiT),依据去噪时序异质性引入时间门控机制,在低噪声阶段动态激活像素通路与分层子像素预测头,实现全局低频结构与局部高频细节的时序解耦生成。
方法详解¶
整体框架¶
Hi-DiT 的核心思想是在同一个参数共享的 Transformer 骨干网络中协调潜空间流(Latent Stream)与像素空间流(Pixel Stream)。输入包含扩散时间步 \(t\)、条件信息 \(c\)(类别或文本提示词),以及对应的潜变量 \(z_t\) 与加噪像素图像 \(x_t\)。在时序演进上,去噪过程被划分为两个明确阶段:在大噪声阶段(\(t \ge \tau\)),模型仅运行潜空间流,通过 VAE 潜变量以低计算开销建立图像的全局骨架与语义轮廓;在跨过时间门控阈值进入低噪声阶段(\(t < \tau\))后,像素分支被动态激活,通过解码当前的潜变量状态作为像素流脚手架,与潜变量 token 协同输入共享的 Transformer 中进行联合多模态特征交互。最终,潜空间预测头估计干净潜变量经由 VAE 解码出低频主体结构,像素预测头输出高频精细分量,二者线性相加合成最终的高分辨率图像。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入: 噪声潜变量 z_t 与条件 c"] --> B["双流共享骨干与时序分工<br/>潜空间流全时序构建全局语义骨架"]
B --> C{"时间门控注入机制<br/>判断去噪阶段 t < tau"}
C -->|高噪声 t >= tau| D["纯潜空间流动速度预测"]
C -->|低噪声 t < tau| E["激活像素分支与加噪图像 Patch 编码"]
E --> F["渐进式子像素预测头<br/>PixelShuffle 两次上采样解算高频残差"]
D --> G["输出: VAE 解码结构"]
F --> H["融合层: VAE 结构底模 + 高频细节残差"]
H --> I["输出: 最终高保真图像 x_final"]
关键设计¶
1. 双流共享骨干与时序分工:解耦宏观语义与微观纹理
为了在不额外成倍增加模型参数量的前提下实现多域特征协同,Hi-DiT 采用了参数完全共享的 Diffusion Transformer 骨干结构。潜变量序列通过 Latent Embedder 投影为特征向量 \(h_z^0\),在全时序范围内学习整图的低频宏观几何形态;当像素流开启时,加噪像素图像 \(x_t\) 通过瓶颈 Patch 编码器(Bottleneck Patch Embedding)以大小为 \(p \times p\) 的图像块划分为像素特征向量 \(h_x^0\)。为了保证空间严格对齐并降低 Transformer 交互的复杂度,Patch 大小 \(p\) 精确设置为与 VAE 的下采样率一致(如 \(p=16\)),使得像素序列与潜变量序列保持完全相同的序列长度。在骨干内部,共享的自注意力与前馈层自然完成了全局结构与局部纹理之间的跨域上下文交互,既免除了双塔独立模型带来的显存翻倍与参数冗余,又赋予了模型兼顾结构先验与像素保真度的能力。
2. 时间门控注入机制:消除全时序容量竞争
常规直接拼接像素与潜变量的朴素方案会在高噪声阶段让网络承受巨大的高频拟合压力,重蹈纯像素空间扩散收敛困难的覆辙。Hi-DiT 提出了硬阈值调度的时间门控机制(Time-Gated Injection),设计门控函数 \(\mathcal{G}(t) = \mathbf{1}(t < \tau)\),将输入 token 构造为: $\(h_{zx}^0 = h_z^0 + \mathcal{G}(t) \cdot h_x^0\)$ 在训练与推理中,阈值固定为 \(\tau = 0.3\)。在 \(t \ge \tau\) 的大噪声阶段,\(\mathcal{G}(t) = 0\),像素流完全静默,算力百分之百聚焦于流匹配速度场 \(v_\theta^z(z_t, t, c)\) 的稳定学习;当 \(t\) 越过阈值(\(t < \tau\))时,门控开启,网络在第 \(t_0\) 步首先调用 VAE 解码器解算当前中间态 \(x^{\text{dec}}_{t_0} = \mathcal{D}(z_{t_0})\) 初始化像素分支,此后每一步通过当前解码图像叠加高频预测量回传更新。这种机制在概念上契合课程学习(Curriculum Learning),先学宏观构图再雕琢微观质感,彻底化解了去噪前期的目标冲突。
3. 渐进式子像素预测头:降低高分辨率回归负担
纯像素扩散模型以往多使用单个线性层从单个 token 直接回归 \(p \times p \times 3\) 的大尺寸像素块,这种极端的单步高维映射欠缺空间局部归纳偏置,容易产生棋盘格伪影。Hi-DiT 受超分辨率网络和 VAE 解码器结构的启发,构建了基于层次化子像素解码的高频像素预测器(High-Frequency Pixel Predictor)。该模块将 Transformer 输出的隐藏状态 \(h\) 重排为 2D 特征图 \(\mathbf{G}_0\),连续穿过两个包含卷积、PixelShuffle(上采样因子 2)、GroupNorm 与 SiLU 的细化块: $\(\mathbf{G}_i = \text{SiLU}(\text{GroupNorm}(\mathcal{PS}(\text{Conv}(\mathbf{G}_{i-1}))))\)$ 两阶段上采样使空间分辨率扩展 4 倍得到稠密网格 \(\mathbf{G}_2\)。经过包含时间步与条件向量的 AdaLN 调制后,末端线性层仅需在每个局部位置回归微小的 \(\frac{p}{4} \times \frac{p}{4}\) 子块。这一级联设计显著减轻了高频拟合负担,以仅 7M 的极轻量参数增量换取了对锐利边缘和真实质感的高精度恢复。
一个完整示例¶
以 ImageNet \(256 \times 256\) 条件图像生成为例,输入图像先通过下采样倍率 16 的 VAE 压缩为 \(16 \times 16 \times 32\) 的潜变量 \(z\)。 - 高噪声阶段(\(t=1.0 \to 0.3\)):系统在潜空间运行 flow matching ODE 采样。门控 \(\mathcal{G}(t)=0\),只有 256 个潜变量 token 通过 SiT 骨干网络,预测潜空间速度场,逐步勾勒出物体的大体朝向、颜色区域与粗粒度轮廓。 - 门控触发瞬间(\(t = 0.3\) 临界点):当前带有一定高斯噪声与微弱瑕疵的潜变量 \(z_{0.3}\) 被临时送入 VAE 解码器,输出初步脚手架图像 \(x^{\text{dec}}_{0.3} \in \mathbb{R}^{256 \times 256 \times 3}\)。 - 低噪声阶段(\(t=0.3 \to 0.0\)):门控打开(\(\mathcal{G}(t)=1\))。\(x_t\) 被划分为 \(16 \times 16\) 个 \(16 \times 16\) 大小的 patch,经瓶颈嵌入得到 256 个像素 token,与 256 个潜变量 token 在共享 Transformer 中进行联合自注意力计算。在输出端,潜空间头预估整体结构并经 VAE 解码为 \(\mathcal{D}(\hat{z}_0)\),子像素预测头则逐层将 \(16 \times 16\) 的 token 图上采样至 \(64 \times 64\),最终预测每个 \(4 \times 4\) 小块的高频残差 \(x_{\text{HF}}\)。两者叠加形成即时像素更新并迭代至 \(t=0\),输出羽毛、毛发丝丝分明的高分辨率真实图像。
损失函数 / 训练策略¶
Hi-DiT 采用联合多任务优化目标。在全时序范围内,对潜空间流使用 Rectified Flow Matching 损失监督速度场: $\(\mathcal{L}_{\text{lat}} = \mathbb{E}_{t, c, z_0, \epsilon} \left\| v_\theta^z(z_t, t, c) - (\epsilon - z_0) \right\|_2^2\)$ 在低噪声阶段(\(t < \tau\)),激活针对最终合成图像 \(x_{\text{final}} = \mathcal{D}(\hat{z}_0) + x_{\text{HF}}\) 的复合像素级损失 \(\mathcal{L}_{\text{pix}}\): $\(\mathcal{L}_{\text{pix}} = \lambda_{\text{rec}} \left\| x_{\text{final}} - x_0 \right\|_1 + \lambda_{\text{per}} \text{LPIPS}(x_{\text{final}}, x_0)\)$ 整体训练损失函数为: $\(\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{lat}} + \mathbf{1}(t < \tau) \cdot \mathcal{L}_{\text{pix}}\)$ 模型训练分为两阶段:首先在 ImageNet 上使用 DINOv2 特征对齐预训练 SiT 潜空间骨干,随后载入预训练权重并开启双流进行端到端联合微调。
实验关键数据¶
主实验¶
在 ImageNet \(256 \times 256\)、ImageNet \(512 \times 512\) 和 MS-COCO 文本生成图像基准上,Hi-DiT 均展现了领先的指标表现。
| 数据集 | 模型 | 范式类别 | 参数量 | FID ↓ | IS ↑ | Precision ↑ | Recall ↑ |
|---|---|---|---|---|---|---|---|
| ImageNet 256×256 | VAR (Tian et al.) | 自回归 (AR) | 600M | 2.57 | 302.6 | 0.83 | 0.56 |
| ImageNet 256×256 | DiT-XL (Peebles et al.) | 潜空间扩散 | 675M | 2.27 | 278.2 | 0.83 | 0.57 |
| ImageNet 256×256 | SiT-XL (Ma et al.) | 潜空间扩散 | 675M | 2.06 | 270.3 | 0.82 | 0.59 |
| ImageNet 256×256 | VA-VAE (Yao et al.) | 潜空间扩散 | 675M | 1.35 | 295.3 | 0.79 | 0.65 |
| ImageNet 256×256 | REPA (Yu et al.) | 潜空间扩散 | 675M | 1.29 | 306.3 | 0.79 | 0.65 |
| ImageNet 256×256 | PixelFlow (Chen et al.) | 像素空间扩散 | 677M | 1.98 | 282.1 | 0.81 | 0.60 |
| ImageNet 256×256 | Deco (Ma et al.) | 像素空间扩散 | 682M | 1.62 | 301.0 | 0.80 | 0.62 |
| ImageNet 256×256 | JiT-G (Li & He) | 像素空间扩散 | 2B | 1.82 | 292.6 | 0.79 | 0.62 |
| ImageNet 256×256 | Hi-DiT (CFG=2.4) | 混合潜空间+像素 | 689M | 1.06 | 296.6 | 0.79 | 0.66 |
| ImageNet 256×256 | Hi-DiT (CFG=3.0) | 混合潜空间+像素 | 689M | 1.10 | 319.4 | 0.79 | 0.65 |
| ImageNet 512×512 | DiT-XL | 潜空间扩散 | 675M | 3.04 | 240.8 | 0.84 | 0.54 |
| ImageNet 512×512 | REPA | 潜空间扩散 | 675M | 2.08 | 274.6 | 0.83 | 0.58 |
| ImageNet 512×512 | JiT-G | 像素空间扩散 | 2B | 1.78 | 306.8 | - | - |
| ImageNet 512×512 | Hi-DiT (CFG=4.6) | 混合潜空间+像素 | 689M | 1.26 | 295.0 | 0.80 | 0.62 |
| MS-COCO 256×256 | Hi-DiT (U-ViT-S) | 混合潜空间+像素 | - | 5.15 | - | - | - |
注:在无 CFG 引导的原生采样下,Hi-DiT 在 ImageNet 256×256 达到 1.66 FID,远超同配置 DiT-XL(9.62)与 VA-VAE(2.17)。MS-COCO 上对比 baseline U-ViT-S/2 的 5.48 FID 降低 0.33。
消融实验¶
消融实验基于 ImageNet 256×256 统一从头训练 80 个 epoch 进行公平横向对比:
| 配置 | 潜变量输入 | 像素预测头类型 | 参数量 | FID ↓ | 说明 |
|---|---|---|---|---|---|
| 纯像素基线 | ✗ | 朴素 MLP | 682M | 3.38 | 缺乏潜空间宏观语义支撑,收敛困难 |
| 混合流 + 线性预测 | ✓ | 朴素 MLP | 682M | 1.74 | 引入潜变量脚手架后生成质量大幅攀升 |
| Hi-DiT 完整模型 | ✓ | 分层高频像素预测头 | 689M | 1.65 | 分层子像素预测头进一步补足细腻纹理 |
时间门控阈值 \(\tau\) 敏感性对比(80 epochs): - \(\tau = 0.7\):FID 为 1.73,过早激活像素流重新引发大噪声阶段的容量竞争。 - \(\tau = 0.5\):FID 为 1.68,性能平稳但仍有优化冗余。 - \(\tau = 0.3\):FID 达到最优的 1.65(IS 267.7),全局结构与细节构建达到最佳平衡。 - \(\tau = 0.1\):FID 退化至 1.75,像素流介入时间过晚,未能充分提炼高频分量。
关键发现¶
- 潜变量脚手架是像素扩散的催化剂:单纯从 patch 级像素直接去噪(FID 3.38)难度极高,而一旦加入潜变量流(FID 降至 1.74),模型优化难度骤降,证实了高层语义对底层纹理合成具有关键的锚定作用。
- 子像素预测头显著优于单步投影:将 682M 的 MLP 预测头替换为仅增加 7M 参数的分层子像素预测头(689M),FID 在 80 epoch 下从 1.74 进一步改善至 1.65,证明空间递进放大策略极大降低了回归高频信息的难度。
- 微小计算开销换取极致质量:Hi-DiT 单图推理耗时为 1.67 秒(对比 SiT 的 1.52 秒和 VA-VAE 的 1.58 秒),训练耗时仅从 17.7 分钟/epoch 略增至 18.1 分钟/epoch,却将 FID 推进至破纪录的 1.06。
亮点与洞察¶
- 时序分工化解空间与频率矛盾:没有盲目堆叠双套独立网络,而是抓住扩散过程前期重结构、后期重纹理的本质规律,在参数共享骨干内用硬门控激活像素流,设计极度克制而有效。
- 对齐 patch 与潜变量分辨率:巧妙地将像素嵌入的 patch 大小设定为与 VAE 下采样率相同(\(p=16\)),让像素序列与潜变量序列在共享注意力中完美对齐,规避了 cross-attention 的复杂映射。
- 可复用的高频预测范式:分层卷积 + PixelShuffle 的轻量化预测头为所有基于 patch 的视觉生成模型提供了一种即插即用的高频残差恢复范式。
局限与展望¶
- 作者指出的局限:在低噪声阶段引入像素流使得晚期采样步骤的计算复杂度略有增加,且像素分支对显存峰值有小幅提升(从 32.9G 增至 35.5G)。
- 实验与机制局限:硬门控阈值 \(\tau=0.3\) 目前为启发式经验参数,未实现针对不同图像复杂度或文本描述长度的自适应门控;此外,初始像素输入仍依赖 VAE 解码中间态,在极度模糊状态下可能引入微小结构偏差。
- 未来改进方向:可探索平滑连续的时间衰减软门控函数,并将混合双流架构扩展到视频扩散生成与高分辨率多模态自回归统一模型中。
相关工作与启发¶
- vs VA-VAE / REPA (潜空间优化): VA-VAE 和 REPA 试图通过特征对齐(如 DINOv2)把判别性表征强行塞入 VAE 潜空间,但其上限始终受限于 VAE 解码器的固有频带缺失。Hi-DiT 允许潜空间专注宏观语义,直接在外部开辟无损像素旁路补足细节,思路更加彻底。
- vs JiT / PixelFlow (纯像素空间扩散): JiT-G 强行使用 20 亿(2B)参数的大模型在纯像素空间硬拟合去噪,FID 仅有 1.82 且耗费海量算力。Hi-DiT 仅以 689M 参数取得 1.06 FID,证明结合潜空间先验能实现降维打击。
评分¶
- 新颖性: 4.5 / 5.0(精妙利用去噪时序异质性解耦潜空间与像素空间扩散,双流参数共享设计兼顾优雅与高效)
- 实验充分度: 5.0 / 5.0(涵盖 ImageNet 256/512 与 COCO,详实的 CFG 曲线扫掠、跨范式全景对比与耗时显存测量)
- 写作质量: 4.8 / 5.0(逻辑严密,图表清晰,从频域与时序角度剖析动机非常透彻)
- 价值: 4.8 / 5.0(为破解 VAE 模糊瓶颈和像素扩散训练困难提供了高落地可行性的新一代 DiT 架构参考)