[論文レビュー] 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

ひとこと要約

本稿では、O(T)からO(1)へGPUメモリ複雑性を低減するため、勾配計算の再構成により、メモリ効率の高いトレースマッチング手法TESLAを提案する。これにより、1クラスあたり最大50枚の画像を用いたImageNet-1Kへのスケーリングが可能となる。さらに、教師モデル出力に基づくソフトラベル割り当て（SLA）を導入し、大規模クラスデータセットでの収束性を向上させ、50 IPCでImageNet-1Kで27.9%のトップ1精度を達成。これは、先行SOTAから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

研究の動機と目的

• ImageNet-1Kのような大規模データセットにスケーリングする際、トレースマッチングベースのデータセット蒸留（MTT）における過大なメモリコストを克服すること。
• MTTがCIFAR-10/100で優れた性能を示すことが、クラス数がはるかに多いImageNet-1Kに対しても一般化されるかを調査すること。
• 合成サンプルにハードラベルを使用した場合、MTTがImageNet-1Kで収束が悪い理由を解明すること。
• 合成サンプルに情報豊富なラベル信号を統合するためのハイパーパrameterフリーかつトレーニング不要な手法を開発すること。
• 従来の手法よりも顕著に高い圧縮率で、ImageNet-1KでSOTAの性能を達成すること。

提案手法

• トレースマッチング勾配をキャッシュ・再配置することで、O(T)からO(1)へメモリ複雑性を低減する、新たな勾配計算手順を提案。TはアンロールされたSGDステップ数を表す。
• 教師モデルの出力（ログティス）を用いて、合成画像のためのソフトラベルを生成する、ソフトラベル割り当て（SLA）を導入。ラベルの共同最適化は行わない。
• 複数の学習エポックからの教師モデルアンサンブルを活用し、合成画像最適化中に動的かつ高品質なソフトラベルを提供する。
• 定数メモリMTT実装とSLAを統合し、1枚のGPUでImageNet-1Kにおける最大50 IPCのエンドツーエンド学習を可能にする。
• SLAでは、追加のハイパーパrameterを一切使用せず、参照モデルの重みのみを用いるため、既存のMTTフレームワークと互換性がある。
• トレーニング不要な推論ベースのラベルインジェクション戦略を採用し、共同ラベル最適化に起因する不安定性を回避する。

実験結果

リサーチクエスチョン

• RQ1 アンロール最適化に起因する高いメモリコストにもかかわらず、トレースマッチングベースのデータセット蒸留（MTT）はImageNet-1Kにスケーリング可能か？
• RQ2 CIFAR-10/100のような小規模データセットで優れた性能を示すMTTの性能が、クラス数がはるかに多いImageNet-1Kに対しても一般化されるか？
• RQ3 合成サンプルにハードラベルを使用した場合、なぜMTTはImageNet-1Kで効果的に収束しないのか？
• RQ4 教師モデルからのソフトラベル割り当ては、追加のハイパーパrameterを導入せずに、大規模クラスデータセットにおけるMTTの性能を顕著に向上させられるか？
• RQ5 1枚のGPUと定数メモリを用いて、50 IPCでImageNet-1KでSOTAの性能を達成することは可能か？

主な発見

• TESLAは、計算オーバーヘッドを最小限に抑えつつ、MTTのメモリ複雑性をO(T)からO(1)へ低減し、ImageNet-1Kへのスケーリングを可能にした。
• 50 IPCでImageNet-1Kで27.9%のトップ1精度を達成。これは、完全データセット学習と比較して僅か5.9%の低下にとどまり、元のデータの4.2%のみを用いている。
• TESLAは、ImageNet-1Kで50 IPCにまでMTTをスケーリングした初の手法であり、先行SOTAから18.2%の絶対的精度向上を達成した。
• 1 IPCでソフトラベル割り当て（SLA）を適用すると、ハードラベルと比較して性能が7.0ポイント向上し、追加のハイパーパrameterは一切不要だった。
• TESLAで得られた蒸留データセットは、アーキテクチャに依存せず良好に一般化され、CIFAR-10、CIFAR-100、ImageNet-1Kで学習したResNet18およびViTにおいて、先行手法を上回った。
• アブレーションスタディにより、SLA単体で1 IPCで精度がハードラベルの0.7%からソフトラベルの7.7%へ向上したことが確認され、初期エポックのソフトラベルが最終エポックのものよりも効果的であることが示された。