Fusic Tech Blog

Fusicエンジニアによる技術ブログ

知識蒸留(Knowledge Distillation)を使ってResNet18をより賢くしてみよう
2024/02/19

知識蒸留(Knowledge Distillation)を使ってResNet18をより賢くしてみよう

こんにちは、機械学習チームのハンです。

知識蒸留(KD:Knowledge Distillation)というものが気になり、簡単なResNetモデルで色々実験を行ってみたので紹介したいと思います。

Knowledge Distillationとは?

あるモデルが学んだ知識を他のモデルに移すという概念であり、Distilling the Knowledge in a Neural Network(2015)で初めて提案されたものです。 最近の機械学習モデルの精度は段々上がっていますが、一方で、モデルが巨大化しています。そのため、実際にデプロイしたり個人が利用したりするには負担になるものも多いです
下の図で確認できるように、Teacherモデルの知識を蒸留しStudentモデルに移して「より軽くて良い精度を持つモデルを作る」のが、KDの目的だと考えられます。

どのような知識を利用するか?

KDで利用される知識(Knowledge)は大きく下記の3つに分けられます。

  • Response-based Knowledge
  • Feature-based Knowledge
  • Relation-based Knowledge

Response-based Knowledge

モデルのResponse、つまり「モデルのOutput」を知識として扱う方法です。巨大なTeacherモデルのOutput(Logit)には、豊かな情報が入っているという考えられ、Outputの分布を知識として扱います。

上図はResponse-base Knowledgeを用いた一般的なKDモデルの構成ですが、右上のDistillation Lossの部分に注目してください。Pre-trainされた「TeacherモデルのOutputが真似できるように学習」していることがわかります。

Feature-based Knowledge

モデルの中間層のFeature-Map、「中間層のOutput」を知識として扱う方法で、この特徴は「Hint」とも呼ばれます。下図のような構造で、Studentモデルは「Teacherモデルの中間層から知識を得て、同様な中間層Outputを作れる」ように学習を進めます。

この方法はResponse-based Knowledgeと比べ、細くて深い層を持つネットワークの学習に役に立つと思われています。しかし、中間層のOutputの形はモデル・レイヤーによって様々であり、Lossを実装する時に手間がかかる場合もあります。

Relation-based Knowledge

モデル中の「レイヤーの間に存在する関係」を知識として扱う方法です。上記2つの方法では、Teacher・Studentモデルのある出力を直接1:1マッチングして比較していますが、Relation-based Knowledgeでは、個別モデルの中間層Hint間の関係を求め、その関係を比較します。

例えば、中間層出力間の距離を関係として扱い「Teacherモデルと似たような関係をStudentが学習する」ような形になります。

知識をどのように移動させるか?

KDを行う方式として、以下の3つのやり方があげられます。

  • Offline Distillation
  • Online Distillation
  • Self Distillation

Offline Distillation

Pre-trainしたTeacherモデルの知識をStudentモデルに学習させる」方法で、一番簡単で一般的なKD方法になります。

KDプロセス上、TeacherモデルはFreezeされ「入力に対して知識になるOutputを推論する」役割、Studentモデルは「Teacherモデルの知識を用いて学習する」役割を行います。一般にTeacherは「スケールが大きく良い精度を持つモデル」が考えられ、個人的には「大きいモデルの知識を蒸留し軽いモデルに移す」というKDの目的にぴったりの方法だと思います。
特徴としては、Pre-trainされたTeacherモデルが要るので「2Stage方式」になることです。

Online Distillation

KDプロセス上、「Teacher・Studentモデルがお互いに学習を行い、知識を共有する方法」で、上の図では、学習能力のあるTeacherモデルの知識をStudentモデルに蒸留することを表しています。 この概念以外にも、Teacherモデル無しで

  • 「構造の異なる2つのStudentモデルがをお互いに知識を与える
  • 「構造は同じだが、設定が異なる2つのStudentモデル」がをお互いに知識を与える

ような方法もあります。(参考:Deep Mutual Learning)

Offline Distillationと異なり、End-to-End学習ができる・モデルお互いに知識の共有ができるという特徴があります。

Self Distillation

一つのモデルでKDを行う」方法で、ネットワーク中でTeacherになる特徴を選択し、その特徴を同じモデルが学習します。例えば、

  • 同じイメージに異なるAugmentationをかけた結果」を一つのモデルに通し、少し異なるOutputをTeacher・Studentターゲットにする
  • 最終レイヤー・中間レイヤー」のOutputをTeacher・Studentターゲットにする

などの方法が挙げられます。

一つのモデルを用いているので上記2つの方法と比べ、学習時間・パラメータ数の面でメリットになる特徴があります。

3つの方法と一般的なモデル学習の比較

| 基準 | 比較 | | :

Han Beomseok

Han Beomseok

Python, AI Engineering, Natural Language