← 返回 2026-05-25

PiD：用像素扩散实现快速高分辨率潜变量解码 PiD: Fast and High-Resolution Latent Decoding with Pixel Diffusion

Yifan Lu, Qi Wu, Jay Zhangjie Wu, Zian Wang, Huan Ling, Sanja Fidler, Xuanchi Ren 📅 2026-05-22 👍 46 2026-07-13 08:36
像素空间生成 扩散模型 模型蒸馏 潜变量解码 超分辨率
arXiv HuggingFace PDF

PiD把潜变量解码重定义为像素扩散，统一解码与超分，4步蒸馏即可秒级生成4K图像。

前置知识

潜在扩散模型（LDM）

在 VAE 编码器得到的紧凑潜空间上进行扩散去噪，再由解码器 $\mathcal{D}$ 映射回像素的图像生成范式。训练和推理成本远低于像素扩散，是 Stable Diffusion、FLUX 等系统的标准架构。

本文所有讨论都围绕 LDM 的解码阶段展开，必须先理解为什么要在潜空间生成、为什么需要解码器，才能体会 PiD 想改的是什么环节。

VAE 解码器与重建损失

传统 VAE 解码器 $x_{dec}=\mathcal{D}(z)$ 用像素重建损失（如 L1、感知损失、GAN 损失）训练，目标只是让 $x_{dec}\approx x$，因此只能复原 latent 中已有的信息，不会主动合成新细节。

PiD 论证的关键就是 VAE 解码器"过于重建取向"导致高频细节差、artifact 会被放大，这是理解整篇论文动机和方法的出发点。

Rectified Flow（整流流）

一种把数据 $x_0$ 与噪声 $\epsilon$ 用线性插值 $x_t=t x_0+(1-t)\epsilon$ 连起来训练的流匹配方法，模型预测速度场 $v_\theta(x_t,t,c)\approx x_0-\epsilon$。比 DDPM 收敛更快、采样步数更少。

PiD 在像素空间采用 rectified flow 目标训练扩散先验和条件解码器，并且 DMD2 蒸馏的 sigma 调度 $\{0.999,0.866,0.634,0.342\}$ 也建立在 flow matching 的连续时间视角上。

DMD2 蒸馏

Distribution Matching Distillation v2，把多步扩散模型蒸馏为少步学生模型，通过让学生的输出分布逼近教师，并配合 GAN 判别器和 fake-score 网络实现 1–4 步推理。是 SDXL-Turbo 等快速模型背后的关键技术。

PiD 用 DMD2 把原本需要 50 步的去噪压缩到仅 4 步，速度从 1 秒级降到 210ms 量级，没有它就达不到论文宣传的实时解码。

ControlNet 风格的潜变量注入

借鉴 ControlNet 的思路，给预训练扩散模型加一组可训练卷积残差块，把额外的条件信号（如 canny 边缘、latent）映射成 token 后注入到骨干网络特定层。本文用 2D 卷积把低分辨率 latent 升采样到 patch 网格再注入 PixelDiT。

这是 PiD 实现"latent-conditioned 像素扩散"的工程基础，决定了 latent 信息能否在不破坏预训练 T2I 能力的前提下进入像素扩散主网络。

Representation Autoencoder（RAE）

用预训练视觉编码器（如 DINOv2、SigLIP）替代 VAE 来获得潜变量。RAE 潜变量语义丰富但欠定低层外观，需要解码器补全纹理细节。FLUX.2、Scale-RAE 等 RAE 路线扩散模型都依赖这种解码。

PiD 强调对 RAE 潜变量同样有效，并且提升幅度比 VAE 更大（SigLIP 上 Unipercept-IAA 从 59.95 涨到 64.94），这是论文重要的卖点。

研究动机

当前几乎所有实用的高分辨率文生图系统（FLUX、SD3、Z-Image、自回归图像生成）都依赖一个 VAE 或 RAE 解码器把潜变量还原成像素。但作者指出这套解码范式有三个核心痛点：(1) VAE 编码-解码是有损的，encoder-decoder 重建不完美，必然损失细粒度信息；(2) VAE 解码器只为还原而优化，遇到生成式 latent 中的 artifact 时往往"忠实复原"甚至放大，而不是像超分模型那样合成新细节；(3) 高分辨率下 VAE 解码成本急剧上升，例如表 3 中 FLUX.1 VAE 在 1024px 仅占 2.6 GB，到 2048px 就飙到 37 GB，并在约 2500² 像素时 OOM。实际工程上被迫采用 cascaded 流水线：低分辨率 latent → 低分辨率 VAE 解码 → 扩散式超分 → 高分辨率 VAE 解码，这套流水线在 2048² 下需要 700–2100ms 才能完成一张图，吞吐严重受限。

