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

モデルのトレーニングは、データサイエンティストの業務の中で最も時間がかかる工程のように思えるかもしれません。しかし、実際の業務時間の大部分（60%から80%に及ぶこともあります）は、データの準備に費やされています。具体的には、モデリングのためのデータの収集、クリーニング、整理です。データセットが拡大するにつれてこの準備時間も増大し、実験の遅延やイテレーションの停滞を招いています。

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

ここで重要となるのが、データセットを劇的に縮小しても、完全なデータでモデルをトレーニングした場合と同じパフォーマンスを維持できるかという問いです。データセット蒸留は、この問いに対する有望な答えの一つです。

これは、大規模なトレーニングデータセットから、モデルが効果的に学習するために必要な本質的なパターンを保持したコンパクトなバージョンを作成する技術です。これにより、トレーニングの高速化、計算リソースの削減、より効率的な実験が可能になります。これはモデルのための学習用「カンニングペーパー」のようなものと考えられ、フルデータセットのコアパターンを教えるために設計された小さな合成データセットです。

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

Link to this sectionデータセット蒸留の理解#

データセット蒸留は、大規模なトレーニングデータセットを、元のデータセットとほぼ同等の情報をモデルに学習させることができる、はるかに小さなデータセットに凝縮するプロセスです。多くの研究者はこのプロセスを「データセット要約（dataset condensation）」とも呼んでいますが、その目的はデータセット全体に現れる本質的なパターンを捉えることにあります。

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

その代わり、最も重要なパターンを捉えるように意図的に最適化されています。この過程で、ニューラルネットワークが蒸留データでトレーニングされた際、まるでデータセット全体でトレーニングされたかのように学習が進むよう、すべてのピクセルと特徴量が調整・最適化されます。

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

図1. 画像データに対するデータセット蒸留の活用 (出典)

それ以降、研究はさらに進展しており、ICMLやICLRで発表された手法などにより、データセット蒸留はより効率的かつスケーラブルになっています。

Link to this sectionデータセット蒸留の重要性#

データセット蒸留は、トレーニング効率を向上させ、開発サイクルを高速化します。モデルが学習に必要なデータ量を減らすことで、計算要件が低減されます。

これは特に、モデルが時間をかけて更新される継続的学習、多くのモデル設計がテストされるニューラルアーキテクチャ探索、そしてメモリや電力が制限された小型デバイス上でモデルを動かすエッジトレーニングにおいて有用です。全体として、これらの利点はデータセット蒸留を、多くの機械学習ワークフローにおける迅速な初期化、短時間でのファインチューニング、および早期プロトタイプ構築のための優れた選択肢にしています。

Link to this sectionデータセット蒸留の仕組みの概要#

データセット蒸留は、人工的に生成されたトレーニングサンプルを作成します。これらのサンプルは、実データでトレーニングするのと非常によく似た方法でモデルの学習を助けます。仕組みとしては、通常のトレーニング過程で3つの重要な要素を追跡します。

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

3つ目はトレーニング軌跡（training trajectory）で、これはエラーと重みが時間とともに段階的にどのように変化するかを表します。合成サンプルは、モデルがそれらを用いて学習する際に、フルデータセットを用いた場合と同じようにエラーが減少し、重みが更新されるように最適化されます。

Link to this sectionデータセット蒸留のステップ・バイ・ステップ#

データセット蒸留プロセスの詳細を解説します。

• ステップ1 - 合成ピクセルの初期化: このプロセスは、学習可能な入力として機能する合成画像から始まります。当初、これらの画像は構造がほとんどなく、白紙の状態に見えます。時間が経つにつれて、それらは情報豊かな例へと最適化されていきます。
• ステップ2 - 勾配マッチングとバックプロパゲーションによる最適化: モデルがこれらの合成画像を学習する際、各ピクセルが実データのトレーニング動作により近づくためにどのように変化すべきかを示す勾配が生成されます。バックプロパゲーションは、ネットワークがミスから学ぶために使用する手法です。エラーをモデルを通じて逆伝播させることで、どのピクセルと重みがその原因であったかを突き止め、わずかに更新します。それらの勾配を利用して、バックプロパゲーションは合成画像を段階的に調整し、トレーニングにとってより有益なものにしていきます。
• ステップ3 - トレーニングステップ間での動作の照合: この手法はトレーニング軌跡も照合します。これは、モデルが学習中にたどる段階的な変化を意味します。これにより、蒸留されたデータセットは、フルデータセットを使用した場合と同じような学習パスをモデルにガイドします。
• ステップ4 - 検証と汎化: 最後に、蒸留されたデータセットを実際の検証データで評価し、トレーニングされたモデルが新しい例に対してどの程度機能するかを確認します。これにより、合成データがモデルに特定のサンプルを記憶させるのではなく、広範で機能的なパターンを教えていることを確認します。

