# You Only Train Once: Differentiable Subset Selection for Omics Data

arXiv: https://arxiv.org/abs/2512.17678 · HTML: https://mlsanity.com/paper/2512.17678/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 10 of its 39 references that are in the map.
- First posted on arXiv 19 Dec 2025.

## Builds on (most cited first, 10 of 10)

- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Model-Agnostic Interpretability of Machine Learning](https://mlsanity.com/paper/1606.05386/) (2016)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [In Defense of the Unitary Scalarization for Deep Multi-Task Learning](https://mlsanity.com/paper/2201.04122/) (2022)
- [Relief-Based Feature Selection: Introduction and Review](https://mlsanity.com/paper/1711.08421/) (2017)
- [Learning Group Importance using the Differentiable Hypergeometric Distribution](https://mlsanity.com/paper/2203.01629/) (2022)
- [Differentiable Random Partition Models](https://mlsanity.com/paper/2305.16841/) (2023)
- [N-ACT: An Interpretable Deep Learning Model for Automatic Cell Type and Salient Gene Identification](https://mlsanity.com/paper/2206.04047/) (2022)

## Abstract (arXiv)

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.