[論文レビュー] How Can We Tame the Long-Tail of Chest X-ray Datasets?

Arsh Verma|arXiv (Cornell University)|Sep 8, 2023

COVID-19 diagnosis using AIMedicine被引用数 3

ひとこと要約

本論文では、ターゲットデータセットで訓練する前に、胸部X線（CXR）データでビジョンモデルをファインチューニングすることで、複雑な損失関数やハイパーパramータチューニングを必要とせずに、希少で長尾に分布するラベルの性能が著しく向上することを提案している。ImageNetに続く2段階の事前学習戦略（ImageNet → CXR特化データセット）を用いることで、超希少ラベルで最大3ポイントの向上が達成され、モデルアンサンブルによってmAPが6%絶対的に向上し、医療画像における長尾一般化において初期化の質が重要な要因であることが示された。

ABSTRACT

Chest X-rays (CXRs) are a medical imaging modality that is used to infer a large number of abnormalities. While it is hard to define an exhaustive list of these abnormalities, which may co-occur on a chest X-ray, few of them are quite commonly observed and are abundantly represented in CXR datasets used to train deep learning models for automated inference. However, it is challenging for current models to learn independent discriminatory features for labels that are rare but may be of high significance. Prior works focus on the combination of multi-label and long tail problems by introducing novel loss functions or some mechanism of re-sampling or re-weighting the data. Instead, we propose that it is possible to achieve significant performance gains merely by choosing an initialization for a model that is closer to the domain of the target dataset. This method can complement the techniques proposed in existing literature, and can easily be scaled to new labels. Finally, we also examine the veracity of synthetically generated data to augment the tail labels and analyse its contribution to improving model performance.

研究の動機と目的

• 胸部X線（CXR）分類タスクにおける希少で長尾に分布するラベルでの性能が低いという課題に対処すること。
• 損失関数やデータサンプリング戦略を変更せずに、ドメイン特化された事前学習によるモデル初期化が、低頻度ラベルへの一般化を改善できるかどうかを調査すること。
• モデルアンサンブルと合成データ増強が、テイルラベルでの性能向上にどの程度寄与するかを評価すること。
• 最小限の再トレーニングで、新しい希少ラベルに対してモデル性能を向上させる、スケーラブルで即時適用可能な手法を提供すること。

提案手法

• 本手法は2段階の事前学習戦略を採用する：まずImageNetの重みで初期化し、その後3つのオープンソースCXRデータセット（NIH CXR、PadChest、CheXpert）でファインチューニングすることでドメイン適応型初期化を実現する。
• その後、26のラベルが長尾分布を示すMIMIC CXRデータセット上で、モデルをエンドツーエンドでファインチューニングする。このデータセットには、事前学習データセットに存在しない5つのユニークなテイルラベルが含まれる。
• モデルアンサンブルは、異なるバックボーン（DenseNet161とResNeXt101）を用いたモデルの予測をアンサンブルすることで、リソースが限られたラベルでのロバスト性を向上させる。
• 合成データ増強は、MIMIC CXRとRoentGenを用いて生成された小さな合成データセット（5,000サンプル）を同時にファインチューニングすることで検討され、テイルラベルの性能向上を図る。
• すべてのモデルは、シグモイド活性化関数を用いたバイナリクロスエントロピー損失を使用し、448×448解像度で20エポック分トレーニングすることで、微細な解剖的特徴を保持する。
• トレーニングパイプラインは患者単位で実行され、1人の患者に対して利用可能なすべての画像ビューが使用され、データ分割も患者単位で作成され、データ漏洩を回避する。

実験結果

リサーチクエスチョン

• RQ1 多様なCXRデータセットで事前学習することで、損失関数やデータサンプリング戦略を変更せずに、希少で長尾に分布するラベルの性能が向上するか？
• RQ2 2段階の事前学習戦略（ImageNetに続くCXR特化データ）は、標準的なImageNet初期化に比べ、テイルラベルでの一般化性能を向上させるか？
• RQ3 モデルアンサンブルは、特にターゲットデータセット固有の低頻度ラベルにおいて、どの程度性能を向上させるか？
• RQ4 合成データ増強は、テイルラベルのデータ不足を緩和できるか？また、モデル性能に測定可能な向上が見られるか？

主な発見

• チェインド事前学習戦略により、DenseNet161では5つのユニークなテイルラベル（Tail-U）で3ポイントの向上が達成されたが、ResNeXt101ではわずか0.5ポイントにとどまり、超希少クラスでの顕著な向上が示された。
• モデルアンサンブルにより、Tail-Uラベルでの平均平均精度（mAP）が6%絶対的に向上（0.1806 → 0.2430）し、個々のモデルをすべて上回った。
• アンサンブルモデルはバリデーションセットの26ラベル中19ラベルで最高性能を示し、ヘッド、ミディアム、テイルの全カテゴリーレベルで優位性を示した。
• 合成データ増強は、規模が限定的（5,000サンプル）であったが、テイルラベルでの性能向上が測定可能であり、特にテイルカテゴリで顕著な向上が見られた。これは将来的なスケーリングの可能性を示唆している。
• 本手法は非常にスケーラブルである：新しいラベルを追加するにはわずか数エポックのファインチューニングで十分であり、臨床現場での実用的導入に適している。
• 結果から、モデル初期化が長尾一般化において極めて重要であり、かつ過小評価されがちな要因であることが示された。特に、希少な病理が臨床的に有意義な医療画像分野において顕著である。