こんにちは、エンジニアの中村です。
今回から数回にわたって、ディープラーニング技術の実用上の重要な課題を解決する蒸留という手法について紹介したいと思います。

ディープラーニングは非常に魅力的なモデル学習技術ですが、実際に使用する際には計算リソースがボトルネックとなり、利用可能なシーンが限定されてしまうことがあります。せっかくの魅力的な技術も、そうなってしまうと宝の持ち腐れです。蒸留は、精度を損なわずに計算負荷の小さいモデルを訓練することを可能にし、ディープラーニングが実際に利用できるシーンを拡大します。

また、蒸留はこうした計算リソースに起因する ディープラーニングの実用上の課題を解決するだけでなく、モデルの精度向上、訓練の効率化、敵対的攻撃(Adversarial Attack)に対する防御など、その他のさまざまな用途に使用可能な応用領域の広い手法でもあります。

今回の記事では、モデル圧縮という蒸留の基本からはじめて、蒸留が取り扱うさまざまな応用領域について、俯瞰的に見ていきたいと思います。

目次

はじめに

ディープラーニング技術の発展により、画像、自然言語、音声等のさまざまな領域において、旧来手法をはるかに凌ぐ性能を示すモデルを学習することが可能になっています。また、ディープラーニングを実用の場面に応用する取り組みも盛んです。
さまざまな領域で高性能を示す ディープラーニングは非常に有望な技術ですが、実応用する際には課題もあります。その1つに、ディープラーニングモデルを稼働させる本番環境において、十分な計算リソースが得られない場合があることが挙げられます。訓練時には、大規模なデータに対して、大規模なモデルを学習させることで、高性能なモデルを得ることが可能です。特に、モデルをより深く、より大きくすることで、性能が向上することが一般的に知られています。しかしながら、訓練で得た大規模なモデルをそのまま使うには、実応用の本番環境の計算リソースが十分でない場合がしばしばあります。例えば、画像認識モデルをスマートフォンアプリで動かす、画像認識モデルを移動型ロボットに搭載して高度な制御を実現する、音声認識・生成モデルをスマートスピーカに搭載する、などです。いずれの場合でも ディープラーニングモデルが実際に稼働するデバイスは小型であり、計算リソースに制約があります。この課題を克服しない限り、ディープラーニングモデルが実際に活用可能な場面が限定されることになります。

ディープラーニングモデルはしばしばモデルサイズが大きく、そのままでは本番環境で利用しにくい (モデル例は AlexNet [Krizhevsky, 2012])
本番環境における計算リソースの制約という課題を解決するための打ち手の1つとして、本番環境で使えるように軽量で、かつ、十分に高性能なモデルを用意する、ということが考えられます。実際に、モデル圧縮 (Model Compression) という研究領域があり、下記を例とするさまざまな方法でニューラルネットワークを上手くダウンサイズしようと試みられています。

  • 性能を損なわない軽量なモデルアーキテクチャのためのモジュールの工夫 (効率的なマイクロアーキテクチャ)
  • ニューラルネットワークを構成する重みのbit数を削減する量子化
  • 訓練後に寄与度の小さい重みを削除する Pruning

今回紹介する蒸留 (Distillation, [Hinton et al., 2015]) も、モデル圧縮のための1手法として登場しました。

分野軽量化方法
効率的なマイクロアーキテクチャ訓練に先立って、
精度と計算のトレードオフを改善するように、
モデルの構成要素を工夫する
量子化訓練後に、重みの浮動小数点数の精度を下げる
Pruning訓練後に、寄与度の小さい重み、
チャネル、レイヤーを削除する
蒸留一度訓練したモデルが学習した知識を、
別の軽量なモデルに継承させる

蒸留の基本的な枠組み

