[論文レビュー] A Comparative Study of CNN, ResNet, and Vision Transformers for Multi-Classification of Chest Diseases

Ananya Jain, Aviral Bhardwaj+2 more|arXiv (Cornell University)|May 31, 2024

AI in cancer detection被引用数 4

ひとこと要約

本研究では、NIH Chest X-rayデータセットを用いて14種類の胸部疾患の多クラス分類を目的としたCNN、ResNet、Vision Transformer (ViT) モデルの比較を行った。事前学習済みのViTモデルは、ImageNet-21kで微調整された場合に94％のテスト精度を達成し、大規模データで事前学習された注意メカニズムに基づく特徴抽出の優位性を示した。

ABSTRACT

Large language models, notably utilizing Transformer architectures, have emerged as powerful tools due to their scalability and ability to process large amounts of data. Dosovitskiy et al. expanded this architecture to introduce Vision Transformers (ViT), extending its applicability to image processing tasks. Motivated by this advancement, we fine-tuned two variants of ViT models, one pre-trained on ImageNet and another trained from scratch, using the NIH Chest X-ray dataset containing over 100,000 frontal-view X-ray images. Our study evaluates the performance of these models in the multi-label classification of 14 distinct diseases, while using Convolutional Neural Networks (CNNs) and ResNet architectures as baseline models for comparison. Through rigorous assessment based on accuracy metrics, we identify that the pre-trained ViT model surpasses CNNs and ResNet in this multilabel classification task, highlighting its potential for accurate diagnosis of various lung conditions from chest X-ray images.

研究の動機と目的

• X線画像からの14種類の胸部疾患における多ラベル分類において、CNN、ResNet、およびVision Transformer (ViT) アーキテクチャの性能を評価・比較すること。
• ImageNet-21kのような大規模データセットでの事前学習が、医療画像分野におけるViTの性能に与える影響を評価すること。
• 低リソースな医療画像認識の文脈において、ViTと従来のCNNベースのモデルとの間で一般化能力および特徴抽出の効率性に差があるかを調査すること。
• 臨床的意義と意思決定の透明性を考慮し、ViTにおける注意マップとモデルの解釈可能性を分析すること。
• 多クラス胸部X線画像分類におけるデータの不均衡やモデル最適化のギャップといった制限要因を特定すること。

提案手法

• ImageNetで事前学習されたモデルと、NIH Chest X-rayデータセット（前向きX線画像112,120枚）を用いてからくりで学習した2種類のViTバリアントを微調整した。
• 同じデータセットとトレーニングプロトコルを用いて、基準モデルとして標準的なCNNとResNet-50をトレーニングおよび評価した。
• ViTエンコーダーの最終マルチヘッドアテンション層の注意重みを抽出することで、事前学習済みInceptionResNetV2モデルを用いて転移学習を実施し、注意マップを生成した。
• 入力として224×224ピクセルのRGB画像を用い、データオーグメンテーションを適用し、多ラベル分類のためのクロスエントロピー損失とシグモイド活性化関数を用いてモデルを学習した。
• トレーニングおよび検証セットを用いて、精度、AUC-ROC、損失曲線を評価し、性能と一般化能力を比較した。
• ヘッドおよびパッチごとの注意重みを平均化し、リサイズして元の画像に重ね合わせることで、注意マップを可視化し、モデルの予測解釈を行った。

実験結果

リサーチクエスチョン

• RQ1 事前学習済みのVision Transformerは、胸部X線画像の多ラベル分類において、CNNやResNetを上回る性能を示すか？
• RQ2 ImageNet-21kのような大規模データセットでの事前学習は、医療画像分類におけるViTの性能にどのように影響するか？
• RQ3 注意メカニズムは、畳み込み特徴マップと比較して、胸部X線画像における臨床的に関連する領域を特定する役割を果たすか？
• RQ4 データの不均衡や希少疾患の例数が少ないことが、アーキテクチャ全体のモデル性能に与える影響は何か？
• RQ5 ViTモデルは、従来のCNNと比較して、注意可視化によって同等またはより優れた解釈性を達成できるか？

主な発見

• 事前学習済みのVision Transformer (ViT-v2) は、NIH Chest X-rayデータセットで微調整した場合に94％のテスト精度を達成し、CNNおよびResNetモデルを顕著に上回った。
• ImageNet-21kで事前学習されたViTモデルは、からくりで学習したViTやすべての基準モデルよりも優れた性能を示し、大規模データでの事前学習の重要性を裏付けた。
• AUC-ROC分析から、ViTモデルは「No Finding」クラスでは性能が低く（AUC = 0.50）、正常例と異常例の区別が困難であることが判明した一方、アテレクトーシス（AUC = 0.51） や胸膜厚化（AUC = 0.53） に対しては中程度の性能を示した。
• 心臓肥大（AUC = 0.48） やヘルニア（AUC = 0.43） といった希少疾患では、データセット内の不均衡が主な要因となり、ViTモデルが困難をきたした。
• ViTモデルから生成された注意マップは、心臓や肺領域といった臨床的に重要な領域を的確に強調しており、モデルが診断的に有意義な領域に注目していることを確認した。
• 平均的には優れた性能を示したが、「No Finding」クラスでは性能が低く、今後の研究ではより良いデータサンプリングやクラスバランス技術の導入が求められることが示唆された。