本文的目标是本文的核心目标是提出一种名为 PiD 的"像素扩散解码器"，把"latent 解码 + 超分放大"两个阶段合并为一个生成式模块，让模型在潜变量 $z$ 和文本 $c$ 的条件下，直接在目标分辨率像素空间做条件扩散去噪。具体目标有三：(1) 在 RTX 5090 消费级显卡上实现 2048×2048 图像 <1 秒解码、GB200 上 210ms；(2) 在视觉质量上全面超过所有 cascaded baseline（含 Real-ESRGAN、SeedVR2、InvSR、TSD-SR 等 SOTA 超分）；(3) 同时适用于 VAE 潜变量（FLUX.1/2、SD3）和 RAE 语义潜变量（DINOv2、SigLIP），并在 LDM 提前终止时也能稳定工作。

与已有工作不同的是，此前已经有一些"扩散式解码"工作（DiVAE、$\epsilon$-VAE、SSDD、DALL-E 3 的扩散解码器），但它们存在三个空白：(1) 受限于重建目标和网络架构，难以扩展到 megapixel 级高分辨率，主流评测还停留在 1024² 以下；(2) 仍然要搭配独立超分阶段才能达到 2K/4K 输出，本质上没解决级联管线的延迟问题；(3) 没有专门为 RAE 这类欠定语义 latent 设计生成能力强的解码器。PiD 的独特切入点是：用一个已经在 2K/4K 像素空间预训练好的 diffusion transformer（PixelDiT 1.3B）作为"像素扩散先验"，再通过 sigma-aware adapter 注入 latent 条件，使整个系统天然支持高分辨率、跨模态 latent 和提前终止。

核心方法

PiD 的整体思路是"用文本到图像的像素扩散先验当解码器"。作者先在 2.6M 高质量 2K 图像上把 PixelDiT（1.3B 参数 MMDiT 像素扩散 Transformer）从 1024² 扩展到 2048²（用 NTK-aware RoPE、shift=6），得到强大 T2I 像素扩散先验 $p_\theta(x_0|c)$。然后冻结 T2I 先验大部分能力，加 ControlNet 风格潜变量投影模块，把低分辨率 latent $z\in\mathbb{R}^{C\times h\times w}$ 上采样到 patch 网格、做卷积特征提取、展平成 token，每隔两个 DiT block 注入一次。注入强度由 sigma-aware gate $g_i=\text{sigmoid}(\text{Linear}_i([h_i,l_i])-\alpha\sigma)$ 控制，使噪声大的 latent 自动被弱化、干净 latent 提供强布局语义。最后用 DMD2 把多步教师蒸馏到 4 步学生，并支持 base LDM 未完全去噪时用 PiD 接续生成（早期终止）。

和已有方法最本质的区别有三层：(1) **生成式先验替代重建式解码器**——传统 VAE 解码器只能从 latent 还原，而 PiD 把"t2i 像素扩散模型"当解码器，让它有能力在解码的同时合成原本不存在的高频细节；(2) **latent-conditioned 像素扩散统一解码与超分**——$\hat{x}_0\sim p^{(s)}_\theta(x_0|z,c)$ 直接建模目标分辨率像素分布，把"解码"和"放大$s$倍超分"合二为一，不需要独立超分阶段；(3) **sigma-aware noisy latent conditioning**——这是论文最核心的工程创新，让同一个解码器既能吃干净 latent（重建任务），也能吃部分去噪的 latent（配合 base LDM 早期终止），注入强度还随噪声水平单调衰减，避免在 latent 不可靠时被错误信息带偏。相比 SSDD 这种单步扩散解码器，PiD 多了 4 倍上采样的超分能力、多了 RAE latent 适配；相比$\epsilon$-VAE，PiD 基于更强的 PixelDiT 像素先验、可扩展到 4K。

方法步骤详情