蒸留では、大規模なモデルの知識を軽量なモデルに継承させることによってモデル圧縮を実現します。ここで重要となるのは、知識の継承先の軽量なモデルの訓練方法です。
蒸留は、基本的に、以下のステップにしたがった訓練によって実現されます。

  • 通常の教師あり学習の方法でモデルを訓練 (これを教師モデルと言います)
  • 教師モデルとは別のより軽量なモデル (生徒モデル) を、教師モデルの出力を正解ラベルの代わりに使用、あるいは、正解ラベルに追加して訓練

蒸留ではしばしば、教師モデルの訓練等で使用する通常の正解ラベルを “hard target“、生徒モデルの訓練に使用する教師モデルの出力を “soft target” と呼びます。 soft target は、正解クラス以外のクラスに関する情報を含むという点で、hard target と異なります。正解クラス以外のクラスに関する情報 (これは “Dark Knowledge” とも呼ばれます) には、クラス間の類似度といった情報が含まれうるため、生徒モデルの効果的な訓練に有効だと一般に考えられています。例えば、猫の画像を与えられた教師モデルは、(少し似ているところがあるために) もしかしたら犬クラスにも相当の値を持つ出力をするかもしれません。しかしながら、魚クラスの値は非常に小さくなり、車クラスに至っては値はほとんど0となるかもしれません。こうした類似度に関する情報によって、生徒モデルはクラス間で共通の特徴表現を効果的に獲得でき、そうした特徴表現のためにより汎化した性能を示すことが可能であると考えられます。
教師モデルに蓄積された知識は、このような形で出力に表れ、それを利用することで生徒モデルに効果的に知識を継承させるというのが、蒸留の基本的なコンセプトとなります。

蒸留の発展

モデル圧縮の方法の1つとして蒸留が有効であること、また蒸留がどのようになされ、効果を発揮するのかという枠組みを以上で説明しました。この枠組みを基本として、蒸留はモデル圧縮以外の目的にも発展しています。
以下で、その発展の例をいくつか見ていきたいと思います。

ディープラーニングモデルの精度向上

蒸留を利用することで、既存の ディープラーニング モデルの性能を向上させられることが報告されています。
ディープラーニング に限らず、モデルの性能に寄与する要素は、以下の通りに考えられます (どのようなモデルをどのようなデータに対してどのように訓練するか)。

  1. モデル
  2. データ
  3. 訓練方法

ディープラーニングにおいて、1. モデルについては、AlexNet に始まり、VGG、ResNet のようにさまざまな (より深い) アーキテクチャが考案されてきました。これに対して、蒸留は特に 2. のデータと 3. 訓練方法について、精度向上に有用な手法を提供します。

蒸留によるデータ利用の改良。利用可能なデータが増えればモデルの性能を向上させることが可能です。しかし、正解ラベルのある形で大規模にデータを収集することは高コストであるため、正解ラベルなしデータを上手く活用することで性能を向上させたいところです。この課題に対して、正解ラベルなしデータに対する教師モデルの出力を正解ラベルとして扱い、生徒モデルの訓練時に正解ラベルなしデータも利用するデータ蒸留 (Data Distillation[Radosavovic et al., 2017]) という手法が提案されています。複数の変換を施した画像を教師モデルに入力、各変換画像に対して得られる教師モデルの出力を平均することで、生徒モデルを訓練するための正解ラベルを得ます。正解ラベルなしデータを正解ラベルありデータに追加して利用することで、正解ラベルありデータのみでの教師あり学習以上の精度が得られることが報告されており、すなわち、データ蒸留が精度向上に有効であることが示されています。

transform A, B, C という3種類の変換を入力画像に施し、教師モデルA に入力。変換された入力それぞれに対する教師モデルの出力をアンサンブルして、生徒モデルに正解として与える [Radosavovic et al., 2017]
このデータ蒸留という手法は、正解ラベルありデータに加えて正解ラベルなしデータを使用するという点で、半教師あり学習と類似した問題設定を扱っています。実際、半教師あり学習の一手法と位置付けることもできます。

