075、超分前沿探索：从扩散模型到隐空间超分的新范式

075、超分前沿探索从扩散模型到隐空间超分的新范式去年年底调一个老项目客户反馈说“人脸放大后像油画纹理全糊了”。我翻出ESRGAN试了试效果还行但一到真实场景——比如监控截图、老照片修复——就崩得厉害。后来试了试扩散模型第一眼看到输出时我愣了半天那种“画上去”的假感没了取而代之的是真实的纹理细节。但代价是一张512×512的图像在V100上跑了快两分钟。这显然不能落地。于是我开始琢磨能不能把扩散模型的计算量砍掉同时保留它的生成质量答案就是隐空间超分——把超分过程从像素空间搬到压缩后的隐空间里。今天这篇笔记就是我在这个方向上踩过的坑和总结的经验。扩散模型为什么能超分先别急着调代码扩散模型做超分的核心思路不是“学习从低分到高分的映射”而是“从噪声中逐步还原出高分图像”。你给它一个低分图像作为条件它就在每一步去噪时参考这个条件最终生成符合条件的高分结果。我最早试的是SR3它的结构其实很简单一个UNet输入是噪声图低分图上采样后的结果输出是预测的噪声。训练时对高分图加噪声让模型去预测噪声推理时从纯噪声开始一步步去噪每一步都拿低分图来“拉一把”。这里有个坑很多人直接拿DDPM的代码改结果发现生成的结果颜色不对。为什么因为DDPM默认假设数据是归一化到[-1,1]的但你的低分图可能还是[0,255]的整数。别这样写——输入前一定要统一归一化方式否则模型学到的条件分布是错的。# 这里踩过坑低分图和高分图必须用相同的归一化# 我习惯用 (x / 127.5 - 1) 映射到 [-1, 1]lr(lr_tensor/127.5)-1.0hr(hr_tensor/127.5)-1.0# 千万别一个用0-1一个用-1到1模型会学疯SR3在FFHQ人脸数据集上效果惊艳但换到真实监控场景立马露馅——因为训练数据太干净了。扩散模型对训练数据的分布极其敏感你拿高清人像训练它就只能生成“像高清人像”的纹理遇到模糊、噪声、压缩伪影它就懵了。隐空间超分把大象装进冰箱需要几步扩散模型慢就慢在像素空间里做迭代。一张256×256的图像像素点就有65536个每一步去噪都要在这个维度上算。那能不能先压缩再超分这就是隐空间超分的思路。典型代表是LDMLatent Diffusion Models它先用一个预训练的自编码器比如VQGAN或KL-VAE把图像压缩到隐空间——通常是原图尺寸的1/8到1/16。然后在隐空间里跑扩散过程最后用解码器还原回像素空间。我试过在LDM基础上做超分做法是把低分图也编码到隐空间然后作为条件输入扩散UNet。这样扩散过程只在低维隐空间里进行速度直接提升了10倍以上。但别高兴太早这里有个关键问题自编码器的压缩是有损的。如果你的自编码器压缩率太高比如1/32高频纹理细节在编码阶段就丢了后面再怎么扩散也补不回来。我踩过的坑是用了1/64的压缩率结果超分结果全是模糊的色块。# 经验值超分任务建议用1/8或1/16的压缩率# 我用的是KL-f8版本隐空间通道数4空间尺寸1/8# 别为了省计算量用f32纹理细节会哭encoderload_vae_encoder(kl-f8.pt)lr_latentencoder(lr_image)# 形状: [B, 4, H/8, W/8]# 然后在这个latent上跑扩散条件注入的几种姿势我全试过隐空间超分的关键是怎么把低分信息注入到扩散过程中。我试过三种方式第一种是Concat直接把低分隐变量和噪声隐变量拼在一起作为UNet输入。简单粗暴但问题是低分信息会随着扩散步数被“洗掉”。早期步数还能看到低分结构后面就完全放飞了。第二种是Cross-Attention把低分隐变量作为key/valueUNet的中间特征作为query。这种方式更灵活模型可以自适应地选择从低分图里提取什么信息。但计算量大了不少而且训练不稳定——我遇到过loss震荡不收敛的情况后来发现是attention的scale参数没调好。第三种是ControlNet的思路在UNet的每个block旁边加一个可训练的副本专门处理条件输入。这个效果最好但参数量翻倍。我最终在项目里用的是这个方案因为质量优先。# ControlNet风格的注入注意这里要冻结主UNet# 只训练controlnet部分否则容易过拟合classControlNetUnit(nn.Module):def__init__(self,unet_block):super().__init__()self.copycopy.deepcopy(unet_block)# 复制一份# 这里踩过坑必须用deepcopy不能共享参数forparaminself.copy.parameters():param.requires_gradTrue训练策略别一上来就全量数据隐空间超分的训练比像素空间敏感得多。我踩过最大的坑是直接用全量数据训练结果模型学到的全是“平均脸”——生成的人脸都长一个样。后来我改成两阶段训练第一阶段只训练条件注入模块比如ControlNet冻结自编码器和主UNet。第二阶段再联合微调所有模块。这样做的原因是自编码器已经学到了很好的隐空间分布你如果一开始就动它它会把“如何解码”这件事忘掉。另外数据增强很重要。扩散模型对数据量极其贪婪我试过在训练时加入随机模糊、噪声、JPEG压缩模型对真实低分输入的鲁棒性明显提升。但注意增强强度不能太大否则低分图和高分图的对应关系会被破坏。# 我的数据增强pipeline踩过坑才总结出来的defaugment_pair(lr,hr):# 先对lr做退化模拟再对hr做同样的增强ifrandom.random()0.5:lrgaussian_blur(lr,kernel_sizerandom.choice([3,5,7]))hrgaussian_blur(hr,kernel_size3)# 轻微模糊别太过# 别对hr做太强的增强否则模型学不到真实纹理returnlr,hr推理加速从50步到4步的实战技巧扩散模型推理慢这是落地最大的障碍。我试过几种加速方法DDIM采样可以把步数从1000降到50质量损失不大。但50步还是太慢。后来我试了DPM-Solver这是一种基于微分方程求解器的采样方法能在10步以内达到不错的效果。但注意DPM-Solver对模型精度要求高如果你的模型训练得不够好它会放大误差。再后来我试了蒸馏——用训练好的扩散模型作为教师训练一个学生模型让学生学会一步或多步生成。具体做法是用教师模型生成大量“伪标签”然后训练学生模型直接预测最终结果。我试过4步蒸馏效果和50步DDIM差不多但速度提升了12倍。但蒸馏也有坑学生模型容易学到教师的“平均行为”丢失多样性。我的经验是在蒸馏时加入对抗损失让学生模型不仅模仿教师的输出还要让判别器分不清谁是谁。# 4步蒸馏的采样过程别写复杂了defsample_fast(model,lr_latent,steps4):xtorch.randn_like(lr_latent)# 初始噪声# 这里用DPM-Solver的schedule步数少但效果好timestepsget_dpm_schedule(steps)fortintimesteps:pred_noisemodel(x,t,lr_latent)xdpm_step(x,pred_noise,t)returnx个人经验哪些坑你大概率会踩自编码器的选择决定上限。我试过VQGAN和KL-VAE后者在纹理细节上明显更好。但KL-VAE的隐空间是连续的扩散模型更容易学。如果你追求极致速度可以考虑VQGAN但要做好纹理丢失的心理准备。条件注入的强度要适中。太强模型会直接复制低分图不做超分太弱模型会忽略条件生成随机结果。我习惯在训练时随机dropout条件比如10%的概率把条件置零这样模型学会“即使条件不完美也要生成合理结果”。别迷信大模型。我试过把UNet的通道数翻倍参数量从1.5B涨到3B但PSNR只提升了0.1dB。对于超分任务模型容量不是瓶颈数据质量和训练策略才是。评估指标要选对。PSNR和SSIM对扩散模型不公平——它们倾向于平滑的结果而扩散模型生成的是锐利但可能有轻微偏移的纹理。我建议同时看LPIPS和FID这两个指标更符合人眼感知。落地时考虑后处理。扩散模型生成的图像有时会有“伪影”——比如边缘的振铃效应。我习惯在最后加一个轻量的CNN精调网络比如3层卷积专门修复这些伪影。这个网络很小几乎不增加推理时间。最后说一句隐空间超分不是银弹。如果你的低分图质量极差比如16×16的人脸隐空间的压缩会丢失太多信息这时候像素空间的扩散模型反而更好。但如果你需要处理1080p甚至4K的视频隐空间超分几乎是唯一可行的方案——因为像素空间的计算量是O(N2)而隐空间是O((N/8)2)差了64倍。这个方向还在快速演进我最近在关注“一致性模型”和“流匹配”在超分上的应用它们有可能把推理步数降到1步。等有更多实战经验了再和大家分享。