You Only Train Once: Differentiable Subset Selection for Omics Data

AI summary

Overview

Research area: Machine learning for computational biology, specifically end-to-end feature selection for single-cell and spatial transcriptomics data.

Technical level: Intermediate. Readers should be comfortable with basic neural network training, gradient-based optimization, and standard classification metrics such as F1-score and AUROC, though the paper explains its core mechanisms clearly.

Scope: The paper introduces YOTO, a single differentiable architecture that jointly learns discrete gene subsets and prediction tasks, achieving state-of-the-art performance on two single-cell RNA-seq datasets while eliminating the need for separate downstream classifiers.

What This Paper Is About

Selecting a small, informative set of genes from high-dimensional single-cell transcriptomic data is critical for biomarker discovery, interpretability, and cost-effective profiling, yet most existing methods either select genes in a separate stage before training a classifier or apply selection only after training via post-hoc attribution. YOTO addresses this by coupling gene selection and prediction inside a single differentiable model, so that the prediction loss directly shapes which genes are chosen, and the chosen genes in turn shape the predictions. The goal is to produce compact, meaningful gene panels that generalize across biological contexts and multiple prediction tasks without retraining.

Key Contributions

1. A fully differentiable end-to-end framework that learns discrete gene subsets and predictions jointly, using a differentiable ranking operator (Gumbel-Softmax plus the Plackett-Luce model) with a straight-through binary mask so that only selected genes contribute during inference.

2. A multi-task design with a shared encoder and task-specific heads, letting one training run serve several prediction objectives (e.g., cell type, disease status, patient identity) and discover gene subsets that generalize across tasks without retraining.

3. Empirical validation on two biologically distinct single-cell RNA-seq datasets (COVID-PBMC and mouse VISp), where YOTO consistently matches or exceeds state-of-the-art baselines including mRMR, block HSIC Lasso, Seurat v3, and PERSIST.

4. An ablation showing that YOTO's gains come from its sparse, differentiable selection mechanism rather than from model size, isolating the contribution of the selection module from the shared multi-task architecture.

Main Findings

• Consistent performance across panel sizes: On the VISp dataset, YOTO achieves state-of-the-art F1-scores for all subset sizes at or above 32 genes, with the advantage growing as the panel size increases. At the extreme of 16 genes, performance becomes sensitive to individual gene choices.

• Strong results on COVID-PBMC: YOTO ranks second at 16 and 32 genes (where the shallow, relevance-based mRMR-f excels because the coarse disease labels are driven by a few genes), and becomes the best method at 64 genes.

• Multi-task efficiency: A single YOTO training run matches or outperforms baselines that require 3 to 6 separate models (Seurat, Seurat v3, mRMR-f, mRMR-rf, block HSIC Lasso, and PERSIST) on both datasets. The largest gains appear on the harder tasks such as patient identity and fine-grained cell types.

• Cross-task gene panel quality: When the gene subsets selected by multi-task YOTO are given to a separate Random Forest classifier, those classifiers beat all baselines across subtasks and occasionally exceed the end-to-end YOTO on COVID-PBMC, confirming that the selected panels carry genuine cross-task signal.

• Robustness under class imbalance: YOTO leads on F1-score and AUPRC for the imbalanced celltype5 task, matching the best baseline on accuracy and trailing slightly on AUROC, which reflects its focus on the decision region that matters for rare cell populations.

• Ablation isolates the selection module: The advantage over the strongest deep learning baseline, PERSIST, is attributable to YOTO's sparse, fully differentiable discrete selection rather than to parameter count, since PERSIST's soft mask lets non-selected genes still influence predictions.

• Limitation under extreme sparsity: When panels are very small and tasks are jointly optimized, patient-informative genes can crowd out optimal cell-type markers, an effect that fades as the subset size grows.

Methodology in Plain English

The researchers start by assigning each gene a learnable importance score. During training, these scores are used to compute a differentiable ranking of genes using a technique borrowed from probabilistic ranking models, so gradients can flow back to the scores even though selection is inherently a discrete top-k operation. To make inference truly sparse, they use a straight-through trick: in the forward pass a hard binary mask keeps only the top-k genes, while in the backward pass a softened version of the selection is used so the scores remain trainable. Over training, the mask becomes progressively sparser until only the final k genes remain.

