跳转至

PixelU: A U-Shaped Transformer for Efficient End-to-End Pixel Diffusion

会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/gzp6688/PixelU
领域: 图像生成
关键词: 像素空间扩散、干净图像预测、U形Transformer、频率解耦、空间下采样

一句话总结

PixelU 在直接预测干净图像的像素扩散中,用浅深层跳连保留空间细节、用单级恒通道下采样集中建模全局语义,在 ImageNet 256×256 上以 136.9 GFLOPs 达到 FID 1.63,相比 JiT-G 的 383 GFLOPs 与 FID 1.82 改善了质量和计算量的折中。

研究背景与动机

潜空间扩散先让 VAE 压缩图像,再在较短序列上生成,代价是生成器最终受制于有损重建。直接在像素空间生成可以绕开这道瓶颈,但 Transformer 必须同时处理物体布局、轮廓、纹理和大量噪声。DeCo、DiP、PixelDiT 等路线给主干增加像素级解码或细节模块,试图让不同网络部分分担语义与纹理;这些模块也带来了额外计算。

PixelU 的切入点来自 JiT:当网络输出的是干净图像,而不是包含高维噪声成分的速度,学习难度可能已经发生根本变化。作者在同一个 JiT-B/16 基线上比较预测目标与解码器,发现 DeCo 解码器将速度预测的 FID 从 190.11 降到 42.07,但在干净图像预测下只从 43.65 降到 41.20。这提示复杂解码器的收益依赖预测目标,不能脱离训练参数化讨论其必要性。

即使直接预测干净图像,深层 Transformer 仍可能丢失细粒度空间信息,因此问题不是把所有细节建模都删掉,而是找到更轻的保留路径。作者选择让浅层特征绕过语义瓶颈,同时把深层空间分辨率主动降低,减少对高频信号的重复处理。核心 idea:在干净图像预测前提下,以跳连保留细节、以单级恒通道下采样压缩语义主干,让两条路径协同替代昂贵的附加像素解码器。

方法详解

整体框架

PixelU 是类别条件的端到端像素空间生成器,不需要先训练图像自编码器。一次去噪调用接收带噪图像、时间步和类别条件,输出干净图像估计;这个估计再转换为流匹配采样器需要的速度,经过多次调用从噪声得到图像。

网络内部先在较高空间分辨率上提取浅层特征,再通过单级恒通道语义瓶颈降低 token 数,最后上采样回原特征分辨率,并借助浅深层细节跳连恢复局部结构。干净图像预测是整个网络的输出参数化,不是进入网络前把噪声图像变干净的预处理。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["带噪图像、时间步<br/>类别条件"] --> B["浅层特征提取"]
    B --> C["单级恒通道<br/>语义瓶颈"]
    C -->|深层建模后上采样| D["浅深层细节跳连"]
    B -->|绕过瓶颈的空间特征| D
    D --> E["干净图像预测"]
    E --> F["干净图像估计<br/>转换为采样速度"]

图中展示的是每次网络调用的数据流。训练时干净图像仅作为目标参与损失,推理时并不存在一条输送真实干净纹理的旁路;类别 token 与 REPA 的作用分别属于条件注入和训练辅助监督。

关键设计

1. 单级恒通道语义瓶颈:让大部分语义处理在更短序列上完成

PixelU 不沿用平坦 DiT 在所有层保持同一空间分辨率的做法,而是将特征网格的高和宽各减半,使中间段的空间 token 数变为四分之一。 关键不是重复多级压缩,而是只做一次下采样和对应上采样,并在浅层、中间和深层保持相同通道宽度。这样既降低中间段的注意力与逐 token 计算,又避免传统 U-Net 在分辨率下降时不断扩宽通道、抵消部分节省。 作者将这一瓶颈解释为低通偏置:布局与大块形状可以在较低空间分辨率上表达,而高频噪声和部分细节被削弱,深层不必在完整空间网格上同时处理所有频率。

