🤖 AI

You Only Train Once: Differentiable Subset Selection for Omics Data

YOTO is an end-to-end, differentiable framework that jointly performs discrete gene subset selection and prediction in a single training pass, outperforming existing multi-stage approaches by enabling tight coupling between feature selection and task performance for improved biomarker discovery in single-cell transcriptomics.

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

Published 2026-06-04

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive jigsaw puzzle, but instead of 1,000 pieces, you have 10,000. Most of these pieces are just background sky or empty table space, and only a few are the actual picture you need to see.

In the world of biology, scientists often face this exact problem with single-cell data. They have measurements for thousands of genes (the puzzle pieces) for every cell, but they don't know which specific genes are the "important ones" that tell them what kind of cell it is or if a patient is sick.

The Old Way: The Two-Step Dance

Traditionally, scientists did this in two separate steps, like a relay race where the baton is dropped:

Step 1: A statistician looks at all the genes and says, "These 50 look interesting based on some math rules."
Step 2: A different scientist takes those 50 genes and trains a separate computer program to make predictions.

The problem? The first step didn't know what the second step needed. They were "weakly coupled." It's like picking ingredients for a soup based on a cookbook, then handing them to a chef who has to guess what dish you wanted to make. The result is often a messy, inefficient process that requires retraining the chef every time you change the recipe.

The New Way: YOTO (You Only Train Once)

The paper introduces YOTO, a new method that combines the "picking" and the "cooking" into one seamless, continuous process.

Think of YOTO as a smart, self-correcting sous-chef who is learning to cook while simultaneously deciding which ingredients to grab from the pantry.

Here is how it works, using the paper's specific claims:

1. The "You Only Train Once" Magic

In the old way, you had to pick ingredients, train the chef, test the dish, and if it tasted bad, go back and pick different ingredients.
With YOTO, the chef and the ingredient picker are the same person. As the chef tries to predict the dish (e.g., "Is this a healthy cell or a sick one?"), they immediately get feedback. If the prediction is wrong, the system instantly whispers to the ingredient picker: "Hey, that last ingredient you grabbed wasn't helpful; try swapping it for this one."

This creates a closed feedback loop. The prediction task tells the system which genes to select, and the selected genes immediately shape the prediction. They learn together, in real-time, until they get it right. You only have to run this training process once.

2. The "Binary Mask" (The Strict Gatekeeper)

Many other AI methods try to be "soft" about selection. They say, "Maybe this gene is 10% important, and that one is 5%." This means the computer still has to look at every single gene to make a decision, which is slow and messy.

YOTO is strict. It uses a binary mask, which is like a gatekeeper with a "Yes" or "No" switch.

Yes: This gene is in the top 50. It gets to enter the kitchen and help cook.
No: This gene is locked out. It contributes nothing to the final decision.

This is crucial because it means the final model only uses the specific, compact list of genes it chose. You don't need to train a new classifier later to see if those genes actually work; the model has already proven it works by using only those genes to make its predictions.

3. The Multi-Task Learner (The Swiss Army Knife)

Usually, if you want to predict two different things (e.g., "What type of cell is this?" AND "Is the patient sick?"), you need two different models.
YOTO is a Swiss Army Knife. It learns one shared "brain" (a shared encoder) that understands the core biology of the cell. Then, it has different "heads" (specialized tools) for different tasks.

Because it learns from all these tasks at once, the "brain" becomes smarter and more robust.
It discovers a set of genes that is useful for all the tasks simultaneously.
You don't need to retrain the model for each new question; you just ask the model a different question, and it answers using the same learned gene subset.

The Results: Does it Work?

The authors tested YOTO on two real-world biological datasets:

COVID-19 Blood Cells: Trying to identify disease states and cell types.
Mouse Brain Cells: Trying to classify specific types of neurons in the visual cortex.

The findings were clear:

Better Performance: YOTO consistently beat the old "two-step" methods and other modern AI methods, especially when they were forced to use very small, compact lists of genes (like 16 to 256 genes out of thousands).
Stability: Even though the specific genes it picked might vary slightly from run to run (due to randomness), the types of biological functions those genes performed were always the same. It consistently found the right "tools" for the job.
Efficiency: It achieved these results by training just one model once, whereas the old methods required training multiple models or retraining classifiers for every new task.

Summary

YOTO is a new way to find the "needle in the haystack." Instead of guessing which needles are important and then testing them later, YOTO builds a machine that learns to find the needles while it is learning to use them. It is faster, more accurate, and requires less work because it does everything in one go, ensuring that the final list of genes is perfectly tuned for the job at hand.

Technical Summary: YOTO (You Only Train Once)

Problem Statement

Selecting compact and informative gene subsets from high-dimensional single-cell transcriptomic data is critical for biomarker discovery, model interpretability, and cost-effective profiling. However, existing approaches suffer from significant limitations:

Decoupled Pipelines: Traditional methods (e.g., statistical filtering, mRMR, Seurat) operate as multi-stage processes where feature selection is performed independently of the downstream predictive task. Consequently, selected genes may not be optimal for the specific models they support.
Weak Coupling in Deep Learning: Recent deep learning methods often rely on post-hoc feature attribution or soft masking mechanisms (e.g., PERSIST). These approaches do not enforce a discrete subset of genes during training; instead, all genes contribute to the prediction with varying weights. This necessitates retraining separate downstream classifiers to evaluate the utility of the selected subset, breaking the end-to-end optimization loop.
Data Scarcity and Noise: Single-cell data are inherently sparse, noisy, and high-dimensional, with labels often incomplete or available only for subsets of samples.

