[论文解读] Squeeze, Recover and Relabel: Dataset Condensation at ImageNet Scale From A New Perspective

Zeyuan Yin, Eric P. Xing+1 more|arXiv (Cornell University)|Jun 22, 2023

Advanced Neural Network Applications被引用 7

一句话总结

本文提出 SRe2L，一种用于大规模数据集压缩的新型三阶段框架，将模型训练与合成数据优化解耦，实现了高效、高分辨率的 ImageNet-1K 数据集压缩。通过在 Squeeze、Recover 和 Relabel 阶段实现解耦训练——利用批量归一化统计匹配和图像块级重标记——该方法在每类仅 50 张图像的情况下，于 ImageNet-1K 上实现了 60.8% 的验证准确率，较之前的方法高出 32.9%，同时将训练时间减少了最多 52 倍。

ABSTRACT

We present a new dataset condensation framework termed Squeeze, Recover and Relabel (SRe$^2$L) that decouples the bilevel optimization of model and synthetic data during training, to handle varying scales of datasets, model architectures and image resolutions for efficient dataset condensation. The proposed method demonstrates flexibility across diverse dataset scales and exhibits multiple advantages in terms of arbitrary resolutions of synthesized images, low training cost and memory consumption with high-resolution synthesis, and the ability to scale up to arbitrary evaluation network architectures. Extensive experiments are conducted on Tiny-ImageNet and full ImageNet-1K datasets. Under 50 IPC, our approach achieves the highest 42.5% and 60.8% validation accuracy on Tiny-ImageNet and ImageNet-1K, outperforming all previous state-of-the-art methods by margins of 14.5% and 32.9%, respectively. Our approach also surpasses MTT in terms of speed by approximately 52$ imes$ (ConvNet-4) and 16$ imes$ (ResNet-18) faster with less memory consumption of 11.6$ imes$ and 6.4$ imes$ during data synthesis. Our code and condensed datasets of 50, 200 IPC with 4K recovery budget are available at https://github.com/VILA-Lab/SRe2L.

研究动机与目标

• 为解决现有双层优化方法在大规模数据集压缩中，尤其是在完整 ImageNet 分辨率下的计算不可行性问题。
• 将模型与合成数据训练的联合优化解耦，以降低内存和计算成本，同时保持高性能。
• 实现在多种模型架构和数据集规模（包括完整的 ImageNet-1K）下的可扩展、高分辨率数据集压缩。
• 相比现有方法，提升压缩图像的质量与语义保真度，尤其在保留结构和类别特异性细节方面。
• 通过在持续学习和与预训练模型集成中的应用，展示方法的实际适用性。

提出的方法

• 该方法提出一种三阶段范式：Squeeze（在真实数据上微调模型以捕获全局统计特性）、Recover（利用 BN 统计匹配从噪声生成合成图像）、Re-label（通过在裁剪区域上的模型预测为合成图像分配精确类别标签）。
• 通过将合成数据生成与真实数据输入隔离，实现双层优化的解耦，避免在训练过程中同时处理真实与合成数据。
• Recover 阶段通过将真实数据集的全局批量归一化统计匹配到合成数据，实现比逐批特征匹配更鲁棒、更全面的分布对齐。
• Re-label 阶段采用图像块级推理策略，为合成图像分配精确标签，提升标签一致性和语义准确性。
• 该框架兼容任何包含 BN 层的预训练模型，支持即插即用集成，降低训练开销。
• 该方法避免双层优化中的展开迭代，显著降低内存和计算成本，同时支持任意输入分辨率和模型架构。

Figure 1 : Left is data synthesis time vs. accuracy on ImageNet-1K with 10 IPC (Images Per Class). Models include ConvNet-4, ResNet-{18, 50, 101}. † indicates ViT with 10M parameters [ 2 ] . Right is the comparison of widely-used bilevel optimization and our proposed decoupled training scheme.

实验结果

研究问题

• RQ1 解耦训练策略是否能在不损失性能的前提下，有效替代大规模数据集压缩中的双层优化？
• RQ2 在整个数据集范围内匹配 BN 统计特征，是否能比逐批特征匹配提供更优的合成数据生成对齐效果？
• RQ3 所提方法是否能在大规模场景（如 ImageNet-1K）下实现高质量、高分辨率的图像生成，且计算成本极低？
• RQ4 在持续学习等下游任务中，SRe2L 框架相较于现有压缩方法表现如何？
• RQ5 该方法在不同模型架构和图像分辨率下的泛化能力如何？

主要发现

• SRe2L 在每类仅 50 张图像的情况下，于 ImageNet-1K 上实现了 60.8% 的验证准确率，较之前最先进方法高出 32.9 个百分点。
• 在 Tiny-ImageNet 上，该方法在 50 IPC 条件下达到 42.5% 的准确率，较之前方法高出 14.5 个百分点。
• 与 MTT（ConvNet-4）相比，该方法快 52 倍，内存消耗降低 11.6 倍；与 MTT（ResNet-18）相比，快 16 倍，内存消耗降低 6.4 倍。
• SRe2L 生成的合成图像在视觉质量上表现更优，语义和轮廓更清晰，而 MTT 生成的图像则模糊且颜色占主导。
• 在 Tiny-ImageNet 上的类别增量学习中，SRe2L 在 5 步和 10 步设置下，随着内存预算增加，始终优于基线方法。
• 该框架支持直接使用现成的、带有 BN 层的预训练模型，降低训练开销，提升实际部署能力。

Figure 2 : Overview of our framework. It consists of three stages: in the first stage, a model is trained from scratch to accommodate most of the crucial information from the original dataset. In the second stage, a recovery process is performed to synthesize the target data from the Gaussian noise.