The selected genes are passed through a single shared encoder that produces a latent representation, which is then routed to task-specific prediction heads. The total loss is the average of per-task losses (cross-entropy for classification, mean squared error for regression), so the prediction objectives directly drive which genes get selected. Because everything is differentiable and the encoder is shared, one training run handles all tasks simultaneously, and evaluation at the end reflects only the selected genes with no additional classifier training needed.

Why This Matters

Impact on research: YOTO removes a long-standing two-stage bottleneck in omics feature selection, where genes are chosen independently of the models they support. It gives the community a principled way to obtain discrete, interpretable gene panels whose predictive value is measured inside the same optimization that produced them, and it shows that sparse, end-to-end multi-task selection can outperform both classical statistical methods and the current deep learning state of the art.

Real-world applications:

• Targeted gene panel design for clinical diagnostics, where a limited number of genes can be measured per assay and each gene has a cost.
• Spatial transcriptomics and imaging-based profiling, where the number of simultaneously measurable genes is physically constrained.
• Biomarker discovery in infectious disease and immunology, demonstrated on COVID-19 patient PBMCs distinguishing disease severity, cell types, and patient variation.
• Cost-effective single-cell studies, by enabling researchers to profile far fewer genes while retaining predictive power across cell types and conditions.

Industry relevance: The approach is directly applicable to pharmaceutical and diagnostics companies building targeted assays, to biotech firms developing spatial or low-plex profiling platforms where panel size is a limiting economic factor, and to any organization that needs interpretable, compact genomic signatures rather than black-box predictions. The fact that a single model covers multiple tasks also reduces the compute and retraining overhead of deploying such systems at scale.

Future Directions

• Patient-disjoint evaluation: The current cell-level train/test split permits patient-specific signals to leak into predictions. Future work should assess generalization to unseen patients, which would strengthen the biomarker-discovery claims.

• Task-conflict handling: When jointly optimized tasks share little biology or when subset sizes are extremely small, tasks can compete for gene slots. Methods that adaptively balance or decouple task-specific selection could address this trade-off.

• Selecting k automatically: The subset size k is currently a hyperparameter and is annealed on a fixed schedule. Learning it directly from the data would remove a manual tuning step and adapt sparsity to the problem.

• Broader omics coverage: The framework extends naturally beyond transcriptomics to proteomics, metabolomics, and multi-modal assays, where the same joint selection-and-prediction problem appears.

• Sparse multi-task generalization: Exploring how the shared gene subset behaves when tasks are only partially related, and whether task-specific selection heads or gating mechanisms would help, is an open design question.

Target Audience

This paper is most valuable to machine learning researchers working on differentiable feature selection and multi-task learning, computational biologists and bioinformaticians analyzing single-cell and spatial transcriptomics data, and practitioners in diagnostics or biotech who need to design compact gene panels for targeted assays. It is also relevant to readers interested in end-to-end architectures that couple selection with prediction, and to methodologists seeking alternatives to post-hoc attribution or multi-stage selection pipelines.

Authors’ abstract

Selecting compact and informative gene subsets from single-cell transcriptomic data is essential for biomarker discovery, improving interpretability, and cost-effective profiling. However, most existing feature selection approaches either operate as multi-stage pipelines or rely on post hoc feature attribution, making selection and prediction weakly coupled. In this work, we present YOTO (you only train once), an end-to-end framework that jointly identifies discrete gene subsets and performs prediction within a single differentiable architecture. In our model, the prediction task directly guides which genes are selected, while the learned subsets, in turn, shape the predictive representation. This closed feedback loop enables the model to iteratively refine both what it selects and how it predicts during training. Unlike existing approaches, YOTO enforces sparsity so that only the selected genes contribute to inference, eliminating the need to train additional downstream classifiers. Through a multi-task learning design, the model learns shared representations across related objectives, allowing partially labeled datasets to inform one another, and discovering gene subsets that generalize across tasks without additional training steps. We evaluate YOTO on two representative single-cell RNA-seq datasets, showing that it consistently outperforms state-of-the-art baselines. These results demonstrate that sparse, end-to-end, multi-task gene subset selection improves predictive performance and yields compact and meaningful gene subsets, advancing biomarker discovery and single-cell analysis.