PixRestore：基于像素扩散Transformer的统一图像修复 PixRestore: Unified Image Restoration via Pixel Diffusion Transformer
免VAE像素扩散Transformer统一修复八种图像退化，单步推理仅44ms
前置知识
流匹配（Flow Matching）
流匹配是扩散生成建模的连续时间形式：模型学习一个随时间变化的速度场 $v_t$，把高斯噪声分布沿常微分方程（ODE）积分到数据分布。本文采用最简单的线性插值路径 $x_t=(1-t)y_{hq}+t\epsilon$，训练时让网络预测速度 $v_t=(x_t-y_{hq})/t$，推理时用一阶 Euler 求解器从纯噪声积分得到干净图像。相比 DDPM 的随机采样，ODE 视角下采样步数可以大幅压缩。
PixRestore 的训练目标就是像素空间条件流匹配损失 $\mathcal{L}_{flow}=\|\hat v_t-v_t\|_2^2$，单步生成器的构造方式（固定 $t=1$）也直接来自该框架，不理解流匹配就看不懂方法主体。
像素空间扩散 vs 潜空间扩散（VAE）
主流文生图模型（SDXL、FLUX 等）先训练变分自编码器 VAE 把图像压缩成低分辨率潜码，再在潜空间做扩散以节省计算。代价是 VAE 的有损压缩会丢掉细小结构和纹理。像素空间扩散直接在 RGB 像素上建模，本文用 patch 大小 8 把像素网格分词，使 token 序列长度与 f8 的 VAE 相当，从而绕开 VAE 的信息瓶颈。
本文的核心主张就是免 VAE 的像素空间建模对修复任务更好，Table 1 用同一 DiT-S 骨干在两种空间做了严格对照实验，这是理解全文动机与结论的关键背景。
DiT（Diffusion Transformer）
DiT 用 Transformer 替换传统扩散模型的 U-Net 骨干：图像被 patch 化为 token 序列，堆叠带自注意力、交叉注意力和前馈网络的块。本文的块采用 RMSNorm、QK 归一化注意力与旋转位置编码 RoPE，并用一个共享 timestep 块为所有块生成 AdaLN 调制参数。DiT 的优势是扩展性好，性能随参数量与计算量平滑增长。
PixRestore 的骨干是 LightningDiT，S/B/L/XL 四个变体由 DiT 规模决定，论文的 scaling 实验（Figure 6）也建立在 DiT 可扩展这一前提上。
DINOv2 自监督视觉基础模型
DINOv2 通过自蒸馏在大规模无标注图像上预训练，其多层特征从浅到深携带互补信息：浅层保留边缘、纹理等局部结构，深层编码全局语义。本文冻结 DINOv2，从均匀分布的 6 层抽取特征 $U_l$ 作为修复的条件与监督信号，并用 LQ–HQ 特征余弦相似度 $s_l$ 度量每层在不同退化下的可靠程度。
自适应分层视觉引导是本文最核心的创新：层路由器以 $s_l$ 为监督学习逐图权重，实现可靠层做条件、不可靠层做监督的机制，是全文方法的灵魂。
修复评价指标（PSNR/LPIPS/MUSIQ/DR-Score）
PSNR/SSIM 是基于像素与结构相似度的全参考指标，衡量保真度；LPIPS/DISTS 基于深度特征衡量感知质量；MUSIQ、AFINE-NR 是无参考（NR）质量指标，不需要参考图。论文发现 NR 指标无法反映退化是否真的被去除，于是提出 DR-Score：用 VLM（Gemini-3.1 Pro）判断目标退化是否被移除，每张图测 5 次取平均。
读懂实验需要理解各指标含义：PixRestore 在 PSNR/LPIPS/DR-Score 上领先，但在部分 NR 指标上反而吃亏，这是论文反复论证的一个关键现象。
研究动机
统一图像修复（UIR）要求单个模型同时处理噪声、模糊、雨纹、雨滴、雾、雪、低光和超分辨率等多种经常共存的退化。早期方案如 PromptIR 等回归式模型用确定性 $\mathcal{L}_1/\mathcal{L}_2$ 损失训练，倾向于输出条件均值，结果过度平滑、高频细节被抑制，且任务级提示或专家路由在复杂真实退化上泛化差。近年的主流做法转向微调大规模文生图（T2I）潜空间扩散模型：FoundIR-v2 在 SDXL 上加 MoE 路由和 MLLM 描述器，FLUX-IR 微调 FLUX，FAPE-IR 还要借助 Qwen2-VL 和 SigLIP 做辅助适配。论文指出这条路线有三大代价：其一，VAE 潜码压缩是有损的，会丢弃修复最需要保留的小结构、文字笔画和锐利边缘；其二，开放式合成先验与忠实还原输入内容的目标错位，容易生成视觉合理但与输入不一致的内容；其三，计算冗余严重——Table 3 显示 FoundIR-v2 有 16910M 参数、18293ms 延迟，FAPE-IR 达 21575M 参数，Flux-IR 需要 795290 GFLOPs，远超修复任务的实际需要。
本文的目标是本文的目标是构建一个不依赖 VAE、不依赖 T2I 预训练的像素空间扩散 Transformer，用单一模型统一处理去模糊、去雾、去噪、去雨纹、去雨滴、去雪、低光增强和超分辨率八种退化，在保真度、感知质量、退化去除能力和推理效率之间取得最佳平衡。具体量化目标包括：模型保持紧凑（主变体约 50M 参数，含冻结 DINO 编码器共 53.7M），推理压缩到单步（NFE=1，A800 上约 44ms），同时在 15 个公开基准和真实世界无 GT 测试集上达到与数十亿参数 T2I 方法相当或更好的整体性能；并希望架构具备可扩展性，能通过加大骨干或缩小 patch 进一步提升质量。为此论文配套构建了约 283 万张图像、覆盖八类退化的训练语料，并提出用 VLM 评估退化去除程度的 DR-Score 作为辅助诊断指标，弥补现有无参考指标不反映退化去除的缺陷。
与已有工作不同的是，论文的独特切入点来自一次重新审视：文生图模型从文本线索合成全新图像，需要极强的开放式生成能力；而修复任务从 LQ 图像出发，输入本身已包含丰富的视觉线索，输出与输入逐像素对齐，因此需要的不是更开放的生成先验，而是对退化的鲁棒性和对像素对齐细节的忠实重建。据此作者彻底放弃 T2I 预训练和 VAE，从零训练像素空间 DiT，用 patch 化控制 token 长度。第二个独特角度是对 DINO 分层特性的系统性观察：不同层携带互补线索（浅层保结构、深层管语义），且各层对不同退化的敏感度不同——浅层对雨、模糊、雪敏感，深层对噪声敏感，中层对低光、超分、雾敏感，因此固定的单层特征或均匀平均都是次优的。作者据此提出按 LQ–HQ 相似度训练轻量层路由器，实现逐图自适应加权，这在此前的修复原生扩散方法（DiffUIR、DA-CLIP、FoundIR）和 T2I 适配方法中都没有出现过。
核心方法
PixRestore 把 UIR 形式化为像素空间上的条件流匹配问题。直觉上：与其把图像压进有损潜码再扩散，不如直接在 patch 化的像素上建模条件流；与其用单一全局退化表示，不如让不同 DINO 层按其对当前退化的可靠程度自适应地提供条件和监督。输入是 LQ 图像 $y_{lq}$ 与带噪状态 $x_t$ 沿通道拼接的 6 通道张量，经 patch 大小为 8 的 patch embedding 变成 token 序列，送入 $N$ 个 DiT 块（含 RMSNorm、RoPE 与共享 timestep 块生成的 AdaLN 调制），每个块通过交叉注意力注入由层路由器融合的 DINO 特征 $U_{fuse}$，网络直接预测干净图像 $\hat y_{hq}$。训练分两阶段：先在约 283 万张图像的语料上等概率采样八类退化训练多步模型 250K 步，再固定 $t=1$、加 DINO 特征域对抗损失微调成单步生成器 100K 步。最终从纯噪声一步还原图像，主变体 PixRestore-S 含冻结编码器共 53.7M 参数、658 GFLOPs、44ms 延迟。
核心创新有三点。第一，免 VAE、免 T2I 的像素空间流匹配：Table 1 在完全相同的训练/测试设置下对比同一 DiT-S，像素版（patch 8）取得 26.62dB PSNR、0.1593 LPIPS、54.32 MUSIQ，全面超过 SD2VAE、FluxVAE 与 QwenVAE（最好的为 22.80dB）三个潜空间版本，直接证明有损潜码瓶颈是 T2I 适配路线的硬伤。第二，相似度引导的自适应分层视觉引导：训练时用余弦相似度与归一化 $L_2$ 距离相似度的平均度量 LQ–HQ 各层特征相似度 $s_l$，softmax 得到 $q_l=\exp(s_l)/\sum_k\exp(s_k)$ 作为伪标签训练轻量预测器；可靠层特征融合成稠密条件注入每个 DiT 块，不可靠层按互补权重 $r_l=\exp(1-s_l)/\sum_k\exp(1-s_k)$ 接受更强的 HQ 特征监督——条件选可靠内容、监督盯退化重的层。第三，把对抗判别器建在冻结 DINO token 而非原始像素上，与条件编码共享先验、开销极小，是单步化的关键助推。
方法步骤详情
分五步。① 特征提取：冻结 DINOv2 从 LQ 图像均匀抽取 6 层特征 $U_l$。② 路由器：以 LQ–HQ 层相似度 softmax 得伪标签 $q_l$，预测器 $\rho_\psi$ 从 LQ 预测权重 $p_l$，交叉熵监督。③ 融合训练：$U_{fuse}=\sum_l p_l U_l$ 经交叉注意力注入各 DiT 块；$\mathcal{L}_{feat}$ 按权重 $r_l$ 约束 HQ 层特征；总损失 $\mathcal{L}=\mathcal{L}_{flow}+0.5\mathcal{L}_{wpred}+0.5\mathcal{L}_{feat}$。④ 单步微调：固定 $t=1$，在 DINO token 上训逐层判别器 $D$，加对抗项 $0.5\mathcal{L}_{adv}$。⑤ 推理：噪声 $\epsilon$ 与 $y_{lq}$ 拼接，单次前向输出 $\hat y_{hq}$（44ms@A800）。训练 AdamW、lr $1\times10^{-4}$、8×A800，250K+100K 步。
技术新颖性
从技术新颖性看，本文的价值在于几个第一次的系统性组合而非单点炫技。首先，它是把 DiT 从零训练在像素空间做统一修复的早期尝试：既有像素生成工作（如 JiT）面向文生图，既有 UIR 工作或基于回归、或适配 T2I 潜空间模型，本文证明修复这一特定任务反而适合同时抛弃两者。其次，相似度引导的层路由机制是全新的：以 LQ–HQ DINO 相似度 $s_l$ 为免费监督信号，推理时只需 LQ 即可预测权重，把哪层特征可信变成可学习的逐图决策，并通过 $q_l$/$r_l$ 互补加权同时驱动条件融合与特征监督；消融 A6（均匀监督）对比 A7（自适应监督）验证了其价值：LPIPS 从 0.1239 降至 0.1209，PSNR 从 27.36 升至 27.66。第三，判别器建在冻结 DINO token 上，与条件编码共享先验、开销极小，使单步微调把 PSNR 从 A9 的 28.07 提到 A10 的 28.49、MUSIQ 从 53.34 拉回 55.52。此外，DR-Score 作为 VLM 驱动的退化去除度量也是方法学上的附加贡献。
实验结果
空间对比（Table 1）：像素版 26.62dB/0.1593/54.32 胜潜空间最优（QwenVAE）22.80dB/0.2181/51.87。基准（Table 2）：去雨纹 32.28dB/0.0902 LPIPS 胜 FAPE-IR*（31.91dB/0.0903），去噪 34.87dB、去模糊 28.32dB 领先；B 变体多数居前二。效率（Table 3）：53.7M 参数、658 GFLOPs、44ms，快扩散基线 7–23 倍，计算量为 FoundIR-v2 的 1/181、Flux-IR 的 1/1209。消融（Table 4）：A0（26.62dB）加 DINO 条件、分层监督、自适应引导与单步微调后，A10 达 28.49dB/0.1120/55.52；NFE=1 时 PSNR 反升至 28.07。范式（Table 5）：流预训练+单步微调 28.49dB 胜同预算回归 27.00dB。真实世界（Table 6）：PixRestore-B 平均 DR-Score 67.88 最佳，超 FAPE-IR*（67.55）；Figure 6 证实 -0.96 平滑缩放。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 像素空间 vs 潜空间（同设置 DiT-S 对照） | PSNR / LPIPS / MUSIQ | Pixel DiT-S：26.62dB / 0.1593 / 54.32（23.41M 参数，无 VAE） | Latent DiT-S + QwenVAE：22.80dB / 0.2181 / 51.87（SD2VAE 22.10dB、FluxVAE 22.63dB） | PSNR +3.82dB，LPIPS 降低 0.0588，且省去 41–78ms VAE 开销 |
| 去雨纹（RainDS-real / RealRain-1k） | PSNR / LPIPS / DR-Score | PixRestore：32.28dB / 0.0902 / 82.75（PixRestore-B：32.85dB / 0.0767 / 84.72） | FAPE-IR*（重训）：31.91dB / 0.0903 / 84.64 | PSNR +0.37dB，LPIPS 持平略优，DR-Score 与最强基线相当 |
| 去雪（WeatherBench） | PSNR / LPIPS | PixRestore：31.26dB / 0.0853（PixRestore-B：31.93dB / 0.0688） | FAPE-IR*（重训）：30.19dB / 0.1136 | PSNR +1.07dB，LPIPS 降低 0.0283 |
| 去噪（DIV2K / PolyU） | PSNR / LPIPS / DISTS | PixRestore：34.87dB / 0.0624 / 0.0735 | FAPE-IR*（重训）：34.22dB / 0.0741 / 0.0750 | PSNR +0.65dB，LPIPS 降低 0.0117 |
| 去模糊（GoPro / UHD-blur） | PSNR | PixRestore：28.32dB（PixRestore-B：29.23dB） | FAPE-IR*（重训）：27.96dB；FoundIR*：29.14dB | 超 FAPE-IR* +0.36dB，B 变体再 +0.09dB 超过所有基线 |
| 真实世界无 GT 测试（6 类退化各 100 张） | 平均 DR-Score | PixRestore：65.20（PixRestore-B：67.88，全场最佳） | FAPE-IR*（重训）：67.55；FoundIR-v2*：63.05；PromptIR：29.09 | B 变体超最强重训基线 +0.33，且推理成本低两个数量级 |
| 模型效率（512×512，A800） | 参数量 / FLOPs / 延迟 | PixRestore：53.70M / 658G / 44ms（单步 NFE=1） | FAPE-IR：21575M / 64758G / 1011ms；FoundIR-v2：16910M / 119334G / 18293ms | 比多数扩散基线快 7–23 倍，FLOPs 为 FoundIR-v2 的 1/181、Flux-IR 的 1/1209 |
| 训练范式对比（同骨干同迭代数） | PSNR / SSIM / LPIPS / MUSIQ | 流预训练 250K + 单步微调 100K：28.49dB / 0.8589 / 0.1120 / 55.52 | 直接回归训练 350K：27.00dB / 0.8179 / 0.1494 / 52.00 | PSNR +1.49dB，证明收益来自流匹配预训练阶段本身 |
局限与改进
作者明确承认三条局限：一是 DiT 架构与固定 token 化导致分辨率不灵活，某一分辨率训练的模型不能直接推广到更高分辨率，需要重训练或修改架构——而真实照片常为 4K 甚至更高；二是与十亿级 T2I 模型相比，紧凑骨干（S 变体约 50M）在极端信息稀缺的退化下可能力不从心；三是 DR-Score 依赖专有 VLM（Gemini-3.1 Pro），VLM 版本更新会带来得分漂移，指标可比性受限。我的补充观察：其一，patch 大小 8 在像素空间意味着 token 数随分辨率平方增长，512 到 2048 时序列长度增 16 倍，高分辨率下的显存与延迟压力远大于潜空间方法，论文未给出高分辨率推理数据；其二，单步设置下 MUSIQ 从 54.97（A7）降到 53.34（A9），无参考美学分数仍受损，模型在依赖 NR 指标的榜单上可能吃亏；其三，DR-Score 每图测 5 次取平均仍是黑盒评判，与人类主观实验的对齐验证规模有限（仅附录部分展示）；其四，八类退化等概率采样的 283 万训练集对单一退化专家模型是否构成完全公平的比较，论文虽用重训基线缓解，但数据配比的影响未被单独消融。
独立分析的弱点
独立分析几个弱点。第一，分辨率刚性是最实际的工程障碍：用户照片普遍超过 512，patch 8 的全局注意力在 2048×2048 时 token 数约 65536，自注意力开销爆炸，可行改进包括窗口化注意力或多尺度渐进 patch。第二，DR-Score 不可复现：依赖闭源 Gemini-3.1 Pro，他人无法严格复现 Table 6 的数字，改进方向是改用开源 VLM、公开提示词与评分协议，甚至训练一个 IR 专用的退化去除评分器。第三，低光任务的保真度并非最优：PixRestore 低光 PSNR 25.62 低于 FAPE-IR* 的 26.04，说明退化极重、信息极度稀缺时小模型生成先验不足，可考虑轻度蒸馏 T2I 先验或直接放大模型（B 变体已展示扩展收益）。第四，对抗微调引入 GAN 固有的训练不稳定性，100K 步微调的超参敏感性未报告，复现者容易踩坑，应补充稳定性消融。第五，真实世界测试每类退化仅 100 张 LQ 图，DR-Score 的统计功效有限，且真实退化缺少 GT 使保真度无从校验，扩大真实评测规模并加入人类主观实验会更有说服力。
未来方向
作者在结论中给出三个方向：分辨率更灵活的架构（如可变 patch、层级化像素建模）、更强的视觉先验（例如从 T2I 或更大 DINO 蒸馏）、以及专用的 IR 质量度量（把 DR-Score 变成标准化、可复现的社区基准）。基于本文成果还可延伸：一是视频统一修复，DINO 层可靠性的思想可直接迁移到时序维度的特征路由；二是把层路由器与 MoE 专家路由结合，针对极端退化动态扩容骨干；三是利用像素空间输出与输入逐像素对齐的特性做可控修复（指定保留与重建区域）；四是端侧部署——44ms@A800 的单步模型经量化与蒸馏后有手机端实时修复的潜力；五是系统研究 NFE 减少反而提升 PSNR 的现象（A9 的 28.07 高于 A7 的 27.66），这可能揭示流积分误差累积与回归过平滑之间的普遍关系，对少步采样理论有价值；六是把 VLM 评判退化去除的范式推广到去反射、去阴影等更多低层视觉任务，以及研究 DR-Score 与人类感知的系统性对齐。
复现评估
复现条件总体友好但完整训练成本不低。开源情况：论文给出项目页与代码链接，承诺在 GitHub 仓库发布代码和整理好的基准；骨干 LightningDiT、编码器 DINOv2 以及全部评测集（GoPro、UHD-blur、RESIDE-6K、UHD-Haze、DIV2K、PolyU、RainDS-real、RealRain-1k、UAV-Rain1k、UHD-LL、LOL、WeatherBench、RealSR、ScreenSR）均为公开资源。数据方面：283 万张训练语料依赖自建数据管道，完整重建数据集需要相当的合成与爬取工作。算力方面：主实验用 8 张 A800 训 250K 加 100K 步、batch 16、512×512 裁剪，属于中上等训练预算，实验室可行、个人难以完整复现；但推理极其廉价（S 变体单步 44ms/张），用发布权重做下游验证和微调很容易。风险点：DR-Score 依赖专有 Gemini-3.1 Pro，该部分结果只能定性复现；对抗微调阶段的稳定性细节披露有限。综合评估：架构与消融结论可信度高，端到端复现难度中等偏上，权重级复用难度低。
论文图表
左图：低光增强与去雨滴任务中各层 DINO 特征的可视化，浅层保留局部结构与细节，深层编码全局语义。右图：八种退化下 DINOv2-B 各层的 LQ–HQ 特征相似度曲线，显示不同层对不同退化的敏感度不同——浅层对雨、模糊、雪敏感，深层对噪声敏感，中层对低光、超分、雾敏感。
这张图是全文最核心设计（自适应分层引导）的实证动机：它证明了没有任何单一层或均匀平均对所有退化最优，从而直接推出需要逐图自适应层路由的结论。