[论文解读] Scaling Up Dataset Distillation to ImageNet-1K with Constant Memory

Justin Cui, Ruochen Wang+2 more|arXiv (Cornell University)|Nov 19, 2022

Multimodal Machine Learning Applications被引用 14

一句话总结

本文提出 TESLA，一种内存高效的轨迹匹配数据集蒸馏方法，通过重新组织梯度计算，将 GPU 内存复杂度从 O(T) 降低至 O(1)，从而实现将 ImageNet-1K 的规模扩展至每类最多 50 张图像（IPC）。此外，本文引入了基于教师模型输出的软标签分配（SLA），以提升在大规模类别数据集上的收敛性，在 50 IPC 下于 ImageNet-1K 上实现了 27.9% 的 top-1 准确率——相比之前最先进方法绝对提升 18.2%，且仅使用原始数据的 4.2%。

ABSTRACT

Dataset Distillation is a newly emerging area that aims to distill large datasets into much smaller and highly informative synthetic ones to accelerate training and reduce storage. Among various dataset distillation methods, trajectory-matching-based methods (MTT) have achieved SOTA performance in many tasks, e.g., on CIFAR-10/100. However, due to exorbitant memory consumption when unrolling optimization through SGD steps, MTT fails to scale to large-scale datasets such as ImageNet-1K. Can we scale this SOTA method to ImageNet-1K and does its effectiveness on CIFAR transfer to ImageNet-1K? To answer these questions, we first propose a procedure to exactly compute the unrolled gradient with constant memory complexity, which allows us to scale MTT to ImageNet-1K seamlessly with ~6x reduction in memory footprint. We further discover that it is challenging for MTT to handle datasets with a large number of classes, and propose a novel soft label assignment that drastically improves its convergence. The resulting algorithm sets new SOTA on ImageNet-1K: we can scale up to 50 IPCs (Image Per Class) on ImageNet-1K on a single GPU (all previous methods can only scale to 2 IPCs on ImageNet-1K), leading to the best accuracy (only 5.9% accuracy drop against full dataset training) while utilizing only 4.2% of the number of data points - an 18.2% absolute gain over prior SOTA. Our code is available at https://github.com/justincui03/tesla

研究动机与目标

• 为解决基于轨迹匹配的数据集蒸馏（MTT）在扩展至 ImageNet-1K 等大规模数据集时面临的高昂内存成本问题。
• 探究 MTT 在 CIFAR-10/100 上表现出色的性能是否可泛化至具有更多类别的 ImageNet-1K。
• 解决在使用硬标签表示合成样本时，MTT 在 ImageNet-1K 上收敛性差的问题。
• 开发一种无需超参数、无需训练的策略，将信息丰富的标签信号注入合成样本中，以实现蒸馏过程中的标签注入。
• 在显著高于以往方法的压缩率下，实现在 ImageNet-1K 上的最先进性能。

提出的方法

• 提出一种新颖的梯度计算流程，通过缓存并重排轨迹匹配梯度，将内存复杂度从 O(T) 降低至 O(1)，其中 T 为反向传播的 SGD 步数。
• 引入软标签分配（SLA），利用预训练教师模型的输出生成合成图像的软标签，且无需联合学习标签。
• 利用多个训练周期的教师模型集成，在合成图像优化过程中提供动态且高质量的软标签。
• 将常量内存 MTT 实现与 SLA 结合，实现在单张 GPU 上对 ImageNet-1K 进行端到端训练，支持最高 50 IPC。
• SLA 仅使用参考模型权重，无需额外超参数，因此可与现有 MTT 框架无缝兼容。
• 采用无需训练、基于推理的标签注入策略，避免了联合标签优化带来的不稳定性。

实验结果

研究问题

• RQ1 尽管存在因展开优化带来的高内存消耗，基于轨迹匹配的数据集蒸馏（MTT）是否仍可扩展至 ImageNet-1K？
• RQ2 MTT 在 CIFAR-10/100 等小规模数据集上的优异表现是否可泛化至 ImageNet-1K 等大规模数据集？
• RQ3 为何在使用硬标签表示合成样本时，MTT 在 ImageNet-1K 上难以有效收敛？
• RQ4 从教师模型中引入软标签分配是否能显著提升 MTT 在大规模类别数据集上的性能，且无需引入额外超参数？
• RQ5 是否可能仅使用单张 GPU 和常量内存，在 ImageNet-1K 上实现 50 IPC 下的最先进性能？

主要发现

• TESLA 将 MTT 的内存复杂度从 O(T) 降低至 O(1)，计算开销可忽略不计，从而实现向 ImageNet-1K 的可扩展性。
• 在 50 IPC 下，该方法在 ImageNet-1K 上实现了 27.9% 的 top-1 准确率——仅比全量数据训练下降 5.9%，且仅使用原始数据的 4.2%。
• TESLA 是首个将 MTT 扩展至 ImageNet-1K 上 50 IPC 的方法，在准确率上相比之前最先进方法绝对提升 18.2%。
• 在 1 IPC 条件下，软标签分配（SLA）相比硬标签将性能提升 7.0 个百分点，且无需额外超参数。
• TESLA 蒸馏得到的数据集在不同架构间泛化能力良好，在 CIFAR-10、CIFAR-100 和 ImageNet-1K 上训练 ResNet18 和 ViT 时，均优于以往方法。
• 消融实验表明，仅 SLA 本身即可使性能从硬标签下的 0.7% 提升至软标签下的 7.7%（在 1 IPC 时）；且来自早期训练周期的软标签比最终周期的软标签更有效。