図2. データセット蒸留の概観 (出典)

Link to this section主要なデータセット蒸留手法#

すべてのデータセット蒸留手法は、異なるアルゴリズムを使用していても、同じ中心的なアイデアに基づいています。ほとんどのアプローチは、性能マッチング、分布マッチング、パラメータマッチングの3つのカテゴリに分類されます。

次に、それぞれの内容と仕組みについて見ていきましょう。

Link to this section性能マッチング#

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

メタラーニングはこのプロセスを改善するために一般的に使用される手法です。蒸留されたデータセットは、繰り返しトレーニングのエピソードを通じて更新されるため、多くの可能な状況で効果的になります。

これらのエピソード中、この手法は現在の蒸留サンプルから学生モデルがどのように学習するかをシミュレートし、その学生モデルが実データでどの程度うまく機能するかを確認し、蒸留サンプルをより良い教師になるように調整します。時間が経つにつれて、蒸留セットは、学生モデルが異なる初期重みから始まる場合や異なるアーキテクチャを使用する場合でも、高速な学習と強力な汎化をサポートするように学習します。これにより、蒸留データセットは単一のトレーニング実行に縛られず、より信頼性の高いものになります。

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

Link to this section分布マッチング技術#

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

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

Link to this section単一層の分布マッチング#

単一層の分布マッチングは、ニューラルネットワークの単一の層に集中し、実データと合成データに対してその層が生成する特徴を比較します。アクティベーションとも呼ばれるこれらの特徴は、モデルがその時点でネットワーク内で何を学習したかを捉えています。

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

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

Link to this section多層分布マッチング#

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

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

これは特に、画像やビデオを理解することを学習するコンピュータビジョンにおいて、有益なパターンが多くの層に分散しているため非常に有効です。複数の深さで特徴分布がよく一致すると、蒸留されたデータセットは元のトレーニングデータのより強力で信頼できる代用品として機能します。

Link to this sectionパラメータマッチング手法#

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

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

Link to this section単一ステップマッチング#

単一ステップマッチングは、実データでのわずか1回のトレーニングステップの後、モデルの重みに何が起こるかを比較します。蒸留データセットは、1ステップだけ学習したモデルが非常に類似した重み更新を生成するように調整されます。この単一の更新にのみ焦点を当てているため、この手法は直接的であり、実行も高速です。

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

Link to this section多ステップパラメータマッチング#

対照的に、多ステップパラメータマッチングは、1ステップだけでなく、いくつかのトレーニングステップにわたってモデルの重みがどのように変化するかを見ます。この更新のシーケンスがモデルのトレーニング軌跡です。

蒸留データセットは、モデルが合成サンプルを学習する際に、その軌跡が実データでの軌跡と密接に従うように構築されます。学習のより長い区間を照合することで、蒸留セットは元のトレーニングプロセスにおけるより多くの構造を捉えます。

学習が時間とともにどのように展開するかを反映するため、多ステップマッチングは通常、モデルが有用なパターンを拾い上げるために多くの更新を必要とする、より大きくまたはより複雑なデータセットに対してよりよく機能します。複数のステップを追跡する必要があるため計算量は増えますが、単一ステップマッチングよりも汎化性能が高く、優れたパフォーマンスを発揮する蒸留データセットを生成することが多いです。

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

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

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

Link to this section蒸留画像の作成と評価#

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

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

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

蒸留されたデータセットは、それらを使ってモデルをトレーニングし、実データでトレーニングされたモデルの結果と比較することで評価されます。研究者は検証精度を測定し、合成セットがクラスを分離するために必要な識別的特徴（モデルが一方のクラスと他方のクラスを区別するために依存するパターンや信号）を保持しているかどうかを確認します。また、蒸留されたデータが過学習を招かないように、異なる実行やモデル設定間での安定性と汎化性能もテストされます。

Link to this sectionデータ蒸留の実世界のアプリケーション#

次に、データが制限されている場合や非常に専門的な場合でも、どのようにして蒸留データセットがトレーニングを加速し、計算コストを削減しつつ強力なパフォーマンスを維持できるか、その例を詳しく見ていきます。