半教師あり学習では、訓練データに十分な正解ラベルありデータが存在しないため、正解ラベルなしデータを上手く活用したい、という点に問題設定の主眼が置かれています。一方、データ蒸留では、訓練データにある程度十分な正解ラベルありデータがすでに存在するが、正解ラベルなしデータを利用してさらなる精度向上を目指したい、という点に問題設定の主眼が置かれています。

以下に、半教師あり学習とデータ蒸留を比較しつつまとめます。

手法概要正解ラベルありデータ正解ラベルなしデータ精度
半教師あり学習十分な正解ラベルありデータが
ない場合の訓練方法を提供
少数十分正解ラベルありデータが十分得られる場合の精度を上限として議論
データ蒸留正解ラベルありデータに加えて、
正解ラベルなしデータを活用することで、
さらなる精度向上を目指す訓練方法を提供
十分十分十分な正解ラベルありデータで訓練した際の精度を下限として、
そこから正解ラベルなしデータを活用することによる精度向上を議論

一方、半教師あり学習で蒸留の考え方が使用されている例もあります (Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results, [Tarvainen et al., 2017])。

蒸留による訓練方法の改良。蒸留によって性能を損なわずに軽量なモデルが得られることをはじめに見てきました。ディープラーニングでは性能向上のためにしばしばモデルを大規模にしますが、では、蒸留において大規模なモデルに対して知識の継承をしたらどうなるのでしょうか。
教師モデルと同程度のキャパシティを持つ生徒モデルに蒸留することで、教師モデルを凌ぐ性能の生徒モデルが得られることが報告されています (Born Again Neural Networks (BAN), [Furlanello et al., 2018])。蒸留によって、モデルやデータについてさほど工夫することなく、ある意味”お手軽に”モデルの性能を向上させることができます。

複雑なタスクの効果的な訓練

ディープラーニングの有効性を一般に知らしめた ImageNet Challenge では 1,000 クラスの分類をすることが求められますが、現実の問題では、より多くの、例えば10万クラスについて、しかも小さな粒度で分類する必要もあります (例えば、猫と識別するだけでなく、その種類まで識別する、等)。このタスクをナイーブに解こうとすると、単に膨大なクラス数を分類する教師あり学習になりますが、大規模なモデルが必要となり、訓練が難しくなります。

こうした粒度が小さくクラス数が多い複雑なタスクに対するより良いアプローチとして、データセットの意味的な構造を活用して、階層的に分類をおこなうというアプローチは有効そうです。すなわち、大カテゴリに対応するような大雑把な分類をした後に、その大雑把な分類の中で詳細に分類するという2段階分類をおこなうということです。実際に、この2段階の階層構造を持つネットワークアーキテクチャを構成し訓練する手法が提案されています (Hierarchical Deep Convolutional Network, HD-CNN, [Yan et al., 2014])。

一方、Knowledge Concentration[Gao et al., 2017] という蒸留の手法では、大カテゴリごとに問題を分割し、各大カテゴリ内で詳細に分類するのに必要な専門的知識を学習するモデル (専門家モデル) をそれぞれ独立に訓練します。蒸留の段階では、それら各専門家モデルの出力を総合して、また別個の生徒モデルを訓練します。個別の専門的知識を、蒸留のプロセスによって上手くまとめあげていると言えます。

入力画像のカテゴリを専門とする専門家モデル (この図では Teacher 2) を選択し、その出力を生徒モデルの正解ラベルとする [Gao et al., 2017]
いずれも、データの意味的な階層構造を捉えて問題にアプローチしていますが、階層的アプローチは、タスクに合わせて、階層性を持つ複雑なネットワークアーキテクチャを設計しなければなりません。このため、階層的アプローチは、構造の複雑さに起因する困難さを軽減するために、場合によっては訓練や推論の仕方を工夫する必要があります。一方、Knowledge Concentration は、階層的なアーキテクチャを設計する必要がなく、モデルの構造をシンプルに保つことが可能です。生徒モデルは蒸留という特別な方法で訓練しなければなりませんが、専門家モデルの訓練は各大カテゴリにおける通常の教師あり学習です。また、蒸留で最終的に得られる生徒モデルは通常の畳み込みニューラルネットワークの構造をしているため、推論の仕方を工夫する必要は通常はないと言えます。さらに、蒸留によって相対的に軽量な生徒モデルを訓練することもできるため、制約のある本番環境においても効用がある手法と言えるでしょう。

