You Only Train Once: Differentiable Subset Selection for Omics Data

Daphné Chopard, Jorge da Silva Gonçalves, Irene Cannistraci, Thomas M. Sutter, Julia E Vogt

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. 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 infer-
ence, eliminating the need to train additional downstream classifiers. Through a multi-task learning design, the model learns shared representations across related objectives, allowing different tasks 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.

Submission Type: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

- added a paragraph in the discussion of 4.4.2 about the cell- vs patient-level split
- added back the link to the github repo (and changed it to the public one)
- moved 4.4.5 (Isolating the Multi-Task Setting), 4.4.6 (Robustness of Selected Genes Across Seeds) and 4.4.7 (Robustness of Selected Genes Across Gene Subset Sizes) from the appendix to the main text
- removed the titles from the plots

Code: https://github.com/chopardda/yoto

Assigned Action Editor: ~Pan_Xu1

Submission Number: 6804