Dataset Distillation as Pushforward Optimal Quantization¶
会议: ICLR2026
arXiv: 2501.07681
代码: 待确认
领域: 模型压缩
关键词: 数据集蒸馏, 最优量化, Wasserstein 距离, 扩散模型, 潜空间聚类
一句话总结¶
将解耦式数据集蒸馏重新形式化为最优量化问题，证明通过扩散先验的潜空间聚类+权重可收敛逼近真实数据分布，提出 DDOQ 算法在 ImageNet-1K 上以极低额外计算量超越 D4M 等基线。
研究背景与动机¶
数据集蒸馏（DD）旨在找到小型合成训练集，使得在其上训练的模型性能接近完整数据训练。早期双层优化方法计算复杂度高且依赖模型架构。解耦方法（如 SRe2L、D4M）通过匹配数据分布并使用生成技术绕过像素空间优化，但缺乏理论保证——没有先前工作从理论上证明蒸馏数据集是否能合理近似原始数据分布。
关键观察：D4M 等方法在潜空间做 \(k\)-means 聚类再解码，本质上是在做 Wasserstein 重心问题（均匀权重），而经典最优量化理论告诉我们加上（自动学习的）权重可以显著减小 Wasserstein 距离。
方法详解¶
整体框架¶
DDOQ（Dataset Distillation by Optimal Quantization）四步流水线： 1. 编码：用 LDM 编码器将训练样本映射到潜空间 \(Z = \mathcal{E}(\mathcal{T})\) 2. 聚类：对每个类别做 mini-batch \(k\)-means（= CLVQ 算法），得到 \(K\) 个质心 \(z_k^{(L)}\) 和对应权重 \(w_k^{(L)}\) 3. 解码：用 LDM 解码器 + 扩散模型生成蒸馏图像 \(x_k^{(L)} = \mathcal{D} \circ \mathcal{U}_t(z_k^{(L)}, \text{emb})\) 4. 加权训练：新模型训练时使用加权损失 \(\min_\theta \sum_{(x,y,w)} w \cdot \ell(x,y,\theta)\)
关键理论¶
定理 1（一致性）：对 VESDE 或 VPSDE 扩散过程，若潜空间分布 \(\mu_T, \nu_T\) 的 Wasserstein-2 距离为 \(\mathcal{W}_2(\mu_T, \nu_T)\)，则经反向扩散到图像空间后： $\(\|\mathbb{E}_{\mu_\delta}[f] - \mathbb{E}_{\nu_\delta}[f]\| \leq C \cdot L \cdot \mathcal{W}_2(\mu_T, \nu_T)\)$ 即扩散生成保持分布的接近性——潜空间的好近似在图像空间也是好近似。
推论 1（收敛率）：随量化点数 \(K\) 增加，近似误差以 \(\mathcal{O}(K^{-1/d})\) 收敛（\(d\) 为潜空间维度），理论上证明了解耦蒸馏方法的一致性。
核心改进：与 D4M 相比仅多了自动确定的权重（CLVQ 聚类过程中自然产生），Wasserstein-2 距离平均降低 15.7%（IPC=10）和 16.1%（IPC=50）。
实验关键数据¶
ImageNet-1K（UNet backbone，ResNet-18 评估）：
| IPC | SRe2L | D4M | RDED | DDOQ |
|---|---|---|---|---|
| 10 | 21.3% | 27.9% | 42.0% | 33.1% |
| 50 | 46.8% | 55.2% | 56.5% | 56.2% |
| 100 | 52.8% | 59.3% | — | 60.1% |
| 200 | 57.0% | 62.6% | — | 63.4% |
- IPC 200 + ResNet-101：DDOQ 68.6% vs D4M 68.1%，相对全精度 69.8% 的误差缩减 30%
- 跨架构泛化（IPC=50）：DDOQ 在 CNN 学生模型上一致优于 D4M（如 MobileNet-V2: 52.1% vs 47.9%）
DiT backbone（DDOQ-DiT）：
| 数据集 | IPC | Minimax-IGD | DDOQ-DiT |
|---|---|---|---|
| ImageNet-1K | 10 | 46.2% | 53.0% |
| ImageWoof | 10 | 43.3% | 48.8% |
| ImageNette | 10 | 65.3% | 68.2% |
- 更强的 DiT backbone 将 ImageNet-1K IPC=10 准确率从 33.1% 提升至 53.0%（+19.9 点）
跨架构泛化详情（IPC=50, ResNet-18 teacher）： - ResNet-18 student: DDOQ 56.2% vs D4M 55.2% - MobileNet-V2 student: DDOQ 52.1% vs D4M 47.9%（+4.2 点） - EfficientNet-B0 student: DDOQ 58.0% vs D4M 55.4%（+2.6 点） - Swin-T student: DDOQ 57.4% vs D4M 58.1%（略低 0.7 点）
Wasserstein 距离分析：加入权重后，蒸馏潜点与编码训练数据的 \(\mathcal{W}_2\) 距离在 IPC=10 平均降低 15.7%，IPC=50 平均降低 16.1%，证实最优量化优于 Wasserstein 重心。
亮点与洞察¶
- 理论贡献扎实：首次证明解耦式蒸馏方法在扩散先验下的一致性和收敛率，填补了该领域的理论空白
- 改进极其简洁：相比 D4M 仅添加自动学习的权重，几乎无额外计算（权重在 \(k\)-means 过程中自然产出）
- 最优量化视角：揭示 \(k\)-means 等聚类方法本质上在求解最优量化问题，权重是 Voronoi 单元的测度
- 扩散模型的理论保证：定理 1 证明扩散生成保持分布接近性，为在潜空间而非像素空间操作提供理论基础
局限性¶
- 低 IPC 设置下（如 IPC=10）仍落后于 RDED 的 patch-based 方法（RDED 42.0% vs DDOQ 33.1%，UNet backbone）
- Swin-T 等 Transformer 学生架构上 DDOQ 略逊于 D4M（57.4% vs 58.1%），可能需要更精细的超参调优
- 收敛率 \(\mathcal{O}(K^{-1/d})\) 随潜空间维度 \(d\) 增大而变慢，对高维潜空间场景效果可能减弱
- 依赖预训练的 LDM/DiT 质量，生成图像的保真度受限于基础模型能力
- 软标签依赖额外预训练分类器（如 ResNet-18），最大性能受限于该分类器精度（69.8%）
- 未探索与扩散引导方法（如 IGD）的结合可能性，二者可能互补
相关工作与启发¶
- 与 D4M 直接对比：仅添加权重即获得一致提升，说明 Wasserstein 重心 → 最优量化的升级是关键
- 与 RDED 比较：RDED 在低 IPC 强但不可扩展，DDOQ 在高 IPC 更优且内存恒定
- 最优量化理论可扩展到其他需要数据近似的场景（如联邦学习中的数据摘要）
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ (最优量化视角 + 一致性证明，理论贡献突出)
- 实验充分度: ⭐⭐⭐⭐ (ImageNet-1K 多 IPC 多架构，但缺少更多数据集)
- 写作质量: ⭐⭐⭐⭐⭐ (理论推导严谨，算法描述清晰)
- 价值: ⭐⭐⭐⭐⭐ (为数据集蒸馏提供理论基础，方法简洁高效)