Methodology: YOTO Framework

The authors propose YOTO (You Only Train Once), an end-to-end, differentiable framework that jointly learns discrete gene subsets and performs multi-task prediction within a single architecture. The model consists of three primary components:

1. Differentiable Subset Selection

Unlike previous works that treat features independently or use soft masks, YOTO explicitly learns a ranking between features to select a discrete subset.

Learnable Scores: A vector of learnable scores $s$ is assigned to each gene.
Differentiable Permutation: The model employs a differentiable permutation matrix $\pi(s)$ based on the Plackett-Luce (PL) distribution and the Gumbel-Softmax trick. This allows the ranking of genes to be differentiable with respect to the scores.
Binary Masking: To ensure sparsity, a binary mask $\Gamma_k(s)$ is constructed by summing the top $k$ rows of the permutation matrix.
Straight-Through Estimator: During the forward pass, the model uses the hard binary mask to ensure only the selected $k$ genes contribute to inference. During the backward pass, the relaxed (soft) version of the permutation matrix is used to propagate gradients, enabling end-to-end training.
Annealing: A temperature parameter $\tau$ is exponentially annealed during training to transition from a soft, stochastic selection to a deterministic, discrete selection.

2. Shared Encoder

A shared neural network encoder maps the selected subset of $k$ gene values into a latent representation. This shared representation captures biological signals common across different tasks.

3. Multi-Task Learning Heads

The framework supports multiple predictive objectives (e.g., cell type classification, disease state, patient identity) simultaneously.

Joint Optimization: Task-specific heads process the shared latent representation to predict labels for $T$ different tasks.
Loss Function: The total loss is the average of task-specific losses (cross-entropy for classification, MSE for regression).
Handling Missing Labels: The design allows for incomplete supervision; tasks with missing labels for certain samples still contribute to refining the shared representation via the available labeled samples.

Key Contributions

End-to-End Discrete Selection: YOTO is the first framework to enforce a strictly discrete, sparse gene subset during training while remaining fully differentiable. This eliminates the need for post-hoc selection or retraining downstream classifiers.
Closed Feedback Loop: The prediction task directly guides gene selection, and the selected genes shape the predictive representation, creating a closed loop that iteratively refines both selection and prediction.
Multi-Task Efficiency: By learning a shared gene subset across multiple related objectives, YOTO discovers gene panels that generalize across tasks without requiring separate training instances for each task.
Single-Model Evaluation: A single training instance of YOTO can be evaluated on all tasks simultaneously, offering significant computational efficiency compared to multi-stage pipelines.

Experimental Results

The authors evaluated YOTO on two distinct single-cell RNA-seq datasets:

COVID-PBMC: Peripheral blood mononuclear cells from COVID-19 patients (predicting cell types, disease status, and patient ID).
VISp: Mouse primary visual cortex cells (predicting cell types at varying granularities).

Key Findings:

Performance: YOTO consistently outperforms or matches state-of-the-art baselines (including Seurat, mRMR, Block HSIC Lasso, and PERSIST) across various metrics (F1-score, AUPRC, AUROC, Accuracy).
Panel Size Robustness: YOTO demonstrates superior performance as the gene panel size increases (e.g., 32–256 genes), suggesting it effectively captures gene interactions and complementary sets. It remains competitive even at very small subset sizes (e.g., 16 genes), though performance is highly sensitive to individual gene choice in extreme sparsity.
Multi-Task Advantage: In multi-task settings, YOTO achieves performance comparable to or better than methods requiring separate models for each task, while using only one training instance.
Ablation Studies:
- Comparing YOTO to PERSIST under identical single-task architectures confirms that YOTO's gains stem from its sparse, discrete selection mechanism rather than architectural size or multi-task supervision alone.
- Multi-task learning improves generalization on difficult tasks but may introduce optimization trade-offs in highly constrained regimes (very small $k$ ) if model capacity is limited.
Robustness: While the specific genes selected vary across random seeds due to feature redundancy, the selected subsets show high functional consistency (e.g., enriched Gene Ontology terms related to synaptic transmission and axon guidance) and preferentially recover known biological marker genes.

Significance and Claims

The paper claims that YOTO advances the field of single-cell analysis by providing a flexible and data-efficient system that unifies feature selection and prediction.

Biomarker Discovery: By enforcing sparsity and end-to-end optimization, YOTO yields compact, meaningful gene subsets that are directly interpretable and validated by their predictive power, rather than by arbitrary thresholds on continuous scores.
Translational Potential: The ability to select informative subsets from high-dimensional data with limited measurement capacity makes YOTO particularly suitable for targeted profiling and spatial transcriptomics applications.
Generalizability: While focused on genomics, the authors note that the framework is a general strategy for discrete feature selection applicable to other high-dimensional domains (proteomics, epigenomics) and settings where small feature sets are crucial.

The authors maintain a modest stance, acknowledging that while multi-task learning enhances robustness, it is not strictly necessary for all tasks and that the assumption of shared biological structure is critical for the multi-task approach to succeed. They also note that cell-level splits used in evaluation may allow models to leverage patient-specific signals, a trade-off they leave for future investigation.