[論文レビュー] On Compressing U-net Using Knowledge Distillation

Karttikeya Mangalam, Mathieu Salzamann|arXiv (Cornell University)|Dec 1, 2018

Energy Efficient Wireless Sensor Networks参考文献 11被引用数 7

ひとこと要約

この論文は、U-Netアーキテクチャを1000倍以上に圧縮する知識蒸留手法を提案している。パラメータ数を元のサイズの0.1%に削減しながら、ほぼ同一のセグメンテーション精度を維持する。蒸留中にバッチ正則化とクラス再重み付けを学生ネットワークに統合することで、訓練の不安定性とクラス不均衡を克服し、標準的な蒸留では失敗するが信頼性のある圧縮を実現する。

ABSTRACT

We study the use of knowledge distillation to compress the U-net architecture. We show that, while standard distillation is not sufficient to reliably train a compressed U-net, introducing other regularization methods, such as batch normalization and class re-weighting, in knowledge distillation significantly improves the training process. This allows us to compress a U-net by over 1000x, i.e., to 0.1% of its original number of parameters, at a negligible decrease in performance.

研究の動機と目的

• モバイルデバイスなどのリソース制限のある環境に大規模なU-Netモデルを導入する課題に対処すること。
• 知識蒸留が、極めて小さなサイズに圧縮されたU-Netアーキテクチャに対して効果的に機能するかを調査すること。
• 標準的な蒸留が、非常に小さなU-Netバージョン（例：2初期チャネル）に圧縮する際に失敗する理由を同定し、解決すること。
• アーキテクチャ的および損失関数の変更を導入することで、小規模な学生ネットワークにおける蒸留の信頼性を向上させること。

提案手法

• 学生U-Netの収縮パスのすべての畳み込み層にバッチ正則化層を導入し、内部共変量シフトを低減し、訓練を安定化させる。
• 交差エントロピー損失に対してクラス再重み付けを適用し、背景画素対前景画素の比率に基づいて前景クラスに17.8の重みを割り当て、クラス不均衡を是正する。
• 温度スケーリングを用いたソフトラベル損失を含む混合知識蒸留を実施。4-Unet教師ネットワークからT=5でのソフトラベルとハードラベルを組み合わせる。
• 学生ネットワークを、蒸留損失（ソフトラベル）と元の交差エントロピー損失（ハードラベル）の両方で訓練し、最終出力を温度スケーリングでソフトマージン化する。
• ハードラベルで事前学習した4-Unetを教師モデルとして用い、それをパラメータ数が最小限の2-Unetに蒸留する。
• 訓練中にソフトラベル損失をT²でスケーリングし、温度に依存する勾配の大きさを一定に保つ。

実験結果

リサーチクエスチョン

• RQ1 標準的な知識蒸留は、U-Netを元のパラメータ数の1%未満に信頼性を持って圧縮できるか？
• RQ2 なぜ標準的な蒸留は、2初期チャネルのような非常に小さなサイズにU-Netを圧縮する際に失敗するのか？
• RQ3 バッチ正則化とクラス再重み付けは、圧縮されたU-Net学生ネットワークの訓練をどのように改善するか？
• RQ4 正則化技術を強化した蒸留で学習された2-Unetは、64-Unetと同等のセグメンテーション性能を達成できるか？

主な発見

• 122,394パラメータの4-Unetは、テスト損失0.0974、IoU 0.807を達成し、元の64-UnetのIoU 0.804と同等の性能を示しており、蒸留なしでもパラメータ削減が有効であることを示している。
• 標準的な蒸留（バニラまたは混合）では2-Unetの訓練に失敗しており、表2のテスト損失0.505～0.507から、一般化性能の低さと不安定性が示されている。
• バッチ正則化とクラス再重み付けを適用した2-Unetは、混合蒸留を用いてIoU 0.759、交差エントロピー損失0.135を達成し、非蒸留2-Unet（IoU 0.752、損失0.134）を上回る性能を示している。
• 最終的な2-Unetモデルは30,902パラメータを有し、元の64-Unetの3100万パラメータの0.1%に相当しながら、ほぼ同一の性能を維持している。
• 蒸留なしで修正を加えた2-Unetの訓練ではテスト損失が0.307に上昇し、蒸留が性能維持に不可欠であることを確認している。
• バッチ正則化とクラス再重み付けの組み合わせにより、標準的な蒸留が失敗する低パラメータ数の状況でも、成功した蒸留が可能になった。