这不是声称任意下采样算子都是严格抗混叠滤波器。缓存正文没有给出具体采样核、滤波器响应或完整算子实现,因此可确认的是单级压缩与恒通道路由,以及频谱可视化支持的低频偏置。 只加下采样反而让 FID 从 43.65 变成 48.87,说明语义瓶颈单独使用会损失生成所需信息;它必须和后面的细节路径共同理解。 模型 B、L、H 的通道宽度分别为 768、1024、1152,编码段/中间段/解码段深度分别为 [4,4,4]、[8,8,8]、[8,20,8]。恒通道是单个模型内部的约束,不是不同规模模型都使用同一宽度。

2. 浅深层细节跳连:让细节绕过空间压缩,而不是在输出端重新学习整套像素解码

中间瓶颈越关注大尺度结构,越需要一条不经过这次压缩的空间信息路径。PixelU 在浅层编码特征与深层解码特征之间建立长跳连,使输出侧既接收语义建模结果,也接收较高空间分辨率的特征。 这可以减少深层反复保存边缘、位置和纹理线索的负担。所谓去掉解码器,准确说是去掉额外的复杂像素解码模块,而不是删除 U 形网络本身的上采样段与解码段。 正文未明确逐条跳连使用相加还是拼接融合,不应凭 U-Net 经验替作者补上实现。

这里的“细节保留”也不意味着浅层特征已经无噪声。网络输入依然是带噪图像,跳连传递的是从它提取的特征;作者关于“未受破坏的高频细节”的说法,更合理的读法是避免这些线索继续被深层压缩损坏。 表1尤其说明它与预测目标的协同:只加跳连时,干净图像预测的 FID 为 38.53,速度预测却仍为 185.54。跳连不是对任何预测目标都有效的通用补丁。 它也并非严格零成本:对应实验的计算量由 21.99 增至 23.20 GFLOPs,只是小于附加复杂解码器的成本。

3. 干净图像预测:保留流匹配采样,但改变网络直接需要拟合的对象

PixelU 沿用 JiT 的参数化,让输出头直接估计干净图像,而不是输出噪声或速度。训练时先把干净图像与高斯噪声按时间步线性插值形成输入;真实速度是干净图像减去噪声,即 \(v=x-\epsilon\)。 网络的干净图像估计并不直接拿去执行速度积分,而是与当前带噪状态一起转换为预测速度,再计算流匹配误差或交给采样器。 因此,“预测干净图像”与“训练速度误差”可以同时成立:前者决定网络输出的参数化,后者决定优化目标。

这种组合把输出空间约束到图像估计,而非要求网络直接表示包含大量噪声的目标;在作者的消融范围内,复杂像素解码器原本补偿的一部分困难因而减轻。 但最终学习目标不是对所有时间步一视同仁的普通图像重建误差。速度转换会使其等价于带时间相关权重的图像预测误差,这也是复现时不能随意把损失替换为普通 MSE 的原因。 缓存的式(1)至式(3)存在文本提取损坏,这里依据完整段落说明机制,不把残缺公式猜补为作者精确公式。

一个完整示例

以 256×256 的 PixelU-B/16 为例,按 /16 的 patch 尺度理解,空间网格为 16×16,即 256 个空间 token;这不计额外的类别 token。 最初 4 个 Transformer 块提取浅层特征,并为后续跳连保留信息。 一次空间下采样把网格变成 8×8,即 64 个空间 token,中间 4 个块在此分辨率处理语义,通道宽度仍为 768。 随后上采样回 16×16,最后 4 个块结合浅层特征细化输出,最终得到干净图像估计。

这组 token 数是依据 patch 尺度和论文半分辨率设计推得的结构示例,不是额外实测。它展示的是网络内部特征压缩,而不是先生成低分辨率成品、再运行另一个超分辨率扩散模型。 训练时可以用真实干净图像构造监督;采样时只从噪声与类别条件出发,反复进行上述网络调用。 从 256 扩展到 512 分辨率时,表4使用 /32 版本,说明分辨率比较同时改变了 patch 尺度,不能直接解释成固定 token 化下的高分辨率扩展实验。

