[論文レビュー] Efficient Folded Attention for 3D Medical Image Reconstruction and Segmentation

Hang Zhang, Jinwei Zhang+5 more|arXiv (Cornell University)|Sep 13, 2020

Medical Image Segmentation Techniques参考文献 34被引用数 8

ひとこと要約

本稿では、3D医療画像再構成およびセグメンテーションのための新規な折りたたみ注意（FA）機構を提案する。この機構により、計算コストとGPUメモリ使用量を顕著に低減しながら、精度を維持または向上できる。全アフィニティ行列の近似に、4つの並べ替えを伴うテンソルの折りたたみ・展開を適用することで、空間的およびチャネル的依存関係を効率的に統合的モデリング可能となり、定量的感受性マッピング（QSM）および複数硬化症（MS）病変セグメンテーションのベンチマークで最先端の性能を達成。標準的な注意機構と比較して、最大97.9%のメモリ削減と88.9%のFLOPs削減を実現した。

ABSTRACT

Recently, 3D medical image reconstruction (MIR) and segmentation (MIS) based on deep neural networks have been developed with promising results, and attention mechanism has been further designed to capture global contextual information for performance enhancement. However, the large size of 3D volume images poses a great computational challenge to traditional attention methods. In this paper, we propose a folded attention (FA) approach to improve the computational efficiency of traditional attention methods on 3D medical images. The main idea is that we apply tensor folding and unfolding operations with four permutations to build four small sub-affinity matrices to approximate the original affinity matrix. Through four consecutive sub-attention modules of FA, each element in the feature tensor can aggregate spatial-channel information from all other elements. Compared to traditional attention methods, with moderate improvement of accuracy, FA can substantially reduce the computational complexity and GPU memory consumption. We demonstrate the superiority of our method on two challenging tasks for 3D MIR and MIS, which are quantitative susceptibility mapping and multiple sclerosis lesion segmentation.

研究の動機と目的

• 3D医療画像処理における従来の自己注意機構の高い計算コストとメモリ使用量に対処すること。
• 計算複雑度を増加させることなく、空間的およびチャネルワイドな依存関係を両方モデリングできる効率的な注意機構を開発すること。
• 3D U-Netベースのネットワークにおける注意モジュールのメモリおよびFLOPsのオーバーヘッドを低減すること。
• QSM再構成やMS病変セグメンテーションのような大規模3D医療画像タスクにおける注意機構の実用的導入を可能にすること。

提案手法

• 本手法は、3D特徴テンソルを小さなサブテンソルに変換するためのテンソルの折りたたみ・展開操作を導入し、1つの大きなアフィニティ行列の代わりに4つのサブアフィニティ行列を計算可能にする。
• 特徴テンソルの4つの並べ替えを用いて、計算複雑度を低減した状態で全グローバル注意行列を近似するサブアフィニティ行列を生成する。
• 各サブアフィニティ行列は、クエリとキーのドット積により計算され、その後ソフトマックス処理が施され、結果は学習可能な投影により集約される。
• 各特徴要素をすべてのチャネルにわたるベクトルとして扱うことで、空間的およびチャネル的コンテキストを統合的にモデリングする。
• アフィニティ行列にランク1制約を適用することで、一般化性能を向上させ、過学習を低減する。
• FAモジュールはプラグアンドプレイであり、U-Netのような既存の3D CNNアーキテクチャに最小限のアーキテクチャ的変更で統合可能である。

実験結果

リサーチクエスチョン

• RQ1 3D医療画像における空間的およびチャネル的依存関係を両方モデリングできる、計算的に効率的な注意機構を設計可能か？
• RQ2 テンソルの折りたたみ・展開により、3D医療画像タスクにおける自己注意のFLOPsおよびGPUメモリ使用量はどの程度低減可能か？
• RQ3 提案された折りたたみ注意機構は、挑戦的な3D医療画像再構成およびセグメンテーションベンチマークで性能を維持または向上させるか？
• RQ4 臨床状況において真値が入手できない場合でも、折りたたみ注意モジュールは未学習データにどの程度一般化可能か？

主な発見

• QSMでは、FA-NetがRMSE 31.18を達成し、QSMnet、DA-Net、RSA-Netを上回った。
• QSMでは、FA-NetがPSNR 49.06を達成し、再構成の忠実性が高く、一般誤差が低減していることを示した。
• QSMでは、FA-NetがHFEN 32.49を達成し、高周波数成分や構造的類似性の保持が優れていることを示した。
• QSMでは、FA-NetがSSIM 0.9833を達成し、真値との構造的およびコントラスト類似性が優れていることを示した。
• FAモジュールは、標準的な自己注意（SA）と比較して、GPUメモリ使用量を97.9%削減し、FLOPsを88.9%削減したが、精度は同等または上回った。
• 定性的な結果から、FA-Netは、真値のCOSMOSデータが入手できない状況でも、より明確な病変境界とより高信号強度の病変を再構築できていた。