(2026|NVIDIA,图像生成,超分辨率,流匹配,矫直流,蒸馏)PiD:基于像素扩散的快速高分辨率潜在解码
PiD: Fast and High-Resolution Latent Decoding with Pixel Diffusion
论文地址:https://arxiv.org/abs/2605.23902
项目页面:https://github.com/nv-tlabs/PiD
进 Q 学术交流群:922230617 或加 CV_EDPJ 进 W 交流群
目录
1. 引言
2. 方法:像素扩散解码器
2.1 问题形式化
2.2 以潜码为条件的像素解码器
2.3 模型训练
2.4 快速推理:蒸馏与提前终止
3. 实验
3.1 数据与训练细节
3.2 量化评估设置
3.3 量化结果与性能
1. 引言
目前主流的高分辨率文生图系统(包括潜空间扩散模型 / LDM 和自回归模型)大多在紧凑的潜空间(Latent Space)中完成生成,再通过解码器将潜码(Latent Code)映射回像素。
然而,传统的潜码解码器以重建为导向,为了反演编码器而优化,而不是合成更多细节;并且,随着图像达到兆像素级别,这类方法的计算开销急剧上升。这促使我们需要一种更具表现力、更高效率的解码新范式。
受近期可扩展像素扩散模型的启发,本文提出PiD(Pixel diffusion Decoder,像素扩散解码器),一种将潜码解码重构为条件像素扩散模型的方法,将解码和超分统一在一个生成模块中。
- 具体地,PiD 直接在像素空间执行去噪,能以低延迟合成 4 倍甚至 8 倍放大图像。
- 为了注入潜空间条件,设计了一个轻量级的 sigma 感知适配器,向像素扩散主干注入带噪潜码,使 PiD 能够解码部分去噪潜码,从而提前终止潜扩散过程。
- 进一步地,采用 DMD2 技术进行蒸馏,将推理步数压缩至 4 步。
本文实现了以下几方面创新:
统一解码与上采样:将潜码解码和高分辨率超分融合为一个生成模块,避免了复杂的级联架构。
高效的潜码条件注入:设计了一个轻量级的 sigma 感知适配器,允许解码器适应不同噪声水平的潜码。这项设计让 PiD 可以解码部分去噪潜码,从而提前终止基座 LDM 的生成过程。
广泛兼容性:PiD 不仅适用于 VAE 潜码,也能有效扩展到语义潜码(如 DINOv2 特征),为 RAE 类模型提供强大的高分辨率解码能力。
极致效率与高质量:在 RTX 5090(13 GB 显存)或 GB200 GPU(210 ms)上实现高速解码,且输出图像质量超越级联超分方案。
2. 方法:像素扩散解码器
2.1 问题形式化
设 z 为从文本条件下的自编码器潜空间中采样得到的潜码。
- 传统方法是先用解码器得到低分辨率图像,再经过一个独立上采样模块。
- PiD 则直接以条件像素扩散建模目标分辨率图像分布,本工作中取放大系数 s = 4 或 8。潜码提供整体结构和语义线索,像素扩散先验则在目标分辨率下合成精细细节。
2.2 以潜码为条件的像素解码器
像素空间生成先验:选用 PixelDiT 作为基础架构,从预训练的 1024×1024 像素扩散模型出发,扩展至高分辨率。
从像素先验到潜码解码器:采用 ControlNet 风格的轻量级适配器,将潜码信息注入像素扩散模型,同时保留前阶段文本到图像先验作为强力初始值。
带噪潜码条件:不仅使用干净潜码,还让解码器接触不同噪声水平的潜码。这可以避免解码器过度依赖潜码而抑制细节生成,也支持基座 LDM 提前退出。
注入与 sigma 感知门控:
- 将潜码通过缩放、残差块、展平、线性投影等方式对齐为潜码条件 token,每隔两个主干块进行注入。
- 用一个 sigma 感知门控来控制注入强度,随着潜码噪声水平提高,注入强度减弱。
2.3 模型训练
首先,使用标准的矫直流(rectified-flow)训练高分辨率像素扩散先验。
- 构建有噪图像:
- 模型预测矫直流速度场:
- 通过如下流匹配(Flow Matching)损失优化:
【
(2024|ICLR|德州大学,InstaFlow & 整流流 & 回流 & 蒸馏)InstaFlow:一步就足以实现基于扩散的高质量文本到图像生成
在矫直流中,轨迹是直线(可参考如上文章),所以,
- (把 t 代入第一个公式)对于 t = 0 时刻的纯噪声 𝜖 和 t = 1 时刻的干净图像 x0,
- 整体速度场可视为 (x0 - 𝜖) / (1 - 0) = x0 - 𝜖
- 或者更简单的,第一个公式直接对 t 求导也可得相同结果
】
然后基于预训练像素先验,使用相同的整流流目标,联合微调扩散主干和潜码注入模块。
训练中模型学习在不同潜码噪声水平下平衡重建忠实度与生成合成能力。
2.4 快速推理:蒸馏与提前终止
少步蒸馏:利用 DMD2 将教师模型蒸馏为仅需 4 次采样步长的学生模型,同时将无分类器引导(classifier-free guidance,CFG)蒸馏到学生模型中,推理时无需分开执行条件和无条件的前向传播。
基座 LDM 提前终止:基座 LDM 在完成全部去噪步骤之前即可停止,部分去噪后的潜码噪声水平与训练时带噪潜码条件一致,可直接送入 PiD 解码。
3. 实验
3.1 数据与训练细节
训练数据包含 MultiAspect-4K-1M 等高质量图像,经过 Q-Align 过滤后使用 2.4M 张图像。每张图像标注三种粒度的文本描述:长描述、中等描述和短描述。
像素扩散先验基于 PixelDiT 1.3B 参数模型,使用 128 张 H100 训练约 1 天。潜码注入微调阶段使用 64 张 H100 训练约半天。蒸馏阶段使用 128 张 H100 训练约 2 小时。
3.2 量化评估设置
测试了三种 VAE 潜码(FLUX.1、FLUX.2、SD3)和两种视觉编码器潜码(DINOv2、SigLIP)。
对比基线包括 VAE/RAE 解码器 + 上采样模型,以及 LUA 等潜空间上采样器。
评价指标包含 MUSIQ、NIQE、DEQA、MANIQA、Q-Align、Unipercept(IAA/IQA)和 VisualQuality-R1。
3.3 量化结果与性能
PiD 在 6 种不同的潜码设置下,大部分指标均取得了最佳或次佳结果。
- 在 RAE 类语义潜码(如 SigLIP)上,PiD 的领先优势最为明显,
- MUSIQ 从 73.68 提升至 74.03,DEQA 从 4.00 提升至 4.17,Unipercept-IAA 从 59.95 提升至 64.94。
端到端解码延迟方面,
- PiD 在单张 GB200 GPU 配合
torch.compile后约为210 毫秒。 - 相比之下,最先进的扩散式一步超分基线(如 TSD-SR 和 InvSR)在相同硬件上需要 724–1237 毫秒,即 PiD 比它们快约3–6 倍,同时图像质量更优。
- 更轻量的上采样器(如 Real-ESRGAN)延迟虽低(62 毫秒),但视觉质量明显下降。
PiD 天然支持基座 LDM 提前终止。
- 实验表明,即便基座 LDM 只完成部分去噪步骤,PiD 仍能提供高质量、高分辨率的输出图像,
- 这为速度与质量之间的权衡提供了灵活的选择空间。