损失函数 / 训练策略

主要损失是预测速度与真实速度之间的均方误差,网络仍直接输出干净图像。另在网络中间块加入 REPA 表征对齐,权重为 0.01;REPA 是训练期辅助,不是在推理链路中加入图像 VAE。 论文还加入 in-context class tokens,用序列中的类别 token 提供条件。表8把这些通用增强逐项加入,因此主结果不能全部归功于裸 U 形结构。

训练采用 ImageNet-1K,global batch size 为 1024,AdamW 的 beta 为 (0.9, 0.95),恒定学习率为 0.0002,使用 8 张 B200。 时间步采用 logit-normal 采样,其 logit 的均值为 -0.8、标准差为 0.8;EMA 衰减为 0.9999。 评测生成 50K 张图像,报告 FID、sFID、IS 与 Precision/Recall;采样采用 50 步 Heun、classifier-free guidance(CFG)和 CFG interval。

表4的 PixelU NFE 记为 100×2,不能把 50 个 Heun 步误当成 50 次模型评估,也不能把单次 GFLOPs 当作完整生成一张图的成本。 图4样例标注 CFG=4.0,但正文没有据此证明所有定量表格都用同一最优 CFG 值。核心结构消融明确使用 160 epochs、cfg=1,与最终带引导结果分开解释。

实验关键数据

主实验

以下来自表3、表4,均为 ImageNet 类别条件生成并使用 CFG。FID 越低越好,IS 越高越好;GFLOPs 是论文报告的网络计算量,而非实测端到端时延。

分辨率 模型 Epochs GFLOPs FID IS
256×256 JiT-G/16 600 383 1.82 292.6
256×256 DeCo-XL/16 600 未报告 1.69 304
256×256 PixelU-H/16 320 137 1.77 306.86
256×256 PixelU-H/16 600 136.9 1.63 305.88
512×512 JiT-H/32 600 183 1.94 309.1
512×512 PixelU-H/32 600 138 1.92 322.10

256 分辨率的 600 epoch PixelU 使用表3更精细的 136.9 GFLOPs,表4四舍五入为 137;相对 JiT-G 约为 35.7%,即减少约 64.3%,不能等同于实测加速三倍。 512 分辨率相对 JiT-H 的 FID 仅改善 0.02,且模型参数量不同,不应描述成等规模下的大幅优势。

原文内部存在三处值得保留的差异:320 epoch FID 在表4为 1.77,附近正文为 1.76;512 结果在表4叫 H/32,正文写 H/16;H/32 参数量在表2为 1171.8M,在表4为 1152M。上表沿用结果表的 FID 与模型名,不自行统一这些差异。

消融实验

表1比较同一 JiT-B/16 基线上的附加模块,条件为 ImageNet 256×256、160 epochs、cfg=1。PixelDiT 解码器一项是作者复现,不能当作其官方完整模型结果。

附加模块 GFLOPs 干净图像预测 FID 速度预测 FID
无,JiT-B/16 21.99 43.65 190.11
DDT decoder 25.35 43.48 189.50
DeCo decoder 25.91 41.20 42.07
PixelDiT decoder,作者复现 33.78 39.20 44.85
跳连 23.20 38.53 185.54

表5固定干净图像预测,使用相同分辨率、训练轮数和 cfg=1,检验两项结构设计是否互补。

跳连 下采样 GFLOPs 参数量 FID IS
否 否 21.99 131.1M 43.65 34.08
否 是 18.25 153.6M 48.87 30.29
是 否 23.20 135.8M 38.53 38.57
是 是 19.46 158.3M 34.92 43.51

