データセット蒸留とは？簡単な概要

モデルのトレーニングは、データサイエンティストの仕事で最も時間のかかる部分のように思えるかもしれません。しかし実際には、その時間の多く、しばしば60～80%が、データの収集、クリーニング、モデリングに向けた整理などのデータの準備に費やされています。データセットが大規模になるほど、この準備時間も増加し、実験の速度が低下して反復が難しくなります。

この課題に対処するため、研究者たちはトレーニングを効率化する方法を何年も模索してきました。合成データ、データセット圧縮、より優れた最適化手法などのアプローチは、いずれも大規模データセットを扱う際のコストと手間を削減し、機械学習ワークフローを高速化することを目指しています。

ここで重要になるのは、完全なデータでモデルをトレーニングした場合と同じ性能を維持しながら、データセットを大幅に縮小できるかという問いです。データセット蒸留は、その有望な答えの1つです。

データセット蒸留では、大規模なトレーニングデータセットの重要なパターンを保持しながら、コンパクトなバージョンを作成します。これにより、トレーニングの高速化、必要な計算リソースの削減、より効率的な実験が可能になります。これは、モデル向けの学習用チートシートのようなものだと考えられます。完全なデータセットと同じ中核パターンを教えるよう設計された、少数の合成データの例です。

この記事では、データセット蒸留の仕組みと、実世界のアプリケーションにおけるスケーラブルな機械学習およびディープラーニングをどのように支援するかを解説します。それでは始めましょう。

データセット蒸留を理解する#

データセット蒸留とは、大規模なトレーニングデータセットを、元のデータセットとほぼ同じ情報をモデルに学習させられる、はるかに小さなデータセットへ凝縮するプロセスです。多くの研究者は、このプロセスをデータセット凝縮とも呼びます。これは、データセット全体に現れる重要なパターンを捉えることが目的だからです。

蒸留データセットは、ランダムに生成された合成データや、実画像から単純に小さなサブセットを選んだものとは異なります。ランダムな偽データセットでも、元のデータセットを縮小したコピーでもありません。

その代わり、最も重要なパターンを捉えるよう意図的に最適化されます。このプロセスでは、ニューラルネットワークが蒸留データでトレーニングされた際に、データセット全体でトレーニングした場合とほぼ同じように学習できるよう、すべてのピクセルと特徴が調整・最適化されます。

このアイデアは、Tongzhou Wang、Jun-Yan Zhu、Antonio Torralba、Alexei A. Efrosによる2018年のarXiv論文で初めて提案されました。初期のテストでは、MNISTやCIFAR-10のような単純なデータセットが使用され、少数の蒸留サンプルで数千枚の実画像を代替できることを容易に示せました。

図1. 画像データにデータセット蒸留を使用（出典）

それ以来、後続の研究によってデータセット蒸留はさらに発展し、ICMLやICLRで発表された手法などにより、凝縮の効率とスケーラビリティが向上しています。

データセット蒸留の意義#

データセット蒸留はトレーニング効率を高め、開発サイクルを高速化します。モデルが学習する必要のあるデータ量を削減することで、計算要件を低減できます。

これは、モデルを時間とともに更新する継続学習、多数のモデル設計をテストするニューラルアーキテクチャ探索、メモリと電力に制約のある小型デバイス上でモデルを実行するエッジトレーニングに特に有用です。全体として、こうした利点により、データセット蒸留は迅速な初期化、素早いファインチューニング、さまざまな機械学習ワークフローにおける初期プロトタイプの構築に適した選択肢となります。

データセット蒸留の仕組みの概要#

データセット蒸留では、合成、つまり人工的に生成されたトレーニングサンプルを作成します。これらのサンプルは、実データでトレーニングした場合に近い方法でモデルが学習するのを助けます。通常のトレーニング中に、3つの重要な要素を追跡することで機能します。

1つ目は損失関数です。これは、予測がどれだけ間違っているかを示すモデルの誤差スコアです。2つ目はモデルパラメータで、学習に伴って更新されるネットワーク内部の重みです。

3つ目はトレーニング軌跡で、時間の経過に伴って誤差と重みがステップごとにどのように変化するかを表します。その後、モデルが合成サンプルでトレーニングされたときに誤差が低下し、完全なデータセットを使った場合と同じように重みが更新されるよう、合成サンプルが最適化されます。

データセット蒸留をステップごとに見る#

データセット蒸留プロセスの仕組みを詳しく見ていきましょう。

