[論文レビュー] Self-supervised Feature Learning for 3D Medical Images by Playing a Rubik's Cube

Xinrui Zhuang, Yuexiang Li+4 more|arXiv (Cornell University)|Oct 5, 2019

Medical Imaging and Analysis参考文献 12被引用数 15

ひとこと要約

本論文は、3次元ボリュームキューブを再構成することで回転および並進不変性を学ぶ、独創的なプロキシタスク「ルービックキューブの回復」を用いた自己教師付き3次元特徴抽出フレームワークを提案する。この手法は、脳出血分類（83.8％の正確性）および脳腫瘍セグメンテーション（77.3％のmIoU）で最先端の性能を達成し、アノテーション付きデータを追加で使用せずに、学習を初期化する場合や自然動画データセットで事前学習する場合を大きく上回る。

ABSTRACT

Witnessed the development of deep learning, increasing number of studies try to build computer aided diagnosis systems for 3D volumetric medical data. However, as the annotations of 3D medical data are difficult to acquire, the number of annotated 3D medical images is often not enough to well train the deep learning networks. The self-supervised learning deeply exploiting the information of raw data is one of the potential solutions to loose the requirement of training data. In this paper, we propose a self-supervised learning framework for the volumetric medical images. A novel proxy task, i.e., Rubik's cube recovery, is formulated to pre-train 3D neural networks. The proxy task involves two operations, i.e., cube rearrangement and cube rotation, which enforce networks to learn translational and rotational invariant features from raw 3D data. Compared to the train-from-scratch strategy, fine-tuning from the pre-trained network leads to a better accuracy on various tasks, e.g., brain hemorrhage classification and brain tumor segmentation. We show that our self-supervised learning approach can substantially boost the accuracies of 3D deep learning networks on the volumetric medical datasets without using extra data. To our best knowledge, this is the first work focusing on the self-supervised learning of 3D neural networks.

研究の動機と目的

• 深層学習モデルを医用診断に適用する際、アノテーション付き3次元医用画像が限られているという課題に対処すること。
• ラベルなし3次元ボリュームデータを活用して、下流タスクの性能を向上させる自己教師付き事前学習フレームワークを開発すること。
• 3次元畳み込みニューラルネットワーク（3D CNN）が並進および回転不変性を学ぶように促す、新たなプロキシタスク「ルービックキューブの回復」を定式化すること。
• 外部データセットを用いず、医用データのみで事前学習を行うことで、ImageNet や UCF101 で事前学習した場合よりも優れた性能が得られることを実証すること。

提案手法

• 本手法は、3次元医用ボリュームを3次元グリッド（例：2×2×2）に分割し、サブボリューム（キューブ）を生成する。
• 元のボリュームの破損版を作成するために、キューブのランダムな入れ替え（キューブの再配置）とランダムな回転（キューブの回転）を適用する。
• プロキシタスクは、キューブの元の配置（空間的順序）および方向（姿勢）を再構成することである。
• 3次元畳み込みニューラルネットワーク（3D CNN）を、対照学習または分類損失を用いて正しい入れ替えと回転を予測するように訓練する。
• 下流タスク（例：脳出血分類や腫瘍セグメンテーション）において、限られたアノテーション付きデータを用いて、事前学習済みエンコーダーを微調整する。
• 本フレームワークは、学習を初期化する場合やUCF101で事前学習した重みを微調整する場合を含むベースラインと比較評価される。

実験結果

リサーチクエスチョン

• RQ1 3次元空間的変形に基づく自己教師付きプロキシタスクは、ボリューム医用画像の特徴抽出を向上させることができるか？
• RQ2 外部データセットを用いず、医用データのみで事前学習を行うことで、自然動画データ（例：UCF101）で事前学習する場合よりも優れた性能が得られるか？
• RQ3 プロキシタスクにキューブの回転を含めることで、キューブの順序のみを考慮する場合と比較して、学習された特徴の汎化性およびロバスト性は向上するか？
• RQ4 ルービックキューブ回復タスクは、分類やセグメンテーションを含む下流の3次元医用画像タスクにおける性能をどの程度向上させるか？

主な発見

• 提案されたルービックキューブ回復事前学習により、脳出血分類タスクで83.8％の正確性を達成し、学習を初期化した場合（72.6％）と比較して11.2％の向上を示した。
• 脳腫瘍セグメンテーションにおいては、3D DUCを用いて77.3％のmIoUを達成し、キューブ順序のみのベースラインと比較して2.3％の向上を示した。
• UCF101で事前学習した重みから微調整する場合に比べ、ルービックキューブ事前学習から微調整する場合、分類タスクで+2.7％、セグメンテーションタスクで+1.1％の性能向上を示した。
• プロキシタスクにキューブの回転を含めることで、並進および回転不変性を持つよりロバストな特徴が学習され、キューブ順序のみの場合よりも汎化性能が向上した。
• 3D DUCアーキテクチャは、すべての事前学習プロトコルにおいてU-Netを上回り、ランダムなデコーダー初期化による性能劣化が低減されたことを示した。
• 本手法は、外部データセットを一切使用せず、ラベルなし医用データのみで有効な自己教師付き事前学習を3次元畳み込みニューラルネットワークに適用した初の研究である。