[論文レビュー] EGE-UNet: an Efficient Group Enhanced UNet for skin lesion segmentation

Jiacheng Ruan, Mingye Xie+3 more|arXiv (Cornell University)|Jul 17, 2023

Cutaneous Melanoma Detection and Management被引用数 12

ひとこと要約

EGE-UNet は、Group multi-axis Hadamard Product Attention (GHPA) と Group Aggregation Bridge (GAB) を用いた軽量 UNet バリアントで、パラメータ数が極端に少なく (~50KB) かつ計算量が低い状態で、最先端の皮膚病変セグメンテーションを実現します。

ABSTRACT

Transformer and its variants have been widely used for medical image segmentation. However, the large number of parameter and computational load of these models make them unsuitable for mobile health applications. To address this issue, we propose a more efficient approach, the Efficient Group Enhanced UNet (EGE-UNet). We incorporate a Group multi-axis Hadamard Product Attention module (GHPA) and a Group Aggregation Bridge module (GAB) in a lightweight manner. The GHPA groups input features and performs Hadamard Product Attention mechanism (HPA) on different axes to extract pathological information from diverse perspectives. The GAB effectively fuses multi-scale information by grouping low-level features, high-level features, and a mask generated by the decoder at each stage. Comprehensive experiments on the ISIC2017 and ISIC2018 datasets demonstrate that EGE-UNet outperforms existing state-of-the-art methods. In short, compared to the TransFuse, our model achieves superior segmentation performance while reducing parameter and computation costs by 494x and 160x, respectively. Moreover, to our best knowledge, this is the first model with a parameter count limited to just 50KB. Our code is available at https://github.com/JCruan519/EGE-UNet.

研究の動機と目的

• 効率的でモバイル対応の皮膚病変セグメンテーションを推進する。
• 高性能を備えた超軽量アーキテクチャの開発。
• 計算量を最小化しつつ、マルチ視点・マルチスケール特徴を抽出するGHPAとGABの導入。
• ISIC2017/ISIC2018で大規模なアブレーションを用いて最先端の結果を示す。

提案手法

• グループ多軸ハダマード積注意機構（GHPA）を導入し、計算量を二次から線形へ削減し、マルチ視点情報を捉える。
• 低レベルと高レベルのマップおよび補助マスクを用いてマルチスケール特徴を統合するグループ集約ブリッジ（GAB）を導入。
• GHPAとGABを六段階のU-Netフレームワークに統合し、多段階出力の深層監視（Deep supervision）を実現。
• ISIC2017およびISIC2018でデータ拡張、AdamW最適化、コサインアニーリング学習率スケジュールを用いてエンドツーエンドで訓練。
• 階層別深部監視ウェイトを持つ二値交差エントロピーとDice損失を組み合わせた損失を使用。
• パラメータ、FLOPs、mIoU、DSC の観点で最先端手法と比較評価。

実験結果

リサーチクエスチョン

• RQ1 GHPAはモバイルヘルスアプリケーションに適した線形計算量のマルチ視点注意を提供できるか？
• RQ2 軽量フレームワークにマスクガイダンスを組み込んで、マルチスケール特徴の融合を改善するか？
• RQ3 EGE-UNetはISIC2017/ISIC2018において大規模なTransformerベース手法と他の軽量モデルの双方とどのように比較されるか？
• RQ4 極端に低いパラメータ数とセグメンテーション性能のトレードオフは何か？

主な発見

データセット	モデル	Params(M)	GFLOPs	mIoU(%)	DSC(%)
ISIC2017	UNet	7.77	13.76	76.98	86.99
ISIC2017	UTNetV2	12.80	15.50	77.35	87.23
ISIC2017	TransFuse	26.16	11.50	79.21	88.40
ISIC2017	MobileViTv2	1.87	0.70	78.72	88.09
ISIC2017	MobileNetv3	1.19	0.10	77.69	87.44
ISIC2017	UNeXt-S	0.32	0.10	78.26	87.80
ISIC2017	MALUNet	0.177	0.085	78.78	88.13
ISIC2017	EGE-UNet (Ours)	0.053	0.072	79.81 ± 0.10	88.77 ± 0.06
ISIC2018	UNet	7.77	13.76	77.86	87.55
ISIC2018	UNet++	9.16	34.86	78.31	87.83
ISIC2018	Att-UNet	8.73	16.71	78.43	87.91
ISIC2018	UTNetV2	12.80	15.50	78.97	88.25
ISIC2018	SANet	23.90	5.96	79.52	88.59
ISIC2018	TransFuse	26.16	11.50	80.63	89.27
ISIC2018	MobileViTv2	1.87	0.70	79.88	88.81
ISIC2018	MobileNetv3	1.19	0.10	78.55	87.98
ISIC2018	UNeXt-S	0.32	0.10	79.09	88.33
ISIC2018	MALUNet	0.177	0.085	80.25	89.04
ISIC2018	EGE-UNet (Ours)	0.053	0.072	80.94 ± 0.11	89.46 ± 0.07

• EGE-UNetはISIC2017およびISIC2018で最先端の性能を達成しつつ、パラメータと計算量を大幅に削減。
• TransFuse のような大規模モデルと比較して、パラメータを494倍、FLOPsを160倍削減。
• EGE-UNetは約50 KBのパラメータサイズで高いセグメンテーション精度を達成。
• GHPAはマルチ軸グルーピングを備えた線形計算量の注意機構を提供し、特徴を強化。
• GABは低・高レベルの特徴とマスクガイダンスを融合し、効果的なマルチスケール融合を実現。
• アブレーションによりGHPAとGABの双方の必須性、GABにおけるマスク情報と拡張（dilation）の重要性が確認された。