[論文レビュー] Invertible Residual Network with Regularization for Effective Medical Image Segmentation

Kashu Yamazaki, Vidhiwar Singh Rathour+1 more|arXiv (Cornell University)|Mar 16, 2021

Advanced Neural Network Applications参考文献 18被引用数 4

ひとこと要約

本稿では、メモリ効率の良い3次元医療画像セグメンテーションを実現するため、部分的可逆ResNet（Partially-InvRes）および完全可逆ResNet（Fully-InvRes）の2つの変種を提案する。加法的カップリングと学習可能な圧縮操作を活用して、可逆的なダウンサンプリング/アップサンプリングを実現することで、iSeg 2019およびBraTS 2020のデータセットにおいて、トレーニング中のメモリ使用量を削減しながらも、競争力のあるセグメンテーション性能を維持する。特にFully-InvResは、ベースラインの3D Unetと比較して2.12GBのメモリ使用量にまで低減した。

ABSTRACT

Deep Convolutional Neural Networks (CNNs) i.e. Residual Networks (ResNets) have been used successfully for many computer vision tasks, but are difficult to scale to 3D volumetric medical data. Memory is increasingly often the bottleneck when training 3D Convolutional Neural Networks (CNNs). Recently, invertible neural networks have been applied to significantly reduce activation memory footprint when training neural networks with backpropagation thanks to the invertible functions that allow retrieving its input from its output without storing intermediate activations in memory to perform the backpropagation. Among many successful network architectures, 3D Unet has been established as a standard architecture for volumetric medical segmentation. Thus, we choose 3D Unet as a baseline for a non-invertible network and we then extend it with the invertible residual network. In this paper, we proposed two versions of the invertible Residual Network, namely Partially Invertible Residual Network (Partially-InvRes) and Fully Invertible Residual Network (Fully-InvRes). In Partially-InvRes, the invertible residual layer is defined by a technique called additive coupling whereas in Fully-InvRes, both invertible upsampling and downsampling operations are learned based on squeezing (known as pixel shuffle). Furthermore, to avoid the overfitting problem because of less training data, a variational auto-encoder (VAE) branch is added to reconstruct the input volumetric data itself. Our results indicate that by using partially/fully invertible networks as the central workhorse in volumetric segmentation, we not only reduce memory overhead but also achieve compatible segmentation performance compared against the non-invertible 3D Unet. We have demonstrated the proposed networks on various volumetric datasets such as iSeg 2019 and BraTS 2020.

研究の動機と目的

• ボリュメトリック医療画像データを用いたトレーニングにおいて、3次元畳み込みニューラルネットワーク（CNN）の高いメモリ使用量を低減すること。
• ラベル付きデータが限られる医療画像分野の課題を克服するため、過学習を抑える正則化を組み込むこと。
• 制限されたGPUメモリ下で、より深く、より大きな3次元ネットワークのトレーニングを可能にするために、可逆残差ブロックを活用すること。
• 一般化性能とメモリ効率を向上させるために、標準の3D Unetベースラインに可逆性と変分オートエンコーダー（VAE）正則化を統合すること。

提案手法

• 入力を2つの部分に分割し、学習可能な関数を用いて変換する加法的カップリングを用いて、可逆残差ブロックを実装する部分的可逆ResNet（Partially-InvRes）を構築する。
• 学習可能なピクセルシャッフル（圧縮）を用いて、可逆的なダウンサンプリングおよびアップサンプリングを実現し、ネットワーク全体の完全な可逆性を保証する完全可逆ResNet（Fully-InvRes）を設計する。
• 入力ボリュームデータを再構築するVAE正則化ブランチを統合し、小規模な医療画像データセットにおける一般化性能の向上と過学習の低減を図る。
• 3D Unetに基づくネットワークアーキテクチャを構築し、標準の残差ブロックを可逆対応のものに置き換え、自己教師付き事前学習用にVAEブランチを追加する。
• ADAM最適化アルゴリズムを用いてトレーニングを行い、ベースラインの学習率を2e-4に設定する。すべての層でインスタンス正則化とLeaky ReLUを適用する。
• PyTorchにおけるMemCNNライブラリを活用し、中間活性化を保存する必要なく、可逆関数の逆算によってメモリ効率の良いバックプロパゲーションを実現する。

実験結果

リサーチクエスチョン

• RQ1 可逆残差ネットワークは、セグメンテーション精度を損なわせることなく、3次元医療画像セグメンテーションにおけるメモリ使用量を削減できるか？
• RQ2 部分的可逆と完全可逆の残差ネットワークは、ボリュメトリックデータセットにおいて、メモリ効率と性能の面でどのように比較されるか？
• RQ3 VAE正則化ブランチを追加することで、データ量が少ない医療画像分野の状況において、一般化性能の向上と過学習の低減はどの程度達成されるか？
• RQ4 可逆ネットワークは、制限されたGPUメモリ下で、3D Unetアーキテクチャに効果的に統合可能か？
• RQ5 提案された可逆ネットワークは、iSeg 2019およびBraTS 2020といった標準ベンチマークにおいて、非可逆的3D Unetベースラインと比較してどの程度の性能を示すか？

主な発見

• iSeg 2019データセットにおいて、Fully-InvResはトレーニング時メモリ使用量が2.12GBにまで低減され、非可逆的3D Unetベースラインの4.18GBと比較して顕著な削減を達成した。
• iSeg 2019データセットでは、Fully-InvResがCSFで95.09%、GMで91.32%、WMで90.31%のDiceスコアを達成し、すべての構造でベースラインをわずかに上回った。
• BraTS 2020データセットでは、Fully-InvResがWTで88.20%、ETで72.01%、TCで77.74%のDiceスコアを達成し、すべての指標でベースライン3D Unetを上回った。
• 両データセットにおいてハウスドルフ距離が低減された。例えば、BraTS 2020では、Fully-InvResがWTで10.46を記録したのに対し、ベースラインは11.45であった。境界の整合性が向上したことを示している。
• 可逆ブロックとVAE正則化の組み合わせにより、部分的可逆ResNetおよび完全可逆ResNetの両方が、限られたトレーニングデータ下でも過学習を回避し、高い精度を達成した。
• 表4は、両提案モデルがメモリ効率が高く、過学習を回避し、高い精度を達成していることを確認している。これに対して、ベースライン3D Unetは、両者の特徴を欠いている。