UNETR++: 効率的かつ正確な3D医用画像セグメンテーションへの取り組み

トランスフォーマーモデルの成功により、最近の研究では3D医用セグメンテーションタスクへの適用可能性が検討されています。トランスフォーマーモデル内では、自己注意メカニズムが、局所的な畳み込みベースの設計に対して長距離依存関係を捉えようとする主要な構成要素の一つです。しかし、自己注意操作は二次の計算量を持ち、これは計算ボトルネックとなります。特に、入力が3Dで多数のスライスを持つ体積医用画像では特に顕著です。本論文では、UNETR++という名前の3D医用画像セグメンテーション手法を提案し、高品質なセグメンテーションマスクとともに、パラメーター、計算コスト、および推論速度に関する効率性を提供します。私たちのデザインの核心は、空間およびチャネル方向の識別的特徴を学習するための新しい効率的ペア注意（EPA）ブロックの導入です。このブロックは、空間およびチャネル注意に基づく相互依存ブランチのペアを使用して、効率的に空間的およびチャネル方向の識別的特徴を学習します。私たちの空間注意の定式化は効率的で、入力に対して線形の計算量を持っています。空間およびチャネル指向のブランチ間の通信を可能にするために、クエリおよびキーのマッピング関数の重みを共有し、相補的な利点（ペア注意）を提供しつつ、計算の複雑さを減少させています。Synapse、BTCV、ACDC、BraTS、そしてDecathlon-Lungという5つのベンチマークに関する広範な評価により、我々の貢献が効率性と精度の両面で効果的であることが明らかになりました。Synapseでは、UNETR++が87.2%のDiceスコアで新たな最先端を樹立し、文献中の最も優れた手法と比較して、パラメーターとFLOPsを71%超削減しました。私たちのコードとモデルは以下で入手可能です: https://tinyurl.com/2p87x5xn.