[論文レビュー] AffinityNet: semi-supervised few-shot learning for disease type prediction

Tianle Ma, Aidong Zhang|arXiv (Cornell University)|May 22, 2018

Machine Learning in Bioinformatics参考文献 24被引用数 8

ひとこと要約

AffinityNet は、高次元でサンプル数が少ないゲノムデータから強固な表現を学ぶために、スタックされた k-最近傍 (kNN) アテンションプーリング層を用いた半教師あり少数ショット学習モデルであり、腫瘍種予測に特化している。40% の訓練データでの腫瘍マッピングで 0.729 の concordance index を達成し、従来のディープラーニングモデルやベースライン Cox モデルを上回り、低データ環境下でも優れた汎化性能を示している。

ABSTRACT

While deep learning has achieved great success in computer vision and many other fields, currently it does not work very well on patient genomic data with the "big p, small N" problem (i.e., a relatively small number of samples with high-dimensional features). In order to make deep learning work with a small amount of training data, we have to design new models that facilitate few-shot learning. Here we present the Affinity Network Model (AffinityNet), a data efficient deep learning model that can learn from a limited number of training examples and generalize well. The backbone of the AffinityNet model consists of stacked k-Nearest-Neighbor (kNN) attention pooling layers. The kNN attention pooling layer is a generalization of the Graph Attention Model (GAM), and can be applied to not only graphs but also any set of objects regardless of whether a graph is given or not. As a new deep learning module, kNN attention pooling layers can be plugged into any neural network model just like convolutional layers. As a simple special case of kNN attention pooling layer, feature attention layer can directly select important features that are useful for classification tasks. Experiments on both synthetic data and cancer genomic data from TCGA projects show that our AffinityNet model has better generalization power than conventional neural network models with little training data. The code is freely available at https://github.com/BeautyOfWeb/AffinityNet .

研究の動機と目的

• 高次元特徴と少ないサンプル数がディープラーニングの性能を制限する『p が大きく N が小さい』ゲノムデータ問題に対処すること。
• 少数ショット学習のシナリオが一般的な生物医学研究において、良好に汎化するデータ効率の良いディープラーニングモデルを開発すること。
• 事前に定義されたグラフ構造を必要とせずに、患者のような複雑で多様性のあるオブジェクトに対して半教師あり表現学習を可能にすること。
• 入力特徴に重みを割り当てる学習可能な特徴アテンション層を介して、ディープラーニングに自動特徴選択を統合すること。
• 豊富なラベルなしデータと限られたラベル付きデータを活用して、がんゲノム研究における生存予測と疾患サブタイピングを改善すること。

提案手法

• コアとなる構成要素は kNN アテンションプーリング層であり、アテンションメカニズムを用いて動的にアフィニティグラフを構築し、各オブジェクトの k 個の最近傍の表現をプールすることで、クラスタリングに類似した正則化を実現する。
• kNN アテンションプーリング層は、グラフ構造が存在しないデータに対しても一般化可能であり、グラフアテンションモデル (GAM) を非グラフ構造データに拡張することで、任意のオブジェクト集合への適用を可能にする。
• 複数の kNN アテンションプーリング層をスタックすることで、深層階層的表現を学習し、畳み込みニューラルネットワーク (CNN) の畳み込み層に類似たプラグインモジュールとして機能する。
• 特徴アテンション層は、kNN アテンションプーリングの特殊ケースであり、バックプロパゲーションにより正規化され非負の重みを入力特徴に学習させることで、エンドツーエンドの特徴選択を可能にする。
• モデルは半教師ありの方法で訓練され、ラベル付きデータとラベルなしデータの両方を活用して表現学習と汎化性能の向上を図る。
• フレームワークは PyTorch で実装されており、再現性と拡張性を考慮して GitHub で公開されている。

実験結果

リサーチクエスチョン

• RQ1 高次元でサンプル数が少ないゲノムデータにおいて、少数のラベル付きサンプルから効果的に学習できるディープラーニングモデルは存在するか？
• RQ2 事前に定義されたグラフ構造が存在しない状況で、アテンションメカニズムをどのように変更して表現を学習できるか？
• RQ3 kNN を用いたアテンションプーリングは、標準的なニューラルネットワークと比較して、疾患サブタイピングの少数ショット学習における汎化性能を向上させられるか？
• RQ4 モデルは、生存予測タスクの性能向上に、どれほどラベルなしデータを活用できるか？
• RQ5 学習可能な特徴アテンション層の統合は、ゲノム分類におけるモデル性能と解釈可能性を向上させられるか？

主な発見

• AffinityNet は、腫瘍マッピングの 40% のサンプルで訓練された場合、腎がんデータで平均 concordance index 0.729 を達成し、ベースライン Cox モデル (0.618) を顕著に上回った。
• 訓練データをたった 10% にまで減らしても、AffinityNet は concordance index 0.694 を達成し、低データ環境下での強力な汎化性能を示した。
• モデルが予測したハザードグループ（低、中、高）は、ログランク検定の p 値 6.7×10⁻¹⁶ を示し、生存結果の分離において、既知の疾患タイプ（p=1.5×10⁻⁵）を上回った。
• 合成データおよび TCGA がんゲノムデータセットにおいて、限られた訓練データで従来のニューラルネットワークモデルよりも AffinityNet が優れた汎化性能を示した。
• 特徴アテンション層は、生物学的に関連する特徴を効果的に特定し強調し、トレーニング中に自動的で微分可能な特徴選択を実現した。
• kNN アテンションプーリング層は、類似したオブジェクトに類似した表現を強制的に学習させ、小規模データ環境下での過学習を緩和する暗黙の正則化として機能した。