• ステップ1 - 合成ピクセルを初期化する: プロセスは、学習可能な入力として機能する合成画像から始まります。最初、これらの画像には構造がほとんどなく、白紙のように見えます。時間の経過とともに、情報量の多い例へ最適化されます。
• ステップ2 - 勾配マッチングとバックプロパゲーションで最適化する: モデルがこれらの合成画像でトレーニングされると、実データのトレーニング挙動により近づくよう各ピクセルをどのように変更すべきかを示す勾配が生成されます。バックプロパゲーションは、ネットワークが誤りから学習するための手法です。誤差をモデル内で逆方向に伝播させ、どのピクセルと重みが原因となったかを特定してから、それらを少しずつ更新します。これらの勾配を使用して、バックプロパゲーションは合成画像をステップごとに調整し、トレーニングにとってより有益なものにします。
• ステップ3 - トレーニングステップ間の挙動を一致させる: この手法では、トレーニング軌跡、つまりモデルが学習中にたどるステップごとの変化も一致させます。これにより、蒸留データセットが、完全なデータセットを使った場合にたどる学習経路と同様の経路へモデルを導きます。
• ステップ4 - 検証と汎化: 最後に、蒸留データセットを実際の検証データで評価し、トレーニング済みモデルが新しい例に対してどの程度良好に動作するかを確認します。これにより、合成データが特定のサンプルをモデルに暗記させるのではなく、広く機能するパターンを教えていることを検証します。

図2. データセット蒸留の概要（出典）

データセット蒸留の主要な手法#

すべてのデータセット蒸留手法は、そこに至るアルゴリズムが異なっていても、同じ中核的な考え方に基づいています。ほとんどのアプローチは、性能マッチング、分布マッチング、パラメータマッチングの3つに分類されます。

次に、それぞれを見て仕組みを確認しましょう。

性能マッチング#

データセット蒸留における性能マッチングは、完全な元のデータセットでトレーニングした場合とほぼ同じ精度にモデルが到達できる、極めて小さく最適化されたトレーニングセットの作成に焦点を当てます。ランダムなサブセットを選ぶのではなく、蒸留サンプルを最適化し、それらでトレーニングしたモデルが、元のデータセットでトレーニングしたモデルと同様の予測、トレーニング中の同様の損失挙動、または同様の最終精度を得られるようにします。

このプロセスの改善には、メタラーニングがよく使用されます。蒸留データセットは、繰り返し行われるトレーニングエピソードを通じて更新されるため、さまざまな状況に対応できるようになります。

これらのエピソードでは、現在の蒸留サンプルから学生モデルが学習する様子をシミュレートし、その学生モデルが実データでどの程度適切に動作するかを確認してから、蒸留サンプルをより優れた教師となるよう調整します。時間の経過とともに、蒸留セットは、学生モデルが異なる初期重みから開始した場合や異なるアーキテクチャを使用した場合でも、高速な学習と高い汎化性能を支えられるようになります。これにより、蒸留データセットの信頼性が高まり、単一のトレーニング実行に依存しなくなります。

図3. メタラーニングプロセス (出典)

分布マッチング手法#

一方、分布マッチングでは、実データセットの統計的パターンに一致する合成データを生成します。モデルの最終精度だけに焦点を当てるのではなく、学習中にニューラルネットワークが生成する内部特徴に着目します。

次に、分布マッチングを支える2つの手法を見ていきましょう。

単一層分布マッチング#

単一層分布マッチングでは、ニューラルネットワークの1つの層に焦点を当て、実データと合成データに対してその層が生成する特徴を比較します。これらの特徴は活性化とも呼ばれ、ネットワーク内のその時点でモデルが学習した内容を捉えます。

合成データから同様の活性化が生成されるようにすることで、この手法は蒸留データセットが元のデータセットと同じ重要なパターンを反映するよう促します。実際には、選択した層の活性化が実画像の活性化とほぼ一致するまで、合成サンプルを繰り返し更新します。

このアプローチは一度に1つの表現レベルだけを整合させるため、比較的単純です。小規模なデータセットや、深い多段階の特徴階層を一致させる必要がないタスクで特に効果を発揮します。1つの特徴空間を明確に整合させることで、単一層マッチングは蒸留データセットを使った学習に安定した意味のある信号を提供します。

多層分布マッチング#

多層分布マッチングは、実データと合成データの比較を1つの層だけでなく、ニューラルネットワークの複数の層で行うという考え方に基づいています。層ごとに異なる種類の情報を捉え、初期層では単純なエッジやテクスチャを、深い層では形状やより複雑なパターンを捉えます。

これらの層全体で特徴を一致させることで、蒸留データセットはモデルが複数のレベルで学習する内容を反映するようになります。ネットワーク全体で特徴を整合させるため、このアプローチは、モデルがクラスを区別する際に依存する、より豊かな信号を合成データに保持させるのに役立ちます。

