[论文解读] 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 注意力池化层，学习深层层次化表征，作为类似卷积神经网络中卷积层的即插即用模块。
• 特征注意力层是 kNN 注意力池化的特例，通过反向传播学习输入特征的归一化、非负权重，实现端到端特征选择。
• 模型以半监督方式训练，利用标注数据和未标注数据共同提升表征学习与泛化能力。
• 该框架使用 PyTorch 实现，已在 GitHub 上公开，便于复现与扩展。

实验结果

研究问题

• RQ1 深度学习模型能否在高维、小样本量的基因组数据中，有效利用少量标注样本进行学习？
• RQ2 注意力机制如何在缺乏预定义图结构的情况下学习表征？
• RQ3 与标准神经网络相比，基于 kNN 的注意力池化是否能提升少样本学习在疾病亚型分类中的泛化能力？
• RQ4 模型在生存预测任务中，能在多大程度上利用未标注数据提升性能？
• RQ5 可学习的特征注意力层是否能提升基因组分类中模型的性能与可解释性？

主要发现

• 当在 40% 的样本上进行训练时，AffinityNet 在肾癌数据上实现了 0.729 的平均 concordance index，显著优于基线 Cox 模型（0.618）。
• 仅使用 10% 的训练数据时，AffinityNet 的 concordance index 达到 0.694，表明其在低数据场景下具备强大的泛化能力。
• 模型预测的危险分组（低、中、高）在 log-rank 检验中 p 值为 6.7×10⁻¹⁶，优于已知疾病类型的 p 值（1.5×10⁻⁵），在区分生存结局方面表现更优。
• 在合成数据集和 TCGA 癌症基因组数据集上，AffinityNet 在有限训练数据条件下展现出优于传统神经网络模型的泛化能力。
• 特征注意力层成功识别并强调了具有生物学意义的特征，实现在训练过程中自动、可微分的特征选择。
• kNN 注意力池化层有效强化了相似对象的表征一致性，作为隐式正则化手段，缓解了小样本场景下的过拟合问题。