方法分四阶段。阶段 1 像素扩散先验预训练：输入 2.6M 高分辨率图像，Qwen3-VL-8B-Instruct + LMDeploy TurboMind 生成三档 caption；训练 $\mathcal{L}_{FM}=\mathbb{E}\|v_\theta-(x_0-\epsilon)\|_2^2$，batch 128、lr $2\times 10^{-5}$、20k 步，约 1 天 / 128 H100。阶段 2 条件解码器微调：取 latent 加噪 $\tilde{z}_\sigma=(1-\sigma)z+\sigma\xi$，$\sigma\sim\mathcal{U}(0,0.8)$；适配器为 Conv2d(16→512)+4 残差块，零初始化注入头；batch 64、lr $5\times 10^{-5}$、30k 步。阶段 3 DMD2 蒸馏到 4 步：sigma $\{0.999,0.866,0.634,0.342\}$，约 2h / 128 H100。阶段 4 推理：4 步采样，base LDM 可在 24/28 或 45/50 步停止后接 PiD。

技术新颖性

新颖性体现在四个交叉点。第一，把像素扩散模型的角色从"生成器"重定义为"解码器"，首次系统地把 ControlNet-style adapter 嫁接到 PixelDiT 上成为 latent-conditioned 解码器。第二，sigma-aware gating 把 latent 噪声水平作为先验控制注入权重 $\text{sigmoid}(\text{Linear}([h,l])-\alpha\sigma)$，使单模型同时支持重建、部分去噪、纯生成三种场景，比 SSDD 固定强度 gating 灵活得多。第三，首次在 RAE 语义潜变量上系统验证生成式解码——SigLIP/Scale-RAE 上 Unipercept-IAA 从 59.95 涨到 64.94，DINOv2/DiTDH 从 58.50 涨到 69.81，提升幅度是 VAE 场景的两倍以上。第四，给出 4K 解码端到端 recipe（训 4K T2I 先验 → 注入 adapter → DMD2 蒸馏），附 NTK-aware RoPE、shift=6、Gemma-2-2B-it 等具体配置，可复现性强。

Overview of PiD. PiD unifies latent decoding and upsampling as a single latent-conditioned pixel diffusion model that predicts the target-resolution pixel-space velocity field.
Figure 3: Overview of PiD. PiD unifies latent decoding and upsampling as a single latent-conditioned pixel diffusion model that predicts the target-resolution pixel-space velocity field.
VAE decoding and PiD decoding at different LDM termination steps.
Figure 6: VAE decoding and PiD decoding at different LDM termination steps.

实验结果

在 6 个 latent 设置上评估。核心数字（表 1）：FLUX.1[dev] 上 PiD(24/28) NIQE 从 4.04 降到 3.50、Unipercept-IAA 从 64.12 涨到 66.21；SD3 上 PiD IAA=62.57（基线 57.09）。RAE 场景提升更显著：SigLIP/Scale-RAE 上 PiD(50/50) IAA=64.94（vs 59.95）；DINOv2/DiTDH 上 PiD IAA=69.81（vs 58.50）、IQA=76.52。延迟（表 3）：GB200+torch.compile 下 2048² 仅 208.8ms（vs SeedVR2 1237.5ms），4096² 仅 1927.3ms。MLLM 偏好（图 4）：Gemini-3-Flash 对 PiD vs SeedVR2 给出 90.1% 胜率。蒸馏（表 2）：4 步学生 MUSIQ=73.26，超过 50 步教师（71.79）、25 步（71.63）、4 步教师（68.32），LPIPS 从 0.16 降到 0.09。

