[論文レビュー] PMFSNet: Polarized Multi-scale Feature Self-attention Network For Lightweight Medical Image Segmentation

Jiahui Zhong, Wenhong Tian+5 more|arXiv (Cornell University)|Jan 15, 2024

Advanced Neural Network Applications被引用数 9

ひとこと要約

PMFSNet は軽量な UNet-風のセグメンテーションモデル（≤1M パラメータ）を提供し、極性化された多尺度自己注意ブロックを用いてグローバル文脈と局所特徴のバランスを取り、2D/3D 医用画像で計算量を抑えつつ競合的な精度を達成します。

ABSTRACT

Current state-of-the-art medical image segmentation methods prioritize accuracy but often at the expense of increased computational demands and larger model sizes. Applying these large-scale models to the relatively limited scale of medical image datasets tends to induce redundant computation, complicating the process without the necessary benefits. This approach not only adds complexity but also presents challenges for the integration and deployment of lightweight models on edge devices. For instance, recent transformer-based models have excelled in 2D and 3D medical image segmentation due to their extensive receptive fields and high parameter count. However, their effectiveness comes with a risk of overfitting when applied to small datasets and often neglects the vital inductive biases of Convolutional Neural Networks (CNNs), essential for local feature representation. In this work, we propose PMFSNet, a novel medical imaging segmentation model that effectively balances global and local feature processing while avoiding the computational redundancy typical in larger models. PMFSNet streamlines the UNet-based hierarchical structure and simplifies the self-attention mechanism's computational complexity, making it suitable for lightweight applications. It incorporates a plug-and-play PMFS block, a multi-scale feature enhancement module based on attention mechanisms, to capture long-term dependencies. Extensive comprehensive results demonstrate that even with a model (less than 1 million parameters), our method achieves superior performance in various segmentation tasks across different data scales. It achieves (IoU) metrics of 84.68%, 82.02%, and 78.82% on public datasets of teeth CT (CBCT), ovarian tumors ultrasound(MMOTU), and skin lesions dermoscopy images (ISIC 2018), respectively. The source code is available at https://github.com/yykzjh/PMFSNet.

研究の動機と目的

• 高い医用画像分割を低計算量と小さなモデルサイズで実現するニーズに対応する。
• グローバル文脈モデリングと局所特徴表現を軽量アーキテクチャでバランスさせる。
• プラグアンドプレイのアテンションブロックを提案し冗長性を削減しエッジデバイス展開を可能にする。
• 多尺度特徴融合を活用して二次コストなしに長距離依存を捉える。

提案手法

• エンコーダ-デコーダ風の構造を持つ PMFSNet を導入しボトルネックにプラグアンドプレイの PMFS ブロックを置く。
• Adaptive Multi-branch Feature Fusion (AMFF) を用いて異なる段階のマルチブランチ特徴を統合。
• Polarized Multi-scale Channel Self-attention (PMCS) と Polarized Multi-scale Spatial Self-attention (PMSS) を実装しチャネルと空間全体のアテンションポイントを拡大。
• クエリ、キー、値を用いたアテンションを線形化された多尺度の式と要素ごとの特徴強化で計算。
• Weighted Extended Dice Loss (WEDL) を用いて訓練しクラス不均衡に対処。

実験結果

リサーチクエスチョン

• RQ1 PMFSNet は ≤1M パラメータで 2D および 3D 医用画像で競争的なセグメンテーション性能を達成できるか？
• RQ2 PMFS の構成要素は計算コストを削減しつつ SOTA モデルと比較して精度を維持または向上できるか？
• RQ3 適応的マルチブランチ融合と極性多尺度注意はセグメンテーションにおけるグローバルおよび局所特徴表現にどのような影響を与えるか？
• RQ4 WEDL は医用 segmentation のクラス不均衡に対処するのに有効か？

主な発見

手法	FLOPs(G) ↓	Params(M) ↓	HD(mm) ↓	ASSD(mm) ↓	IoU(%) ↑	SO(%) ↑	DSC(%) ↑
UNet3D	2223.03	16.32	113.79	22.40	70.62	70.72	36.67
DenseVNet	23.73	0.87	8.21	1.14	84.57	94.88	91.15
AttentionUNet3D	2720.79	94.48	147.10	61.10	52.52	42.49	64.08
DenseVoxelNet	402.32	1.78	41.18	3.88	81.51	92.50	89.58
MultiResUNet3D	1505.38	17.93	74.06	8.17	76.19	81.70	65.45
UNETR	229.19	93.08	107.89	17.95	74.30	73.14	81.84
SwinUNETR	912.35	62.19	82.71	7.50	83.10	86.80	89.74
TransBTS	306.80	33.15	29.03	4.10	82.94	90.68	39.32
nnFormer	583.49	149.25	51.28	5.08	83.54	90.89	90.66
3D UX-Net	1754.79	53.01	108.52	19.69	75.40	73.48	84.89
PMFSNet3D (Ours)	15.14	0.63	5.57	0.79	84.68	95.10	91.30

• PMFSNet3D は IoU 84.68%、SO 95.10%、DSC 91.30% を 0.63M パラメータと 15.14 GFLOPs で 3D CBCT 歯データセットで達成。
• PMFSNet3D は FLOPs とパラメータを大幅に削減し（いくつかの SOTA モデルと比較して）高い IoU および DSC を維持。
• 3 つのデータセット（CBCT、MMOTU ultrasound、ISIC 2018）に渡って、PMFSNet の変種は競合手法よりはるかに少ないパラメータで高い精度を達成。
• PMFS ブロックは線形アテンション様の複雑さで長距離依存関係を効率的にモデリングし、エッジデバイス展開を支援。