YOTO: Differentiable Subset Selection for Omics Data

This repository contains the implementation of YOTO (You Only Train Once), an end-to-end framework for joint gene subset selection and prediction in high-dimensional omics data.

Overview

Datasets

Currently supported datasets:

• VISp
  corresponding paper

• Covid-PBMC original dataset source corresponding paper

Adding a new dataset

To integrate a new dataset into the pipeline, follow these steps.

1. Define Dataset Configuration

in NDSConfig.py, add a dataclass with the template

@dataclass()
class NEWDataConfig(DataConfig):
	name: str = "new_dataset_name"
	
	dir_data: str = "<PATH/TO/DATADIR>"
	fn_data: str = "<PATH/TO/FILE.h5mu>"
	
	num_features_transcripts: int = <number_of_transcripts_in_dataset>
	num_features_proteins: int = <number_of_proteins_in_dataset>

2. Register Dataset

in main_subset_selection.py:

from config.NDSConfig import NEWDataConfig

Then register it with Hydra:

cs.store(group="data", name="new_dataset_name", node=NEWDataConfig)

3. Select Dataset

in config.yaml, update

defaults:
  - data: new_dataset_name

Running the Model

Example command:

python main_subset_selection.py

Make sure the dataset paths are correctly configured before running.

Citation

If you use our model in your work, please cite us using the following citation

@inproceedings{chopard2026yoto,
  title={You Only Train Once: Differentiable Subset Selection for Omics Data},
  author={Chopard, Daphné and da Silva Gonçalves, Jorge and Cannistraci, Irene and Sutter, Thomas M and Vogt, Julia E},
  year = {2026},
  booktitle = {Transactions of Machine Learning Research},
}

This project builds upon:

• DRPM (Differentiable Random Partition Model)
  https://github.com/thomassutter/drpm

@inproceedings{sutterryser2023drpm,
  title={Differentiable Random Partition Models},
  author={Sutter, Thomas M and Ryser, Alain and Liebeskind, Joram and Vogt, Julia E},
  year = {2023},
  booktitle = {Advances in Neural Information Processing Systems},
}

• MVHG (Multivariate Hypergeometric)
  https://github.com/thomassutter/mvhg

@inproceedings{sutter2023mvhg,
  title={Learning Group Importance using the Differentiable Hypergeometric Distribution},
  author={Sutter, Thomas M and Manduchi, Laura and Ryser, Alain and Vogt, Julia E},
  year = {2023},
  booktitle = {International Conference on Learning Representations},
}