Image quality and latency across decoding pipelines.
Table 1: Image quality and latency across decoding pipelines.
Performance of different inference steps of teacher and student model.
Table 2: Performance of different inference steps of teacher and student model.
Ablation study on FLUX.1 [dev] decoding (left) and small-text reconstruction (right).
Table 4: Ablation study on FLUX.1 [dev] decoding (left) and small-text reconstruction (right).
4K decoding results. PiD synthesize more details at 4k resolution.
Figure 2: 4K decoding results. PiD synthesize more details at 4k resolution.
Pairwise image quality preference judged by closed-source MLLMs.
Figure 4: Pairwise image quality preference judged by closed-source MLLMs.
Image reconstruction comparison. Given a latent encoded from a clean image, PiD reconstructs the image at higher resolution with sharper details than the original VAE / RAE decoder.
Figure 5: Image reconstruction comparison. Given a latent encoded from a clean image, PiD reconstructs the image at higher resolution with sharper details than the original VAE / RAE decoder.
Image quality of PiD at different LDM termination step for FLUX.1 [dev] (28 denoising steps in total).
Figure 8: Image quality of PiD at different LDM termination step for FLUX.1 [dev] (28 denoising steps in total).
Comparison with native 2048 × 2048 px generation.
Figure 9: Comparison with native 2048 × 2048 px generation.
查看结构化数据
任务指标本文基线提升
FLUX.1[dev] 潜变量 4× 解码 (PiD(24/28)) MUSIQ (paq2piq) ↑ 73.26 73.40 (VAE+InvSR-1) -0.14, 与最强 baseline 几乎持平，但 PiD 在 NIQE/IAA/IQA/VQ-R1 上全部领先
FLUX.1[dev] 潜变量 4× 解码 NIQE ↓ 3.50 4.04 (SSDD+InvSR-1) 绝对下降 0.54，相对改善 ~13.4%，是衡量自然度的关键指标
FLUX.1[dev] 潜变量 4× 解码 Unipercept-IAA ↑ 66.21 64.12 (VAE+InvSR-1) +2.09，绝对提升约 3.3%，美学评分显著更高
SD3-medium 潜变量 4× 解码 (PiD(24/28)) Unipercept-IAA ↑ 62.57 57.09 (VAE+InvSR-1) +5.48，相对提升 ~9.6%，PiD 优势在 SD3 上更明显
FLUX.2[dev] 潜变量 4× 解码 (PiD(45/50)) Q-Align ↑ 4.70 4.69 (VAE+Real-ESRGAN) +0.01，绝对值微小但 PiD 在所有美学指标上一致领先
DINOv2/DiTDH 潜变量 4× 解码 Unipercept-IAA ↑ 69.81 58.50 (RAE+TSD-SR) +11.31，相对提升 ~19.3%，是 VAE 提升幅度的两倍以上
SigLIP/Scale-RAE 潜变量 8× 解码 Unipercept-IAA ↑ 64.94 59.95 (RAE+TSD-SR) +4.99，相对提升 ~8.3%，验证 PiD 对欠定 RAE latent 的有效性
FLUX.1[dev] 2048² 解码延迟 (GB200+torch.compile) Latency (ms) ↓ 211.2 1237.5 (VAE+SeedVR2) 约 5.86× 加速；同时 PiD 视觉质量更好
RTX 5090 2048² 解码延迟 (eager) Latency (ms) ↓ 1388.8 OOM (FLUX.1 VAE) PiD 仍可在 13 GB 显存内完成，VAE 在该规模直接 OOM
4 步学生 vs 50 步教师 (FLUX.1[dev] 重建) MUSIQ ↑ 73.26 71.79 (Teacher 50 steps) +1.47，蒸馏后少 12× 步数但 IQA 更优

局限与改进

作者承认的局限：(1) 小字重建像素精度不及多步教师——4 步学生 PSNR=24.19 低于 50 步教师的 24.96，作者明确指出学生优先保证感知相似度（LPIPS 从 0.16 降到 0.09），不适合 OCR 类精确字符恢复；(2) 训练算力门槛极高——仅 T2I 先验就需 128 H100 × 1 天，蒸馏又要 128 H100 × 2h，4K 模型更需 96 GB200；(3) 依赖 PixelDiT 这类成熟像素扩散先验，对小团队不友好。我观察到的额外局限：(a) 表 1 延迟数据全部基于 NVIDIA GPU（GB200/H100/RTX 5090），无跨厂商对比；(b) 没有公开权重和训练代码的可下载链接，项目页 https://research.nvidia.com/labs/sil/projects/pid/ 尚未上线权重；(c) 评测完全依赖无参考 IQA 和闭源 MLLM 偏好，没有 FID/IS 等有参考指标对照；(d) 训练数据 2.6M 中 "internally procured images" 是闭源数据，难以 1:1 复现。

独立分析的弱点

独立分析后我认为 PiD 还有几个值得改进的弱点。第一，单模型跨 6 种 latent 需要重新微调——对 DINOv2/SigLIP 必须冻结 PixelDiT 骨干以避免颜色漂移，改进方向是引入 latent 类型 embedding 或 IP-Adapter 式条件归一化。第二，sigma-aware gate 的 $\alpha$ 是全局标量，无法对不同 latent 区域差异化调控，理想做法是 spatial-aware 门控。第三，4 步蒸馏在小字重建上 PSNR 偏低（24.19 vs 24.96），纯 GAN + DMD 损失牺牲高频细节保真度，可加轻量超分 refiner 或 OCR-aware 蒸馏损失。第四，显存占用仍不算轻量——PiD 在 RTX 5090 4096² 需要 9238ms + 22.5 GB，比 Real-ESRGAN 高约一个数量级，可探索 INT8 量化或蒸馏到更小 backbone。第五，PiD 必须和特定 PixelDiT 权重配套，若 LDM 换成 consistency model 或 EDM hybrid 能否无缝衔接还需要验证。

