Stable Diffusion 中的 VAE 起什么作用?
简化版
Stable Diffusion 的 VAE 负责在像素图像和低维潜变量之间转换。编码器把图片压缩成 latent,扩散模型在 latent 上加噪和去噪,解码器再把最终 latent 还原为 RGB 图片。这样能显著降低注意力和卷积计算量。
VAE 不是无损压缩:过强压缩会丢失小字、细线和高频纹理;解码器也可能引入色偏、块状伪影。更换 VAE 时要同步 latent 通道、缩放因子和数值范围,否则图像可能发灰、过曝或完全失真。
扩散模型决定“潜空间里生成什么”,VAE 决定“这些潜变量最终能被还原成怎样的像素细节”。
详细版
典型流程是:
RGB image x: [B, 3, H, W]
-> VAE Encoder qφ(z|x)
latent z: [B, 4, H/8, W/8]
-> scale -> diffusion / denoising
-> unscale -> VAE Decoder pψ(x|z)
reconstructed RGB image
对 512×512 图像,像素张量有 3×512×512=786,432 个值;4 通道的 64×64 latent 只有 4×64×64=16,384 个值,元素数约缩小 48 倍。虽然网络计算量不只由元素数决定,这仍是潜扩散比像素扩散高效的重要原因。
VAE 训练常用近似目标:
L = L_reconstruction(x, x_hat)
+ λ_perc L_perceptual(x, x_hat)
+ β KL(qφ(z|x) || N(0,I))
+ λ_adv L_adversarial
重建项保留内容,感知或对抗项提升视觉锐度,KL 项让 latent 分布更规则。KL 太强会造成后验坍缩或模糊,太弱则 latent 分布不利于扩散模型学习。
| 问题 | 常见表现 | 优先检查 |
|---|---|---|
| scaling factor 错 | 对比度、曝光异常 | encode/decode 是否各缩放一次 |
| 数值范围错 | 整体偏黑或偏白 | [0,1] 与 [-1,1] 转换 |
| 精度溢出 | NaN、彩色斑点 | FP16 解码、强制 upcast |
| VAE 不匹配 | 色偏、纹理崩坏 | 模型配置、latent 通道和版本 |
完整版教学
一、为什么不直接在像素空间扩散
像素空间分辨率高,冗余也高。对 1024×1024 图像直接做大规模去噪,需要在一百多万个空间位置上处理特征;先压缩到 128×128 latent,可把空间位置降到原来的 1/64。
潜扩散用 VAE 学习感知压缩,把肉眼不敏感的像素级冗余丢掉,让主生成模型把容量集中在构图、语义和纹理上。
二、编码器输出为什么是分布
VAE 编码器通常输出均值 μ 和对数方差 log σ²,再用重参数化采样:
ε ~ N(0, I)
z = μ(x) + σ(x) ⊙ ε
训练时采样让相邻图像映射到较平滑的潜空间。推理或确定性重建时也可以使用分布均值,但这会改变随机性和重建特征,必须明确实现选择。
三、KL 正则在约束什么
KL 项让编码后的后验接近标准正态先验:
KL = -0.5 × Σ(1 + log σ² - μ² - σ²)
若完全没有分布约束,latent 可能形成不连续、尺度悬殊的区域,扩散过程更难学习。若 KL 权重过大,编码器会忽略输入细节,所有后验都靠近先验,重建变模糊。
四、为什么还需要感知和对抗损失
单纯像素 MSE 会对多个合理纹理取平均,容易产生平滑结果。感知损失比较预训练网络中的特征,更关注结构与语义;PatchGAN 类对抗损失鼓励局部纹理真实。
| 损失 | 优点 | 代价 |
|---|---|---|
| L1/L2 重建 | 训练稳定,保留整体像素 | 容易模糊 |
| 感知损失 | 更符合视觉相似度 | 依赖特征网络偏好 |
| 对抗损失 | 纹理锐利 | 可能产生虚假细节,训练不稳 |
| KL 正则 | 潜空间连续、尺度规则 | 权重大时损伤重建 |
VAE 的最终目标不是追求某个单项最低,而是在可扩散的 latent 与可接受重建之间平衡。
五、Latent Scaling Factor 为什么容易出错
原始 VAE latent 的标准差未必接近 1。Stable Diffusion 系列常使用配置中的 scaling factor,把编码结果调整到扩散模型训练时的尺度:
z_for_unet = vae.encode(image).sample() * scaling_factor
image = vae.decode(z_for_unet / scaling_factor)
不同模型的因子可能不同,不能把某个常见值硬编码为通用常量。漏乘、漏除或重复处理都会让 U-Net/DiT 看到错误幅度。
scaling factor 是 VAE 与去噪网络之间的接口契约,不是任意的图像增强参数。
六、VAE 的压缩瓶颈决定哪些细节难生成
空间下采样 8 倍意味着 latent 网格的相邻位置间距对应约 8 个像素,但卷积感受野通常大于单个 8×8 小块。极小文字、发丝、网格线和细密图案可能在编码阶段丢失;生成骨干能依据先验补出合理纹理,却无法确认输入中已丢失的真实细节。
增加 latent 通道、降低压缩倍数或使用更强解码器能改善细节,但会提升显存和计算成本。选择应结合目标分辨率和任务:摄影生成、图标、OCR 文字图对 VAE 的要求并不相同。
七、VAE 如何参与训练与图生图
训练潜扩散时,数据图片先经 VAE 编码为 latent,再在 latent 上加入噪声。很多系统会冻结 VAE,避免 latent 分布在扩散训练中漂移。
图生图也先编码输入图,再按 strength 对应的时间步加噪。若 VAE 重建本身已有明显色偏或文字损伤,即使 strength 接近 0,输出也不可能与原图像素完全相同。
八、精度、切片与平铺解码的工程问题
高分辨率解码可能占用大量显存。常见优化包括切片处理 Batch、按空间 Tile 解码和在注意力处使用节省显存的实现。但平铺边界处理不当会出现接缝。
部分 VAE 在 FP16 下容易溢出,因此配置会在解码时临时 upcast 到 FP32。排查 NaN 时应记录 latent 最小值、最大值、均值和 dtype,而不是只看最终黑图。
九、如何评测和选择 VAE
至少要分别评估“重建能力”和“对生成模型的影响”:
固定真实图片 -> encode/decode -> PSNR、SSIM、LPIPS、色差与人评
固定生成 latent -> 不同 VAE decode -> 纹理、色彩、伪影对比
固定完整 Pipeline -> FID/偏好、文字与人脸切片、延迟与显存
例如在 1000 张目标域图片上比较 LPIPS、平均 ΔE 色差和人脸细节失败率。只凭几张精选图判断“某 VAE 更鲜艳”容易把饱和度偏移误当成质量提升。
十、VAE 异常的排查顺序
可以从接口到模型逐层检查:
- 输入 RGB 是否正确归一化到训练范围。
- 通道顺序、颜色空间和 Alpha 处理是否正确。
- latent 通道数和下采样倍数是否匹配。
- scaling factor 是否仅在接口两端各处理一次。
- dtype 是否出现溢出,是否需要 force upcast。
- 模型权重与配置是否来自同一版本。
- 单独 encode/decode 的重建基线是否正常。
十一、常见误区与追问
- 误区:VAE 只影响速度,不影响画质。 压缩瓶颈、损失函数和解码器直接决定颜色与高频细节上限。
- 误区:VAE 是无损编解码器。 即便不做扩散,重建也可能损失文字、细线和纹理。
- 误区:任何 Stable Diffusion VAE 都可以无条件互换。 latent 通道、缩放、分布和训练域必须兼容。
- 误区:输出发灰就提高 CFG。 这可能是 scaling factor 或数值范围错误,CFG 不能修复接口错配。
- 误区:PSNR 最高的 VAE 一定视觉最好。 像素指标可能偏好平滑结果,需要结合 LPIPS 和人评。
- 追问:为什么图生图 strength=0 也可能变化? 输入仍经过有损 VAE 重建,且实现可能存在最小去噪步。
- 追问:为什么冻结 VAE? 固定 latent 分布可以让去噪网络面对稳定训练目标,也节省训练资源。
十二、加强记忆
两端转换:Encoder 把像素压到 latent,Decoder 把 latent 还原为像素;一大收益:空间缩小约 8 倍,主扩散网络计算显著下降;三类约束:重建保内容,感知/对抗保观感,KL 规整分布。关键接口:通道、归一化、scaling factor 和 dtype 必须匹配;能力上限:VAE 丢掉的细节,扩散模型无法准确找回。