ImageNet-1K 圧縮 20 倍、Top-1 精度が初めて 60% を超える: 大規模データセット蒸留の転換点

データの圧縮や蒸留というタスクは、ここ数年大きな注目を集めています。データ圧縮方法は、大規模なデータセットを代表的なコンパクトなサブセットに圧縮することにより、元のデータセットの重要な情報を保持しながら、高速なモデルトレーニングと効率的なデータストレージを実現します。データ圧縮は大量のデータを処理するプロセスで重要な役割を果たすため、研究やアプリケーションにおけるデータ圧縮の重要性は過小評価できません。データ圧縮は、高度なアルゴリズムを採用することで大きな進歩を遂げました。しかし、既存のソリューションは主に、低解像度の小さなデータセットを圧縮するのに優れています。この制限は、2 レベルの最適化プロセス中に多数の展開反復を実行することで発生する膨大な計算オーバーヘッドによるものです。

MBZUAI と CMU チームによる最新の研究SRe2L は、この問題に対処することを目的としています。この研究は、現在、大規模高解像度データセットの蒸留を実現した唯一のフレームワークです。Imagenet-1Kの元の120万データサンプルを0.05Mに圧縮し（圧縮率1:20 ）、蒸留に一般的に使用される224x224の解像度を使用します。ImageNet-1K標準検証セット（val set）で現在最高の60.8%のTop-1精度を達成し、TESLA（ICML'23）の27.9%の精度など、これまでのすべてのSOTA手法をはるかに上回りました。

この作業は、蒸留されたデータ、蒸留プロセス、トレーニング コードを含め、完全にオープン ソースになりました。

論文: https://arxiv.org/abs/2306.13092

コード: https://github.com/VILA-Lab/SRe2L

データセットの蒸留/圧縮タスクの定義と難易度

従来のモデル蒸留は、モデルのパフォーマンスを可能な限り高く保ちながら、よりコンパクトなモデルを取得することを目的としています。対照的に、データセット蒸留タスクは、よりコンパクトで表現力豊かな圧縮データセットを取得する方法に重点を置いています。データ サンプルは元のデータセットよりもはるかに小さくなります (モデルを最初からトレーニングする際の計算オーバーヘッドを節約します)。同時に、モデルは圧縮されたデータセットでトレーニングされ、元のデータ検証セットでテストされて、依然として良好な精度を実現します。

データセット蒸留タスクの主な難しさは、必要なサンプルを効率的かつ実現可能に生成するための生成アルゴリズムをどのように設計するかにあります。生成されたサンプルには、元のデータセットのコア情報が含まれている必要があります。現在、最も一般的に使用されている方法には、勾配マッチング、特徴マッチング、軌跡マッチングなどがありますが、これらの方法の共通の欠点は、大規模なデータセットに拡張できないことです。たとえば、計算と GPU メモリの制限により、標準の ImageNet-1K またはより大きなデータセットを抽出することはできません。必要な計算量と GPU メモリが大きすぎる主な理由は、これらの方法では生成プロセス中に大量の情報を一致させて保存する必要があるためです。現在、多くの GPU メモリは一致させる必要のあるすべてのデータ情報を収容することができないため、これらの方法のほとんどは小さなデータ セットにしか適用できません。

これらの問題に対処するために、新しい論文では、データ生成とモデルトレーニングの 2 つのステップを切り離すことで、3 段階のデータセット蒸留アルゴリズムを提案しています。新しいデータを抽出するプロセスは、元のデータセットで事前トレーニングされたモデルのみに依存するため、計算量とビデオメモリの要件が大幅に削減されます。

ソリューションの核となるアイデア

これまでのデータセット蒸留方法の多くは、サンプル生成とモデルトレーニングの 2 レベル最適化を中心に展開したり、モデルパラメータの軌跡マッチングに基づいて圧縮データを生成するものでした。これらの方法の最大の制限は、スケーラビリティがあまり高くないことです。大量のビデオ メモリと計算が必要であり、ImageNet-1K 全体またはより大きなデータセットに拡張することはできません。

これらの問題に対処するために、著者らはデータ生成とモデルトレーニングを分離し、元のデータの情報抽出プロセスとデータ生成プロセスを互いに独立させる方法を提案しています。これにより、メモリの追加が不要になるだけでなく、元のデータと生成されたデータを同時に処理した場合に、元のデータ内のノイズが生成されたデータにバイアスを与えることも防げます。

