[論文レビュー] Partial transfusion: on the expressive influence of trainable batch norm parameters for transfer learning

Fahdi Kanavati, Masayuki Tsuneki|arXiv (Cornell University)|Feb 10, 2021

COVID-19 diagnosis using AI参考文献 25被引用数 14

ひとこと要約

この論文では、医療画像の転移学習において、トレーニング可能なバッチ正規化（BN）パラメータ—スケールγとシフトβ—のみをファインチューニングすることを提案し、このアプローチがフルモデルファインチューニングと同等の性能を達成することを示している。この手法により、パラメータ数が少なく、学習率を高く設定できるため、収束が速くなり、4つのアーキテクチャと9つの医療画像データセットで一貫した結果が得られた。

ABSTRACT

Transfer learning from ImageNet is the go-to approach when applying deep learning to medical images. The approach is either to fine-tune a pre-trained model or use it as a feature extractor. Most modern architecture contain batch normalisation layers, and fine-tuning a model with such layers requires taking a few precautions as they consist of trainable and non-trainable weights and have two operating modes: training and inference. Attention is primarily given to the non-trainable weights used during inference, as they are the primary source of unexpected behaviour or degradation in performance during transfer learning. It is typically recommended to fine-tune the model with the batch normalisation layers kept in inference mode during both training and inference. In this paper, we pay closer attention instead to the trainable weights of the batch normalisation layers, and we explore their expressive influence in the context of transfer learning. We find that only fine-tuning the trainable weights (scale and centre) of the batch normalisation layers leads to similar performance as to fine-tuning all of the weights, with the added benefit of faster convergence. We demonstrate this on a variety of seven publicly available medical imaging datasets, using four different model architectures.

研究の動機と目的

• 医療画像の転移学習におけるトレーニング可能なバッチ正規化（BN）パラメータの表現的影響を調査すること。
• 最適な性能を達成するにはフルモデルファインチューニングが必要であるという一般的な仮定に反論すること。
• ファインチューニングの対象をBNのアフィンパラメータ（γとβ）に限定した場合、フルファインチューニングと同等の性能を達成できるかどうかを調査すること。
• 転移学習中にBNパラメータを更新するか、固定された移動平均を使用するかの影響を評価すること。
• ランダム初期化された重みを用い、BNパラメータのみをトレーニングする方法の妥当性を検討し、BN層の内在的表現力の理解を深めること。

提案手法

• 他のすべての重みを固定したまま、バッチ正規化層のトレーニング可能なスケール（γ）およびシフト（β）パラメータのみをファインチューニングする。
• 5つのトレーニング戦略—特徴抽出、フルファインチューニング、BNパラメータのみのファインチューニング、移動平均更新（BMA）、ランダム重みにBNパラメータのみをトレーニングする—の性能を比較する。
• 9つの公的医療画像データセットで、4つの事前学習済みCNNアーキテクチャ—DenseNet121、ResNet50V2、InceptionV3、EfficientNetB3—を用いる。
• 最終全結合層を最初に訓練した後、BNパラメータをファインチューニングするか、両者を同時に訓練する。
• トレーニング中にBN層を推論モードに保ち、移動平均の更新による不安定化を回避する。
• マルチラベルおよびバイナリ分類タスクのテストセットで、受信器操作特性曲線（AUC）を用いてモデルを評価する。

実験結果

リサーチクエスチョン

• RQ1 バッチ正規化アフィンパラメータ（γとβ）のみをファインチューニングすることで、医療画像分類においてフルモデルファインチューニングと同等の性能を達成できるか？
• RQ2 ファインチューニング中にバッチ正規化の移動平均統計量を更新すると、固定したままにした場合と比較して性能が低下するか？
• RQ3 ランダム初期化された重みを用い、BNパラメータのみをトレーニングした場合、性能はどの程度か？
• RQ4 BNパラメータのみのファインチューニングが収束速度およびモデル安定性に与える影響は？
• RQ5 アーキテクチャやデータセット固有の要因が、BNパラメータのみのファインチューニングの有効性に影響を与えるか？

主な発見

• BNパラメータのみをファインチューニングすることで、4つのアーキテクチャと9つのデータセットすべてで、フルモデルファインチューニングと同等の性能が達成された。
• DenseNet121では、BNパラメータのみのファインチューニングが、すべてのデータセットでフルファインチューニングと1〜2%以内のROC AUCスコアを達成し、どのタスクでも顕著な性能低下が見られなかった。
• BNパラメータのみのファインチューニングにより、トレーニング可能なパラメータ数が少なく、より高い学習率が使用可能になったため、収束が速くなった。
• 事前学習済みモデルを特徴抽出器として使用する（ファインチューニングなし）場合、フルファインチューニングおよびBNパラメータのみのファインチューニングと比較して一貫して性能が低かった。
• 移動平均統計量の更新（BMA）をファインチューニング中に実行した場合、大多数のデータセットで性能が低下し、固定したままにした場合の利点が確認された。
• ランダム初期化された重みにBNパラメータのみをトレーニングした場合、Chest X-ray や Malaria などの単純なデータセットでは妥当な性能が得られたが、EfficientNetB3では勾配飽和のためトレーニングに失敗し、アーキテクチャ依存性が顕著に現れた。