[論文レビュー] UNeXt: MLP-based Rapid Medical Image Segmentation Network

Jeya Maria Jose Valanarasu, Vishal M. Patel|arXiv (Cornell University)|Mar 9, 2022

Advanced Neural Network Applications被引用数 54

ひとこと要約

UNeXtは医用画像分割のための畳み込みステムに続くトークン化されたMLPブロックで、TransUNetおよびUNet系よりはるかに少ないパラメータとはるかに低い計算量で最先端の性能を実現し、現場でのCPU推論を高速化します。

ABSTRACT

UNet and its latest extensions like TransUNet have been the leading medical image segmentation methods in recent years. However, these networks cannot be effectively adopted for rapid image segmentation in point-of-care applications as they are parameter-heavy, computationally complex and slow to use. To this end, we propose UNeXt which is a Convolutional multilayer perceptron (MLP) based network for image segmentation. We design UNeXt in an effective way with an early convolutional stage and a MLP stage in the latent stage. We propose a tokenized MLP block where we efficiently tokenize and project the convolutional features and use MLPs to model the representation. To further boost the performance, we propose shifting the channels of the inputs while feeding in to MLPs so as to focus on learning local dependencies. Using tokenized MLPs in latent space reduces the number of parameters and computational complexity while being able to result in a better representation to help segmentation. The network also consists of skip connections between various levels of encoder and decoder. We test UNeXt on multiple medical image segmentation datasets and show that we reduce the number of parameters by 72x, decrease the computational complexity by 68x, and improve the inference speed by 10x while also obtaining better segmentation performance over the state-of-the-art medical image segmentation architectures. Code is available at https://github.com/jeya-maria-jose/UNeXt-pytorch

研究の動機と目的

• 現実世界の点検用途での計算資源を制約した状態での画像セグメンテーションを動機づける。
• 畳み込みとトークン化されたMLPコンポーネントを組み合わせた軽量なエンコーダ-デコーダアーキテクチャを開発する。
• 潜在表現を効率的にモデル化するための軸方向シフティングを備えたトークン化MLPブロックを導入する。
• パラメータとFLOPsを大幅に削減しつつ、セグメンテーション精度を維持・向上させることを示す。

提案手法

• 初期の畳み込みステージに続く2段階構造と、Tokenized MLPステージ。
• Tokenized MLPブロックは畳み込み特徴をトークンに射影し、ローカル依存性をモデル化するためにシフトしたMLPを適用する。
• 軸方向シフティング（WとH）はトークン化の前に局所性を誘導し、Tok-MLPブロックには深さ方向畳み込みとGELU活性化を使用する。
• Tokenized MLPブロック内で残差接続と層正規化を使用する。
• エンコーダとデコーダ間のスキップ接続はUNetを模倣し、デコーダはTokenized MLPブロックに続く畳み込みブロックを使用する。

実験結果

リサーチクエスチョン

• RQ1 潜在空間での畳み込みステムとトークン化MLPがパラメータと計算量を削減しつつセグメンテーション精度を維持できるか。
• RQ2 Tok-MLPブロックにおける軸方向シフトは、競合的な医用画像セグメンテーションに十分な局所性を提供するか。
• RQ3 UNeXtは精度、パラメータ、FLOPs、CPU推論時間の点でUNet、UNet++、ResUNet、MedT、TransUNetと比較してどうか。

主な発見

ネットワーク	パラメータ（M）	推論速度（ms）	GFLOPs	ISIC F1	ISIC IoU	BUSI F1	BUSI IoU
UNet	31.13	223	55.84	84.03 ± 0.87	74.55 ± 0.96	76.35 ± 0.89	63.85 ± 1.12
UNet++	9.16	173	34.65	84.96 ± 0.71	75.12 ± 0.65	77.54 ± 0.74	64.33 ± 0.75
ResUNet	62.74	333	94.56	85.60 ± 0.68	75.62 ± 1.11	78.25 ± 0.74	64.89 ± 0.83
MedT	1.60	751	21.24	87.35 ± 0.18	79.54 ± 0.26	76.93 ± 0.11	63.89 ± 0.55
TransUNet	105.32	246	38.52	88.91 ± 0.63	80.51 ± 0.72	79.30 ± 0.37	66.92 ± 0.75
UNeXt	1.47	25	0.57	89.70 ± 0.96	81.70 ± 1.53	79.37 ± 0.57	66.95 ± 1.22

• UNeXtはISICおよびBUSIデータセットで競合的または優れたF1およびIoUスコアを達成する。
• UNeXtは1.47Mパラメータおよび0.57 GFLOPsを使用し、TransUNet（105.32M、38.52 GFLOPs）よりもはるかに少ない。
• UNeXtはISICで89.70 F1および81.70 IoU、BUSIで79.37 IoUを達成し、CPU推論時間は25 ms。
• ベースライン全体で、UNeXtは精度と効率の最良のバランスを提供し、計算量とパラメータ数の面で注意機構ベースのモデルを上回る。
• アブレーション研究は、Tok-MLPのシフトと畳み込みとMLPステージの組み合わせが、複雑さの増加を最小限に抑えつつ最良の性能を示すことを示す。