[論文レビュー] Multi-task UNet: Jointly Boosting Saliency Prediction and Disease Classification on Chest X-ray Images

Hongzhi Zhu, Robert Rohling+1 more|arXiv (Cornell University)|Feb 15, 2022

Visual Attention and Saliency Detection被引用数 6

ひとこと要約

本稿では、深層学習を用いて、胸部X線（CXR）画像において同時に注目度予測と疾患分類を実行するマルチタスクUNetアーキテクチャを提案する。共有特徴学習と過学習を軽減する最適化されたマルチタスク学習方式を活用することで、注目度予測や分類専用のモデルよりも優れた性能を両タスクで達成し、最先端の性能を発揮した。

ABSTRACT

Human visual attention has recently shown its distinct capability in boosting machine learning models. However, studies that aim to facilitate medical tasks with human visual attention are still scarce. To support the use of visual attention, this paper describes a novel deep learning model for visual saliency prediction on chest X-ray (CXR) images. To cope with data deficiency, we exploit the multi-task learning method and tackles disease classification on CXR simultaneously. For a more robust training process, we propose a further optimized multi-task learning scheme to better handle model overfitting. Experiments show our proposed deep learning model with our new learning scheme can outperform existing methods dedicated either for saliency prediction or image classification. The code used in this paper is available at https://github.com/hz-zhu/MT-UNet.

研究の動機と目的

• マルチタスク学習を活用して表現学習を向上させることで、医療画像におけるデータ不足の問題に対処する。
• CXR画像のヒトの視覚的注目度と疾患分類を同時に予測する統合的深層学習フレームワークを構築する。
• マルチタスク学習の最適化された訓練方式を用いて過学習を低減することで、モデルの頑健性と一般化性能を向上させる。
• CXRデータセットにおいて、注目度予測または疾患分類に特化した既存の最先端モデルを上回る性能を発揮する。
• ヒトの視覚的注目を深層学習パイプラインに統合することで、医療画像タスクにおける移譲可能で効率的な学習を実現する。

提案手法

• 注目度予測と疾患分類のヘッド間でエンコーダー特徴を共有するマルチタスクUNetアーキテクチャを設計する。
• CXR画像からの階層的特徴を抽出するために、UNetアーキテクチャに基づく共有エンコーダーを活用する。
• 2つの並列ディコーダーヘッドを実装する：1つは2次元の注目度ヒートマップの予測用、もう1つは多クラス疾患分類用。
• 勾配のバランスを最適化し、訓練中に過学習を軽減するための動的損失重み付けを導入した最適化されたマルチタスク学習方式を採用する。
• 一般化性能を向上させるために、データ増強と正則化技術を適用する。
• 注目度予測にはバイナリクロスエントロピー、分類にはカテゴリカルクロスエントロピーを組み合わせた統合損失関数を用いて、エンドツーエンドでモデルを訓練する。

実験結果

リサーチクエスチョン

• RQ1 注目度予測と疾患分類の共同学習は、単一タスクモデルと比較して、両タスクの性能を向上させることができるか？
• RQ2 共有特徴を用いたマルチタスク学習は、データが少ない医療画像環境における表現学習をどのように向上させるか？
• RQ3 提案された最適化されたマルチタスク訓練方式は、小規模なCXRデータセットにおいて過学習をどの程度軽減するか？
• RQ4 注目度予測によるヒトの視覚的注目を統合することで、疾患分類モデルの解釈性と性能が向上するか？
• RQ5 統合アーキテクチャは、CXR画像における注目度予測と疾患分類の両方で、専用モデルを上回る性能を発揮できるか？

主な発見

• 提案されたマルチタスクUNetは、注目度予測と疾患分類の両タスクで、最先端のモデルを上回る性能を発揮した。
• 最適化されたマルチタスク学習方式は、特に小規模な医療データセットにおいて過学習を顕著に低減し、一般化性能を向上させた。
• 共同訓練により、単一タスクベースラインと比較して、注目度予測および分類の両方の指標でより優れた特徴学習が実現した。
• 疾患分類ではより高いAUCスコアを達成し、人間がアノテートした注目度マップとの相関性も、既存の手法を上回った。
• アブレーションスタディの結果、共有特徴学習と動的損失重み付けが性能向上に不可欠であることが確認された。
• コードとトレーニング済みモデルはGitHubで公開されており、医療分野におけるマルチタスク学習の再現性とさらなる研究を可能にしている。