手法概要訓練時の特徴推論時の特徴
HD-CNN階層的なネットワークアーキテクチャを設計階層的アーキテクチャの各構成要素を事前学習、
その後 全体を fine-tune
モデルが階層性を持ち、複雑であるため、
推論の工夫が必要な場合がある
Knowledge
Concentration
問題を分割し専門家モデルを訓練
蒸留で専門知識を総合
専門家モデルを大カテゴリごとに、
個別に通常の教師あり学習
専門家モデルを生徒モデルに蒸留
推論の工夫が不要
軽量な推論モデルを得ることも可能

蒸留によって複雑なタスクをこなすモデルを効率的に訓練するというコンセプトは、強化学習においても見られます。方策蒸留 (Policy Distillation[Rusu et al. 2015]) という手法が提案されています。この手法では、強化学習タスクの達成に必要な多様なスキルを、教師方策から生徒方策に蒸留します。

敵対的攻撃 (Adversarial Attack) に対する防御

ディープラーニングは幅広い分野で高い性能を示していますが、課題もあります。その一つとして、画像認識領域において、敵対的攻撃(Adversarial Attack)という外部からの攻撃に脆弱であることが報告されています。すなわち、画像認識モデルに対して敵対的に生成された摂動が画像に加えられることによって、誤識別が誘発されるという問題です。実は、蒸留には、その訓練方法で敵対的攻撃に対して頑健なモデルが得られる、という性質があることが知られています (Defensive Distillation[Papernot et al., 2015])。蒸留によって、入力に対する出力の勾配を抑えられ、入力の近傍における出力を頑健にできるため、と考えられています。

蒸留が対象とする課題の整理

以上述べてきた蒸留について、対象とする課題を主軸に、蒸留以外の手法も含めて整理してみます。

課題アプローチ打ち手
計算リソースの制約がある本番環境でも
ディープラーニングを利用したい

本番環境における
計算リソースの制約を緩和
小型かつ計算力の高い計算機の開発
モデルサイズを軽量化マイクロアーキテクチャ、量子化、pruning、
蒸留
モデルの精度を向上させたいモデルアーキテクチャを工夫(深層な) アーキテクチャのデザイン
モデルの訓練方法を工夫最適化手法の改良、
蒸留 (BAN)
データの利用可能性を拡大、
利用可能なデータを増やす
Data Augmentation、半教師あり学習、データ蒸留
粒度が小さく、クラス数の多い分類タスクを
うまく学習させたい
大規模モデルを訓練教師あり学習
問題の階層性をアーキテクチャで捉える階層的ネットワーク
問題の階層性に応じて問題を分割蒸留 (Knowledge Concentration)
敵対的攻撃 (Adversarial Attack) に対して
頑健なモデルを開発したい
攻撃に対して頑健な
モデルアーキテクチャを考案
分類器の頑健化 (Classifier Robustifying)
Adversarial Example を加えて訓練Adversarial Training
攻撃者が入力に加える摂動に対して
ロバストなモデルを訓練
蒸留 (Defensive Distillation)

このように、蒸留ではモデル圧縮をはじめとしてさまざまな課題を対象としていることが改めて見てとれます。

今回の記事では、蒸留について、その起源であるモデル圧縮から始めて蒸留の基本的な考え方をおさえ、その後 蒸留の発展として、ディープラーニングモデルの精度向上、専門的知識の効果的な訓練、敵対的生成に対する防御といったさまざまな応用領域を俯瞰的に見てきました。

次回以降、蒸留の各手法について、具体的な研究を参照して、詳細に紹介する予定です。