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

モデルのトレーニングは、データサイエンティストの仕事の中で最も時間がかかる部分のように思えるかもしれません。しかし、彼らの時間のほとんど（多くの場合60%から80%）は、実際にはデータの準備（収集、クレンジング、モデリングのための整理）に費やされています。データセットが大きくなるにつれてその準備時間も増え、実験が遅くなり、イテレーションが困難になります。

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

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

これは、モデルが効果的に学習するために必要な重要なパターンを保持しながら、大規模なトレーニングデータセットのコンパクトなバージョンを作成します。これにより、より高速なトレーニング、低い計算要件、およびより効率的な実験への道が開かれます。これは、モデルのためのスタディ用チートシートのようなものであり、フルデータセットと同じコアパターンを教えるために設計された少数の合成データサンプルのセットです。

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

データセット蒸留の理解#

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

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

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

このアイデアは、2018年にTongzhou Wang、Jun-Yan Zhu、Antonio Torralba、Alexei A. EfrosによるarXivの論文で初めて発表されました。初期のテストではMNISTやCIFAR-10などのシンプルなデータセットが使用され、少数の蒸留サンプルが何千もの実画像のかわりになることが簡単に示されました。

図 1. 画像データに対するデータセット蒸留の使用（ソース）

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

データセット蒸留の重要性#

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

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

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

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

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

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

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

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

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

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

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

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

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

性能マッチング#

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

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

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

図 3. メタラーニングプロセス（ソース）

分布マッチング技術#

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

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

単一層の分布マッチング#

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

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

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

多層分布マッチング#

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

蒸留画像の作成と評価#

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

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

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

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

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

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

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

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

図 4. データセット蒸留の例（ソース）

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

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

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

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

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

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

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

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

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

データ蒸留の長所と短所#

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

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

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

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

重要なポイント#

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

コミュニティに参加し、GitHubリポジトリをチェックして、AIの詳細を発見してください。独自のビジョンAIプロジェクトの構築を検討している場合は、ライセンスオプションをご確認ください。ソリューションページにアクセスして、医療分野のAIやリテール分野のビジョンAIなどのアプリケーションの詳細を探索してください。