Link to this sectionコンピュータビジョンアプリケーションへのデータセット蒸留の活用#

コンピュータビジョンにおいては、画像やビデオなどの視覚データを理解するようにモデルをトレーニングすることが目標です。これらのモデルはエッジ、テクスチャ、形状、オブジェクトといったパターンを学習し、画像分類、オブジェクト検出、セグメンテーションといったタスクに利用します。ビジョンの問題は、照明、背景、視点の変化が非常に激しいため、コンピュータビジョンシステムは通常、汎化するために大規模なデータセットを必要とし、それがトレーニングをコストのかかる遅いものにしています。

図4. データセット蒸留の例 (出典)

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

データセット蒸留は、元のトレーニングセットを、分類器にとって最も重要な視覚的手がかりを含む少数の合成画像に圧縮できます。ImageNetのような大規模なベンチマークにおいて、元の約4.2%の画像しか使用しない蒸留セットでも、強力な分類精度を維持できることが示されています。つまり、小さな合成プロキシが、はるかに少ない計算リソースで何百万もの実サンプルを代替できることを意味します。

Link to this sectionニューラルアーキテクチャ探索#

ニューラルアーキテクチャ探索（NAS）は、タスクに最適な設計を見つけるために、多くの可能なニューラルネットワーク設計を自動的に探索する手法です。NASは多数の候補モデルをトレーニングおよび評価する必要があるため、フルデータセット上で実行すると遅く、非常に計算負荷が高くなる可能性があります。

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

Link to this section継続的学習とエッジデプロイメント#

一度トレーニングして終わるのではなく、新しいデータが到着するたびに更新され続ける継続的学習システムは、高速でメモリ効率の高い更新を必要とします。カメラ、携帯電話、センサーなどのエッジデバイスも、計算リソースやストレージの制限が厳しいため、同様の制限に直面しています。

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

データセット蒸留は、大規模言語モデルのための知識蒸留と並行して機能することもあります。小さな蒸留データセットは、教師モデルからの最も重要なタスク信号を保持できるため、圧縮された学生モデルがパフォーマンスを大きく損なうことなく、より効率的にトレーニングまたはリフレッシュできます。これらのデータセットは非常に小さいため、ストレージや計算リソースが制限されているが、更新後もモデルの精度を維持したいエッジやデバイス上の使用には特に役立ちます。

Link to this sectionデータ蒸留の長所と短所#

データセット蒸留を使用する利点は以下の通りです。

• 迅速な実験に最適。 毎回巨大なデータセットで再トレーニングすることなく、新しいアーキテクチャ、損失関数、ハイパーパラメータをテストできます。
• 潜在的なプライバシー上の利点。 生データが直接公開されないため、蒸留された合成サンプルを共有する方が、実際のユーザーデータポイントを共有するよりも安全な場合があります。
• 単純なサブセット選択よりも強力な場合が多い。 単にサンプルを選択するのではなく、蒸留は情報を最大限に引き出せるようにそれらを積極的に最適化します。

データセット蒸留にはいくつかの利点がありますが、留意すべき制限事項もあります。

• 過学習の可能性： 蒸留データは通常、蒸留中に使用されたアーキテクチャに対して最もよく機能し、大きく異なるモデルへはうまく移行できない可能性があります。
• ハイパーパラメータに敏感。 結果は、学習率、初期化、蒸留ステップ数などの要因に大きく依存する可能性があります。
• 実世界の複雑さへのスケーリングが難しい。 ベンチマークでうまく機能する手法であっても、大規模で煩雑なデータセットや高解像度データセットでは精度が低下する可能性があります。

Link to this section重要なポイント#

データセット蒸留により、少数の合成サンプルが完全なデータセットとほぼ同等の効果でモデルを学習させることが可能になります。これにより、機械学習はより高速で効率的になり、スケーリングも容易になります。モデルが成長し、より多くのデータを必要とするようになるにつれて、蒸留データセットは精度を犠牲にすることなく計算コストを削減する実用的な方法を提供します。

私たちのコミュニティに参加し、GitHubリポジトリをチェックしてAIについてさらに学びましょう。独自のビジョンAIプロジェクトの構築をお考えの場合は、ライセンスオプションをご確認ください。ヘルスケアにおけるAIや小売におけるビジョンAIといったアプリケーションについて、ソリューションページでさらに詳しく解説しています。