DyWeight: Dynamic Gradient Weighting for Few-Step Diffusion Sampling¶
日期: 2026-03-12
arXiv: 2603.11607
代码: GitHub
机构: 浙江大学、西湖大学 AGI Lab、同济大学、中国科学技术大学
领域: 图像生成 / 扩散模型加速
关键词: diffusion sampling, ODE solver, dynamic weighting, few-step, distillation-free
一句话总结¶
提出 DyWeight，一种基于学习的多步 ODE 求解器，通过放松经典数值约束（∑w≠1）实现动态梯度加权 + 隐式时间校准（time shifting & scaling），将梯度聚合与步长调节隐式耦合，仅需末端监督即可单轮优化——CIFAR-10 上 5 NFE 达到 3.02 FID（iPNDM: 7.77, S4S-Alt: 3.73），FLUX.1-dev 上全面超越 DPM-Solver++ 和 iPNDM。
研究背景与动机¶
-
领域现状: 扩散模型生成质量优异但采样慢（需要数百次函数评估）。多步 ODE 求解器（DPM-Solver++, iPNDM）通过复用历史梯度加速采样，但其系数是手工设计的固定值。
-
现有痛点:
- 经典求解器的系数基于平稳性假设（相邻时间步梯度统计相似），但扩散采样动态高度非平稳——固定权重无法适应变化
- 少步采样时大步长导致严重截断误差，经典小步长假设失效
- 现有学习型求解器（S4S, DLMS, DSS）采用显式解耦范式——独立参数化时间表和系数，导致优化景观复杂（如 S4S-Alt 需要 K=8 的双层循环优化）
-
核心 idea: 放松 ∑w=1 约束，让权重既做梯度聚合（归一化部分）又做隐式时间平移（非归一化部分），将困难的联合优化问题转化为高效的单轮学习。
方法详解¶
1. 动态梯度加权（Dynamic Gradient Weighting）¶
- 传统多步求解器用固定系数聚合历史梯度，约束 ∑w=1 以保证数值精度
- DyWeight 放松此约束，允许权重自由学习——归一化部分自适应聚合梯度，非归一化的权重和隐式调整宏观积分边界（time shifting）
2. 时间校准机制（Time Calibration）¶
- 隐式时间平移（Time Shifting）: 通过放松权重约束实现，动态调整有效步长；实验表明（Fig. 4）time shifting 将训练梯度方差稳定在 1.0 附近，显著改善收敛
- 时间缩放（Time Scaling）: 在微观层面调整送入神经网络的查询时间，修正大步长导致的噪声预期不对齐
- 两者共同消除数值时间线与网络内部去噪动态的错位
3. 无数据蒸馏优化¶
- 教师-学生框架：35 步 iPNDM 作为教师，生成 10k 样本对
- 末端监督（endpoint supervision）：只约束最终输出匹配，允许中间轨迹自由偏离——比路径监督（DLMS, DSS）更灵活
- 像素空间用 Inception 特征空间 L2 距离，潜空间直接用 L2 距离
- 用 Adams-Bashforth 系数初始化提供强归纳偏置，仅 ~600 张图即可获得显著增益
- 单轮优化，无需复杂的交替策略
实验关键数据¶
Table 1: 像素空间 FID↓（50k 样本）¶
| 数据集 | NFE | DyWeight | S4S-Alt | iPNDM | DPM-Solver++ |
|---|---|---|---|---|---|
| CIFAR-10 | 3 | 8.16 | 16.95 | 24.55 | 55.76 |
| CIFAR-10 | 5 | 3.02 | 3.73 | 7.77 | 9.94 |
| CIFAR-10 | 9 | 2.13 | 2.31 | 2.83 | 2.99 |
| FFHQ | 5 | 5.85 | 6.25 | 13.80 | 13.47 |
| AFHQv2 | 3 | 9.16 | 14.71 | — | 35.05 |
| ImageNet64 | 5 | 6.30 | — | 15.54 | 16.87 |
Table 2: FLUX.1-dev 文本到图像（MS-COCO 512×512）¶
| NFE | 指标 | DyWeight | DPM-Solver++ | iPNDM |
|---|---|---|---|---|
| 5 | FID↓ | 21.02 | 21.81 | 25.03 |
| 7 | FID↓ | 20.69 | 22.18 | 22.22 |
| 5 | CLIP↑ | 26.73 | 26.38 | 26.37 |
| 7 | CLIP↑ | 26.49 | 26.20 | 26.17 |
Table 3: FLUX.1-dev DrawBench（NFE=7）¶
- ImageReward↑: DyWeight 0.921 vs DPM-Solver++ 0.863 vs iPNDM 0.852
- HPSv2.1↑: DyWeight 0.288 vs DPM-Solver++ 0.286
消融实验关键发现¶
- 时间校准缺一不可: 去掉 time shifting 或 time scaling 均导致性能下降，同时去掉则大幅退化
- 最优历史阶数 K=3: K 从 1→3 持续提升，K=N（用全部历史）反而因旧梯度噪声积累而退化
- 训练极轻量: CIFAR-10 约 3 分钟（2×RTX 4090），ImageNet-64 约 5 分钟（8×RTX 4090），LSUN-Bedroom 约 9 分钟（8×H100）
- 对教师求解器不敏感: 无论用 iPNDM/DPM-Solver-2/Heun 作教师，最终性能相当
亮点与洞察¶
- 核心洞察：经典求解器的 ∑w=1 约束在少步采样下并非最优，放松后隐式获得时间调节能力——一个简洁优雅的设计
- 末端监督 > 路径监督：允许中间轨迹自由偏离，反而能找到更高效的采样路径
- 在 FLUX.1-dev 等大规模 flow-matching 模型上也有效，且在 DrawBench 分布外 prompt 上泛化良好——说明学到的是通用求解器动态而非过拟合训练分布
- 与蒸馏方法不同，DyWeight 不修改底层模型，推理时无额外开销——即插即用
- 隐式耦合 vs 显式解耦：一个 relaxation 同时解决了梯度聚合和步长调节两个问题，避免了 S4S 那样的多轮双层优化
评分¶
- 新颖性: ⭐⭐⭐⭐ 放松权重约束实现隐式耦合的 idea 简洁深刻
- 实验充分度: ⭐⭐⭐⭐⭐ 6 个像素/潜空间数据集 + FLUX.1 + DrawBench 泛化 + 完整消融
- 价值: ⭐⭐⭐⭐ 训练代价极低的即插即用少步采样加速，实用性强