Factorizer: A Scalable Interpretable Approach to Context Modeling for Medical Image Segmentation

要約

U 字型アーキテクチャを備えた畳み込みニューラル ネットワーク (CNN) は、さまざまな臨床目的に不可欠な医療画像セグメンテーションを支配しています。
ただし、畳み込みの固有の局所性により、CNN は、脳病変などの一部の構造をより適切に認識するために不可欠なグローバル コンテキストを十分に活用できません。
トランスフォーマーは、主に長期的な依存関係をモデル化する機能により、セマンティック セグメンテーションを含むビジョン タスクで有望なパフォーマンスを発揮することが最近証明されました。
それにもかかわらず、注意の二次的な複雑さにより、既存の Transformer ベースのモデルは、何らかの方法で画像解像度を下げた後にのみ自己注意レイヤーを使用します。これにより、より高い解像度で存在するグローバル コンテキストをキャプチャする機能が制限されます。
したがって、この作業では、Factorizer と呼ばれるモデルのファミリを導入します。これは、エンド ツー エンドのセグメンテーション モデルを構築するために低ランクの行列因数分解の力を活用します。
具体的には、コンテキスト モデリングへの線形にスケーラブルなアプローチを提案し、Nonnegative Matrix Factorization (NMF) を U 字型アーキテクチャに統合された微分可能なレイヤーとして定式化します。
シフト ウィンドウ手法は、ローカル情報を効果的に集約するために、NMF と組み合わせて使用​​することもできます。
ファクタライザーは、精度、スケーラビリティ、および解釈可能性の点で CNN およびトランスフォーマーと有利に競合し、脳腫瘍セグメンテーション用の BraTS データセットおよび脳卒中病変セグメンテーション用の ISLES’22 データセットで最先端の結果を達成します。
非常に意味のある NMF コンポーネントは、CNN やトランスフォーマーよりも、Factorizer に追加の解釈可能性を提供します。
さらに、私たちのアブレーション研究では、トレーニング済みの Factorizer の推論を大幅に高速化することを可能にする、Factorizer の特徴的な機能が明らかになりました。
コードとモデルは、https://github.com/pashtari/factorizer で公開されています。

要約(オリジナル)

Convolutional Neural Networks (CNNs) with U-shaped architectures have dominated medical image segmentation, which is crucial for various clinical purposes. However, the inherent locality of convolution makes CNNs fail to fully exploit global context, essential for better recognition of some structures, e.g., brain lesions. Transformers have recently proven promising performance on vision tasks, including semantic segmentation, mainly due to their capability of modeling long-range dependencies. Nevertheless, the quadratic complexity of attention makes existing Transformer-based models use self-attention layers only after somehow reducing the image resolution, which limits the ability to capture global contexts present at higher resolutions. Therefore, this work introduces a family of models, dubbed Factorizer, which leverages the power of low-rank matrix factorization for constructing an end-to-end segmentation model. Specifically, we propose a linearly scalable approach to context modeling, formulating Nonnegative Matrix Factorization (NMF) as a differentiable layer integrated into a U-shaped architecture. The shifted window technique is also utilized in combination with NMF to effectively aggregate local information. Factorizers compete favorably with CNNs and Transformers in terms of accuracy, scalability, and interpretability, achieving state-of-the-art results on the BraTS dataset for brain tumor segmentation and ISLES’22 dataset for stroke lesion segmentation. Highly meaningful NMF components give an additional interpretability advantage to Factorizers over CNNs and Transformers. Moreover, our ablation studies reveal a distinctive feature of Factorizers that enables a significant speed-up in inference for a trained Factorizer without any extra steps and without sacrificing much accuracy. The code and models are publicly available at https://github.com/pashtari/factorizer.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google