Flow-Based Density Ratio Estimation for Intractable Distributions with Applications in Genomics
Authors: Egor Antipov, Alessandro Palma, Lorenzo Consoli, Stephan Günnemann, Andrea Dittadi, Fabian J. Theis
Organizations: Institute of Computational Biology, Helmholtz Munich · School of Computation Information and Technology, Technical University of Munich · Munich Center for Machine Learning (MCML) · Munich Data Science Institute (MDSI) · TUM School of Life Sciences Weihenstephan, Technical University of Munich
Estimating density ratios between pairs of intractable data distributions is a core problem in probabilistic modeling, enabling principled comparisons of sample likelihoods under different data-generating processes across conditions. While exact-likelihood models such as normalizing flows offer a promising approach to density ratio estimation, naive evaluations are computationally expensive and prone to discretization errors because they require simulating each distribution's likelihood independently. In this work, we leverage condition-aware flow matching to derive a single dynamical formulation for tracking density ratios along generative trajectories. We demonstrate competitive performance on simulated benchmarks for closed-form ratio estimation, and show that our method supports versatile tasks in single-cell genomics data analysis, where likelihood-based comparisons of cellular states across experimental conditions enable treatment effect estimation and batch correction evaluation.
Figures & tables
Figure 1: Naively approximating the likelihood ratio q(x1∣y)/q(x1∣y′) requires estimating the likelihood of x1 separately under two conditional generative models ptθ(⋅∣y) and ptθ(⋅∣y′) , incurring two independent integral evaluations with compounding discretization errors. scRatio instead estimates the log-ratio directly as the solution of a single ODE tracked along trajectories simulated from data to noise under a vector field bt . Here, we illustrate the case where bt:=utθ(⋅∣y) , i.e., trajectories are simulated under the numerator field.
Figure 2: a. The average MSE across 5 training runs between true and estimated likelihood ratios across models and simulation parameters. SB paths stand for Schrödinger Bridge probability paths. b. Runtime in seconds for likelihood ratio estimation comparing scRatio with the naive approach across simulation parameters. Dynamical ratios are simulated with the adaptive dopri5 solver.
Model
d
20
40
80
160
320
Classifier baseline
0.27 ± 0.04
9.99 ± 0.00
19.51 ± 0.18
0.83 ± 0.11
26.68 ± 3.70
TSM with SB paths
0.03 ± 0.01
0.38 ± 0.02
0.25 ± 0.10
0.89 ± 0.11
3.55 ± 0.12
CTSM with SB paths
0.06 ± 0.01
0.09 ± 0.04
0.23 ± 0.03
0.87 ± 0.10
2.15 ± 0.08
scRatio σmin=0,λ=0
0.05 ± 0.01
0.05 ± 0.02
0.05 ± 0.03
0.16 ± 0.05
21.18 ± 0.14
scRatio σmin=0.1,λ=0
0.04 ± 0.02
0.03 ± 0.02
0.02 ± 0.01
0.16 ± 0.07
21.00 ± 0.26
Table 1: MAE ( ↓ ) between true and predicted MI across data dimensions. Results are averaged over three training runs and reported as mean ± standard error. The best and second-best results per dimension are highlighted in bold and underlined, respectively.
Model
ρ(AUC,a) ( ↑ )
AUC ( a≥0.3 ) ( ↑ )
ρ(NAR,a) ( ↑ )
NAR ( a≥0.3 ) ( ↑ )
ρ(CSP,a) ( ↑ )
CSP ( a≥0.3 ) ( ↑ )
MrVI
0.71 ± 0.09
0.89 ± 0.00
0.83 ± 0.01
7.33 ± 0.27
0.67 ± 0.06
0.97 ± 0.00
MELD
0.48 ± 0.00
0.93 ± 0.00
0.17 ± 0.00
6.45 ± 0.00
0.27 ± 0.00
0.99 ± 0.00
scRatio ( σmin=0,λ=0 )
1.00 ± 0.00
0 .95 ± 0.00
1.00 ± 0.00
9 .64 ± 0.45
0 .95 ± 0.01
0 .98 ± 0.00
scRatio ( σmin=0.1,λ=0 )
0 .99 ± 0.01
0.96 ± 0.00
1.00 ± 0.00
11.39 ± 0.30
0 .95 ± 0.03
0 .98 ± 0.00
scRatio ( σmin=0,λ=0.01 )
0 .99 ± 0.01
0 .95 ± 0.00
0 .98 ± 0.02
9.39 ± 0.19
0.98 ± 0.01
0 .98 ± 0.00
Table 2: Comparison of scRatio with three distinct scheduling settings and baseline methods across increasing levels of DA controlled by the parameter a (Eq. 14 ). We report (I) the Spearman correlation ρ(⋅,a) between each metric and the ground-truth DA level a , and (II) the average metric value over high-DA regimes ( a≥0.3 ). Results are averaged over three independent training runs and reported as mean ± standard error. In bold and underlined are the best and second-best results per metric, respectively.
Figure 3: Ratio-based batch correction evaluation on the NeurIPS 2021 (panel a ) and C. Elegans (panel b ) datasets. Left: Distribution of the absolute estimated log -ratio before and after batch correction for each of the measured batch and cell type combinations in the two datasets, estimated by training scRatio alternatively on the original and batch-corrected representations. Right: UMAP plots calculated before and after batch correction with scVI. To compute the UMAP, we use 50-dimensional representations obtained by PCA and scVI for uncorrected and corrected data, respectively. Batch names are annotated below each plot.
Figure 4: On the x-axis, the mean absolute log -likelihood ratio logrθ(x∣(d1,d2),d1) evaluated on cells perturbed with both (d1,d2) and only d1 . On the y-axis, the mean log-odds of a classifier trained to discriminate cells treated with (d1,d2) from cells treated with d1 only. Each point is a combination. UMAP plots show extreme cases of the combinatorial effect and lack thereof.
Figure 5: Patient-specific differential response to perturbation. ( a ) Depiction of the experimental setup. Following Oesinghaus et al. (2025) , we divide donors into two groups responding differently to distinct cytokine treatments. (b) Absolute log -likelihood ratios between treated and control distributions evaluated on perturbed cells from different donor groups. A higher ratio indicates evidence of a strong response to a cytokine by a donor group, while log -ratios around 0 refer to a lower shift from controls.
Appendix figures & tables17 assets
Supplementary material from the paper’s appendix.
Appendix
αt
σt
ut(xt∣x1)
∇xtlogpt(xt∣x1)
t
1−t
(x1−xt)/(1−t)
(tx1−xt)/(1−t)2
t
1−(1−σmin)t
(x1−(1−σmin)xt)/(1−(1−σmin)t)
(tx1−xt)/(1−t+tσmin)2
t
C=[(1−λ)t2+(λ−2)t+1]21
(((λ−2)t+2)x1+(λ+2(1−λ)t−2)xt)/(2C2)
(tx1−xt)/((1−λ)t2+(λ−2)t+1)
Appendix
Table 3: Closed-form conditional vector field ut(xt∣x1) and score ∇xtlogpt(xt∣x1) for different conditional probability path schemes. The tuple (αt,σt) represents the schedule of the path between noise and data. A more detailed description of the schedules is in Sec. B.4 .
αt
σt
α˙t
σ˙t
∇xtlogptθ(xt)
utθ(xt)
t
1−t
1
−1
(tutθ(xt)−xt)/(1−t)
(xt+(1−t)∇xtlogptθ(xt))/t
t
1−(1−σmin)t
1
σmin−1
(tutθ(xt)−xt)/(1−(1−σmin)t)
(xt+(1−(1−σmin)t)∇xtlogptθ(xt))/t
t
C=[(1−λ)t2+(λ−2)t+1]21
1
(λ−2(λ−1)t−2)/2C
(2(tutθ(xt)−xt))/((λ−2)t+2)
(2xt+((λ−2)t+2)∇xtlogptθ(xt))/2t
Appendix
Table 4: Closed-form relation between score and vector field for different noise schedules. The tuple (αt,σt) represents the schedule of the path between noise and data, while (α˙t,σ˙t) is the respective time derivative. The relationship between score and vector field is obtained from Eq. 12 .
Hyperparameters
Values
Latent space size
32, 64, 128, 256, 512
λ
0, 0.1, 0.25, 0.5, 0.75, 1
σmin
0, 1e-4, 1e-3, 0.01, 0.1
Use sinusoidal time embedding
True, False
Sinusoidal time embedding size
32, 64, 128
Condition embedding size
32, 64, 128
Appendix
Table 5: The hyperparameter grid explored for the Gaussian simulation dataset. Each model with a different configuration of hyperparameters is compared based on the test-set ratio estimation performance measured as MSE.
d
MI
20
5
40
10
80
20
160
40
320
80
Appendix
Table 6: The real MI values for different data dimensions.
Hyperparameters
Values
λ
0, 0.1, 0.25, 0.5
σmin
0, 0.01, 0.1
Appendix
Table 7: The hyperparameter grid explored for the MI estimation experiment.
Hyperparameters
Grid
λ
0, 0.01, 0.1, 0.25, 0.5
σmin
0, 0.01, 0.1
Batch size
256, 512,1024
Top-k PCs
10, 20, 30, 40, 50
Appendix
Table 8: The hyperparameter grid explored for the DA estimation experiment on the PBMC68k dataset.
Figure 6: a. UMAP plot of the PBMC68K dataset colored by cluster labels. The clusters are assigned to each cell through the Leiden algorithm ( Traag et al., 2019 ) . b. The proportion of treated cells per cluster for different values of a (see Eq. 14 for a description of a ). c. UMAP plot colored by the treatment label sampled for each cluster using different values of a .
Figure 7: Mean Squared Error (MSE) of different methods across varying data dimensionalities and values of s (mean shifts from the origin). TSM and CTSM denote baseline models, and SB paths indicate Schrödinger Bridge probability paths.
Figure 8: a. Comparison between scRatio and a standard classifier baseline on the task of predicting likelihood ratios between Gaussians (see Sec. 5.1 ). Performance is evaluated for different numbers of features (x-axis) as the MSE between ground truth and predicted ratios. b. Overlap between 2D Gaussian distributions as a function of the distance between their means. In all plots, one of the distributions is fixed as a standard Gaussian.
Figure 9: scRatio’s MSE results for different levels of sample complexity on the Gaussian log-likelihood ratio estimation task. From Sec. 5.1 we use the harder setting with s=2 and d∈{2,5,10,20,30} ( Yu et al., 2025 ) . For the best hyperparameter configuration, we retrain the model with different numbers of training points and report the performance on the test set.
Figure 10: scRatio’s MSE results for neural parameterization and closed-form reparameterization of the scores on the Gaussian log-likelihood ratio estimation task. From Sec. 5.1 we use the harder setting with s=2 and d∈{2,5,10,20,30} ( Yu et al., 2025 ) . For the best hyperparameter configuration in Sec. 5.1 , we compare the performance of the model across dimensions between closed-form reparameterization and neural parameterization of the score and vector field in Prop. 4.1 .
Figure 11: Mean Squared Error (MSE) as a function of the absolute value of the log-ratio, shown for different locations and dimensionalities.
Figure 12: Test MSE as a function of dimensions and path variance for CTSM and scRatio across hyperparameter configurations. Boxplots represent the distribution of MSE (linear scale for CTSM, log-scale for scRatio) as a function of input dimensions. Boxes for both models are colored by the level of path stochasticity in the interpolants (Schrödinger Bridge and flow matching path variance, respectively).
Figure 13: Heatmap of Mean Squared Error per number of training steps for velocity utθ and score stψ fields. On the x -axis the number of steps used to train the velocity field utθ is presented, the y -axes shows the number of steps for the score field stψ , while the hue value encodes the MSE attained by the combination of the two.
Figure 15: Right: Four UMAP plots showing the gradual increase of combinatorial effects across several single perturbations and combinations thereof. Left: Corresponding box plots for each UMAP, displaying the absolute values of the estimated log-likelihood ratios.
d
λ
σmin
MSE ( s=1.0 )
MSE ( s=2.0 )
2
0.0
0.0
0.01
0.05
0.0001
0.01
0.05
0.001
0.01
0.13
0.01
0.01
0.05
0.1
0.02
0.17
0.10
0.0
0.02
0.07
Appendix
Table 9: Sweep results on Gaussian simulation experiments for different noise schedules. λ controls the stochasticity along conditional probability paths from noise to data, while σmin is the variance of Gaussian noise around data points. The αt parameter is kept fixed to αt=t , with t∈[0,1] . MSE results for different shifts from the origin ( s ) are averaged across three training replicates. Best results are highlighted per dimensionality.
Figure 16: Performance as a function of number of features for different model hyperparameters (number of condition embedding dimensions, learning rate, time feature size, number of latent dimensions). Evaluations are conducted on Gaussian data from Sec. 5.1 with s=2 . The performance quality is measured through the MSE between ground-truth and estimated likelihood ratios.
Flow matching has emerged as a powerful generative family for simulation-based inference (SBI). However, sampling and density evaluation require numerical integration at inference time. This makes posterior sampling in neural posterior estimation costly. More importantly, it can make repeated likelihood evaluation prohibitively expensive and impractical for Markov chain Monte Carlo (MCMC) sampling, which may require thousands of repeated density evaluations. Here, we develop conditional mean flows for posterior and likelihood approximation, combining few-step sampling with tractable density evaluation. We show that the mean flow objective identifies conditional flow maps and that independent masking enables direct estimation of marginal likelihoods. We evaluate both posterior and likelihood estimation across several tasks and show that conditional mean flows retain the flexibility of flow-based density estimators while significantly reducing inference time. In particular, single-step conditional mean flow density evaluation yields estimation error comparable to 200-step flow matching while enabling about 200-fold faster downstream MCMC sampling.
Jonas Arruda, Jan Hasenauer, Stefan T. Radev
University of Bonn · Rensselaer Polytechnic Institute
High-dimensional count data arise in applications such as single-cell RNA sequencing and neural spike trains, where mappings between distributions across successive batches or time points form critical components of data analysis. The recent success of diffusion- and flow-based deep generative models for images, video, and text motivates extending these ideas to count-valued settings, but many existing methods either treat each count as a categorical state or transform counts into a continuous space, neither of which is natural or efficient when the count range is large. We propose count-FM, a flow-matching framework for count data based on a continuous-time birth-death process with local unit jumps. Count-FM learns marginal transitions efficiently in count space through simulation-free training of conditional transition rates, allowing transport between arbitrary count-distributed source and target populations. In simulation, count-FM variants achieve strong sample quality while using substantially fewer parameters. We further apply count-FM to scRNA-seq and neural spike-train data for unconditional generation, transport, and conditional generation. Across these tasks, count-FM yields improved sample quality, greater modeling efficiency, and interpretable transport paths.
Ganchao Wei, John Pearson
Department of Neurobiology Department of Statistical Science Duke University Durham, NC, USA · Department of Neurobiology Department of Electrical and Computer Engineering Duke University Durham, NC, USA
Density ratios quantify distribution shift from a probability-mass point of view, whereas displacement fields describe, from a dynamical point of view, how one distribution is transported onto another. Although both offer complementary insights, they are usually estimated separately, and converting one into the other requires post-processing. In this paper, we estimate the density ratio between a target and a base distribution by parametrizing it through a displacement field acting on the base: the log-ratio is modeled as minus the Stein operator of the base applied to the field, up to a normalizing constant. This gives both statistical and dynamical descriptions of the distribution shift through a single convex optimization problem. Iterating this estimate-and-move step gives two inference algorithms: push-forward moves the model and corrects a pretrained sampler without retraining it, whereas pull-back moves the data closer to the base and fits a transformation model one layer at a time. Applications to distribution shift in simulation-based inference and to nonlinear independent component analysis illustrate the benefits and limitations of the approach.