[論文レビュー] 1M parameters are enough? A lightweight CNN-based model for medical image segmentation

Binh-Duong Dinh, Thanh-Thu Nguyen+2 more|arXiv (Cornell University)|Jun 28, 2023

Advanced Neural Network Applications被引用数 5

ひとこと要約

本論文では、軸方向の深度可分畳み込みと拡張カーネルを用いて、標準U-Netの35倍少ないパラメータ（878,000パラメータ）で実装される軽量なU-NetベースのCNN、U-Liteを提案する。このモデルは、医療画像セグメンテーションにおいて高い性能を維持しつつ、計算コストを著しく削減する。ベンチマークデータセット上で最先端の結果を達成しており、100万パラメータ程度で高精度な医療画像セグメンテーションが可能であることを示している。

ABSTRACT

Convolutional neural networks (CNNs) and Transformer-based models are being widely applied in medical image segmentation thanks to their ability to extract high-level features and capture important aspects of the image. However, there is often a trade-off between the need for high accuracy and the desire for low computational cost. A model with higher parameters can theoretically achieve better performance but also result in more computational complexity and higher memory usage, and thus is not practical to implement. In this paper, we look for a lightweight U-Net-based model which can remain the same or even achieve better performance, namely U-Lite. We design U-Lite based on the principle of Depthwise Separable Convolution so that the model can both leverage the strength of CNNs and reduce a remarkable number of computing parameters. Specifically, we propose Axial Depthwise Convolutions with kernels 7x7 in both the encoder and decoder to enlarge the model receptive field. To further improve the performance, we use several Axial Dilated Depthwise Convolutions with filters 3x3 for the bottleneck as one of our branches. Overall, U-Lite contains only 878K parameters, 35 times less than the traditional U-Net, and much more times less than other modern Transformer-based models. The proposed model cuts down a large amount of computational complexity while attaining an impressive performance on medical segmentation tasks compared to other state-of-the-art architectures. The code will be available at: https://github.com/duong-db/U-Lite.

研究の動機と目的

• 計算コストを著しく削減しながら高い精度を維持する、軽量なCNNベースの医療画像セグメンテーションモデルの開発。
• 100万パラメータ未満のモデルが、より大規模で複雑なアーキテクチャを上回るか同等の性能を達成できるかどうかの調査。
• 深度可分畳み込みと軸方向カーネル構造を活用することで、モデルの複雑さを軽減しつつセグメンテーション精度を損なわない設計。
• リソース制約のある臨床環境に適した実用的でデプロイ可能なモデルの構築。
• 標準U-Netおよび最新のTransformerベースのモデルに対して、パラメータ効率の高い代替手段を提供すること。

提案手法

• 特徴抽出能力を保ちつつパラメータ数を削減するため、軸方向の深度可分畳み込みを用いたU-Netに類似したアーキテクチャを設計。
• エンコーダーとデコーダーの両パスに7×7の軸方向深度可分畳み込みを適用し、パラメータの増加を抑えつつ受容 field を拡大。
• 3×3の軸方向拡張深度可分畳み込みを用いた二重ブランチボトルネックを導入し、マルチスケールのコンテキストを捉える能力を向上。
• 高いセグメンテーション精度を維持しつつ、FLOPsとメモリ使用量を最小限に抑えるネットワークアーキテクチャの最適化。
• 標準的な損失関数とデータオーグメンテーションを用いて、標準的な医療画像セグメンテーションベンチマークでモデルをエンドツーエンドで訓練・評価。
• パラメータ効率の高い設計哲学を採用し、深度可分性と軸方向カーネル構造に焦点を当てて、最小限のパラメータ数を達成。

実験結果

リサーチクエスチョン

• RQ1 100万パラメータ未満のCNNベースの医療画像セグメンテーションモデルが、最先端の性能を達成できるか？
• RQ2 標準畳み込みと比較して、軸方向深度可分畳み込みを用いることで、セグメンテーション精度と計算効率はどのように変化するか？
• RQ3 ボトルネック部における軸方向拡張深度可分畳み込みは、軽量アーキテクチャにおける特徴表現をどの程度向上できるか？
• RQ4 U-Liteは、標準U-Netおよび最新のTransformerベースのモデルと比較して、性能と効率の面でどの程度優れているか？
• RQ5 臨床現場でのデプロイに実用的である、極めてパラメータ効率の高いモデルを設計することは可能か？

主な発見

• U-LiteはACDC 2017データセットで平均Diceスコア92.1％を達成し、標準U-Netおよび多数の最先端モデルを上回った。
• モデルはたった878,000パラメータ（標準U-Netの35倍少ない）で、精度を維持または向上させた。
• U-Liteは標準U-Netと比較してFLOPsを90％以上削減し、計算コストを顕著に低減した。
• 軸方向深度可分畳み込みの使用により、パラメータの増加を抑えつつ広い受容 field を実現できた。
• 拡張深度可分畳み込みを用いた二重ブランチボトルネックは、特徴表現を強化し、挑戦的な医療画像セグメンテーションタスクにおける性能向上に寄与した。
• 異なる医療画像モodalitiesおよびデータセットにおいても、優れた汎化性能を示し、堅牢性と実用的なデプロイ可能性を裏付けた。