[論文レビュー] iUNets: Fully invertible U-Nets with Learnable Up- and Downsampling

Christian Etmann, Rihuan Ke+1 more|arXiv (Cornell University)|May 11, 2020

Parallel Computing and Optimization Techniques参考文献 25被引用数 13

ひとこと要約

この論文では、指数的歪対称行列に基づく学習可能で可逆なアップサンプリングおよびダウンサンプリング操作を備えた、完全に可逆なU-Netアーキテクチャ、iUNetsを紹介する。可逆性によりバックプロパゲーションのメモリ効率が向上し、GPUメモリ制限下でもより深く、より大きなネットワークの学習が可能となり、3D医療画像セグメンテーションおよびCT後処理において標準的なU-Netを上回る性能を発揮する。また、競争力のある尤度性能を示すノーマライジングフローとしても利用可能である。

ABSTRACT

U-Nets have been established as a standard architecture for image-to-image learning problems such as segmentation and inverse problems in imaging. For large-scale data, as it for example appears in 3D medical imaging, the U-Net however has prohibitive memory requirements. Here, we present a new fully-invertible U-Net-based architecture called the iUNet, which employs novel learnable and invertible up- and downsampling operations, thereby making the use of memory-efficient backpropagation possible. This allows us to train deeper and larger networks in practice, under the same GPU memory restrictions. Due to its invertibility, the iUNet can furthermore be used for constructing normalizing flows.

研究の動機と目的

• 大規模3D医療画像処理における標準的U-Netの高メモリ消費を軽減し、メモリ効率の良い学習を可能にする。
• 中間活性化を保存せずにバックプロパゲーションを可能にする完全可逆U-Netアーキテクチャの開発。
• 固定されたGPUメモリ制約下で、より大きな、より深いネットワークの使用を可能にする。
• このアーキテクチャが生成的モデリングのためのノーマライジングフローとして使用可能であることを実証する。
• ボリュメトリックセグメンテーションやCT再構成を含む3D画像対画像タスクにおける性能向上。

提案手法

• 指数的歪対称畳み込みカーネルを用いた学習可能な可逆ダウンサンプリングを提案。これは双方向かつ微分可能である。
• 可逆アップサンプリングを学習済みダウンサンプリングの逆演算として定義し、完全な可逆性を確保する。
• 歪対称行列から導かれる直交畳み込み演算子を用い、訓練中の可逆性と安定性を維持する。
• 空間的ダウンサンプリングをチャネル拡張にマッピングする構造的カーネル再配置機構を採用し、情報の保持を図る。
• 活性化再構築を用いたメモリ効率の良いバックプロパゲーションを活用し、3D医療画像処理タスクにこのアーキテクチャを適用する。
• アフィンカップリング層を用いてiUNetブロックをスタックし、尤度ベースの生成的モデリングが可能なノーマライジングフローを構築する。

実験結果

リサーチクエスチョン

• RQ1 U-Netアーキテクチャにおける非可逆操作（例：マックスプーリングや転置畳み込み）に置き換えることのできる学習可能で可逆なアップ/ダウンサンプリング操作を設計可能か？
• RQ2 得られた完全可逆U-Net（iUNet）は、固定されたGPUメモリ制約下で、より深く、より大きなネットワークのメモリ効率の良い学習を可能にするか？
• RQ3 iUNetは、標準的U-Netと比較して3D医療画像セグメンテーションおよび後処理において優れた性能を発揮するか？
• RQ4 iUNetアーキテクチャはノーマライジングフローとして使用可能であり、最先端手法と比較して尤度性能はどの程度か？
• RQ5 学習された可逆操作は、アーチファクトを生成せずに多様な画像対画像タスクに一般化可能か？

主な発見

• 64チャネルのiUNetは、BraTS2018の検証セットでDiceスコア0.850を達成し、標準的U-Net（0.833）およびより大きなiUNets（例：iUNet-64: 0.850 vs. U-Net: 0.833）を上回った。
• 3D CT後処理において、iUNetはベースラインのU-Netを顕著に上回り、メモリ効率がもたらすモデル容量の拡大の利点を示した。
• iUNetを用いたメモリ効率の良いバックプロパゲーションは、標準的バックプロパゲーションと同等の損失曲線を生成し、学習の安定性を確認した。
• iUNetベースのノーマライジングフローは、CIFAR10で3.60ビット/次元の負の対数尤度を達成し、Real NVP（3.49ビット/次元）をわずかに下回った。これは改善の余地があることを示唆している。
• 可逆性に寄与する行列指数級数の切断近似により、尤度推定に潜在的なバイアスが生じ、性能差の要因である可能性がある。
• iUNet構造により、GPUメモリ制限下でもより大きなネットワークの学習が可能であることが実証された。特に、チャネル数の増加（例：iUNet-64）が性能向上に寄与している。