Hotdry.
归档
mlops

Linum 图像-视频 VAE 优化：架构设计与训练管道实战参数

基于 Linum 4 个月实验，详解图像-视频统一 VAE 的压缩策略、联合训练 Loss 平衡、稳定性修复及多分辨率 Curriculum，提供工程化参数与监控清单。

在图像到视频生成系统中，变分自编码器（VAE）负责将像素空间压缩至连续潜在空间，以降低扩散 Transformer 的计算负担。Linum 通过 4 个月实验优化了图像 - 视频统一 VAE 架构与训练管道，实现高保真重建的同时确保下游生成质量，避免了过度重建噪声导致的潜在空间扭曲。核心观点是：压缩率需平衡内存与保真度，联合训练需精确 Loss 归一化，稳定性依赖规范化层选择，多分辨率需渐进 Curriculum。本文提炼可落地参数与清单，帮助复现类似系统。

VAE 架构设计：3D CNN 与压缩策略

Linum VAE 采用经典 Encoder-Decoder 结构，但全用 Conv3D 处理时空维度，支持图像（填充为 4 帧静态视频，经时下采样塌缩至单帧潜在）和视频统一编码。初始实验从 FLUX.1 启发，尝试 4× 空间 + 4× 时间下采样（有效压缩 12×），但 360p 1s 视频单样本即 OOM 80GB H100。经网格搜索，最终选 8× 空间 + 4× 时间（有效 48× 压缩，考虑 RGB/16 通道）：180p 视频重建无明显 artifact，高运动场景偶现瑕疵。

可落地参数：

  • 空间下采样：Encoder 中 3 层 ×2 下采样（总 8×），Mid-block 前尽早下采样减 Attn 内存。
  • 时间下采样：4×（2 层 ×2），图像 pad 至 4 帧。
  • 潜在通道：16（RGB 3→16 上采样对称）。
  • Attn：仅 Mid-block 单头，因全时空 Attn 内存爆炸。

此配置下，720p 5s@24FPS 像素 110M tokens 压缩至可训规模，下游 DiT 序列长大幅缩减。

联合训练 Loss 制定：模态与尺寸平衡

纯视频 VAE 易训，但联合图像时，重建 Loss（sum over tensor /batch_size）随视频 tensor 尺寸线性放大：180p 2s@24FPS 视频～10× 图像（256×256 pad4 帧）。优化器忽略图像，导致单图模糊 smear faces。

解决方案：改 sum-based Loss 相对固定参考形状归一化，再显式加权模态：

L_recon = (sum(|x - recon|) / ref_shape_elements) * modality_weight
total_L = L_image * w_img + L_video * w_video

参考形状统一为 180p 基准。等权 NaN，图像权重降至 0.25 仍爆；引入 FiLM 层（CNN 隐藏特征经 image/video embedding 调制 scale/shift）辅助区分，但梯度爆炸，后移除。

完整 Loss：

  • 核心：Laplacian NLL（学共享 scale，退化为自适应 L1）。
  • KL：权重 1e-6（仅轻正则）。
  • Perceptual：VGG 特征距离。
  • Adversarial：GAN 判别器锐化细节。
  • 权重示例：w_img=0.1~0.25, w_video=1.0（调优防 NaN）。

重建忽略 pad 重复帧，仅首帧 perceptual/adv。

稳定性修复：Splotches 与规范化

高分辨率（360p+）现 “黑斑”（splotches），类似 LiteVAE。钩子追踪激活 L2-norm，发现 Encoder Mid-block Attn 前均匀，后爆出高 norm 像素簇。

修复链：

  1. Conv3D 块：GroupNorm+Conv → Self-Modulating Convolution (SMC)：通道 - wise 权重 L2-norm 归一 + 学 scale，独立调制防组内主导。
  2. Mid-block Attn：GroupNorm → PixelNorm（像素级 L2 /sqrt (dim+eps)），防 outlier 劫持 Attn map。

加 Adaptive Gradient Clipping (AGC)：EMA 追踪 grad_norm /weight_norm，clip 超阈通道。

监控清单：

指标 阈值 异常行动
激活 L2-norm (各层) <10 绘图定位 splotch 源头
梯度 norm <1e3 启用 AGC，clip ratio>1.5
Loss NaN 0 降 w_img，检查 scale param
重建 FID (rFID) <0.5 (180p) 过滤低质样本

多分辨率 Curriculum 与下游注意

顺序 fine-tune 180p→360p→720p 导致低分辨 “灾难性遗忘”。改混合采样：持续低分辨 + 渐增高分辨，超参扫 Loss 权重：

  • 180p: 1.1
  • 360p: 0.1
  • 720p: 0.01

高分辨视觉冗余多，易学，故低权重。最终虽重建优，但 Linum 发现更好 rFID 不增下游 gFID：过度拟合 JPEG 噪声扭曲语义潜在空间。建议：数据集过滤压缩 artifact，VAE 加 alignment 正则（如 DINO 特征对齐）。

训练管道清单：

  1. 数据：视频 180p/360p/720p 混合，图像 pad4 帧；预滤低质（blocky faces, shaky）。
  2. 预训：视频 only 1 周基线。
  3. 联合：3 个月，curriculum 2 周。
  4. 嵌入：离线全数据集 → latents，存 HDF5。
  5. 硬件：H100 80GB，FSDP 分片。
  6. 回滚：若 NaN，降 batch_size 至 1，w_img=0.05。

Linum 最终用 Wan 2.1 VAE 生产（更小快），但开源自研 VAE 供复用。此优化适用于任何 image-to-video 系统，确保 VAE 不仅是压缩器，更是 “学得好的” 潜在桥接。

资料来源： [1] Linum Field Notes: Better Reconstruction ≠ Better Generation, https://linum.ai/field-notes/vae-reconstruction-vs-generation
[2] Linum.ai 主页 & Field Notes, https://linum.ai/field-notes

查看归档