これは、モデルが画像や動画を理解するタスクを意味するコンピュータビジョンで特に有用です。役立つパターンが多くの層に分散しているためです。複数の深さで特徴分布が適切に一致すると、蒸留データセットは元のトレーニングデータの、より強力で信頼性の高い代替として機能します。

パラメータマッチング手法#

データセット蒸留におけるもう1つの主要カテゴリは、パラメータマッチングです。精度や特徴分布を一致させるのではなく、トレーニング中にモデルの重みがどのように変化するかを一致させます。蒸留データセットでのトレーニングによって、実データでのトレーニングと同様のパラメータ更新が生じるようにすることで、モデルはほぼ同一の学習経路をたどります。

次に、2つの主要なパラメータマッチング手法を説明します。

単一ステップマッチング#

単一ステップマッチングでは、実データで1回だけトレーニングステップを実行した後、モデルの重みに何が起こるかを比較します。その後、蒸留データセットを調整し、それで1ステップだけトレーニングしたモデルが非常に似た重み更新を行うようにします。単一の更新だけに焦点を当てるため、この手法は単純で高速に実行できます。

欠点は、特にモデルがより豊かな特徴を構築するために多数の更新を必要とする難しいタスクでは、1ステップでは完全な学習プロセスを反映できないことです。そのため、単一ステップマッチングは、有用なパターンをすばやく学習できる単純な問題や小規模なデータセットで最も適切に機能する傾向があります。

多ステップパラメータマッチング#

これに対して、多ステップパラメータマッチングでは、1ステップだけでなく、複数のトレーニングステップにわたってモデルの重みがどのように変化するかを調べます。この更新の連続が、モデルのトレーニング軌跡です。

蒸留データセットは、モデルが合成サンプルでトレーニングされたときに、実データでたどる軌跡に近い軌跡をたどるように構築されます。より長い学習過程を一致させることで、蒸留セットは元のトレーニングプロセスの構造をより多く捉えます。

学習が時間とともに進む様子を反映するため、多ステップマッチングは通常、モデルが有用なパターンを学習するために多くの更新を必要とする、より大規模または複雑なデータセットで優れた性能を発揮します。複数のステップを追跡する必要があるため計算量は増えますが、単一ステップマッチングよりも汎化性能と精度に優れた蒸留データセットを生成できることが多くあります。

合成データセットの生成と最適化の仕組み#

主要な蒸留アプローチを理解したところで、次に合成データがどのように作られるかを見ていきましょう。データセット蒸留では、最も重要な学習信号を捉えるよう合成サンプルを最適化するため、小さなセットで、はるかに大規模なデータセットを置き換えられます。

次に、この蒸留データがどのように生成・評価されるかを見ていきます。

蒸留画像の作成と評価#

データセット蒸留では、合成ピクセルが多数のトレーニングステップにわたって更新されます。ニューラルネットワークは現在の合成画像から学習し、実データセットのパターンにより適合するよう各ピクセルをどのように変更すべきかを示す、勾配ベースのフィードバックを送ります。

これは、プロセスが微分可能（つまり、すべてのステップが滑らかで明確に定義された勾配を持つため、小さなピクセルの変化が損失の予測可能な変化につながる）だからこそ可能であり、勾配降下法によってモデルが合成データを滑らかに調整できます。

最適化が進むにつれて、合成画像はモデルが認識する形状やテクスチャなど、意味のある構造を形成し始めます。これらの洗練された合成画像は、分類器が学習に必要とする主要な視覚的手がかりを捉えるため、画像分類タスクでよく使用されます。

蒸留データセットは、それらでモデルをトレーニングし、実データでトレーニングしたモデルと結果を比較することで評価されます。研究者は検証精度を測定し、クラスを区別するためにモデルが依存する識別特徴（あるクラスと別のクラスを区別するためのパターンや信号）が保持されているかを確認します。また、蒸留データが過学習を引き起こさないことを確認するため、異なる実行やモデル構成に対する安定性と汎化性能もテストします。

データ蒸留の実世界での応用#

次に、データが限られている場合や非常に専門的な場合でも高い性能を維持しながら、蒸留データセットがトレーニングを高速化し、計算コストを削減する例を詳しく見ていきます。

コンピュータビジョンアプリケーションでのデータセット蒸留の利用#

コンピュータビジョンでは、画像や動画などの視覚データを理解できるようモデルをトレーニングすることが目標です。モデルはエッジ、テクスチャ、形状、物体などのパターンを学習し、それらを画像分類、物体検出、セグメンテーションなどのタスクに利用します。ビジョンタスクでは照明、背景、視点に大きなばらつきがあることが多いため、コンピュータビジョンシステムは通常、十分に汎化するために大規模なデータセットを必要とし、トレーニングにコストと時間がかかります。

図4. データセット蒸留の例（出典）

