冻结的像素空间扩散模型可用自身样本自我引导 A Frozen Pixel-Space Diffusion Model Can Guide Itself with Its Own Samples
冻结像素扩散主干,用合成样本训练轻量适配器,靠中间与最终预测之差自我引导采样。
前置知识
像素空间扩散模型 (Pixel-Space Diffusion)
直接在原始像素空间(如 256×256 的 RGB 图像)上做去噪的扩散/流匹配模型,不引入 VAE 把图像压缩到潜空间。代表作为 JiT,它采用 plain ViT 架构、用大 patch 直接预测干净图像 $\hat{x}=f_\theta(x_t,t)$。优点是端到端、避免 VAE 瓶颈与重建失真;难点是单个模型必须同时拟合全局语义与高频局部纹理,而高频分量通常欠拟合,导致纹理模糊。JiT、PixelREPA、DeCo、PixelDiT 等均属此类。
SSG 的全部动机与实验对象都围绕像素空间扩散模型展开,必须先理解它相对潜空间扩散(LDM)的高维建模难题,才能理解为何作者强调高频细节难以拟合、为何自引导能针对性地改善它。
流匹配与速度预测 (Flow Matching)
JiT 采用的生成目标。在干净图像 $x$ 与高斯噪声 $\epsilon$ 之间线性插值 $x_t=tx+(1-t)\epsilon$,对应目标速度 $v_t=x-\epsilon$;模型学习速度场,训练损失为 $\mathcal{L}_{FM}=\mathbb{E}_{x,\epsilon,t}\|v_\theta(x_t,t)-v_t\|_2^2$。由于 JiT 直接预测 $\hat{x}$,再由 $v_\theta=(\hat{x}-x_t)/(1-t)$ 转成速度。采样时用 Heun 积分器沿 ODE 从 $t=1$ 积分到 $0$。
SSG 的适配器训练沿用同一流匹配目标(JiT/PixelREPA 预测 clean image,DeCo 预测 velocity),自引导公式在 clean 空间和 velocity 空间等价。不懂流匹配就无法读懂 $x_{guided}=x_{inter}+w(x_{final}-x_{inter})$ 为何同样作用于速度场。
引导方法 CFG / Autoguidance / Internal Guidance
采样期改善生成质量的手段。分类器无关引导 CFG 用条件预测 $v_c$ 与无条件预测 $v_u$ 外推:$v_{guided}=v_u+w(v_c-v_u)$。Autoguidance (AG) 用一个单独的弱模型替代无条件分支;Internal Guidance (IG) 在主干中间层联合训练一个输出头作为弱预测,需训练 600 个 epoch。它们都需要“弱预测+强预测”一对。
SSG 本质上也是一种“弱+强”引导,但它从同一冻结主干的中间层和最终层各取一个预测,无需独立弱模型也无需联合训练。把 SSG 放在 CFG/AG/IG 的脉络里看,才能理解其即插即用的本质区别(表 8 直接对比 IG)。
分类器无关引导 CFG 与引导尺度 w
CFG 通过放大条件方向提升样本质量但牺牲多样性,由引导尺度控制强度;SSG 引入自己的尺度 $w$,对 $x_{guided}=x_{inter}+w(x_{final}-x_{inter})$ 而言 $w=1$ 退化为原始输出,$w>1$ 沿模型自身精修方向外推、放大高频。SSG 可单独使用,也可与 CFG 串联:先各自做 self-guidance 再做 CFG。
论文表 1-3 分别报告了 with/without CFG 的结果,且 SSG 倾向于配合更低的 CFG 尺度(IS 略降)。理解 CFG 与 SSG 尺度的交互(图 5b 的尺度扫描)是读懂实验结论的前提。
研究动机
像素空间扩散模型(如 JiT)直接在原始像素上做去噪,架构简洁、避免了 VAE 两阶段管线的瓶颈,但代价是单个端到端模型必须在同一高维空间里同时拟合全局语义和局部高频纹理,而高频分量尤其难拟合——这是像素扩散生成图像纹理模糊、细节缺失的根源。现有改进手段(改预测目标为 x-prediction、加感知监督、表示对齐 PixelREPA、解耦高低频的层级架构 DeCo/PixelDiT 等)无一例外都要从头设计和训练一个全新模型,通常在 600M–1B 参数量级上训练约 600 个 epoch,算力开销巨大,且没有充分利用已预训练模型里已有的先验知识。即便如此,JiT-H/16 在 ImageNet 256×256 上配合 CFG 仍只有 FID 1.86,高频细节明显欠拟合。能否在不重训主干的前提下,廉价地“榨取”一个已经训好的像素扩散模型,把它变成更强的生成器,是一个未被回答的问题。
本文的目标是作者的目标是提出一种廉价、即插即用的方法,直接增强任意预训练像素扩散 Transformer,做到三点:(1) 完全冻结主干、不重训,只训练一个挂在中间层的轻量适配器,算力不到完整模型训练量的 1%;(2) 训练适配器时不需要任何真实图像,只用模型自身生成的合成样本;(3) 在 JiT、PixelREPA、DeCo 等多种像素扩散模型、ImageNet 256 和 512 两种分辨率、有/无 CFG 两种设置下都稳定降低 FID。作者希望借此证明:一个冻结的像素扩散模型可以用自身的中间表示引导自身的采样,从而把欠拟合的高频细节“补”回来。
与已有工作不同的是,本文的独特切入角度是“同一冻结主干内的自我引导”。此前引导类工作要么需要额外的弱模型(Autoguidance),要么需要把中间输出头和主干联合训练 600 个 epoch(Internal Guidance);自改进类工作(SIMS、Neon)则要在权重空间微调或外推。SSG 不走任何一条老路:它发现 plain 像素扩散 Transformer 的中间层本身就解码出只含低频结构的粗预测(Layer 4 的低频带功率已达最终预测的 99.5%,而高频带仅 87.7%),最终层才逐步精修高频,于是中间与最终预测天然构成一对“弱+强”预测,可直接用作引导信号。更反直觉的是,作者证明训练这个中间适配器时,模型自己生成的合成样本比真实 ImageNet 图像更有效——因为适配器只需要提供一个低频参考而非复现真实图像的全频谱。
核心方法
核心直觉是:预训练像素扩散 Transformer 的深度方向上存在“由粗到精”的内在结构——浅层已编码好全局低频结构,深层才精修局部高频细节。SSG 把这一过程显式化:在某个浅/中层挂轻量适配器,把该层表示解码成粗预测 $x_{inter}$,主干完全冻结;采样时把最终预测 $x_{final}$ 沿模型自身精修方向外推,得 $x_{guided}=x_{inter}+w(x_{final}-x_{inter})$,$w>1$ 为 SSG 尺度。适配器训练沿用流匹配目标 $\mathcal{L}_{FM}=\mathbb{E}_{x,\epsilon,t}\|v_\theta(x_t,t)-v_t\|_2^2$,干净预测 $\hat{x}=f_\theta(x_t,t)$、速度 $v_\theta=(\hat{x}-x_t)/(1-t)$;但训练数据全部来自冻结模型生成的 1M 张合成图(每类 1K 张)。整个适配器训练量不足完整模型训练量的 1%(base/large 训 30 epoch、huge 训 50 epoch),可单独使用,也可与 CFG 串联。
本质创新有两点,且都反直觉。第一,引导信号来自同一冻结主干内部的“粗→精”层级差异,无需独立弱模型(区别于 Autoguidance)、无需与主干联合训练(区别于 Internal Guidance 那种 600 epoch 联训头)。作者用频域分析(图 3)证实:中间预测的低频带功率接近最终层(99.5%),高频带明显偏低(87.7%),二者之差恰好集中在模型欠拟合的高频分量上。第二,训练中间适配器时应使用模型自己的合成样本而非真实图像——因为适配器只需提供低频参考,而合成样本天然被限制在模型能表达的频带内,给出更干净的低频对应物、更大的残差 $(x_{final}-x_{inter})$,正好把外推能量集中到该补的高频上。表 4 印证:合成训练在几乎所有设置上超过真实训练(如 JiT-H/16 无 CFG:FID 2.26 vs 2.51)。这使得 SSG 成为真正零真实数据、即插即用的方法。
方法步骤详情
(1) 合成数据生成:用预训练模型按其论文最佳 CFG 生成 1M 张图(每类 1K 张)作干净目标。(2) 适配器构造:在浅/中层(B/L 取 Layer 6,H/PixelREPA-H/DeCo-XL 取 Layer 8)挂轻量适配器——B/L 用 1 个 transformer block、H 用 2 个,接线性层映回像素;DeCo 挂 encoder 并复用 pixel decoder;主干全程冻结。(3) 训练:Adam($\beta_1{=}0.9,\beta_2{=}0.95$)、lr $2\times10^{-4}$(DeCo $5\times10^{-5}$)、batch 1024(DeCo 256),训 30 epoch(B/L)或 50 epoch(H);JiT/PixelREPA 预测 clean image、DeCo 预测 velocity。(4) 自引导采样:每步算 $x_{inter}$、$x_{final}$,得 $x_{guided}=x_{inter}+w(x_{final}-x_{inter})$;无 CFG 全程 $t\in[0,1]$ 应用,有 CFG 时两者都在 $t\in[0.1,1.0]$ 应用、Heun 50 步,对条件/无条件预测分别先做 self-guidance 再做 CFG。
技术新颖性
技术新颖性体现在:(a) 用频域分析(图 3)定量刻画中间层 vs 最终层的低/高频带功率差,把“粗→精”从经验直觉变成可测的弱/强预测对,并据此选适配层(图 5a 显示 Layer 4–6 最佳、深层退化回基线)。(b) 与 CFG 用“无条件 vs 条件”不同,SSG 用同一条件预测内部的“粗 vs 精”,因此与 CFG 正交、可叠加,且互补于而非替代 CFG。(c) 合成数据优于真实数据这一结论由径向功率谱(图 3b)支撑:合成训练的引导使生成图的高频谱最接近真实图像。(d) 真正即插即用:同一主干 checkpoint、零真实数据、不足 1% 算力,并可迁移到 DeCo 这种层级架构(适配器挂 encoder、复用 pixel decoder)。(e) 表 8 进一步证明 SSG(FID 3.29)还优于需联训主干的 IG(3.41)。
实验结果
核心结论:不足 1% 算力换显著 FID 下降。带 CFG 的 256×256(表 1):JiT-H/16 1.86→1.67(−10.2%)、PixelREPA-H/16 1.81→1.59(−12.2%,像素空间新最优)、DeCo-XL/16 1.69→1.63;512×512(表 2)JiT-H/32 1.94→1.84。无 CFG(表 3)降幅超 50%:JiT-B/16 25.42→9.47、JiT-L/16 13.85→4.47、JiT-H/16 7.15→2.26。关键消融:(1) 合成 > 真实数据(表 4),如 JiT-H/16 无 CFG 2.26 vs 2.51、带 CFG 1.67 vs 1.78;(2) 仅 1 万张合成样本即够,1M→10K 时 FID 变化 <0.05(表 5);(3) Layer 4–6 最佳、深层退化(图 5a);(4) B/L 用 1 个 block、H 用 2 个(表 6);(5) 对生成 CFG 鲁棒(表 7);(6) SSG(3.29)优于联训 IG(3.41);固定 CFG 扫描(图 5b)证明增益不全来自更大搜索空间。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet 256×256 类条件生成(带 CFG) | FID↓ | JiT-H/16+SSG 1.67;PixelREPA-H/16+SSG 1.59(像素空间最优) | JiT-H/16 1.86;PixelREPA-H/16 1.81 | 分别 −10.2% 与 −12.2% |
| ImageNet 256×256 类条件生成(带 CFG) | FID↓ | DeCo-XL/16+SSG 1.63 | DeCo-XL/16 1.69 | −3.6%,证明可迁移到层级架构 |
| ImageNet 512×512 类条件生成(带 CFG) | FID↓ | JiT-H/32+SSG 1.84 | JiT-H/32 1.94 | −5.2% |
| ImageNet 256×256 类条件生成(无 CFG) | FID↓ | JiT-H/16+SSG 2.26 | JiT-H/16 7.15 | −68%,无 CFG 下降均超 50% |
| ImageNet 256×256(无 CFG)适配器训练数据对比 | FID↓ | 合成训练 JiT-H/16 2.26 | 真实图像训练 2.51 | 合成数据反超真实数据 |
| ImageNet 256×256(带 CFG)与 IG 对比 | FID↓ | JiT-B/16+SSG 3.29 | JiT+IG(联训主干)3.41 | 冻结+合成即优于联训 IG |
局限与改进
作者承认的局限:实验仅覆盖类条件 ImageNet 生成(256 与 512)且仅针对像素空间扩散模型,中间→最终的精修结构能否推广到更大的文本条件模型或其他图像分布尚不清楚。我的补充观察:(1) 方法依赖 plain ViT 式像素扩散主干清晰的“由粗到精”深度结构,频域分析(图 3)只在 JiT-B/16 上做过,其他主干的最优层是经验挑选,结构差异大的模型未必能给出有用的弱/强对;(2) 表 1 中 IS 普遍略降,因为 SSG 在 FID 最优点倾向更低的 CFG 尺度,说明存在分布/多样性上的权衡;(3) 适配器仍额外引入 12–66M 参数和采样开销,且需对 SSG 尺度 $w$ 与 CFG 尺度联合调参(表 S1);(4) 外推 $w$ 过大时可能放大噪声;(5) Recall 等多样性指标提升有限,方法主要在精修质量上发力。
独立分析的弱点
弱点 1:依赖清晰的粗→精层级结构——频域分析仅对 JiT-B/16 做了,其他主干的最优适配层靠经验挑选,结构差异大时可能失效,改进方向是用同一带功率指标自动选层。弱点 2:合成数据生成本身仍需 1M 张图(虽然表 5 证明 1 万张就够),默认 1M 浪费了算力,可改用渐进式/主动采样按需扩充。弱点 3:在已带 CFG 的强基线上增益相对温和(1.86→1.67),而无 CFG 时增益巨大,提示在强 CFG 下边际收益递减,可与其他高频专用损失组合。弱点 4:仅类条件验证,文本到图像的像素扩散尚未测试,适配条件维度更高的 T2I 模型是关键缺口。弱点 5:IS 回退提示模式/分布权衡未被充分刻画,建议补充多样性敏感指标与样本多样性分析。
未来方向
作者提出的方向:把 SSG 扩展到文本到图像生成、更广的数据集以及其它扩散架构。基于成果可延伸的方向:(1) 将自引导原则推广到潜空间扩散,验证其中间层是否也有粗→精细修,从而把 SSG 变成跨架构通用引导;(2) 用多个中间层的多尺度引导替代单层适配器,自适应组合不同频段;(3) 把“合成数据更优”这一洞见推广到其它辅助头训练,例如 SIMS 式负引导头也用合成样本训;(4) 直接在适配器上加谱损失,针对性强化目标频段的残差 $(x_{final}-x_{inter})$;(5) 与 Autoguidance 式弱模型组合,探索复合引导的叠加效应;(6) 把自引导用在视频/3D 像素扩散上,检验时序一致性是否也能受益。
复现评估
可复现性很强。代码开源于 https://github.com/zfu006/SSG。论文在表 S2 中给出完整配置(主干深度、隐藏维度、patch 大小、适配器深度/参数、学习率、batch、采样器),表 S1 给出每个模型的精确引导尺度。JiT、PixelREPA、DeCo 均用公开发布的 checkpoint。适配器训练廉价(<1% 算力,30–50 epoch,1M 合成样本在中量级硬件即可生成)。附录 C 详细描述了频域分析协议(正交 2D FFT、径向频带、0.5 处切分低/高频)。主要复现门槛在于 50K 样本的 FID 评测需要标准 ImageNet 评测流水线,并对 SSG 与 CFG 尺度做仔细扫描,但流程本身清晰直接。整体复现评级:高。
论文图表
并排展示 No Guidance、CFG、CFG+SSG 三种引导策略下 JiT-H/32 的生成样例对比,直观呈现加入 SSG 后局部纹理与精细细节的明显改善。
在进入技术细节前用一张图直观建立“SSG 能补回高频细节”的感性认知,是理解全文价值的最佳入口。