未来方向

作者提到三条未来方向：(1) 把 PiD 推广到 4K 以上的超高分辨率（已展示 4K 能力，主体实验在 2K）；(2) 支持视频 latent 解码，结合时间一致性约束；(3) 与更新的 LDM 路线（如 Consistency Models）结合。基于本文结果我建议的延伸方向：(a) 统一的 multi-latent 解码器——通过 latent-type embedding 让单个 PiD 同时支持 VAE、DINOv2、SigLIP、MAE 等多种 latent，无需重训；(b) 与 autoregressive 图像生成结合——LlamaGen、VAR、MAR 等 AR 模型同样需要解码器，PiD 框架理论上可兼容；(c) 加入多模态控制——扩展为 ControlNet-style 的多条件（depth/edge/pose）输入；(d) 用于图像编辑——把待编辑图的 VAE latent 当条件输入 PiD，按新 prompt 重新生成细节；(e) 蒸馏到更小 backbone（600M 以下）和 INT8/FP8 量化，争取把 2048² 延迟压到 200ms 以下。

复现评估

复现评估整体偏难但比同档工作更友好。积极因素：4.2 节完整列出超参（batch 128/64/16、lr $2\times 10^{-5}/5\times 10^{-5}/1\times 10^{-5}$、sigma $\{0.999,0.866,0.634,0.342\}$、shift=6、$\sigma_{max}=0.8$）；1M 数据来自公开 MultiAspect-4K-1M；caption 用 Qwen3-VL-8B-Instruct + LMDeploy TurboMind、text encoder Gemma-2-2B-it，全部公开。消极因素：未公开权重（项目页 https://research.nvidia.com/labs/sil/projects/pid/ 尚未上线）；"internally procured images" 闭源，无法 1:1 复现训练集；算力门槛极高。思路可借鉴并用小规模数据复现 demo，但要复现 SOTA 数字需 NVIDIA 级别算力和数据资源。

论文图表

PiD directly decodes latents from VAE or vision encoders into higher-resolution images, replacing the decode–then–upsample cascade while achieving lower latency and higher visual quality.
Figure 1: PiD directly decodes latents from VAE or vision encoders into higher-resolution images, replacing the decode–then–upsample cascade while achieving lower latency and higher visual quality.

上半部分横向对比 6 个 1024² 解码结果：FLUX.1[dev] 的 VAE Decode、RAE Decode、PiD Decode（SD3 latent / DINOv2 latent / FLUX.1[dev] latent）、PiD Decode（Z-Image / FLUX.2[dev]），下方对应 VAE Decode 灰度图作对照。下半部分是流水线对比：左侧 Conventional Cascaded Pipeline 为 Low-Res → VAE Dec. → Upsampler DM → VAE Dec. → High-Res Latent → Image；右侧 PiD Pipeline 为 Latent of 512² → PiD 直接输出 2048² Image。

这是论文的"封面图"，同时展示 PiD 的多 latent 通用性、视觉优势、级联 vs 端到端架构差异，是理解全文 motivation 和 overview 的第一张必看图。

Figure 7: PiD vs. cascaded super-resolution.

四行四列对比，每行一个 prompt，每列为 VAE Decode (512²) → +SeedVR2-3B (2048²) → +InvSR (2048²) → +TSD-SR (2048²) → Ours (2048²)。每张图左下角标注延迟：18.25ms / 1237.4ms / 1017.7ms / 724.7ms / 211.2ms。PiD 在保持细节锐度（如发丝、树叶、小字）的同时延迟最低。

和级联 SR baseline 的横向视觉对比，是表 1 数字的图像化补充，直接证明 PiD 不仅延迟低、视觉也好。

Table 3: PiD latency and decoder memory usage across output resolutions, compared with the original VAE decoder.

上半部分是 Latency（RTX 5090 / H100 / GB200 × eager / torch.compile × 256/512/1024/2048/4096 px）；下半部分是 Memory（FLUX.1 VAE vs PiD 在同样硬件和编译模式下）。FLUX.1 VAE 在 4096px OOM，PiD 在 22.5 GB 内完成；H100+compile 下 PiD 4096² 仅 3754.6ms。

提供分辨率维度的延迟/显存扩展曲线，是论文核心卖点"快且省显存"的量化基础。