医療スキャン、野生動物のモニタリング、工場での欠陥検出などの画像分類ユースケースでは、モデルは精度とトレーニングコストの難しいトレードオフに直面することがよくあります。これらのタスクでは通常、大規模なデータセットが必要です。

データセット蒸留では、元のトレーニングセットを、分類器にとって最も重要な視覚的手がかりを含む少数の合成画像へ圧縮できます。ImageNetのような大規模ベンチマークでは、元の画像のわずか4.2%しか使用しない蒸留セットでも、高い分類精度を維持できることが示されています。つまり、ごく小さな合成プロキシで、はるかに少ない計算量で数百万件の実サンプルを置き換えられます。

ニューラルアーキテクチャ探索#

ニューラルアーキテクチャ探索（NAS）は、さまざまなニューラルネットワーク設計を自動的に探索し、タスクに最適なものを見つける手法です。NASでは多数の候補モデルをトレーニングして評価する必要があるため、完全なデータセットで実行すると時間がかかり、計算負荷も非常に高くなります。

データセット蒸留は、元データの主要な学習信号を保持した小さな合成トレーニングセットを作成することで、この課題を解決します。そのため、各候補アーキテクチャをはるかに高速にテストできます。これにより、NASは設計を効率的に比較しながら、良いアーキテクチャと悪いアーキテクチャの順位を十分な信頼性で維持でき、最終的なモデル品質を大きく損なわずに探索コストを削減できます。

継続学習とエッジデプロイメント#

継続学習システム、つまり一度だけトレーニングするのではなく新しいデータが届くたびに更新を続けるモデルには、高速でメモリ効率の高い更新が必要です。エッジデバイスであるカメラ、スマートフォン、センサーも、計算資源とストレージの予算が限られているため、同様の制約に直面します。

データセット蒸留は、どちらの場合にも、大規模なトレーニングセットを小さな合成セットへ圧縮することで役立ちます。これにより、モデルは完全なデータセットではなく、小規模なリプレイセットを使って適応や再トレーニングを行えます。たとえば、カーネルベースのメタラーニングの研究では、わずか10個の蒸留サンプルで、標準的な画像分類ベンチマークであるCIFAR-10において64%を超える精度を達成できることが示されました。リプレイセットが非常にコンパクトであるため、特にモデルを頻繁に更新する必要がある場合に、更新がはるかに高速で実用的になります。

データセット蒸留は、大規模言語モデルの知識蒸留と組み合わせて使用することもできます。小規模な蒸留データセットに教師モデルから最も重要なタスク信号を保持させることで、圧縮された学生モデルを、性能を大きく損なわずにより効率的にトレーニングまたは更新できます。これらのデータセットは非常に小さいため、ストレージと計算能力が限られていても、更新後のモデル精度を維持したいエッジやオンデバイスの用途に特に有用です。

データ蒸留のメリットとデメリット#

データセット蒸留を使用するメリットは次のとおりです。

• 迅速な実験に最適です。 新しいアーキテクチャ、損失関数、ハイパーパラメータを、毎回大規模なデータセットで再トレーニングすることなくテストできます。
• プライバシー面で有利な可能性があります。 蒸留された合成サンプルの共有は、実際のユーザーデータを共有するより安全な場合があります。生の例が直接公開されないためです。
• 単純なサブセット選択より優れていることがよくあります。 例を選ぶだけでなく、最大限に有益になるよう積極的に最適化するためです。

データセット蒸留にはいくつかの利点がある一方、次のような制限にも注意が必要です。

• 過学習: 蒸留データは、蒸留に使用したアーキテクチャで最もよく機能することが多く、大きく異なるモデルにはうまく転移しない場合があります。
• ハイパーパラメータの影響を受けやすいです。 結果は、学習率、初期化、蒸留ステップ数などの要因に大きく左右される可能性があります。
• 実世界の複雑さへのスケールが困難です。 ベンチマークで良好に機能する手法でも、大規模で雑多なデータセットや高解像度のデータセットでは精度が低下する可能性があります。

主なポイント#

データセット蒸留により、少数の合成サンプルで完全なデータセットとほぼ同じ効果でモデルを学習させられます。これにより、機械学習はより高速で効率的になり、スケールしやすくなります。モデルが大規模化してより多くのデータを必要とする中、蒸留データセットは精度を犠牲にせず計算コストを削減する実用的な方法を提供します。

コミュニティに参加し、GitHubリポジトリを確認して、AIについてさらに詳しく学びましょう。独自のビジョンAIプロジェクトの構築を検討している場合は、ライセンスオプションをご確認ください。ヘルスケアにおけるAIや小売業におけるビジョンAIなどのアプリケーションについては、ソリューションページをご覧ください。