[論文レビュー] MT-TransUNet: Mediating Multi-Task Tokens in Transformers for Skin Lesion Segmentation and Classification

Jingye Chen, Jieneng Chen+4 more|arXiv (Cornell University)|Dec 3, 2021

Cutaneous Melanoma Detection and Management被引用数 15

ひとこと要約

MT-TransUNet は、タスク固有のトークンを介して統合し、アテンションマップと予測の整合性を強制することで、皮膚腫瘍のセグメンテーションと分類を統合的に実行するトランスフォーマーベースのモデルを提案する。4800万パラメータと0.17秒の推論時間で、ISIC-2017およびPH2データセットにおいて最先端の性能を達成し、精度と効率性の両面で先行手法を上回り、ピクセル単位のアノテーションと画像単位のアノテーションの両方を活用する。

ABSTRACT

Recent advances in automated skin cancer diagnosis have yielded performance on par with board-certified dermatologists. However, these approaches formulated skin cancer diagnosis as a simple classification task, dismissing the potential benefit from lesion segmentation. We argue that an accurate lesion segmentation can supplement the classification task with additive lesion information, such as asymmetry, border, intensity, and physical size; in turn, a faithful lesion classification can support the segmentation task with discriminant lesion features. To this end, this paper proposes a new multi-task framework, named MT-TransUNet, which is capable of segmenting and classifying skin lesions collaboratively by mediating multi-task tokens in Transformers. Furthermore, we have introduced dual-task and attended region consistency losses to take advantage of those images without pixel-level annotation, ensuring the model's robustness when it encounters the same image with an account of augmentation. Our MT-TransUNet exceeds the previous state of the art for lesion segmentation and classification tasks in ISIC-2017 and PH2; more importantly, it preserves compelling computational efficiency regarding model parameters (48M~vs.~130M) and inference speed (0.17s~vs.~2.02s per image). Code will be available at https://github.com/JingyeChen/MT-TransUNet.

研究の動機と目的

• 皮膚腫瘍分類を独立したタスクとして扱う従来手法の限界を是正し、腫瘍セグメンテーションの利点を無視することを防ぐ。
• セグメンテーションと分類の間で共有される特徴を活用する共同学習フレームワークを構築する。
• ピクセル単位のアノテーションが存在しない画像に対してもモデルの頑健性を向上させるために、二重タスクおよび注目領域の整合性損失を導入する。
• 従来のマルチタスクモデルと比較して、高い計算効率を維持しながら最先端の性能を達成する。
• 両タスクに共通する1つのエンコーダーを共有することで、データ効率的かつモデル効率的な学習と推論を実現する。

提案手法

• トランスフォーマーエンコーダ・デコーダアーキテクチャ内にタスク固有の分類用およびセグメンテーション用トークンを導入し、タスク間の特徴の調整を可能にする。
• 分類用アテンションマップとセグメンテーションマップの整合性を図るために、二重タスク整合性損失を採用し、弱教師ありデータにおける一般化性能を向上させる。
• 分類ヘッドとセグメンテーションヘッドの間の整合性を強制するために、注目領域の整合性損失を適用し、特にピクセル単位のアノテーションが欠落している場合に顕著な利点を発揮する。
• 両タスクに共通するエンコーダーを採用することで、モデルパラメータ数を削減し、推論速度を向上させる。
• 半教師あり設定において、セグメンテーションの監督信号の代わりにレベルセット関数を補助的に用いることで、真値マスクのない画像に対しても対応可能にする。
• セグメンテーションと分類が共有表現と整合性制約を通じてお互いに補完し合うように、マルチタスク学習パイプラインを設計する。

実験結果

リサーチクエスチョン

• RQ1 皮膚腫瘍セグメンテーションと分類の共同学習は、単一タスクアプローチと比較して、両タスクの性能を向上させることができるか？
• RQ2 分類とセグメンテーションの予測の整合性をどのように強制することで、特に画像単位のラベルしか利用できない状況でも頑健性を向上させられるか？
• RQ3 共通のトランスフォーマーベースアーキテクチャは、従来のマルチタスクモデルと比較して、顕著に小型化され、高速な推論を達成できるか？
• RQ4 提案手法は、微調整なしで異なるデータセットにどの程度一般化可能か？
• RQ5 二重タスク損失および注目領域整合性損失の使用は、弱教師ありデータにおける性能にどのような影響を与えるか？

主な発見

• MT-TransUNet は、ISIC-2017およびPH2データセットにおいて最先端の性能を達成し、セグメンテーションのジャッカードスコアが88.5/92.2、デイクスコアが93.6/95.9、分類精度が97.1%を記録した。
• 1枚あたりの推論時間を0.17秒に短縮し、MB-DCNNの2.02秒と比較して91.6%の高速化を達成した。
• 4800万パラメータというわずかなサイズで、MB-DCNNの1億3000万パラメータと比較して63%のモデルサイズ削減を実現した。
• ISIC-2017で事前学習した後、PH2でテストした際、認識精度が3.1%、ジャッカードスコアが2.8%向上し、一般化性能が向上した。
• 二重タスク損失および注目領域整合性損失により、データオーグメンテーションに対する耐性が向上し、ピクセル単位のアノテーションが欠落している画像でも性能が向上した。
• セグメンテーションと分類の共同学習により、相互に性能向上が見られ、皮膚画像診断における共有表現学習の利点を実証した。