关键发现

  • 下采样单独使用损害 FID,但与跳连组合后,相对基线降低 8.73 FID,并减少约 11.5% GFLOPs;参数量却增加,因此“计算更少”不等于“参数更少”。
  • 表6的 [4,4,4] 达到 FID 34.92,[5,2,5] 为 35.47,[3,6,3] 为 38.08。平衡分配是该 B 级消融的结论,不意味着更大 H 模型也应采用等深三段。
  • 表7中 0、1、2 次下采样的 FID 分别为 38.53、34.92、42.70。恒通道条件下继续压缩会形成过窄瓶颈,多级 U-Net 不一定更优。
  • 表8的 H/16 在 160 epochs 时,裸模型 FID 12.41,加入类别 token 后为 8.93,再加 REPA 为 6.95,使用 CFG 与 CFG interval 后为 2.42;训练到 600 epochs 才达到 1.63。

亮点与洞察

  • 先改变预测目标,再判断架构必要性,比孤立比较模块更有解释力。表1揭示某些“必需的细节模块”可能是在补偿特定输出参数化的困难。
  • 降低空间 token 数与增加细节旁路是一对相互依赖的操作。表5中单独下采样变差、联合后变好,是比单张漂亮频谱图更直接的设计依据。
  • 图5对 10 个 ImageNet 类别、每类 100 个样本做 t-SNE,图6分析中间特征的二维 FFT。二者分别从类别组织和频率分布支持语义瓶颈解释,但仍不是完整的因果证明。

局限与展望

  • 实验集中于 ImageNet 类别条件生成,尚未证明同样设计在开放域文本到图像、复杂排版或视频生成中仍能保持优势。
  • “像素空间最优”需要限定比较范围:正文提到未列入表4的 PixelDiT 与 SiD2 的参考 FID 分别为 1.61、1.38;潜空间模型 RAE-XL 的表4 FID 为 1.13。不能把 1.63 称为所有已报告方法中的绝对最佳。
  • 全局 FID 改善不代表各项指标全面领先。256 分辨率下 PixelU 的 sFID 为 5.04,DeCo 为 4.59;512 分辨率下 PixelU 的 Recall 为 0.58,DeCo 为 0.60。
  • 主模型包含 REPA,且不同模型的参数量、训练预算与采样设置并非完全一致。GFLOPs 节省也尚未通过统一硬件上的吞吐、显存峰值和总训练成本验证。
  • 低频不等于全部语义,细粒度识别也可能依赖纹理。后续可以在高频决定类别或包含小物体的分布上测试单级压缩的失效边界,并公开完整算子与训练配置。

相关工作与启发

  • vs JiT:继承干净图像预测,改变的是内部特征路由。PixelU 不让全部层在相同空间分辨率上建模,而是引入显式语义瓶颈和细节跳连。
  • vs U-ViT / U-DiT:U 形拓扑和跳连并不是新发明;本文的贡献在于像素空间、干净图像预测、单级下采样与恒通道的组合及其针对性消融。
  • vs DeCo / PixelDiT:它们用专门像素解码或细节建模模块划分职责,PixelU 改用浅深层特征传递降低这部分成本。结论适用于本文测试设置,不等于复杂解码器在所有任务中都无用。
  • vs 潜空间扩散:PixelU 不依赖 VAE 图像压缩,但内部仍压缩特征。可迁移的启发是区分“生成空间是否为像素”与“主干是否必须全程保持完整空间分辨率”。

代码链接来自论文中的发布声明;本笔记未联网验证仓库当前是否已公开实现。

评分

  • 新颖性: 4/5。组件经典,但对预测目标与架构必要性的交互分析具有清晰价值。
  • 实验充分度: 4/5。包含多规模、双分辨率与针对性消融,仍缺开放域任务和实际运行成本对照。
  • 写作质量: 3/5。论证主线清楚,但若干数值与型号不一致,且“零成本”“纯净语义”等表述偏强。
  • 价值: 4/5。为直接像素生成提供了易理解的轻量结构路线,复现仍依赖更完整的实现细节。