具体的には、本論文では、下図に示すように、Squeeze、Recovery、および Relabel (SRe2L) と呼ばれる新しいデータセット圧縮フレームワークを提案しています。このフレームワークは、トレーニング中にモデルと合成データの二重層最適化を 2 つの独立した操作に分離し、さまざまなサイズ、さまざまなモデル アーキテクチャ、および高解像度の画像のデータセットを処理して、効果的なデータセット圧縮を実現します。

私たちのアプローチは、さまざまなデータセット サイズで柔軟性を発揮し、1) 合成画像の任意の解像度、2) 高解像度でのトレーニング コストとメモリ消費量の低さ、3) 任意の評価ネットワーク構造に拡張できる機能など、いくつかの点で複数の利点を示します。この論文では、Tiny-ImageNet および ImageNet-1K データセットに対して多数の実験を実施し、非常に優れたパフォーマンスを実証しています。

3段階データセット蒸留フレームワーク

この論文では、3 段階のデータセット蒸留フレームワークを提案します。

• 最初のステップは、通常実行するモデルトレーニングと同様に、データセット全体のコア情報をモデルに圧縮し、モデルパラメータを通じて元のデータセットに情報を保存することです。
• 2 番目のステップは、トレーニングされたモデル パラメータからこの非常に抽象的な情報を回復することです。この論文では、さまざまな損失と正規化関数が復元された画像の品質とデータセットの蒸留タスクに与える影響について説明します。
• 3 番目のステップは、生成されたデータのカテゴリ ラベルを再調整するという、最も大きな改善が見られるステップでもあります。ここで著者は FKD メソッドを使用して、各作物に対応するソフト ラベルを生成し、それをデータセットの新しいラベルとして保存します。

3 段階のプロセスを次の図に示します。

写真

パフォーマンスとコンピューティングのエネルギー効率

50 IPC（クラスあたり50枚の画像）では、本論文で提案された方法は、Tiny-ImageNetとImageNet-1Kで現在最高のTop-1精度である42.5%と60.8%を達成しており、これはそれぞれ以前の最良の方法よりも14.5%と32.9%高い値です。

さらに、提案された方法は、MTT よりも約 52 倍 (ConvNet-4)、約 16 倍 (ResNet-18) 高速であり、データ合成中に必要なメモリも少なく、MTT 方法と比較してそれぞれ 11.6 倍 (ConvNet-4)、6.4 倍 (ResNet-18) 削減されます。具体的な比較は次の表に示されています。

写真

実験結果

実験のセットアップ

この研究は主に大規模なデータセットの蒸留に焦点を当てているため、ImageNet-Tiny と ImageNet-1K という比較的大きな 2 つのデータセットが実験用に選択されました。バックボーンネットワークについては、本論文ではResNet-{18, 50, 101}、ViT-Tiny、および独自に構築したBN-ViT-Tinyをターゲットモデル構造として使用します。テストフェーズでは、以前の研究と同様に、モデルを最初からトレーニングして圧縮データセットの品質を評価し、ImageNet-Tiny と ImageNet-1K の元の検証セットでのテスト精度を報告します。

ImageNet-1Kデータセット全体の結果

写真

同じ IPC 条件下では、この論文の実験結果は以前の方法 TESLA をはるかに上回っていることがわかります。同時に、この方法で抽出されたデータセットでは、モデル構造が大きいほどトレーニング精度が高くなり、一貫性とスケーラビリティが優れていることが反映されます。

下の図は、パフォーマンスの比較を視覚化したものです。以前の方法である TESLA 蒸留によって得られたデータセットでは、モデルが大きくなるほどパフォーマンスが低下し、大規模なデータセットの蒸留には不利な状況であることがわかります。対照的に、本論文で提案された方法は、モデルが大きいほど精度が高くなり、常識と実際のアプリケーションのニーズにより一致しています。

圧縮データの視覚化

上の図からわかるように、MTT によって生成されたデータ (1 行目と 3 行目) と比較すると、この論文によって生成されたデータ (2 行目と 4 行目) は、品質、明瞭性、意味情報の点で大幅に優れています。

蒸留プロセス画像生成のアニメーション

さらに、50、200 IPC（4K リカバリ バジェット付き）を含む圧縮データセット ファイルは、次のリンクから入手できます: https://zeyuanyin.github.io/projects/SRe2L/

このアプローチを継続学習課題に拡張した結果

写真

上の図は、200 個のカテゴリ (Tiny-ImageNet) を 5 または 10 個の学習ステップに分割し、各ステップでそれぞれ 40 個と 20 個のカテゴリを収容する、5 ステップと 10 ステップの増分学習戦略を示しています。この論文の結果はベースラインのパフォーマンスよりも大幅に優れていることがわかります。

詳細については、元の論文とコードをお読みください。