[Paper Review] Select-ProtoNet: Learning to Select for Few-Shot Disease Subtype Prediction

Ziyi Yang, Jun Shu+3 more|arXiv (Cornell University)|Sep 2, 2020

Gene expression and cancer classification48 references4 citations

TL;DR

This paper proposes Select-ProtoNet, a meta-learning framework that enhances Prototypical Networks with feature selection and sample reweighting to improve few-shot disease subtype prediction on noisy, high-dimensional gene expression data. By filtering irrelevant genes and suppressing noisy samples, the model achieves state-of-the-art accuracy, outperforming its backbone ProtoNet by over 20% on lung cancer subtype prediction and demonstrating robustness under high noise levels.

ABSTRACT

Current machine learning has made great progress on computer vision and many other fields attributed to the large amount of high-quality training samples, while it does not work very well on genomic data analysis, since they are notoriously known as small data. In our work, we focus on few-shot disease subtype prediction problem, identifying subgroups of similar patients that can guide treatment decisions for a specific individual through training on small data. In fact, doctors and clinicians always address this problem by studying several interrelated clinical variables simultaneously. We attempt to simulate such clinical perspective, and introduce meta learning techniques to develop a new model, which can extract the common experience or knowledge from interrelated clinical tasks and transfer it to help address new tasks. Our new model is built upon a carefully designed meta-learner, called Prototypical Network, that is a simple yet effective meta learning machine for few-shot image classification. Observing that gene expression data have specifically high dimensionality and high noise properties compared with image data, we proposed a new extension of it by appending two modules to address these issues. Concretely, we append a feature selection layer to automatically filter out the disease-irrelated genes and incorporate a sample reweighting strategy to adaptively remove noisy data, and meanwhile the extended model is capable of learning from a limited number of training examples and generalize well. Simulations and real gene expression data experiments substantiate the superiority of the proposed method for predicting the subtypes of disease and identifying potential disease-related genes.

Motivation & Objective

• To address the challenge of few-shot disease subtype prediction in genomics, where training data are limited and noisy.
• To simulate clinical reasoning by modeling interrelated clinical variables simultaneously, rather than treating each disease subtype in isolation.
• To enhance the Prototypical Network for genomic data by addressing high dimensionality and noise inherent in gene expression profiles.
• To develop a data-efficient model that generalizes well from few examples by leveraging shared knowledge across related clinical tasks.
• To identify biologically relevant genes through interpretability, validated via survival analysis on TCGA data.

Proposed method

• Extends Prototypical Networks with a learnable feature selection layer that automatically identifies disease-relevant genes from high-dimensional expression data.
• Introduces a sample reweighting strategy that adaptively downweights noisy or outlier samples during support set embedding.
• Uses a weighted prototype computation where each class's prototype is derived from a weighted average of support samples, with weights learned via a neural network.
• Employs a meta-learning paradigm to train on multiple related clinical tasks (e.g., disease subtypes from TCGA), enabling knowledge transfer to new few-shot tasks.
• Trains the model end-to-end using episodic training, where each episode simulates a few-shot classification task with support and query sets.
• Incorporates a multi-task learning setup using a curated subset of TCGA tasks, forming the mini-TCGA Meta-Dataset for benchmarking.

Experimental results

Research questions

• RQ1 Can meta-learning techniques improve few-shot disease subtype prediction in genomic data with limited and noisy samples?
• RQ2 How does incorporating feature selection and sample reweighting enhance model robustness to high-dimensional and noisy gene expression data?
• RQ3 To what extent can shared knowledge from interrelated clinical tasks improve generalization in few-shot genomic classification?
• RQ4 Can the model identify biologically meaningful genes that correlate with patient survival outcomes?
• RQ5 How does Select-ProtoNet compare to conventional supervised methods and its backbone ProtoNet in accuracy and robustness under varying noise levels?

Key findings

• Select-ProtoNet achieves 97.78% accuracy on the lung cancer subtype prediction task from the TCGA Meta-Dataset, outperforming its backbone ProtoNet by over 20 percentage points.
• On the mini-TCGA Meta-Dataset, Select-ProtoNet maintains high performance under increasing noise, achieving 78.15% accuracy at 50% noise, compared to 57.06% for Logistic Regression.
• The model improves upon the baseline ProtoNet by 12.7 percentage points in accuracy on the lung cancer task, demonstrating strong generalization from meta-learned knowledge.
• The 20 top-ranked genes selected by Select-ProtoNet show significant survival differences (logrank p-value = 5.801e-4), with patients lacking alterations in these genes exhibiting longer survival.
• The feature selection layer effectively filters out irrelevant genes, and the sample reweighting mechanism reduces the impact of noisy samples, improving model stability.
• Survival analysis confirms that genes identified by Select-ProtoNet are biologically relevant and clinically meaningful in the context of lung cancer.