医用画像では位置関係が重要だった！！医用画像セグメンテーションのためのContrastive Learning

Contrastive Learning 2022年05月31日

3つの要点
✔️ 医用画像セグメンテーションのためのContrastive Learning
✔️ 画像同士の位置関係をもとに正例・負例ペアを決定
✔️ 半教師あり学習・転移学習タスクで既存手法を超えるセグメンテーション精度を実現

Positional Contrastive Learning for Volumetric Medical Image Segmentation
written by Dewen Zeng, Yawen Wu, Xinrong Hu, Xiaowei Xu, Haiyun Yuan, Meiping Huang, Jian Zhuang, Jingtong Hu, Yiyu Shi
(Submitted on 16 Jun 2021 (v1), last revised 28 Sep 2021 (this version, v3))
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

はじめに

自己教師あり学習はラベルなしの画像データを用いた学習を行うことができるため、アノテーションがついたデータを得ることが難しい医用画像タスクにおいては非常に有効な手法とみなされています。

自己教師あり学習の中でもSimCLRやMoCoなどの手法が有名なContrastive Learningは、コンピュータビジョンの分野で特に大きな成功を収めています (詳しい内容はこちら)。

Constrastive Learningを簡単に説明すると、「正例ペアの特徴量を近づけて、負例ペアの特徴量を遠ざけること」を目的とした自己教師あり学習です。学習後に得られる特徴量は、下流タスク (画像分類、物体検出、セグメンテーションなど)で、精度を向上させるために使用されます。

Contrastive Learningでは、正例・負例ペアの決定方法が重要になっています。既存の手法では、正例ペアは通常の画像とDAが施された画像とで作られ、負例ペアは異なる画像同士で作られます。

一方で、既存の手法をそのまま医用画像タスクに対して使おうとすると、正例ペア・負例ペアに関して問題が生じます。その結果として、下流タスクでのセグメンテーション精度の低下が引き起こされてしまいます。

その問題とは具体的に、不適切な負例ペアを作ってしまうという問題です。つまり、特徴量を遠ざけるべきではない画像同士が負例ペアとなり、正しい学習を行うことができなくなるということです。

この問題は、データセット内に同じ組織や臓器が、複数の画像にまたがって存在することが原因となっています。つまり、見た目がほとんど同じ画像であるにも関わらず、負例ペアとして特徴量を遠ざけるように学習してしまうのです。

本論文では、この問題を解決するべく下図に示す医用画像セグメンテーションのための新たなConstrative Learning：Positional Contrastive Learning (PCL)を提案しています。PCLは医用画像の位置関係を効果的に利用して正例・負例ペアを決定することで上記の問題を解決しています。

医用画像のセグメンテーションにおいて、既存の手法とPCLを事前学習として用いた半教師あり学習・転移学習で比較すると、PCLは既存の手法を超えるセグメンテーション精度を達成しました。さらに、PCLはアノテーションのついた画像が少ない場合により大きな効果を発揮することがわかりました。

本記事では、PCLの解説及び、PCLの有用性を示す実験結果について紹介します。

提案手法：Positional Contrastive Learning

まずは、PCLによる自己教師あり学習について紹介します。

上の図は、PCLの概要を表しています。図に示すように入力データとしては、3D画像からz方向に沿って切り出した2D画像 (x-y平面画像)を使用します。

PCLの肝である、正例・負例ペアの決め方について紹介します。まず、2D医用画像の位置を定義するためにpositionを定義します。positionは、3D画像から切り出される2D画像の数をn、2D画像が切り出されたz方向の座標をn (0 < m < n)して、2D画像のをm/n (値は0~1の間)と定義されます。

このpositionをもとに、正例・負例ペアは以下のように決められます。

正例ペア：positionの差がある一定の閾値内である2D画像同士
負例ペア：positionの差が閾値を超える画像同士

ここで重要なことは、positionの差のみでペアが決定していることです。同じ3D画像から切り取られた2D画像はもちろん異なる3D画像から切り取られた2D画像同士でも、positionの差が閾値を超えない限り正例ペアとして扱われます。こうすることによって、同じ組織や臓器が複数の画像にまたがって存在していても、それらの画像の特徴量を近づけることを可能としています。

PCLでの損失関数は以下のようになります。

通常のContrastive Learningとほとんど同じですが、PCLでの損失関数はバッチ内の画像に対するすべての正例ペアでContrastive lossの平均をとったものになります。式 (1)は、式 (2)の損失をデータ全体に関して和を取ったものになります。また、通常のContrastive Learningと同様に各画像に対してデータ拡張を行っています。

次に、ファインチューニングの方法について紹介します。