Python
機械学習
MachineLearning
DeepLearning
TensorFlow

ニューラルネットを蒸留にする "Distilling the knowledge in a neural network" の実装 TensorFlowのサンプルコード付

はじめに

読んだ論文 : Hinton, Geoffrey, Oriol Vinyals, and Jeff Dean. "Distilling the knowledge in a neural network." arXiv preprint arXiv:1503.02531 (2015).

a.png

気になっていたDistillation(蒸留)の論文で読んで実装してみました。1パラグラフごとに一言で要約しながらまとめたかったのですが、単なる翻訳になってしまった箇所があります。Hinton先生が提案した手法ということは知っていたのですが、Jeff Dean氏も共著だったのはちょっと驚きでした。

参考 : 全盛期のJeff Dean伝説

蒸留に関しては、TFUG Utsunomiya #04 でもLT発表をしました。

thumbnail


要約

ワンパラ要約

複雑な機械学習モデルのソフトマック関数が出力する誤った答えに対する確率は有用な情報である。これに着目した「Distillation(蒸留、抽出)」と呼ばれる手法により、複雑なモデルの知識を小さいニューラルネットに移すことができた。

先行研究

Buciluǎ, Cristian, Rich Caruana, and Alexandru Niculescu-Mizil. "Model compression." Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining. ACM, 2006.

アンサンブル学習器を1つのモデルに圧縮することが可能であることを示した。

動機・目的

高いパフォーマンスを誇るアンサンブル学習だがいくつかのモデルを使うので計算コストが多くかかってしまう。複雑なモデルを単純なモデルに圧縮したい。

提案手法

目的関数に教師ラベルによるハードターゲットと複雑なモデルによるソフトターゲットを用いる。

ベンチマーク・実験の対象

  • MNIST
  • 音声認識

Abstract

機械学習アルゴリズムのパフォーマンスを向上させる単純な方法として、同じデータセットを異なるモデルで学習させて予測の平均化すること(アンサンブル学習)が挙げられるが、アンサンブル学習器や大規模なニューラルネットを多くのユーザに向けた配備には多くの計算量を必要とする。Caruanaらはアンサンブル学習器を1つのモデルに圧縮することが可能であることを示したが、私たちは彼らの手法を発展させてMNISTや音声認識において成果を挙げた。

1 Introduction

一般的に大規模な機械学習においては学習段階と配備段階ではほぼ同じモデルが使用される。しかし、音声認識やオブジェクト認識などのタスクにはリアルタイム性が必要とされ、加えて多数のユーザに対するモデルの配備には膨大な計算量が必要となる。複雑なモデル(cumbersome model)とはモデルのアンサンブル(アンサンブル学習器)やドロップアウトを用いたモデルを指し、これらのモデルの知識をより配備に適した小さな小さなモデルに移すことを「蒸留(distillation)」と呼ぶ。

学習されたモデルの知識と学習されたパラメータの値を同一視する傾向がこれまで有望なアプローチを妨げた概念的な障壁となっており、同じ知識を保ちながらモデルの形式を変えられるかを知ることを難しくしていた。多クラス分類を行う複雑なモデルでは通常、平均対数確率(average log probability)を最大化することであるが、副作用として複雑なモデルは不正解となる答えに対しても確率を割り当ててしまうが、この誤った答えに対する確率により複雑なモデルがどのようにして一般化されるかが分かる。

大きなモデルから小さなモデルへ蒸留する場合、大きなモデルと同様の方法で小さなモデルを学習させることができる。複雑なモデルを一般化させるならば、同じ方法で汎化能力を向上させるために学習された小さなモデルは、同じ訓練データセットで普通に学習させる場合よりも、典型的にはテストデータに対してパフォーマンスがよいだろう。

複雑なモデルの汎化能力を小さなモデルに移す明白な方法は、複雑なモデルが出すクラス確率をソフトターゲット(soft targets)として小さなモデルを学習させるために使うことである。ソフトターゲットが高いエントロピーを有する場合、ハードターゲットによりも多くの情報かつ訓練データ間の勾配におけるより小さな分散を提供する。その結果より高い学習率を使うことができるのでより服内データで学習できる。

補足:ハードターゲットは教師(正解)ラベルのことを指す。

多クラス分類では例えば2が3や7に似ているというような有用な情報が確率の割合として表されるが、MNISTのように常に高い信頼度で正解を出すタスクではソフトターゲットの値が非常に小さくなってしまうため、クロスエントロピーコスト関数への影響が非常に小さいものとなってしまう。そこで、CaraanaらはSoftmaxへ入力されるロジット(logits)を使い複雑なモデルによるロジットと小さなモデルによるロジットの自乗差を最小化することで問題を解決した。そこで、我々はより一般的な解決方法として蒸留と呼ばれるSoftmaxの温度を上げる手法を提案する。

オリジナルの訓練データセットを使うとうまくいくことがわかったが、ソフトターゲットだけではなく真のターゲットも小さなモデルに予測させるようにする項を追加すると特に効果を発揮する。

メモ

以下のlogitに対するソフトマックス関数の出力をグラフにしてみた(横軸:入力値、縦軸:尤度) :

x.png

y.png

以下が温度付きのソフトマックス関数 :

T=2の場合

z1.png

T=4の場合

z2.png

