Improved Training Technique for Shortcut Models (iSM)¶
会议: NeurIPS 2025
arXiv: 2510.21250
代码: 未开源
领域: 图像生成
关键词: Shortcut Models, Flow Matching, 少步生成, CFG引导, 小波损失, 最优传输, EMA策略
一句话总结¶
针对 Shortcut Models 的五大性能瓶颈（指导累积、固定引导、频率偏差、自一致性偏离、弯曲轨迹），提出 iSM 统一训练框架，通过内禀引导、多级小波损失、缩放最优传输和双 EMA 策略，在 ImageNet 256×256 上实现单步 FID 5.27、四步 FID 2.05 的大幅提升。
研究背景与动机¶
领域现状：基于 Flow Matching 的扩散模型在图像生成上取得了卓越成果，但采样需要大量迭代步数，限制了部署效率。加速采样是当前关键研究方向。
Shortcut Models 的优势：Shortcut Models (SM) 是一种优雅的加速方案——网络同时以噪声水平 \(t\) 和目标步长 \(d\) 为条件输入，通过自一致性损失训练，使单个网络支持单步/少步/多步采样。
SM 的困境：尽管框架设计优美，SM 的实际性能却远落后于同类方法（单步 FID 10.60 vs. IMM 7.77），阻碍了广泛采用。
核心痛点：作者识别出五个系统性问题——(1) CFG 引导在大步长下的指数级累积放大效应；(2) 训练时固定引导强度导致推理灵活性丧失；(3) 像素级损失带来的低频偏差使生成图像模糊；(4) EMA 目标网络的时间滞后与自一致性目标的矛盾；(5) 随机噪声-数据配对导致的弯曲生成轨迹。
切入思路：不修改 SM 的核心框架，而是系统性地解决上述五个瓶颈，使 SM 成为真正具有竞争力的生成范式。
核心问题分析¶
问题一：CFG 引导累积（Compounding Guidance）¶
这是本文最重要的发现之一。SM 在训练时使用固定 CFG 强度 \(w=1.5\) 构造自一致性目标。当模型进行单步生成 \(s_\theta(x_0, 0, c, 1)\) 时，它隐式聚合了 \(N=128\) 个中间步的效果。作者首次形式化证明：每个隐式中间步上实际作用的引导强度不是 \(w\)，而是 \(w' = w^{\log_2(N)}\)。当 \(N=128, w=1.5\) 时，\(w' = 1.5^7 \approx 17\)，导致严重的颜色过饱和和模糊伪影。
问题二：固定引导不灵活¶
原始 SM 在训练时将 CFG 强度硬编码为 \(w=1.5\)，推理时无法调整多样性-保真度权衡，且最优 \(w\) 值依赖于推理步数，固定值无法适应不同场景。
问题三：频率偏差（Frequency Bias）¶
SM 使用像素级 \(\ell_2\) 损失优化直接域预测，神经网络天然倾向于先学低频特征，导致生成图像缺乏高频纹理细节，整体偏模糊。
问题四：自一致性偏离¶
EMA 目标网络使用慢衰减率以保持训练稳定性，但这导致目标网络代表的是在线网络的"历史状态"。在线网络同时需要优化当前轨迹（流匹配）和对齐历史目标（自一致性），形成矛盾学习信号。
问题五：弯曲流轨迹¶
标准流匹配中随机的噪声-数据配对导致前向轨迹频繁交叉，迫使逆向生成过程走弯曲路径，增加了大步长预测的难度。
方法详解¶
整体框架¶
iSM 保留 SM 的核心双损失结构（流匹配 + 自一致性），引入四个关键改进组件，最终训练目标为： $\(\mathcal{L}_{\text{total}}(\theta) = \alpha \mathcal{L}_{\text{velocity}}(\theta) + \beta \mathcal{L}_{\text{guidance}}(\theta) + \gamma \mathcal{L}_{\text{consistency}}(\theta)\)$ 其中 \(\alpha = \beta = \gamma = 1\)。
改进一：内禀引导（Intrinsic Guidance）¶
- 功能：将引导强度 \(w\) 作为网络的显式条件输入，训练模型直接输出 CFG 调制后的速度 \(s_\theta(x_t, t, c, d, w)\)
- 核心设计：
- 流匹配目标：在 \(d=0, w=0\) 时训练基础速度场，用标准 dropout 随机加入空条件
- 内禀引导目标：在 \(d=0, w>0\) 时训练模型直接学习 CFG 的缩放行为。目标为 \(s_{\text{velocity}} + w \cdot \text{sg}(s_{\text{guidance}})\)，其中 stop-gradient 防止干扰基础预测
- 引导自一致性目标：在 \(d>0, w \geq 0\) 时保持任意步长和引导强度下的自一致性
- 效果：消除引导累积、支持推理时灵活调节 \(w\)、单步即可使用 CFG、推理时间减半（无需额外的无条件前向传播）
- 区间引导：在高噪声区域（\(t < t_{\text{interval}} = 0.3\)）不施加引导，避免过早模式坍缩
改进二：多级小波损失（Multi-Level Wavelet Loss）¶
- 功能：用离散小波变换（DWT）将预测和目标分解为多频段表示，在小波域计算损失
- 核心设计：递归分解小波子带至 \(L=5\) 级（\(32 \times 32\) 潜空间的最大分解深度），在每个频段独立计算误差
- 效果：引入频率感知的误差信号，迫使模型恢复被 \(\ell_2\) 损失忽视的高频细节，生成更锐利的纹理
改进三：缩放最优传输（Scaling Optimal Transport, sOT）¶
- 功能：解耦 OT 计算的批大小与训练批大小，实现大规模 OT 匹配
- 核心设计：每 \(K\) 个训练 batch（大小 \(M\)），汇聚所有 \(K \times M\) 个噪声-图像样本计算一次全局 OT 方案，然后拆分回 \(K\) 个 mini-batch 进行训练
- 实现细节：\(K=32\)，额外训练时间开销仅约 4%
- 效果：大幅减少前向轨迹交叉，生成更直的逆向路径，降低自一致性和流匹配损失
改进四：双 EMA 策略（Twin EMA）¶
- 功能：维护两套 EMA 参数替代传统单 EMA
- 核心设计：
- 推理参数 \(\theta_{\text{infer}}^-\)：慢衰减率（标准），仅用于推理，保证稳定高质量生成
- 目标参数 \(\theta_{\text{target}}^-\)：快衰减率（\(\rho = 0.95\)），用于生成自一致性目标，紧跟在线网络当前状态
- 效果：在线网络面对的自一致性目标更贴近当前分布，消除"对齐历史版本"的矛盾，同时推理仍享受慢衰减的稳定性
训练与推理细节¶
训练配置¶
- 骨架网络：SiT-XL/2（675M 参数），潜空间大小 \(32 \times 32\)
- 数据集：ImageNet 256×256（类条件生成），使用 sd-vae-ft-mse 编码到潜空间
- 训练迭代：800K（主实验），250K（消融实验）
- 引导尺度采样：\(w\) 从 \([0, w_{\max}=3.5]\) 中以 0.25 步长离散采样；\(w_{\max}\) 过高（5.0）会引入不必要的复杂度，过低（2.0）则缺乏高质量引导
- 区间引导阈值：\(t_{\text{interval}} = 0.3\)，低于此阈值时 \(w=0\)（不施加引导）
- 小波损失：\(L=5\) 级 DWT 分解（\(32 \times 32\) 潜空间的最大可行深度 = \(\log_2(32) = 5\)）
- sOT 参数：\(K=32\)，每 32 个 batch 汇聚样本计算一次全局 OT 方案，额外训练时间约 4%
- Twin EMA：推理参数 \(\theta_{\text{infer}}^-\) 衰减率 0.9999（标准慢衰减），目标参数 \(\theta_{\text{target}}^-\) 衰减率 0.95（快衰减）
- 条件 dropout：训练时随机以一定概率将条件替换为空条件 \(\varnothing\)，遵循标准 CFG 训练范式
推理流程¶
- 推理时使用慢衰减 EMA 参数 \(\theta_{\text{infer}}^-\) 生成样本
- 可灵活选择采样步数 NFE \(\in \{1, 2, 4, 8, 128\}\)，无需重新训练
- 引导强度 \(w\) 在推理时可自由调节，网络直接输出 CFG 调制后的速度——无需额外的无条件前向传播，因此推理时间相比标准 CFG 减少约 50%
- 单步生成时直接 \(x_1 = x_0 + s_\theta(x_0, 0, c, 1, w)\)；多步生成按 Euler 积分迭代
实验关键数据¶
主实验（ImageNet 256×256，SiT-XL/2，800K 迭代）¶
| 模型 | NFE | FID-50K ↓ |
|---|---|---|
| SM（原版） | 1 | 10.60 |
| IMM | 1 | 7.77 |
| iSM（本文） | 1 | 5.27 |
| SM | 2 | — |
| IMM | 2 | 3.99 |
| iSM | 2 | 2.44 |
| SM | 4 | 7.80 |
| IMM | 4 | 2.51 |
| iSM | 4 | 2.05 |
| SM | 8 | — |
| IMM | 8 | 1.99 |
| iSM | 8 | 1.93 |
| SM | 128 | 3.80 |
| iSM | 128 | 1.88 |
额外指标验证（FD-DINOv2 / IS）¶
| 模型 | NFE | FD-DINOv2 ↓ | IS ↑ |
|---|---|---|---|
| SM | 1 | 500.92 | 102.66 |
| IMM | 1 | 247.78 | 128.87 |
| iSM | 1 | 232.31 | 223.52 |
| SM | 2 | 329.53 | 125.66 |
| IMM | 2 | 152.08 | 173.66 |
| iSM | 2 | 107.63 | 302.29 |
| SM | 4 | 265.90 | 136.79 |
| IMM | 4 | 110.88 | 204.95 |
| iSM | 4 | 83.70 | 298.23 |
值得注意的是，iSM 在 FD-DINOv2 指标上的提升更为显著：4 步时 iSM (83.70) 相比 SM (265.90) 降低超过 3 倍，表明改进不局限于 Inception 特征空间。IS 指标同样展现出巨大优势，2 步时 iSM (302.29) 是 SM (125.66) 的 2.4 倍。
消融实验（250K 迭代，逐步叠加，超参数逐块搜索）¶
每个组件的最优超参数在其对应的消融块中确定后，carry forward 到下一组件的搜索中。
| 组件 | FID (1步) | FID (4步) |
|---|---|---|
| Intrinsic Guidance (\(w_{\max}=3.5\)) | 9.62 | 3.17 |
| + Interval Guidance (\(t=0.3\)) | 8.49 | 2.81 |
| + Wavelet Loss (\(L=5\)) | 8.12 | 2.64 |
| + sOT (\(K=32\)) | 7.97 | 2.23 |
| + Twin EMA (\(\rho=0.95\)) | 6.56 | 2.16 |
超参数敏感性要点¶
- \(w_{\max}\)：2.0 → 3.5 → 5.0 对应 1 步 FID 10.10 → 9.62 → 10.38，存在 sweet spot
- \(t_{\text{interval}}\)：0.0 → 0.1 → 0.3 → 0.5 对应 4 步 FID 3.17 → 3.14 → 2.81 → 2.84，0.5 时 1 步 FID 骤升至 19.22（引导覆盖范围过窄）
- DWT 层数：0 → 1 → 3 → 5 层，效果单调递增，5 层为潜空间最大分解限制
- sOT \(K\) 值：0 → 1 → 8 → 32，单步 FID 从 8.12 → 8.07 → 8.03 → 7.97，4 步 FID 从 2.64 → 2.51 → 2.28 → 2.23
- 目标 EMA 衰减率：0.9999 → 0.999 → 0.95，单步 FID 从 7.97 → 7.43 → 6.56，快衰减对单步提升幅度最大
高分辨率泛化（ImageNet 512×512，FlowDCN 架构，300K 迭代）¶
| 模型 | NFE | FID ↓ | Precision ↑ | Recall ↑ |
|---|---|---|---|---|
| SM | 1 | 43.81 | 0.56 | 0.11 |
| iSM | 1 | 37.05 | 0.60 | 0.55 |
| SM | 4 | 12.16 | 0.86 | 0.19 |
| iSM | 4 | 9.94 | 0.78 | 0.62 |
Recall 的巨大提升（0.11 → 0.55, 0.19 → 0.62）说明 iSM 极大改善了样本多样性——这与内禀引导消除引导累积造成的模式坍缩直接相关。4 步时 Precision 略有下降（0.86 → 0.78）符合多样性-保真度权衡的预期。
亮点¶
- 引导累积的形式化：首次严格证明 SM 中 CFG 的指数级累积问题（\(w' = w^{\log_2 N}\)），解释了单步生成伪影的根源，是重要的理论洞察
- 问题诊断全面且系统：不是"头痛医头"，而是一次性识别五个瓶颈并给出统一解法框架
- 每个组件都有明确贡献：消融实验清晰展示各改进的增量收益，从 9.62 → 6.56（单步 FID），设计决策有据可查
- 架构与分辨率泛化：在 FlowDCN + 512×512 上验证了框架通用性，Recall 从 0.11 提升到 0.55 尤其显著
- 训练开销可控：sOT 仅增加 ~4% 训练时间，Twin EMA 额外参数量可忽略，整体改进不引入严重的计算负担
局限与展望¶
- 仅在 ImageNet 上验证：缺少文本到图像（如 T2I）或更大规模数据集上的实验，实际应用价值需进一步确认
- 与蒸馏方法的对比有限：DMD、DMD2 等蒸馏方法在少步生成上有更激进的压缩能力，论文未充分对比
- 单步 FID 仍有差距：与 GAN（StyleGAN-XL 2.30）相比，iSM 单步 5.27 仍有明显差距，SM 框架在极端一步场景的天花板待探索
- sOT 的 \(K\) 值上限：论文止步于 \(K=32\)，更大的 \(K\) 是否会持续收益未知，且随着 \(K\) 增大批间延迟增加
- 小波损失的通用性：多级 DWT 损失写死在潜空间 \(32 \times 32\)，对不同分辨率的潜空间需要重新调整分解层数
与相关工作的对比¶
| 维度 | SM（原版） | IMM | iSM（本文） |
|---|---|---|---|
| 可变步数 | ✅ 单网络多步 | ✅ 单网络多步 | ✅ 单网络多步 |
| CFG 灵活性 | ❌ 训练时固定 | 部分 | ✅ 推理时可调 |
| 单步 FID | 10.60 | 7.77 | 5.27 |
| 四步 FID | 7.80 | 2.51 | 2.05 |
| 频率感知 | ❌ | ❌ | ✅ 小波损失 |
| OT 匹配 | ❌ | ❌ | ✅ 缩放 OT |
| EMA 策略 | 单 EMA | 单 EMA | 双 EMA |
启发与关联¶
- 内禀引导可推广：将 CFG 强度作为网络条件的做法可以直接迁移到其他需要 CFG 的生成模型（如一致性模型、Flow Matching），是一个通用的涨点技巧
- 小波域损失的启示：在潜空间而非像素空间做频率分解是高效且有效的，值得在其他潜空间扩散模型中尝试
- 双 EMA 思想：训练和推理使用不同 EMA 参数的设计思路简洁且有效，可能适用于其他自一致性/自蒸馏框架
- sOT 的跨 batch 池化：以极低代价实现大规模 OT 的做法，可以作为通用的 Flow Matching 训练增强手段
- 与 VeCoR 互补：VeCoR 从速度场的对比正则化角度改进 FM，iSM 从 SM 的训练策略角度改进，两者正交——SM + VeCoR + iSM 可能是值得探索的组合
评分¶
- 新颖性: ⭐⭐⭐⭐ （引导累积的理论分析和双 EMA 设计新颖，但各组件单独看并非全新）
- 实验充分度: ⭐⭐⭐⭐ （消融全面、指标多样，但缺少 T2I 和更大规模验证）
- 写作质量: ⭐⭐⭐⭐⭐ （问题-方案结构清晰，图表精美，公式推导严谨）
- 价值: ⭐⭐⭐⭐ （系统性提升 SM 竞争力，为可变步数生成范式提供了实用指南）
相关论文¶
- [NeurIPS 2025] Scaling Offline RL via Efficient and Expressive Shortcut Models
- [NeurIPS 2025] Show-o2: Improved Native Unified Multimodal Models
- [ICCV 2025] Improved Noise Schedule for Diffusion Training
- [NeurIPS 2025] Training-Free Constrained Generation with Stable Diffusion Models
- [NeurIPS 2025] RLVR-World: Training World Models with Reinforcement Learning