[論文レビュー] Model compression for faster structural separation of macromolecules captured by Cellular Electron Cryo-Tomography

Jialiang Guo, Bo Zhou+3 more|arXiv (Cornell University)|Jan 31, 2018

Computational Physics and Python Applications参考文献 8被引用数 3

ひとこと要約

本稿では、細胞内電子コートトモグラフィー（CECT）におけるマクロマolecule分類の高速化を目的として、知識蒸留に基づくモデル圧縮手法を提案する。深層教師ネットワーク（DSRF3D-v2）から得られるソフトラベルを用いて、コンパクトな学生ネットワークを訓練することで、モデルサイズを最大113倍まで縮小し、推論時間を2.85倍高速化するが、元の精度の93.8％を維持する。これにより、大規模なCECTデータセットにおける効率的なデプロイが可能になる。

ABSTRACT

Electron Cryo-Tomography (ECT) enables 3D visualization of macromolecule structure inside single cells. Macromolecule classification approaches based on convolutional neural networks (CNN) were developed to separate millions of macromolecules captured from ECT systematically. However, given the fast accumulation of ECT data, it will soon become necessary to use CNN models to efficiently and accurately separate substantially more macromolecules at the prediction stage, which requires additional computational costs. To speed up the prediction, we compress classification models into compact neural networks with little in accuracy for deployment. Specifically, we propose to perform model compression through knowledge distillation. Firstly, a complex teacher network is trained to generate soft labels with better classification feasibility followed by training of customized student networks with simple architectures using the soft label to compress model complexity. Our tests demonstrate that our compressed models significantly reduce the number of parameters and time cost while maintaining similar classification accuracy.

研究の動機と目的

• 高スループットのCECTデータから数百万個のサブトモグラムを処理する際の計算上のボトル neck を軽減すること。
• マクロマolecule分類に用いられる大規模で高精度な3D-CNNモデルの高い推論コストを低減すること。
• 計算リソースが限られた環境でも、大規模なCECTデータセットに効率的に深層学習モデルをデプロイ可能にする。
• 実世界の生物学的応用においても、顕著なモデル圧縮に対しても高い分類精度を維持すること。

提案手法

• シミュレートされたCECTサブトモグラム上で複雑なDSRF3D-v2教師ネットワークを訓練し、高品質なソフトラベルを生成する。
• 構造を簡素化した3つの段階的な小型学生ネットワーク（DSRF3D-v2-s1, s2, s3）を設計する。
• 教師ネットワークのソフトラベルが学習を導く知識蒸留を用いて学生ネットワークを訓練する。
• 教師から学生への知識伝達を強化するために、温度調整付きソフトマックス関数（T=5）を用いる。
• 確率的勾配降下法とネステロフモーメンタムを用いたカテゴリカル交差エントロピー損失関数で学生モデルを最適化する。
• 検証損失に基づく早期停止を適用し、過学習を防ぎ、一般化性能を確保する。

実験結果

リサーチクエスチョン

• RQ1 知識蒸留を用いることで、顕著な精度損失なしに、CECTサブトモグラム分類のための大規模3D-CNNモデルを効果的に圧縮可能か？
• RQ2 3D生物学的画像解析において、モデル圧縮率、推論速度、分類精度の間にはどのようなトレードオフがあるか？
• RQ3 教師ネットワークからの知識蒸留は、ハードラベルのみで学習する場合と比較して、小型学生ネットワークの性能をどのように向上させるか？
• RQ4 この文脈において、圧縮、速度、精度のバランスが最も優れた学生ネットワークアーキテクチャはどれか？
• RQ5 モデル圧縮は、大規模なCECTデータセットにおける深層学習モデルの実用的デプロイをどの程度可能にするか？

主な発見

• DSRF3D-v2-s1学生モデルは18.1倍の圧縮率を達成し、パラメータを18.3Mから1.0Mに削減しながら、教師モデルの93.8％の精度を維持した。
• DSRF3D-v2-s1モデルは、100万個のサブトモグラムあたりの推論時間を15.91秒から9.41秒に短縮し、1.69倍の高速化を達成した。
• 教師モデルのソフトラベルを用いて学生ネットワークを訓練することで、ハードラベルのみで学習する場合と比較して、最大8.16ポイントの精度向上が達成された。
• DSRF3D-v2-s3モデルは最高の圧縮率（113.3倍）を達成したが、精度の低下も最大であり、教師モデルの精度の89.83％しか維持しなかった。
• 知識蒸留により、ハードラベルのみで学習したモデルと比較して、学生モデルがより良好に一般化することが示された。これにより、ソフトラベルのガイドラインの価値が裏付けられた。
• 結果から、3D-CNNにおける蒸留によるモデル圧縮は、大規模な生物学的データセットの効率的処理に非常に有効であることが明らかになった。