転移学習手順を改良して医療用モデルを少量データでも開発しやすくする(2/2)

１．転移学習手順を改良して医療用モデルを少量データでも開発しやすくする(2/2)まとめ

・事前学習を３ステップにする事で最大で600倍少ないデータで高品質なモデルを学習可能
・この手法はモデルのアーキテクチャやデータセットに依存しないことが確認されている
・45枚の画像embeddingsデータで学習したモデルは放射線科医に非劣性を達成した

２．3ステップの転移学習

以下、ai.googleblog.comより「Simplified Transfer Learning for Chest Radiography Model Development」の意訳です。元記事は2022年7月19日、Akib UddinさんとAndrew Sellergrenさんによる投稿です。

アイキャッチ画像はlatent diffusionでプロンプトはChest X-Ray Network。

CXRの事前学習による効果

このembeddings層を、肺の空隙の不透明度を例にして、各ステップで可視化しました(下図参照)。SupConによる事前学習前では、CXRの正常なembeddingsと異常なembeddingsがあまりうまく分離されていない事がわかります。また、SupConによる事前学習後は、ポジティブな例とネガティブな例がより密接にグループ化され、それぞれのカテゴリーの画像が似ていることをモデルが識別したことを示しています。

一般的なネットワークのembeddingsとCXRに特化したネットワークのembeddingsをt-SNEを使って視覚化。embeddingsは情報量の多い数値ベクトルであり、それだけでクラスを区別することができます。(この場合は空隙の不透明度がポジティブかネガティブか)

本研究では、2段階目の事前学習を追加することで、一般的な非医療データセットで事前学習したモデルを活用する従来の転移学習アプローチと比較して、最大で600倍少ないデータで高品質なモデルを学習できることを示唆しています。

これは、モデルのアーキテクチャ(ResNetやEfficientNetなど)や、自然画像の事前学習に用いたデータセット(ImageNetやJFT-300Mなど)に依存しないことが確認されました。このアプローチにより、研究者や開発者は、必要なデータセットサイズを大幅に削減することができます。

上：ディープラーニングモデルでは、ニューラルネットワークには複数の人工ニューロンの層があり、最初の層がCXR画像を入力とし、中間層が追加の計算を行い、最後の層が分類(例えば、空隙の不透明度：あり vs なし)を行います。embeddings層は通常、最後の層の1つです。
左下：従来の転移学習では、2段階の学習が行われ、事前に学習した一般的なネットワークが、関心のある予測タスクに対して直接最適化されます。私達の提案する3段階学習は、SupCon ML技術を用いたCXRネットワークの生成(ステップ2)の後、目的の予測タスクに最適化(ステップ3)します。
右下：embeddingsを利用する場合、より小さなモデルを学習するか(最初の2つのストラテジー)、十分なデータがある場合はネットワーク全体を微調整します。(3つめのストラテジー)

結果

初期モデルの学習後、CXR embeddingsに適用した線形モデルと非線形モデル、およびネットワーク全体を微調整して作成した非線形モデルについて、曲線下面積(AUC)指標を使用して性能を測定しました。

ChestX-ray14 や CheXpert などの公開データセットでは、様々な学習データセットサイズといくつかの知見に基づいて開発されたモデルについて、データと精度のトレードオフを実質的かつ一貫して改善することができました。例えば、結核モデルの開発能力を評価する場合、データ効率の向上はより顕著で、わずか45枚の画像のembeddingsデータで学習したモデルは、外部検証データセットにおいて結核の検出で放射線科医に非劣性を達成しました。結核と重症COVID-19の両方において、凍結したembeddingsで学習した非線形分類器は、データセット全体で微調整したモデルよりも性能が優れていることを示します。

CXR特化ネットワークの転移学習(赤)とベースライン転移学習アプローチ(青)を様々なCXR異常(左上)、結核(左下)、COVID-19状態(右下)で比較

このアプローチにより、同じデータセットサイズで性能が向上するか、同じ性能に到達するために必要なデータセットサイズが小さくなります。興味深いことに、\(8^5\)画像までのデータセットサイズでは、CXRネットワークと、より高速に学習できる単純なMLモデル(赤)を使用した方が、フルネットワーク(黒)の学習よりも性能が良いです。

結論と今後の課題

少ないデータ量と計算量でCXRモデリング作業を加速するために、私達はCXR Foundationツールと線形および非線形分類器を学習するスクリプトをリリースすることにしました。

これらのembeddingsにより、研究者はより単純な転移学習法を用いたCXRモデリング作業を迅速に開始することができます。この手法は、小規模なデータセットを使用した予測モデリングや、患者集団の分布が変化した場合(時間の経過や異なる施設間など)にCXRモデルを適応させる場合に特に有効です。

私たちは、Northwestern MedicineやApollo Hospitalsなどのパートナーと協力し、この技術の効果をさらに追求していきたいと考えています。限られたデータと計算機でCXRモデルを開発できるようにすることで、より多くの開発者が集団にとって最もインパクトのある問題を解決できるようになることを期待しています。

謝辞

このプロジェクトに貢献した Google の主要メンバーには、Christina Chen, Yun Liu, Dilip Krishnan, Zaid Nabulsi, Atilla Kiraly, Arnav Agharwal, Eric Wu, Yuanzhen Li, Aaron Maschinot, Aaron Sarna, Jenny Huang, Marilyn Zhang, Charles Lau, Neeral Beladia, Daniel Tse, Krish Eswaran および Shravya Shetty が含まれます。また、Sreenivasa Raju Kalidindi, Mozziyar Etemadi, Florencia Garcia-Vicente, そして David Melnickも多大な貢献と意見を寄せてくれました。

ChestX-ray14 のデータセットについては，NIH Clinical Center の公開に感謝します。また、著者らはGoogle Health Radiology とラベリングソフトウェアチームの多くのメンバーに感謝します。また、研究期間中、画像の解釈と注釈付けを行い、この研究を可能にしてくれた放射線科医、画像の注釈付け作業を調整してくれたJonny Wong、原稿にフィードバックをくれたCraig MermelとAkinori Mitani、ブログ記事にフィードバックをくれたNicole LintonとLauren Winer、アニメーション作成のTom Smallに心からの感謝を捧げます。