[論文レビュー] MALUNet: A Multi-Attention and Light-weight UNet for Skin Lesion Segmentation

Jiacheng Ruan, Suncheng Xiang+3 more|arXiv (Cornell University)|Nov 3, 2022

Cutaneous Melanoma Detection and Management被引用数 6

ひとこと要約

本稿では、特徴抽出とサンプル間関係性の向上を目的とした4つの新規アテンションモジュール（DGA、IEA、CAB、SAB）を備えた軽量U-Netアーキテクチャ、MALUNetを提案する。標準U-Netと比較してパラメータ数を44倍、FLOPsを166倍削減しながら、ISIC2017およびISIC2018で最先端の性能を達成し、mIoUを2.39%、DSCを1.49%向上させた。

ABSTRACT

Recently, some pioneering works have preferred applying more complex modules to improve segmentation performances. However, it is not friendly for actual clinical environments due to limited computing resources. To address this challenge, we propose a light-weight model to achieve competitive performances for skin lesion segmentation at the lowest cost of parameters and computational complexity so far. Briefly, we propose four modules: (1) DGA consists of dilated convolution and gated attention mechanisms to extract global and local feature information; (2) IEA, which is based on external attention to characterize the overall datasets and enhance the connection between samples; (3) CAB is composed of 1D convolution and fully connected layers to perform a global and local fusion of multi-stage features to generate attention maps at channel axis; (4) SAB, which operates on multi-stage features by a shared 2D convolution to generate attention maps at spatial axis. We combine four modules with our U-shape architecture and obtain a light-weight medical image segmentation model dubbed as MALUNet. Compared with UNet, our model improves the mIoU and DSC metrics by 2.39% and 1.49%, respectively, with a 44x and 166x reduction in the number of parameters and computational complexity. In addition, we conduct comparison experiments on two skin lesion segmentation datasets (ISIC2017 and ISIC2018). Experimental results show that our model achieves state-of-the-art in balancing the number of parameters, computational complexity and segmentation performances. Code is available at https://github.com/JCruan519/MALUNet.

研究の動機と目的

• 高コストな計算リソースとメモリ要件により、リソース制限のある臨床環境への複雑なディーブラーニングモデルの導入が困難であるという課題に対処する。
• モデルの複雑さを増さずに、局所的およびグローバルな特徴表現を強化することで、皮膚腫瘍のセグメンテーション性能を向上させる。
• 自己アテンションベースのモデルでしばしば無視されがちな、医療画像セグメンテーションにおけるサンプル間関係性を確立する。
• チャネル軸および空間軸の両方で、マルチスケールかつマルチステージの特徴を効果的に統合し、パラメータの増加を最小限に抑えてセグメンテーション精度を向上させる。

提案手法

• 拡張畳み込みとゲート付きアテンションを組み合わせたDilated Gated Attention (DGA)ブロックを導入し、パラメータを削減しつつグローバルおよびローカルなコンテキストを捉える。
• データセット全体レベルの表現を学習することで、サンプル間関係性をモデル化するためのInverted External Attention (IEA)ブロックを提案する。
• チャネル軸と空間軸のそれぞれで、マルチステージ特徴を統合し、アテンションマップを生成するためのChannel Attention Bridge (CAB)およびSpatial Attention Bridge (SAB)モジュールを設計する。
• 6段階のU-Netアーキテクチャに4つのモジュールを統合し、MALUNetを構築することで、モデルサイズを最小限に抑えつつ性能を最大化する。
• DGAで深度分離畳み込みを用いることで、計算複雑性を低減し、効率性を維持する。
• 各段階のチャネル数を8、16、24、32、48、64に削減することで、アーキテクチャを最適化し、コンactかつ強力なバックボーンを実現する。

実験結果

リサーチクエスチョン

• RQ1 パラメータ数とFLOPsを大幅に削減した軽量U-Netアーキテクチャが、顕著に低いリソースで最先端のセグメンテーション性能を達成できるか？
• RQ2 アテンション機構を設計することで、医療画像セグメンテーションにおけるサンプル内およびサンプル間の特徴関係性を向上させられるか？
• RQ3 チャネルおよび空間レベルでのマルチステージ特徴統合機構は、セグメンテーション精度の向上にどの程度効果的か？
• RQ4 グローバル（IEA経由）とローカル（DGA経由）のアテンションを統合することで、モデルの複雑さを増さずに性能向上が達成できるか？
• RQ5 提案されたアーキテクチャは、ISIC2017およびISIC2018のような多様な皮膚腫瘍データセットにおいても高い性能を維持できるか？

主な発見

• MALUNetは、標準U-Netと比較してパラメータ数を44倍、計算複雑度（GFLOPs）を166倍削減しながら、ISIC2018でmIoUを2.39%、DSCを1.49%向上させた。
• DGAブロックは、FLOPsとパラメータ数を削減した状態で、ベースラインU-Net（BasicUNet）と比較してmIoUを0.90%、DSCを0.66%向上させた。
• IEAブロックは、標準EAブロックを上回り、パラメータ数の増加がわずかであるにもかかわらず、mIoUを0.46%、DSCを0.29%向上させた。
• DGAとIEAに加え、CABとSABを組み合わせた場合が最良の性能を示し、ISIC2018でmIoU 80.23%、DSC 89.14%を達成した。
• アブレーションスタディの結果、各モジュールが独立して性能向上に寄与しており、CABとSABは段階間の特徴統合を強化していることが確認された。
• MALUNetは、モデル効率性とセグメンテーション精度の両面で最先端のバランスを達成しており、ISIC2017およびISIC2018の両データセットで既存の軽量モデルを上回った。