確かに温度を導入すると、割り当てられてたクラスの確率が高い場合は低く、低い場合は高くなっていることが分かる。


2 Distillation

ニューラルネットはロジットを変換してSoftmax出力層によりクラス確率を出力する。温度付きのソフトマックス関数は以下のように記述される。

$$q_{i}=\frac{\exp \left( \frac{z_{i}}{T} \right)}{\sum_{j=1}^{K}{\exp \left( \frac{z_{j}}{T} \right)}}$$

  • $z_i$ : ロジット(logit)
  • $q_i$ : クラス確率
  • $T$ : 温度(通常は1)

もっとも単純な蒸留の方法は、複雑なモデルに高い温度のSoftmaxで生成したソフトターゲットの転移セットでモデルを学習させることである。学習させたら、温度は1とする。転移セットと教師ラベルを使ったより方法では、2つの目的関数を用いる。一つ目は、ソフトターゲットによるクロスエントロピー。二つ目は通常の教師ラベルによるクロスエントロピーを用いる。非常に小さい重みを二つ目の目的関数に使うことでベストな結果が得られた。勾配のスケールが$1/T^2$となるので、両方の項に$T^2$を掛けることが重要。

2.1 Matching logits is a special case of distillation

3 Preliminary experiments on MNIST

それぞれの層に1200個のReLuユニットである2つの隠れ層を持つニューラルネットを使って、単一のニューラルネットを学習させる。ニューラルネットは、ドロップアウトと重み制限により制限される。

  • 学習データセット:60,000
  • もととなるネット : テストデータセットに対して、64件のテストエラーがあった
  • 2つのReluユニットを持つ2つの隠れ層:146件のテストエラー
  • 上記にソフトターゲットを設けた場合+T=20:74件のテストエラー

以上の結果から、ソフトターゲットは汎用化するかということを含んだ知識を小さなモデルに移したと言える。

300個以上のユニットでは、8以上のすべての温度ではすべて同じ結果となったが、それぞれの層を30まで減らした場合温度は2.5から4の値がよりよく機能した。

転移セットから3の数字を排除し、蒸留されたモデルからは3は今までみたことがない数字となる。蒸留されたモデルは206件のテストエラーのみを提示し、そのうちテストセットの1010件の数字3の133件がテストエラーだった。たいていのエラーは学習された数字3のクラスに対するバイアスがあまりにも低いために引き起こされており、テストデータセットでの全体の最適化によりバイアスが3.5ほど上げた場合、蒸留モデルは109件のエラーを提示しそのうち数字3に関するエラーは14件となる。したがって、正しいバイアスにより、蒸留モデルは学習時に数字3をまったくみたことがないにもかかわらず数字3に対して98.6%の精度となった。


追試検証

TensorFlowによる蒸留手法の実装は以下のレポジトリにまとめました。

諸元

Hinton先生の検証実験とは諸元が異なります。

  • データセット : 訓練データ:60000、テストデータ:10000
  • 検証方法 : ホールドアウト法
  • 性能指標 : 精度(accuracy)
  • ニューラルネット : 共通として同じノード数の2層の多層パーセプトロンを用いる
    • (1) 教師モデル:ドロップアウト(p=0.5) 784 → Relu(1200) → RELU(1200) → Softmax
    • (2) 生徒モデル:ドロップアウト(P=0.5) 784 → relu(600) → relu(600) → softmax
    • 温度 : T=2, T=4
    • (3) 蒸留なしのMLP : ドロップアウト(P=0.5) ) 784 → relu(600) → relu(600) → softmax
  • エポック数:50
  • バッチサイズ:50
  • 学習アルゴリズム : SGD(学習率:0.001)

実験結果

精度 :

result01.jpg

result02.jpg

  • 教師モデルよりも生徒モデルの精度が0.0141ほど向上した
  • 蒸留なしのMLPよりも生徒モデルの精度が0.0216ほど向上した
  • 2層のMLPに蒸留を適用した結果、蒸留された小さいネットワークでは性能が向上した

TensorFlowのコード

温度付きのSoftmax関数 :

import tensorflow as tf

def softmax_with_temperature(logits, temp=1.0, axis=1, name=None):
    logits_with_temp = logits / temp
    _softmax = tf.exp(logits_with_temp) / tf.reduce_sum(tf.exp(logits_with_temp), axis=axis, keep_dims=True)
    return _softmax

誤差関数 :

import tensorflow as tf

# ...
T = tf.placeholder(tf.float32)

hard_target = tf.placeholder(tf.float32, [None, 10])
soft_target = tf.placeholder(tf.float32, [None, 10])
logit = tf.matmul(h, W) + b

y_hard_target = tf.nn.softmax(logit)
y_soft_target = softmax_with_temperature(logit, temp=T)

loss_hard_target = -tf.reduce_sum(hard_target * tf.log(y_hard_target), reduction_indices=[1])
loss_soft_target = -tf.reduce_sum(soft_target * tf.log(y_soft_target), reduction_indices=[1])

loss = tf.square(T) * loss_hard_target + tf.square(T) * loss_soft_target

おわりに

2回読んでようやく大体のヘソがつかめた(気がします)。ツッコミ歓迎です。よろしくお願いします。暗黒知識(Dark knowledge)に触れていないのであとで追記します。

次は以下あたり読みたい。