GENOT: Entropic (Gromov) Wasserstein Flow Matching with Applications to Single-Cell Genomics

Authors: Dominik Klein, Théo Uscidda, Fabian Theis, Marco Cuturi

NeurIPS 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental 5 Experiments
Researcher Affiliation Collaboration Dominik Klein Helmholtz Munich dominik.klein@helmholtz-munich.de Théo Uscidda CREST-ENSAE theo.uscidda@ensae.fr Fabian Theis Helmholtz Munich fabian.theis@helmholtz-munich.de Marco Cuturi Apple cuturi@apple.com
Pseudocode Yes Algorithm 1 U-GENOT. Skip teal steps for GENOT.
Open Source Code Yes The GENOT model along with the code to reproduce the experiments can be found at https: //github.com/MUCDK/genot, while a more modular implementation can be found in OTT-JAX [14]. Additionally, we implement applications in moscot [39].
Open Datasets Yes The dataset of the developing mouse pancreas was published in Bastidas-Ponce et al. [2] and can be downloaded following the guidelines on https://www.ncbi.nlm.nih.gov/geo/query/acc. cgi?acc=GSE132188.
Dataset Splits No For each drug, we project the singlecell RNA-seq readout of the unperturbed and perturbed cells to a 50-dimensional PCA embedding. Subsequently, we split the data randomly to obtain a train and test set with a ratio of 60%/40%.
Hardware Specification No We only perform single-GPU training, thus we assume there is no limitation to reproduce single experiments / there is no environmental/societal effect due to a single experimental run.
Software Dependencies No The GENOT framework is implemented in JAX [6]. We use the discrete OT solvers provided by OTT-JAX [14].
Experiment Setup Yes Batch size: n = 1024. Entropic regularization strength: ε = 10 2. By default, we do not scale the cost matrices passed to discrete OT solvers. Unbalancedness parameter: τ = (1, 1). This means that by default, we impose the hard marginal constraints. Number of training iterations: Titer = 10, 000. Optimizer: Adam W with learning rate lr = 10 4, and weight decay λ = 10 10.