Count Bridges enable Modeling and Deconvolving Transcriptomic Data
Authors: Nic Fishman
Gokul Gowri
Tanush Kumar
Jiaqi Lu
Valentin De Bortoli
Jonathan S. Gootenberg
Omar Abudayyeh
Venue: ICLR 2026 Poster (2026)
Keywords: ordinal data
diffusion
schrodinger bridge
flow matching
single cell genomics
spatial transcriptomics
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We extend diffusion models to count data.

Abstract: Many modern biological assays, including RNA sequencing, yield integer-valued counts that reflect the number of molecules detected. These measurements are often not at the desired resolution: while the unit of interest is typically a single cell, many measurement technologies produce counts aggregated over sets of cells. 
Although recent generative frameworks such as diffusion and flow matching have been extended to non-Euclidean and discrete settings, it remains unclear how best to model integer-valued data or how to systematically deconvolve aggregated observations.
We introduce Count Bridges, a stochastic bridge process on the integers that provides an exact, tractable analogue of diffusion-style models for count data, with closed-form conditionals for efficient training and sampling. We extend this framework to enable direct training from aggregated measurements via an Expectation-Maximization approach that treats unit-level counts as latent variables.
We demonstrate state-of-the-art performance on integer distribution matching benchmarks, comparing against flow matching and discrete flow matching baselines across various metrics. We then apply Count Bridges to two large-scale problems in biology: modeling single-cell gene expression data at the nucleotide resolution, with applications to deconvolving bulk RNA-seq, and resolving multicellular spatial transcriptomic spots into single-cell count profiles. Our methods offer a principled foundation for generative modeling and deconvolution of biological count data across scales and modalities.

Published as a conference paper at ICLR 2026
COUNT BRIDGES ENABLE MODELING AND DECONVOLVING TRANSCRIPTOMIC DATA
Nie Fishman Gokul Gowri?, Tanush Kumar', Jiaqi Lu', Valentin de Bortoli", Jonathan S. Gootenberg & Omar Abudayyeh 'Harvard University: 'MIT, *CNRS;
'Beth Israel Deaconess Medical Center; 'Brigham and Women's 's Hospital; HHarvard Medical School "Corresponding author: n se com
ABSTRACT
Many modern biological assays, including RNA sequencing, yield integer-valued counts that reflect the number of molecules detected. These measurements are often not at the desired resolution: while the unit of interest is typically sin- gle cell, many measurement technologies produce counts aggregated over sets of cells. Although recent generative frameworks such as diffusion and flow matching have been extended to Ann-uulidean and discrete settings, it remains unclear how best to model integer-valued data or how to systematically deconvolve aggregated observations. We introduce Count Bridges, stochastic bridge process on the integers that provides an exact, tractable analogue of diffusion-styl models for count data, with fosed-form conditionals for efficient training and sampling, We extend this framework to enable direct training from aggregated measurements via Expectation-Ne satrriion-Maiiiiiatiinn-ittt latent variables. We demonstrate state-of-the-art performance on integer distribu- tion matching benchmarks, comparing against flow matching and discrete flow matching baselines across various metrics. We then apply Count Bridges to two large-scale problems in biology: modeling single-cell gene expression data at the nucleotide resolution, with applications to deconvolving bulk RNA-seq, and resolv- ing multicellular spatial transcriptomic spots into single-cell count profiles. Our methods offer principled foundation for generative modeling and deconvolution biological count data across scales and modalities.
approach that treats unit-level counts
INTRODUCTION
Integer-valued counts are fundamental product of scientific measurements because of the discrete nature of molecules. Modern biological assays yield massive streams count data: RNA-seq read counts, fluorescence imaging molecule counts, and mass cytometry ion counts (Klein et al. 2015 Raj et al. 2008 Bendall et al.) 2011). However, these measurements are often aggregated over multiple individual units, obseuring the fine-grained patterns underlying these natural phenomena. Transcriptomics technologies exemplify this challenge, with technologies such as Visium capturing 10-50 cells per spot (Stahl al. (0166) and bulk RNA-seq aggregating thousands millions of cells per readout, yielding averages rather than high-resolution details. Deconvolving these aggregates into single-cell profiles is critical for the precise mapping of cellular heterogeneity, cell-cell interactions, and tissue architecture (Moses & Pachter 2022; Armingol et al.) 2021). l). The challenge is twofold: building generative models that respect the integer nature counts and extending these models to infer leit- profiles from aggregated observations.
Recent developments generative modelling only partially addresss the problem. Discrete diffusion models (Austin et al. 2021 Lou et al.) 2023) treat counts unordered categories through masking or uniform noise. Blackout Diffusion (Santos et al.) 2023), the only count-specific approach, uses pure-death processes that cannot transport between arbitrary distributions. The biological decon- volution literature on the other hand focuses on deconvolving cell-type (cluster-level) proportions (Kleshchevnikov et al. 2022 Cable et al.) 2022 Li al. 2023), rather than unit-level count profiles. Thus, there is need for framework that respects the integer and ordinal structure of counts, enables transport between arbitrary distributions, and systematically deconvolye aggregated observations.

---

Published as a conference paper at ICLR 2026
We introduce Count Bridges: stochastic bridge process Zd using Poisson birth-death dynamics. This yields closed-form conditionals for exact sampling and extends naturally to deconvolution via an EM algorithm treating unit-level counts latent. The birth- mechanism allows transport between arbitrary integer-valued distributions while preserving the ordinal structure, as both increments and decrements respect the natural ordering of counts. We show that Count Bridges outperform existing methods on synthetic benchmark datasets and scale more favorably to high-dimensional settings. We then showease Count Bridges on two real-world biological applications centered on deconvolution: nucleotide-reson resolution single-cell RNA-sequence modeling for bulk RNA-seq deconvolution and reference-Fie frreence--eeeeeeee spatial transcriptomic deconvolution. The codebase available here.
BACKGROUND ON DIFFUSION MODELS
Diffusion models specify a time-indexed family of bridge kernels connecting Xo ~ to a simple
source distribution X1 P1 (often Gaussian). There are two layers of structure: an unconditional forward process (X1))e[o.1] with kernels (T Iо) = Law(X, Xo zo); (ii) for any
0 S s S t < 1, family of bridge kernels K Io, Ie) = Law(X, Xo = xo. X₁ = xz).
Diffusion models require two consistency properties. First we require bridge consistency identity.
For any S su Kaju(T's x.)) = Kaje(ar, wx) ( w.) dac.
Thus multi-step sampling along any grid u->-~> matches the single-step u - bridge. Second, the kernel must have projective posterior:
K Kaje(ar 22))= 20 2x) Kg 0x zo xze dæa.
(2)
where (xn .x)) = Law X, 2²)). This identity expresses Kaje as mixture over the posterior of the po data. It is essential for denoising: during sampling, each predicted X, changes the posterior qoje, so the reverse kernels must be projective under this posterior update.
Together, equation 1 and equation 2 lets define general diffusion approach. First we train denoiser that approximates the posterior, Xo go(- wx, I) as Law(Xo Xx=0;;), using tuples (t, Xx, Xo) drawn from the "global" bridge: sample zo ~ Po. 21 ~ P1, Unif[0, 1] and then X,
K- 20, z1).
20,
For sampling, pick a grid tk
to = O, draw X1 ~ Pl, set Xtx + -1,, sampling
x ~ go( Xte44 tx+1), Xi. ~ Ke, (-
X (k+1) (tt..))
(3)
By our consistency properties, this multi-step procedure is equivalent to sampling directly from the (0, 1) bridge, so the model cannot drift out of the training distribution.
2.1 DIFFUSION AS BRIDGE BETWERN NOISE AND DATA
Let us consider the unconditional Kojo process (X))eejo1) of the following form
X, = (1) + Br,
where (Bi) (B))) is a d-dimensional Gaussian process with non-decreasing standard deviation (f), and (t) a non-increasing function. Note that (0))= and (0))=
We want to define process that interpolates smoothly between Xo ~ po and X, given by another
distribution as in |Peluchetti (2023); Albergo et al. (202): Delbracio & Milanfar (2023) Liu et al. 2002 2023). We have the following proposition defining the global and local bridge.
Proposition 2.1. Let (Xt)ee(o,1) be given by equation 4 For 0 < < L. S 1, consider (Xs)sejo.) conditioned on X, = It and Xo Then the conditional law Kyjo,( xo, xx) is Gaussian and can be written
a(s) r(s. ((s, +o(ss r(s,t)) a(t)
X, d a()((1 r(s,
(5)
where ~ (O, Id) independent of (Xo, and r(s,t) t) a()))(s))
(K o5s defined by equation satisfies equations 7 and Folzz In particular, the family
2

---

Published as a conference paper at ICLR 2026
Note that if Xg (V( Id), a(1) = 0 and (1) we have X; _ (t u + (t) Furthermore,
our equation recovers the interpolation described in Albergo et al. (023) with the identification (tt + a(t) r()),, atf(() - B(t) and ott) r(t)) 7
2.2
SAMPLING THE POSTERIOR
In this paradigm the bridge is only the first of two choices that define the model. We also have to choose how to model the posterior Xol f. There are two core options: we can use differential equations model the posterior in the limit of small steps or we can focus more directly on modeling the posterior. In Euclidean space, the former lets us learn simple conditional expectation, whereas the latter always requires a distribution model.
Infinitesimal. Consider small backward step of size 5 > O. The local bridge between times and t is Gaussian, so conditioned on X; = x we can write to first order in 8 5 b(z, t) + Vō
X X₁
&. ~ (0, E(x, t)),
where b is the reverse-time drift and Σ is the diffusion covariance of the bridge.
The conditional law X X, is Gaussian and can be computed in closed form:
b(x, 1) (1) x + B2(1) E[Xo Xx = x] + bo()),
E(x,() Eo(().
The difffusion covariance depends only on f (from the Brownian increment), and the drift depends on the posterior Law(XX, Xx,) only through its mean. This justifies learning the mean qo (z. t) sz E[Xa| X, = z]] (equivalently, score or velocity) as in standard diffusion models (Song et al. 2020).
Distributional. Following De Bortoli al. 2025
can learn the conditional law go( ze, t) z LXw X2=rz), using any distribution learning approach. We can then sample and directly plug into the bridge
x (1++ qo [ (x+1) X1, Kr, 10 (- x (k+ .(.41)
to sample the posterior. The distributional perspective is particularly powerful when the infinitesimal perspective fails to admit a simplification to the conditional expectation, which motivates our use of the distributional approach for Count Bridges (see Sec. 3.2). In categorical discrete settings, all approaches are distributional since they are based on cross-entropy losses, see Campbell et al. (2022)
Austin et al. (2021) Shi (2024); ) Sahoo al. (2024).
COUNT BRIDGES
3.1 AN INTEGER BRIDGE BETWEEN DISTRIBUTIONS
Mirroring Sec. 2 we seek a bridge for integer-valued data. Instead of a Gaussian process, we use pair
of independent Poisson birth/death processes (B+)tejo.1j and ( (Dx)eeo.11 that increment//decrement
the counts. We define an increasing "jump-intensis" function w : [0, + R>0 with w()) w(1 (1) = 1, and then write the cumulative birth/death intensities as (t) a w(t) for some d > 0
so B. Poi(A (f)) and D, Poi(A-(()) From here we can define the unconditional kernel Kijo:
X₁ = Xo + B; Dr.
(6)
Denoting the displacement dN = X, - Xy, the total number of jumps N₂ = B, + De, and the slack variable M, min (B, D,). Any two of these variables determine the third:
N₂ dd] + 2M.
B. = (N, + de). D.
Br.
From the (Ne, B,) perspective, Poisson superposition and thinning imply that, conditional on (N₁, B,) at time i, the earlier counts (N., B.) for < can be sampled by Binomial draw for N, and a Hypergeometric draw for B,. Switching to (Me, de). Poisson change of variables yields the slack posterior M, de whose pmf has Bessel form (see Prop. A.6 in App. A). These two ingredients together give a count analogue of Proposition 2.1 the full derivation
App. A
Proposition 3.1. Let (Xx)uejo.1) be given by equation 6 Now, consider (Xx) conditioned by X, = Z2 and Xn To- Then, we have the Poisson Birth-Death bridge K .o...
X, d Xo B. Dsy
(8)
3

---

Published as a conference paper at ICLR 2026
Troeccaries
04-
03
12
15
Recursive trajectoiies
Posterior at FPM d]
stop of
3.0
03-
43
00
10.3
15.5
Figure 1: Left: Sample paths for several endpoint gaps d1 (top). Fixing the prefix [0, resample (t, I] by the recursive kernel (bottom). Middle: Bessel slack posteriors at initial and intermediate times. The slack M, concentrates near 0 as |d] grows. Right: ECDFs of X, from one-step kernel (1-+s) and a two-step kernel (1 (1 s) are indistinguishable, confirming composition.
where we condition on d, X, Xo and sample M, d, Bes([de]; A (t), A_ (t)), changing variables to N₂ and B, to sample Ba, and D2 which can plug into equation
N. N₂ Bin No, w(s)
B, (NN, B.) B,) ~ Hyp(N., B,, N.), D. = N. B..
(9)
w(t)
The family {K Josasesi defined by equation 8 satisfies equations andd 2
We visualize this process in Fig. [ where we show the trajectories for the one- and two-step models along with the core composition property that drives bridge models. This setup enables training and sampling from a Count Bridge, see Algorithms [nd These results leverage our custom CUDA kernel implementing the fast Bessel sampler of Devroye (002) to enable sampling at scale.
In Fig. 1 we also see that as d, grows the slack Mt concentrates near zero, so there is no slack. This means that Count Bridges are an instance of the static Schrödinger bridge problem (Léonard 2013): they solve entropy-regularic
be the jump intensity and (T =1 = (o ( K [x1 (zo) be the joint law of (Xo+ X2) induced by the kernel. Over the
optimal transport. Let K = V ======
space of couplings C(po. P1) = {C on X X C(, Po. C(X,-) pı}. Count Bridges solve
Cx E arg min CeC(po.P1)
Letting + oo yields the independent coupling Po & Pr. but as 0 we obtain
(I ~ log 2 E HC)),
so 0 recovers discrete OT with cost | rol (see App. A.2
This echoes the Gaussian case (Sec.) 2h where we define o = (11) and phet- and as a $
KL PI R 2 Ec X -H()),
so σ + 0 recovers quadratic OT, while + oc again gives @ (Shi et al. 2023). Thus the bridge parameters x (count) and a (Gaussian) play the same role as entropy mntropy-regularizztion strengths.
3.2 DISTRIBUTIONAL SCORING LOSS FOR THE DENOISER
Training requires distributional loss due to the discrete nature the space. As shown by Holderrieth et al. (2024) the ELBO for discrete generators (e.g.. jemgp processes) is distributional and cannot be reduced to expectations over point estimates. This mirrors the need for cross-entie cross-entropy in discrete diffusion and flow models. We can cross-entropy with Count Bridges (we test this, App. D.I), but it has two issues: first, it does not incorporate the lattice structure; second, cross-entropy cannot model the joint of X, X, without exponential cost in dimension, cross entropy is usually
4

---

Published as a conference paper at ICLR 2026
Require: dataset (xo, x1), (+)),+++))
Require: Ttx = 21, model go. ((-),A+())
1: for each minibatch do
1: for k = K, K 1,.... ..,, do
2:
sample (Ko, 21 (xo: X1)
sample tu go( Feaz (i)
3:
~ 1n
3:
di t Lta - 20
4:
dj + 21 20
Mr. Bes(( |: A4 (tx), (tx), (tx))
5:
Mi [ess([d]:: (1),
5:
Ni + Ide. + 2M
6:
N1 < (d1 11|+
6:
Bi + [NVc + dea)
7:
B1 + 5(N11 + d1)
r f w(tkk1)//o((
8:
N (in,Ni, w())
8:
No Bin (onNN
9:
B₁ Hyp(N1, (N, B1, N₂)
9:
B Hyp New: Bea Noa-1)
It G I1 -2(B B,)+ (N1-NN)
10:
Tta-1 Ita 2( Bir Bia- )+ (Nox Non-1)
1I:
update 0 on (())
11:
end for
end for
12:
return 2to
Algorithm 1: Training Poisson-BD Bridge
Algorithm 2: Sampling Poisson-BD Bridge
factorized, modeling each coordinate of X, X, independently or autoregressivell. Specializing to count data we can go beyond cross-entropy by using proper scoring rule that (i) incorporates the geometry and (ii) enables modeling of the joint.
Zg, be our Formally, let (Xo, Xx) denote a training pair from Kojo,1 time E [o, 1]. and let (
denoiser. We train using strictly proper distributional scoring rule (Gneiting & Raftery, [2007] De Bortoli al. 2025). Fix negative-type semimetric p on 2 (all our experiments focus on the pl
with B = 1). For any distribution and outcome the energy score is
Sp(p,)) v) = Ex,x''
Ex- [p(x, and (()) = Ex..x., (d9s( (x,, Xo)]
which is strictly proper when p characteristic. Taking m ii.d. samples z(6) ~g( (2,,)) we can use the plugin estimator: Sp = m(m-1) Ejjy , m EØ= (( (0)))
DECONVOLUTION WITH COUNT BRIDGES
We extend Count Bridges to handle unit-level generation when we only observe aggregates. Consider G units in the one-dimensional case where the group-level state time is vector X, ZG with entries X, for unit at time f. Each entry evolves independently according to the bridge in Section The key challenge: we observe the unit-level endpoint x but only the aggregate at time
4g = >g Ziz E 2, not the unit-level vector Xg. Our goal is to learn a count bridge go (wo We, 1, that generates unit-level endpoints given start data time
and side information
We formulate this as generalized EM problem, similar to Rozet et al. 2024), where X, is latent and ag Eg Xgo is observed. Let z + be a linear aggregate map (e.g., sums across units, block sums). For (z,, 2). the denoiser q- Ir, f, defines an i.i.d. product prior over X, = (Xio...., Xoo). Conditioning on the aggregate yields
G
Qo(Xo au, F2,t,2)
[ g=1
go go 2,, 2) 1(A(Xo) = au}.
In the E-step we will generate "latent" 2 using the model and in the M-step we will use these x to train the model at the aggregate level. We summarize the overall procedure in Algorithms 3 and
E-Step The ideal step would sample from the exact aggregate- conditional law
X Qo do, te, 2).
We could then the sampled x latent variables to sample 3e between (xj. x1) using the unit-level kernel Kaju from Prop. 3.1 Unfortunately, Qe is generally intractable to sample from,
'The same method described here can be used with distributional diffusion on continuous space, but we focus counts since most often when observe aggregates helieve they based discrete underlying data.
5

---

Published as a conference paper at ICLR 2026
Require: (x1, ao, 2), w(-), A= (.), qo.
Require: (x1, do, ao, w(-), (-), go.
1: for k == K, K - 1,..., , do
1:
for each minibatch do
2: Sample X0.f, ~ qo(: Xta tk. 2)
2: E-step: Sample latent x from
3: Xo.te t [I au, 2]
conditional on do via Alg 4: M-step: t Unif[0. 1]
3:
4: Update Xt, by running the reverse step
5:
using steps 4-10 of Alg.
2 with Xo.ta
6: Sample x via the forward bridge
5:
6: end for
(x x1) using steps 4-10 of Alg.
7:
x sample and project X0,4
7: Update using the gradient of Luge
8:
return x
8:
end for
Algorithm 3: Guided Sampling to for z
Algorithm 4: Training with Aggregate Supervision
CFM
DFM
DCB
Figure 2: A scaled and rounded variant of the classic 8 gaussian to two moons task. Here we compare the trajecto- ries of continuous flow matching, dis- crete flow matching, and count bridges. CB achieves the lowest W2, MMD, and EMD, see Table
given just a unit-level model, so we approximate it through the diffusion sampling process itself. Starting from X1, run the sampling process in Algorithm 2 but at each timestep tx we: (1) predict & ~ (- X tk. 2), (2) project & to satisfy the aggregate constraint (see Sec. 4). yiclding No, and (3) perform the sampling step using Xy the predicted endpoint. This projection-guiee diffusion ensures the aggregate constraint is incorporated throughout the denoising trajectory (see
Alg, 3). This process produces latent 2 samples that are consistent with the aggregate constraints, which we can then use in the M-step to train the model.
M-Step With these unit-level samples in hand, the M- M-step runs the bridge process as in Section 3
But instead of computing the loss the unit-level latents, we compute the loss with respect to the aggregates. Given the ground-truth aggregate do, lift the same strictly proper score aggregates:
(p, [p(A((), A(X')) - Epp [p( ((X)(X)))))) and Lagz BA x. (/ X,, 2),, Ao)
with the plug-in in obtained by sampling X qo t, and forming ((4) A(X
Approximate Sampling from the conditional distribution Given a predicted endpoint &o from our diffusion model and target aggregate do, we need to sample from the conditional distribution Q+(- A(Xo) = ao). While this is intractable, we can derive principled approximation.
Proposition 4.1 (First-orde aggregate projection). Let A A(Xa) be the aggregate, and let po be
the prior law of Xo- Under the regularity conditions in App. B.I] the aggregate-conditione law Qo(- Ao an) admits a first-order exponential tilt. The corresponding generalized KL projection
T(xa)) arg min DKI(yol|xo))
yo: A(yo)=oo
gives a kind of first-order approximation to (- Ao = aa). For an elementwise sum A(x,) = Eg go this projection the simple scaling ( Og (Eg Z.go))
The proposition shows that the natural rescaling operation is not ad hoc, but be justified kind first-order approximation to the true conditional distribution a large sample regime (see Appendix B.I. When unit-level training data exist, we can leam projection IX 2, og) that actually enables sampling conditional on the mean. See Sec. where we show how to learn such projection.
RELATED WORKS
Stochastic interpolants. Our formulation allows us to transport any integer-valued distribution
Pi another integer-valued distribution Po- In the case of Fuclidean state space early works such

---

Published as a conference paper at ICLR 2026
as (De Bortoli et al. 2021: Vargas et al. 2021 Chen et al. 2021) have shown how to perform such an interpolation leveraging (Entropic) Optimal transport and the concept of Schrödinger Bridges. In more recent works, ignoring the Optimal Transport constraints, several works have proposed to bridge distributions a more relaxed formulation leveraging the concept of Markov projection, see Peluchetti (2023); Allergo al.] (2023); Delbracio & Milanfar (2023): [Liu et al. (2022 (023) for instance. Those frameworks can be shown to be strictly equivalent to diffusion models in the case
where one of the end distribution is a unit Gaussian, see Gao al. (2025), However, those works are limited to the Euclidean setting, and extension to the integer-valued valued setting is required.
Discrete diffusion models. Recently, with the advent of language diffusion models such as Ye et al. 2025); Song et al. (2025); [Sahoo et al. (2024); (Shi et al. (2024); Ou et al. ( [Ariola et al. (2025 ; Nie et al. (2024); Zheng et al. (2023). ), discrete diffusion models have gained considerable traction. Most works rely on discrete equivalents of the original formulation of diffusson models, explicitly or implicitly replacing the continuous Gaussian noising process by a Continuous- Time
Markov Chain (CTMC) (Austin et al.. 2021
Campbell et al. [0022] Lou et al. [2023] Campbell et al.,
2024; Kitouni et al. 2024; Sun et al. 2023 Other approaches include relying on some Euclidean
relaxation (Chen al. (022) modelling the space of probability (Avdeyev et al. 2023 Stark et al.
2024). Similarly, flow matching techniques have been extended to cover this paradigm (Gat et al.
2024). Most of these works focus on categorical data and therefore consider uninformed forward process such as uniform or masking process. In contrast, in this work, we focus on ordinal data. To the best of our knowledge, the only existing work that also deals with such a process Blackout Diffusion (Santos et al., 2023), which considers a pure-death process where an image is taken to the all-zero limit, as opposed to an endpoint conditioned bridge. Our approach generalizes this setup in two ways: first, we allow births and deaths at every time, recovering their pure birth construction in the limit as K- + 0; second, we generalize the process to a bridge which can transport X1 to Xo-
Finally, we highlight that diffusion models have been extended to the very general setting where only an infinitesimal generator is available Benton et al. (2024); Holderricth et al. (2024). While our work can be seen give any information regarding the design of the forward process for integer-valued data, the specific parameretization in terms of slack variables and the necessity of the distributional diffusion loss.
instanciation of this general framework, these general frameworks do not
Distributional Diffusion Models. In De Bortoli et al. (2025); Shen et al. (2025). the authors learn the conditional distribution Poje (2ol2z)) through the use of scoring rules, going beyond the classical training framework of diffusion, which approximates the conditional mean B[X,
2z). The
importance of approximating the covariance was already noted by Nichol & Dhariwal (2021) and
further analyzed in (Ho et al. 2020 Nichol & Dhariwal 2021 Bao al.) 2022a Ou al. 2024b).
In similar flavor (Xiao et al.
2022 uses GAN to approximate Poja (To (T)).
Sequence-to--xppeesion models An ambitious goal in biology is to predict gene expression from DNA sequence information. There have been several attempts to train deep learning models for sequence-to-expresion prediction tasks (Barbadila-Marttnnzz et al. 2025), including Enformer (Avsec) et al. 202T), a state-of-the-art transformer-ban like the vast majority of sequence-te teree-ttt--e---ooooeee models, was trained on bulk gene expression data and is not able to predict single-cell expression profiles, missing the cellular heterogencity and fine-grained regulatory patterns that shape tissue function.
DNA sequence model. While powerful, Enformer,
Spatial transcriptomic deconvolution Spatial transcriptomics encompasses a family of recently developed techniques which measure gene expression and spatial location in tissues. The majority of these techniques are not capable of resolving individual cells, instead providing aggregate information over small neighborhoods consisting of on the order of tens of cells (Moses & Pachter 2022). To address this limitation, number of deconvolution methods have been developed infer single- cell level information (Li et al. 2023). The majority of these methods, including cell2location (Kleshchevnikov et al. (022) and RCTD (Cable et al., (022), require a paired non-spatially resolved scRNA-seq atlas, and output cluster-level mixture proportions rather than single cell counts. The ideal deconvolution would recover full single-cell count profiles directly from spatial data without requiring external reference atlases. DestVI (Lopez et al., 2022). which outputs count profiles but requires a reference, and STDeconvolve (Miller et al., (002) which does not require a reference but outputs cluster-level predictions, both take steps toward this goal.
7

---

Published as a conference paper at ICLR 2026
Method
Steps NFE 32
CFM
DFM
Count Bridge
NFE-
NFE 128
4
0
4
8
16
32
128
256
512
Dimension d
Figure 3: CFM, DFM, and CB on our low-rank mixture of Gaussians transport experiment across dimensions and NFE. See App. D.2 for full details.
6 APPLICATIONS
We evaluate with three distributional metrics: the Energy score, the Wasserstein-2 distance, and the MMD (RBF). For deconvolution, we evaluate cell-type proportion predictions using RMSE, the Jensen-Shannon Divergence (ISD), and Spearman correlation following Li al. (2023). Synthetic tasks have std. errors over 3 training seeds; main applications have std. errors 3 over inference seeds.
6.1
SYNTHETIC DISTRIBUTIONS
Here, we benchmark count bridges (CB) against continuous flow matching (CFM) Lipman et al.) 2022 and discrete flow matching (DFM) (Gat et al. 2024) across range of synthetic experiments.
Discrete 8-Gaussians 2-Moons. We adapt this classic task the integers. We plot the learned trajectories in Fig Qualitatively CB achieves the best performance. DFM is much more competitive in this experiment than CFM, but DFM trajectories are decoupled from the underlying geometry. whereas CB produces OT-like trajectories similar to CFM. These qualitative evaluations are confirmed quantitatively: CB achieves the best performance across W2, Energy, and MMD (see App. D.I)
Scaling in Low-Rank Gaussian Mixtures. To test scalability to higher dimensions, construct integer-valuee datasets with fixed intrinsic dimensionality while ambient dimension increases in powers of two from 4 512. Each dataset is 5-component Gaussian mixture with latent rank projected to 2d In Fig. see that CB has the best scaling in dimensionality (see App. D.2 for more).
Deconvolution of Gaussian Mixtures. We extend the low- rank mixture task to evaluate deconvolution capabilities. In this experiment, each observation is an aggregate constructed by summing group of G samples. For each group, the G samples are drawn from group-specific Gaussian mixture whose com- ponent weights are sampled from Dirichlet distribution with concentration parameters (a),..., ). The labels of the G source components are provided as unit-level side information. We then vary the size of the group G and the extent of variation between groups by changing the concentration parameter (see Appendix D.3 for details). In Fig. 4 we see performance degrades as groups become more uniform and larger. We ex- plore the theoretical limits to deconvolution in Apps. Bn2 B.3 which confirm that deconvolution requires between-group which is inherently lost as groups become large. Despite these limits, we demonstrate practical deconvolution on moderately-sized groups in our spatial transcriptomies application (Section
Disched o-1008
12-9
Size
Figure 4: Deconvolution of the low- rank Gaussian mixture across dif- ferent group sizes and levels of between-group heterogeneity.
heterogeneity to enable identification,
6.2 MODELLING GENE EXPRESSION AT SINGLE-CELL AND SINGLE-NUCLEOTIDE RESOLUTION
central goal in biology is to understand the relationship between DNA sequence and gene expression. Many models relate sequence and expression, the most prominent of which, such Enformer (Avsec)
8

---

Published as a conference paper at ICLR 2026
Method
Bulk MSE
CT MSE
Comparison
MMD
W2
Energy
Fine-tuned Enformer Count Bridge
2.590
3,142
Bulk mean
0,515
0.208
56,800
0.601
1.410
Count Bridge
0.446
0.182
28.583
+0,000
-0.002
+0.001
+0.003
Table Nucleotide-Ievel MSE for bulk and Table 2: Gene expression count profile deconvo-
bulked cell-type (CT) specific predictions.
lution error for bulk RNA sequencing data.
et al. 2021), are Transformer-based models that predict expression from sequence. More recent work has explored fine-tuning Enformer on single-cell data (Hingerl et al. 2024). On the other hand, there is a mature literature on deconvolving bulk RNA-seq (Newman et al., 2019 Wang et al. 2019). These methods operate the gene (rather than nucleotide) level. leveraging bulk cell-type profiles or single-cell references to deconvolve bulk profiles into cell-type proportions (not count profiles).
We use CBs to jointly model sequence and single-cell expression counts in scRNA-seq data, and to enable nucleotide-eevel decon- volution of bulk profiles. To validate CBs in this setting, we demonstrate two key re- sults. First, show that CBs trained on single-cell data produce meaningful count profiles and outperform fine-tuned En- former model on sequence-tor toaicceto-exppessiion prediction. Second, we show that conditioning CBs on bulk profiles enables deconvolution of bulk gene expression into inferred single-cell gene expres- sion profiles. We validate these deconvolved profiles distributionally and show that they achieve state-of-the-art performance relative to cell-type proportion deconvolution models.
Method
JSD
RMSE
Spearman
CIBERSORTX
0.194
0.109
0.079
MuSiC
0.313
0.140
0.186
Count Bridge
0.113
0.073
0.267
+0.001
+0.000
+0.005
Table 3: Cell-type proportion deconvolution error for nucleotide level bulk RNA sequencing data.
Modeling sequence and single-cell counts We train CBs on PBMC ScRNA-seq counts at nucleotide resolution using 10 cells across 10x donors Yazar et al. 2022). Each training example corresponds to a nucleotide position in single cell, and is represented by the noisy count 2, and diffusion time t from the CB forward process, a cell-type embedding, local genomic context 2 obtained by encoding the surrounding DNA sequence with Enformer, and i.i.d. noise Ç for the distributional loss. These features are concatenated and passed through residual multi-head attention blocks and a final softplus head that parameterizes the conditional count distribution Xo) I, z. The model is trained directly on unit-level (single-cell) expression profiles rather than only on aggregated counts. During training randomly mask cell-type labels so that the model supports both unconditional and cell-type-condition sampling at test time.
Learned projection for deconvolution Since we have unit-level data can learn a better projection operator than the simple rescaling function in Prop. 4.1 We augment the CB with a small projection module HIg attention block operating batch. Given an initial CB prediction in. an observed aggregate dg. and the noisy state Tt, the module outputs Xo Thp (&o, 40, Xx, z), we train this using the distributional loss to learn sample Xo A(Xo)=oo. Xz, t. To support both unconditional and aggregate- aggregate-condditonee conditioned inference, we apply the projection module only on a random 10% of training examples where ag provided.
each nucleotide (represented by 2) across cells in the
Bulk gene expression We first evaluate the ability of our model to predict expression from sequence, both unconditionally and conditional cell type. As baseline, use an Enformer model fine- tuned directly on the PBMC dataset. We find that Count Bridge predictions outperform fine-tuned Enformer (Table for results by cell type and further details see App.
Deconvolved profiles We can use this unit-level model for deconvolution tasks: we can condition on an aggregate (bulk profile) to sample single-cell profiles from the model while matching that aggregate. We next evaluate the ability of CBs to deconvolve mixtures of cell types from held- out individuals. We held out 10% of patients from our training set and synthetically bulked these patients. Since we have the ground truth data, we can then evaluate deconvolution quality. We first evaluate the distributional quality of these predictions against the bulk mean, further validating compare to CIBERSORTx
the CB count profiles (Table 2). As more robust set of baselines, (Newman et al.) 2019) and MuSiC (Wang et al., 2019). To facilitate comparison, we aggregate our nucleotide-leve predictions into gene counts and assign each of our deconvolved cells to the closest

---

Published as a conference paper at ICLR 2026
cell type. CBs achieve better performance on JSD, RMSE, and Spearman correlation while providing nucleotide- counts (Table 3). In App. E we plot the UMAP for qualitative comparison.
6.3 DECONVOLVING SPATIAL TRANSCRIPTOMIC SPOTS INTO SINGLE-CELL COUNTS
Next, we show how CBs can be used to infer single cell gene expression profiles from spot-level aggregates in spatial transcriptomic data. In spatial transcriptomic data generated by Visium (Stähl
et al. 2016), it is common to have access to side information beyond the spot-level count aggregates. In particular, many datasets include images of the cells with a nuclear stain (Palla et al. 2022). CBs
provide natural way to leverage this cell-level side information to deconvolve aggregate count data.
Modeling spatial aggregates We train CBs on a MERFISH mouse brain dataset (Vizgen 2021), which is resolved at the single-cell level, and artificially aggregate neighborhoods of cells to simulate spot-level Visium data. This synthetic dataset gives us access to spot-level aggregates and their corresponding single-cell ground truth, well as single-cell nuclear images. Following the notation in Sec. the spot-level counts can be treated as aggregates Og. and single-cell images can be treated unit-level side information z. In this application, spot-level aggregates and the single-cell images. We leverage a UViT (Bao et al., 2023) extended to incorporate count and noise patches (see App. We use simple source distribution X Poi(10).
never observe single-cell count profiles, only
Cell type proportions We benchmark CBs against STDeconvolve (Miller et al. 2022). )
Method
JSD
RMSE
Spearman
widely used spatial transcriptomic deconvolu- tion method which is state-of-the-art among reference- free approaches Li et al. (2023) (see
STDeconvolve
0.288
0.177
0.255
Count Bridge
0.231
0.110
0.332
#0.002
+0.001
+0.001
Appendix forr comparisons to reference- methods). STDeconvolve outputs cell type (clus-
Table 4: Cell-type proportion deconvolution error for spatial transcriptomics.
ter identity) proportions for each spot rather than single cell counts. As such, we evalute the quality of deconvolution by comparing the predicted cell type proportions to the true cell type proportions per spot. For CBs, which provide single-cell count profile predictions rather than cell type proportions, we assign each predicted count profile its nearest neighbor cell type in order compare against STTDecovvlvve. CBs outperforms STDDecovolvee on both the JSD and the RMSE (Table
Count profiles We next evaluate the quality of the count profiles inferred by CBs. Here, because STDe- convolve does not provide these predictions, we instead consider a simple baseline: predicting the spot-level mean (oo/G) for each cell. This baseline, while seemingly naive, is actually biologically well- motivated. spatial transcriptomics, cells within a spot represent local tissue organization where neigh- boring cells coordinate their functions (Armingol et al., 202I). As such, we expect cells in spatial neighborhoods to have correlated gene expression profiles, making the spot mean reasonable baseline. Nonetheless, CBs outperform the spot-level mean baseline (see Table showing CBs can learn meaningful unit-level distributions from real-world aggregate data. In App. P we provide more detailed biological evaluation of the cell types and pathways in our generated data, alongside the UMAP facilitate qualitative comparison.
Comparison
MMD
W2
Energy
Spot mean 0.409
0.030
41.717
Count Bridge
0.203
0.017
8.903
-0.000
+0.000
+0.014
Table 5: Gene expression count profile decon- volution error for spatial transcriptomics
7 CONCLUSION
Count Bridges offer a tractable, discrete-natio alternative to continuous diffusion models, unifying direct count generation with deconvolution from aggregates. We demonstrate the power of Count Bridges for mucleotide-level deconvolution of bulk RNA-seq and spatial transcriptomic deconvolution.
Limitations (i) When counts well-approximmate continuous, Euclidean models may match or exceed performance. (ii) Identifiability in pure deconvolution degrades as group sizes grow or between-group heterogeneity shrinks, so our EM procedure is most reliable moderate aggregation. (iii) The projection step we use is a first-order surrogate and lacks serious theoretical support.
Despite these caveats, Count Bridges lay the groundwork for rigorous discrete generative model- ing and invite future work on a deeper understanding of the projection-guided sampler, sharper identifiability bounds, and generally stronger guarantees for projection-guided EM.
10

---

Published as a conference paper at ICLR 2026
Ethics Statement. This study uses publicly released, de-identified single-cell and spatial tran- scriptomics datasets under their respective licenses; no new human subject data were collected, and institutional review board (IRB) approval was therefore not required. We do not foresee serious ethical implications to Count Bridges beyond the risks already posed by standard diffusion/flow matching models. Our deconvolution methods could possibly pose some additional privacy risks. We used LLMs to help draft portions of the code used in our experiments and to edit portions of this manuscript. All our models are intended for research use only, not clinical use. LLMs were not used in any way significantly outside the current norms of academic research.
Reproducibility Statement. We have taken significant steps to ensure that all results presented in this work are reproducible. An anonymous source code repository is provided here containing complete implementations of the Count Bridge framework, including model architectures, training procedures, projection-base based deconvolution, and evaluation pipelines. The appendix includes full mathematical derivations and proofs of all theoretical claims. We also provide descriptions of all data preprocessing steps for synthetic benchmarks, PBMC scquence-to-expresiie transcriptomic aggregation, as well as architectural and hyperparameter specifications. Together, these materials are intended to allow independent researchers to fully reproduce our theoretical and empirical findings.
prediction, and spatial
REFERENCES
Nist digital library of mathematical functions. https://dlmf http:://ddmff..is....om .avt.goov/// 2025. See (10.4 Accessed 2025-09-24.
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.087997 2023.
Erick Armingol, Adam Officer, Olivier Harismendy, and Nathan E Lewis. Deciphering cell-cell interactions and communication from gene expression. Nature reviews. Genetics, 22(2):71-88, February 2021. ISSN 1111-0006,,47110000 doi: 10. 10.108//441376000000022222
Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autoregres- sive and diffusion language models. arXiv preprint DoXiv:2503.0952. 2025.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:1444:9 177791999999 2021.
Pavel Avdeyev, Chenlai Shi, Yuhao Tan, Kseniia Dudnyk, and Jian Zhou. Dirichlet diffusion score model for biological sequence generation. In International Conference on Machine Learning, pp. 1276-1301. PMLR, 2023.
Ziga Avsec, Vikram Agarwal, Daniel Visentin, Joseph R Ledsam, Agnieszka Grabska-Barn Barwinska, Kyle R Taylor, Yannis Assael, John Jumper, Pushmeet Kohli, and David R Kelley. Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, 18 116):1196-1203, 4 October 2021. ISSN 1548-7091, 15537709,,1548-10000 1548-7105. doi: 10. 00.1008//41122-02--001
Fan Bao, Chongxuan Li, Jiacheng Sun, Jun Zhu, and Bo Zhang. Estimating the optimal covariance with imperfect mean in diffusion probabilistic models. In International Conference Machine Learning, 2022a.
Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-DPM: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. In International Conference on Learning Representations, 2022b.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 22669-22679, 2023.
Lucía Gaabadilla-Martiiee Noud Klaassen, Bas Steensel, and Jeroen de Ridder. Predicting gene expression from DNA sequence using deep learning models. Nature reviews. Genetics, 26(10): 666-680, 13 May 2025. ISSN 1471-0056, 1011-117166666 1-0064. doi: 10. 10.108884415776600500081111
11

---

Published as a conference paper at ICLR 2026
Sean C Bendall, Erin F Simonds, Peng Qiu, El-Ad D Amir, Peter O Krutzik, Rachel Finck, Robert V Bruggner, Rachel Melamed, Angelica Trejo, Olga Omatsky, Robert Balderas, Sylvia K Plevritis, Karen Sachs, Dana Pe' er. Scott D Tanner, and Garry P Nolan. Single-cell mass eytometry of differential immune and drug responses across a human hematopoietic continuum. Science (New York, N.Y.), 332(6030):687-663) 6 May 2011. ISSN 1056-8075,1095-9220 doi: 10.1 2.1126/science. 1198704.
Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Amaad Doucet. From denoising diffusions to denoising markov models. Journal of the Royal Statistical Society Series Statistical Methodology, 86(2))28-331, 2024.
Dylan M Cable, Evan Murray, Luli S Zou, Aleksandrina Goeva, Evan Z Macosko, Fei Chen, and Rafael A Irizarry. Robust decomposition of cell type mixtures in spatial transcriptomics. Nature biotechnology. 40(4):517-526, April 2022. ISSN 154--116,,146166990 doi: 10.1038/ 541587-021-00830-w
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Amnad Doucet. A continuous time framework for discrete denoising models. Advances Neural Information Processing Systems, 35:28266-28277 2022.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXhv:2402.04997 2024.
Tianrong Chen, Guan-Horng Liu, and Evangelos A Theodorou. Likelihood training of schr\)" odinger bridge using forward-backe backward sdes theory. arXiv preprint 1221:
Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning. arXiv preprint urXiv:2208.04202. 2022.
Valentin De Bortoli, James Thortoon, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in neural information processing systems, 34: 17695-17709, 2021.
Valentin De Bortoli, Alexandre Galashov, Swaroop Guntupalli, Guangyao Zhou, Kevin Murphy, Arthur Gretton, and Arnaud Doucet. Distributional diffusion models with scoring rules. arXiv preprint arXiv:2502.02488; 2025.
Mauricio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising diffusion for image restoration. arXiv preprint arXiv:2303.1144 2023.
Luc Devroye. Simulating bessel random variables. Statistics & probability leteers 57(3):249-22
2002.
C Domínguez Conde, Xu, B Jarvis, D Rainbow, S B Wells, T Gomes, S K Howlett, O Suchanek, K Polanski, H W King. L Mamanova, N Huang, P A Szabo, L Richardson, L Bolt, E S Fasouli, K T Mahbubani, S Pritchard, T Li, R Elmentaite, J Park, E Rahmani, D Chen, D K Menon, O A Bayraktar, I. K Sims, D L. Farber, Saeb-Parsy,
Prete, L Tuck, N Richoz, Z K Tuong, L Campos, H S Mousa, J Needham,
James, B Meyer, N Yosef, M R Clatworthy, and S A Teichmann. Cross-tissue immune cell analysis reveals tissuc-specific features in humans. Science (New York, N.).), 37666594) 13 May 2022. ISSN 0036-8075,1095-9203, doi: 10.1 12//sciencee.ab51197.
Jones,
Ruiqi Gao, Emiel Hoogeboom, Jonathan Heek, Valentin De Bortoli, Kevin Patrick Murphy, and Tim Salimans, Diffusion models and gaussian flow matching: Two sides of the same coin. In The Fourth Blogpost Track ICLR 2025, 2025.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. Advances in Neural Information Processing Systems, 37: 133345-133385, 2024.
Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477):35--378 2007.
12

---

Published as a conference paper at ICLR 2026
Johannes C Hingerl, Laura D Martens, Alexander Karollus, Trevor Manz, Jason D Buenrostro, Fabian Theis, and Julien Gagneur. scooby: Modeling multi multi-modal genomic profiles from DNA sequence single-cell resolution. bioRxiv.org: the preprint server for biology, pp. 2024.09. 2024.09.19.613755 19.61374, 23 September 2024, doi: 10.1 00.110//2024.09.19.111117
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models, In Advances in Neural Information Processing Systems, 2020.
Peter Holderrieth, Marton Havasi, Jason Yim, Neta Shaul, Itai Gat, Tommi Jaakkola, Brian Karrer, Ricky TQ Chen, and Yaron Lipman. Generator matching: Generative modeling with arbitrary markov processes. arXiv preprint arXiv:2410.20587, 2024.
Ouail Kitouni, Niklas Nolte, James Hensman, and Bhaskar Mitra. Disk: A diffusion model for structured knowledge, 2024, hRttps::://axxt /awxiv.vvivv.//v..//eee/aaaaavva
05253
Allon M Klein, Linas Mazutis, like Akartuna, Naren Tallapragada, Adrian Veres, Victor Li, Leonid Peshkin, David A Weitz, and Mare W Kirschner. Droplet barcoding for single-cell transcriptomics applied to embryonic stem cells. Cell, 161(5):1187-100) 21 May 2015. ISSN 0092-8674,1097- 4172. doi: 211.166/.eel..221..444..
Vitalii Kleshchevnikov, Artem Shmatko, Emma Dann, Alexander Aivazidis, Hamish W King, Tong Li, Rasa Elmentaite, Artem Lomakin, Veronika Kedlian, Adam Gayoso, Mika Sarkin Jain, Jun Sung Park, Lauma Ramona, Elizabeth Tuck, Anna Arutyunyan, Roser Vento-Tormo, Moritz Gerstung, Louisa James, Oliver Stegle, and Omer Ali Bayraktar. Cell2location maps fine-grained cell types in spatial transcriptomics. Nature biotechnology, 40(5):661-671, 13 May 2022. ISSN 1187-0156,1546-16663 doi: 10.1038/<415877021--00
Christian Léonard. A survey of the schr odinger problem and some of its connections with optimal transport. arXiv preprint arXiv: 0215, 2013.
Haoyang Li, Juexiao Zhou, Zhongxiao Li, Siyuan Chen, Xingyu Liao, Bin Zhang, Ruochi Zhang, Yu Wang, Shiwei Sun, and Xin Gao. A comprehensive benchmarking with practical guidelines for cellular deconvolution of spatial transcriptomics. Nature communications, 14(1):1548, 21 March 2023. ISSN 1141-1773,2041-1772
doi: 10. 10-0008//41446--03333716666
Yaron Lipman, Ricky TQ Chen, Heli Gen-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Guan-Homg Liu, Arash Vahdat, De-An Huang, Evangelos Theodorou, Weili Nie, and Anima Anandkumar. I 2 sb: image-to-image schr\"" odinger bridge. arXiv preprint arXiv:2302.0587. 2023.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300) 2022.
Romain Lopez, Baoguo Li, Hadas Keren-Shaul, Pierre Boyeau, Merav Kedmi, David Pilzer, Adam Jelinski, Ido Yofe, Eyal David, Allon Wagner, Can Ergen, Yoseph Addadi, Ofra Golani, Franca Ronchese, Michael Jordan, Ido Amit, and Nir Yosef. DestVI identifies continuums of cell types in spatial transcriptomics data. Nature biotechnology, doi: 00.10388/411577-022-0111
September 2022. ISSN
0087-0156,15461111
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310. 16834, 2023.
Brendan F Miller, Feiyang Huang, Lyla Atta, Arpan Sahoo, and Jean Fan. Reference-fres cell type deconvolution of multi-cellular pixel-resolution spatially resolved transcriptomies data. Nature communications, 13(1):2339, 29 April 2022. ISSN 2041-1722,204111722 doi: 10. 10.10388641147--022-30333:
Lambda Moses and Lior Pachter. Museum of spatial transcriptomics. Nature methods, 19(5):534-546, 10 May 2022. ISSN 1548-7091, 1548-7705. doi: 10. 10122-4-4419222222-440--2222
13

---

Published as a conference paper at ICLR 2026
Aaron M Newman, Chloé B Steen, Chih Long Liu, Andrew J Gentles, Aadel A Chaudhuri, Florian Scherer, Michael S Khodadoust, Mohammad S Esfahani, Bogdan A Luca, David Steiner, et al. Determining cell type abundance and expression from bulk tissues with digital cytometry. Nature biotechnology, 37(7):773-782, 2019.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, 2021.
Shen Nie, Fengqi Zhu, Chao Du, Tinyu Pang, Qian Liu, Guangtao Zeng, Min Lin, and Chongxuan Li. Scaling up masked diffusion models text. arXiv preprint arXiv:2410.18514, 2024.
Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint arXiv:2400.033755 2024a.
Zijing Ou, Mingtian Zhang, Andi Zhang, Tim Z Xiao, Yingzhen Li, and David Barber. Diffusion model with optimal covariance matching. arXiv preprint arXiv:2406. 100000880 2024b.
Giovanni Palla, Hannah Spitzer, Michal Klein, David Fischer, Anna Christina Schaar, Louis Benedikt
Kuemmerle, Sergei Rybakov, Ignacio L Ibarra, Olle Holmberg, Isaac Virshup, Mohammad Lot- follahi, Sabrina Richter, and Fabian Theis. Squidpy: a scalable framework for spatial omics analysis. Nature methods, 19(2):771177, 28 February 2022. ISSN 1548-7091,1548-7155 doi: 10. 10.10388441592-00--0133333
Stefano Peluchetti. Non-denoising forward-time diffusions. arXiv preprint arXiv:2312.1458) 2023.
A Raj, P V D den Bogaard, Scott A Rifkin, Oudenaarden, and Sanjay Tyagi. Imaging individual mRNA molecules using multiple singly labeled probes. Nature methods, 5(10):877-879 21 September 2008. ISSN 1548-709,, doi: 10. 00.1038/mmehh.1555
François Rozet, Gérôme Andry, François Lanusse, and Gilles Louppe. Learning diffusion priors from observations by expectation maximization. Advances Neural Information Processing Systems, 37:87647-87682, 2024.
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems,
36-13018444 2024.
Javier Santos, Zachary Fox, Nicholas Lubbers, and Yen Ting Lin. Blackout diffusion: generative diffusion models in discrete-state spaces. In International Conference on Machine Leurning, pp. 9034-9059. PMLR, 2023.
Xinwei Shen, Nicolai Meinshausen, and Tong Zhang, Reverse markov learning: Multi-step generative models for complex distributions. arXiv preprint arXiv:2502. 13747, 2025.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnad Doucet, and Michalis Titsias. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37: 103131-103167. 2024.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Difffusion schrödinger bridge matching. Advances in Neural Information Processing Systems,
2023.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 2020.
Yuxuan Song, Zheng Zhang, Cheng Luo, Pengyang Gao, Fan Xia, Hao Luo, Zheng Li, Yuehang Yang. Hongli Yu, Xingwei Qu, Yuwei Fu, Jing Su, Ge Zhang, Wenhao Huang, Mingxuan Wang. Lin Yan. Xiaoying Jia, Jingjing Liu, Wei-Ye Ying Ma, Ya-Qin Zhang, Yonghui Wu, and Hao Zhou. Seed diffusion: A large-scale scale diffusion language model with high-speed inference, 2025. URL
https::/erxxls www.siivoooooooooooooooooo
org abss///////// /508 02193
14

---

Published as a conference paper at ICLR 2026
Patrik L Ståhl, Fredrik Salmén, Sanja Vickovic, Anna Lundmark, José Fernández Navarro, Jens Magnusson, Stefania Giacomello, Michaela Asp. Jakub O Westholm, Mikael Huss, Annelie Mollbrink, Sten Linnarsson, Simone Codeluppi, Åke Borg, Fredrik Pontén, Paul Igor Costea, Pelin Sahlén, Jan Mulder, Olaf Bergmann, Joakim I Lundeberg, and Jonas Frisén. Visualization and analysis of gene expression in tissue sections by spatial transcriptomics. Science (New York, N.Y.), 353(6294):78-82, July 2016. ISSN 0036-8075,1095-200
doi: 10.1 126///iincc..aaa240000
Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to dna sequence design. arXiv preprint arXiv:2402.005811, 2024.
Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models, 2023. Jtttts:s:///axxi http:///aawi...ooa//////a///////
org httpa//////..1111/.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger
bridges via maximum likelihood. Entroyy, 23(9): 1334, 2021.
Vizgen. Vizgen Data Release V1.0, May 2021. Title of the publication associated with this dataset:
Mouse Brain Receptor Map.
Xuran Wang, Jihwan Park, Katalin Susztak, Nancy R Zhang, and Mingyao Li. Bulk tissue cell type deconvolution with multi-subject single-cell expression reference. Nature communications, 10(1): 380, 2019.
Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion GANs. In International Conference on Learning Representations, 2022.
Seyhan Yazar, Jose Hernandez, Kristof Wing, Anne Senabouth, M Grace Gordon, Stacey Andersen, Qinyi Lu, Antonia Rowson, Thomas R P Payllor, Linda Clarke, Katia Maccora, Christine Chen, Anthony L. Cook, Chun Jimmie Ye, Kirsten A Fairfax, Alex W Hewitt, and Joseph E Powell. Single-cell eQTL mapping identifies cell type-specific genetic control of autoimmune disease. Science (New York, N.Y.). 3e666519) 8 April 2022. ISSN 0036-8075,1095-9200 doi: 10.1 1el1///ciincc.abf34444
Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Dream 7b: Diffusion large language models, 2025. URL 2508. 15487
org, abs/
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffussio model for text generation. arXiv preprint arXiv:2302.05737, 2023.
A COUNT BRIDGES
A.1 COUNT BRIDGES
To review the setup, we have Xo Po. X1 ~ Pl, where Xo. X1 lie on a discrete lattice (Z+). We derive an unconditional forward process kernel Koj = P(),(Xo) and a family of bridge kernels K Kajo.1 P(X, X,, Xo), S S such that these kernels satisfy the consistency property described [
In this section, we define the model and the associated kernels formally, and show how the Binomial and Hypergeometric distributions emerge in our framework. We prove they satisfy the consistency property thus establishing their correctness as diffusion kernels for count data.
A.1.1 SAMPLING PROCESS
Forward Kernel: Let w [0. - R>o be an increasing cuump-intensigee shape function with w() = 0 and w(1) 1. Intuitively, w(t) describes the proportion of "jumps" that have occurred by time ł for t. E [0, 1]. Further, let X A_ > be intensities for the birth and death processes respectively, We then define the cumulative birth (+) death (-) intensities
A+) X+ w()),
A (t) X_ w()).
15

---

Published as a conference paper at ICLR 2026
We use these jump intensities define the independent (non-homogeneous) Birth/Death Poisson processes, (Bt)kejo.1) and (Di)xe(o.1) respectively, such that:
B, ~ Poi(A- Dr ~ (oi(A-(t))
These Random Variables define counts at t, X, = Xx= + B, D.. giving us the unconditional forward kernel Kejo-
Bridge Kernels: To derive Kajo,r, we first define three more random variables that arise through the Poisson process.
Let d, denote the net displacement at time L, i.e:
de = X₁ - Xo
e [0,
Let N, denote the total number of jumps that have occurred until time t, i.e:
N₂ B, Dr
t [0,1]
Finally, we define M, denoting the number of "slack" jumps that cancel each other out during the process. For example, birth followed by a death results in no net displacement, but result 1 "slack"
jump. By definition:
M, = min( B,, D:)
E [0,
follows from the definitions that any two of the variables define the third, as shown in equation 7
While de is observed, N₁ and M, known and fixed before deriving a tractable posterior distribution for Mr in the next section. We now condition on the points Xo To. X₁ Iz and M,- to obtain:
For the following lemmas, we first assume that M, is
1 (N, (N, + de),
d; 24
N₂
|d + 2m
B₁ =
De = B..
2
We first characterize two properties of (Ne)te(o.1)
Proposition A.1 (Total jumps process and time-rescaling invariance). (No)celo.) is (non- homogeneous) Poisson process with cumulative intensity A(t) (A+ + A-) w((). Moreover, conditional on N₁ = n, the n unordered jump times are ii.d. with cdf
P(T S N, 2 = w(t), EE [0,
hence w is precisely the time-rescaling that makes the jump times uniform on [0, 1].
Proof. Note that N, B, D, is the of two independent Poisson processes, and by Poisson superposition principle, is Poisson process itself with cumulative intensity equal to the sum of the cumulative intensities of the individual processes. The second statement direct result of the Probability Integral Transform.
To compute Kojo. = P(X, = + X, X,), we essentially have to compute (N B.) or (B,, D.). both being equivalent to computing the state X,, given the state X, and (N,, Bt). For now, consider the (0, 1) bridge, i.e K-o.1-
This problem can be visualized in the one one-dimensional case as follows: Given (N1, B,). view the total N1 jumps as a set of N1 points on [o, 1), each labeled +1 denoting birth or -1 denoting death, with the density governed by the function w(-). Note that these labels exchangcable, for both processes have the same intensity function. Thus the goal is to find the number of points to the left of t (N₂), and the number of those that +1s (B.).
From this perspective, counting how many of the N₁ points fall to the left of is analogous to sampling coloured balls from probability w())). giving a Binomial draw. Determining how many of those are births is equivalent to drawing subset of size N, from the original population with B, successes without replacement, i.e a Hypergeometric draw. We formalize this:
with replacement (each jump independently lands before with
Lemma A.2 (Binomial Hyymmiilm--yeeeeemettiiiec structure of Count Bridges). Fix t E (0, 1] and condition on N₁ n. and B, bj. Then:
16

---

Published as a conference paper at ICLR 2026
1. Binomial structure. Each of the nı jumps independently falls before time with probability w((t (Proposition A.I hence
No Ni - Bin(ni, w()).
2. Hypergeometric structure. Since the jump labels are exchangeable, the nt jumps in (0, 8 form a uniformly random subset of the nı total jumps, hence
Be (N) = (11,b1,7e) ~ Hyp(m1, b1, ne).
We can generalize this arbitrary time points. For 0 <
t 1 define 7 wws)/w(t).
Conditional on (N,, B,) = (n4. be), the ne jumps in (0, t] form a non-homogeneous Poisson process with cumulative intensity (t))/((t) so the N₂ jumps in (0, s) satisfy
N₂ Ni = ne ~ Bin( ne, 7),
and, by the same exchangeability argument applied within (0, 4).
B, (N,, B,, N.) (14, ns)) Hyp(me, be, ns),
D, N₂ B..
A. A.1 COMPOSITION CONDITIONAL ON THE SLACK
In this section we prove two elementary lemmas showing that the Binomial and Hypergeometric distributions have compositional property, and then use these results to show that the bridge satisfies the consistency equation under a fixed and known slack. In the next sections, we derive the posterior distribution of the slack and show that the result maintains for the general case when we "mix" the bridge over the slack posterior.
Binomial composition arises from the observation that each of the N N objects in binomial draw are treated completely independent of each other. It states that the following two-step process:
1. Given N balls in an urn, classify each ball as "success" or "failure" independently with probability 8. Let K balls be marked "success".
2. Of those K balls, further classify each ball as "success" "failure" independently again with probability Let L balls be marked "success".
equivalent to the one-step process of classifying each of the original N balls as "success" or "failure" with probability On. We prove this formally in the following lemma.
Lemma A.3 (Binomial composition). Let N E N and 0.n [0, Let K Bin(N, 0) and conditional on K, let L ~ Bin(K,) 7). Then ~ Bin(N, Or).
Proof. We begin by writing the joint distribution directly from the model:
( 0t (1
[N'
Nx~~ (k) (4) (1 ()-
P(K k. L = () = P(K =P((KK))PPPLCCC
())
Apply the standard combinatorial identity
= (M) (N -
N
to obtain the factorization
P(K k, L =) To get the marginal law of L, the joint pmf over all h: 2 l.
(
N
(0r) (1 (1- Nr)
(M--))
(N (N- WW-1)
(0(1-))
k-1
1-8
(N-K)
T-8n
1-0n
N Σ
(7) (6n)
01) N-6 N (N -1) (#(1-)) (
k-t /-ve( (- (N-R)
P(L=0==
on)
11-00
1-07
k
k=t
17

---

Published as a conference paper at ICLR 2026
Let then ranges from to N - L, and
N-t P[L -)
Σ (on') -07) N- (N. (N-6)
m ( (
(011-1)
m
(1-0n)
(N (N-6)-m
1-60
Let
a 0(1 17)
1-
b
On
1 0nj
and factor out all terms not depending on m.
N-f
P - (7) (0n) 0n) N-C
Σ
N
a amp(N-0))0
m
m==
Since a + b 1, the inner sum is exactly the Binomial theorem,
--t N
a m 1 (a b)
m==
Thus
P(L P(L=E)= (7) (1))
which is the pmf of Bin(N,Or)
We now show hypergcometric composition:
Lemma A.4 (Hypergeometriic composition). Let N denote a population of N E N balls an urn, with M N marked "success". Suppose we perform two-stage sampling without replacement:
1. Draw K < N balls from the population without replacement and keep them aside. Amongst those balls, let B S K be the number of balls marked "success". l.e. B ~ Hyp(N, M,
2. From those K balls, draw another L < K balls, and let A A L be the number of successes observed. Le,, A ~ Hyp( Sop(K, B,
Then the marginal distribution of A
ALL M,,((,,,,L
Proof. We prove this by computing the probability observing given L-subset after the two-step procedure, and show that it equals the pmf of a Hypergeometric draw.
To do so, we first show that all L-subsets are equally likely. The problem then simply amounts to calculating the proportion of L-subsets with A successes and L A failures.
Let IC be the subset of size selected after the first draw, and C be the set of size selected after the second draw.
Let Sg denote all possible sets of size K that could have formed C. Formally, S. = {U C c N NUI K). Note that Sc cannot be empty for any given L. Moreover, under sampling without replacement, each set in Sc is equally likely. Therefore, the probability of selecting any sized subset is:
P(C) Σ
Σ11 UES2 K
1
Note that the number of such sets Sc is equal to the number of ways to select the remaining K items (since L is already determined by selecting L) from the remaining N L balls. Therefore:
(N
P(C)
18

---

Published as a conference paper at ICLR 2026
The right-hand side does not depend L, conditional on L every L L-subeet of [1....,NN equally likely.
N3 is
There are
(*) such subsets, among which exactly
N M (M---1) A-A)
subsets contain A successes and L A failures. Hence
(N-M) (2) ( (-a)
P(A - a L ()
which is the pmf of the hypergeometric distribution (op(NN,,D))
We now state composition of the birth-death bridge for fixed slack using the above two lemmas. We show transitioning through the fixed-slack bridge kernel to timepoint & from the starting timepoint 1 yields the same state X, in distribution regardless of the choice of the number of steps taken to reach
Theorem A.5 (Bridge composition conditional on slack). Fix 0 < & < and endpoints X, = To.
X1 31, and slack M, (equivalently, fix (N1, B₁) via equation 7 Let [0, [0, 1] be the time-rescaling function from the Poisson birth-death reference.
Define the fixed-slack bridge kernels
alo,b (Ta (o, P(Xe Fa Xu To. X, Ib, M1 =)),
0 Osasbsl.
Then the two woostage transition (1 + + s) yields the same law for (N.. B&) as the single-stage transition (1 + 8), and the fixed-slack kernels satisfy the bridge-consistency identity
k((m) (xx x0,
Σ 2t
k(() s00,3 (xxx (o, K k((m) (are Fo, T1).
(10)
1|0,1
Proof. We show that the two-step transition (1 - s) is equivalent to the one-step transition (1 + s) for (N,, B.). or equivalently, XX,,, 2a + 2B, N.).
Conditional on (N,, B,) (equivalently on (zo, 11 M1 = m)), we have N₁ jumps with B, "births" and N1 - B, "deaths".
The Binomial structure Lemma A.2 yields
N₂ N1 ~ Bin N1, w()),
N₂ Nr Bin(N, w(s)/w(t)).
Applying Lemma
A.3 with 0 = (f) and r = w(s)/w(t) yields
N₂ N1 Bin N1, w(s))
which is equivalent to the direct (1 + 8) transition.
Now conditioning on the number of jumps at I, No. the Hypergeometric structure described in Lemma A.2 yields
B, Hyp(N,, BI,
(N,, Br, N.) ~ Hyp(N,, Br, N.)
Further, drawing N₂ using the binomial draw and applying Lemma A yields:
(N B1, N.) ~ Hyp(N,, B,, B1, N.)
which is equivalent to the direct (1 + 8) transition.
Thus, at fixed (æo, 21, Mi m). the two-step transition (1 - - 8) and the one-step transition
(1 s) agree law for (N,, B.), and hence for X,- Equivalently, the conditional distributions
<(m)) (x, Zo. 21) = IP(X. = Xo = zo, Xi Iı, Mi m) obey the kernel identity equation 10 which is exactly the bridge consistency with fixed slack.
19

---

Published as a conference paper at ICLR 2026
A.1.3 SLACK POSTERIOR
In the previous section, we established that, conditional on a fixed slack value, the bridge satisfies exact composition.
Note that Nt, Br, Dr, MM are not observed in the sampling process. To obtain observable Markov bridge, we resample the sluck at each time from its posterior distribution, conditional only on observed dz
mm de) = P(M, m de = - z0).
Once we sample Mc. we compute (N,, B,) and apply the conditional BH transition on (Ne, BB)
to obtain (Ng, Bx), as described in Lemma A.2
This subsection derives the explicit form of Tee (e di) and, erucially, proves its closure under the timerrescaling wtt) of the birth-death process.
Proposition A.6 (Bessel slack posterior). Under the Poisson birth-death reference described in section A.I.I fix E (0, 1] and let d, Xt- Xo B, Dt. Conditional d,, the pair (B,, D:)
lies on the lattice
(m (m+ + de, m),
d, 2 0,
(B,, D.) =
m E N,
(m, m (del);
de 0,
and the induced posterior pmf of the slack Mr = min (Br, D:) is (A+ (t)A (t)) (m + [del))m!
Tm(m di) PPMM m de)
0,
(11)
with normalizer
Zi(d,) = (A+( lde 12 I|a,1 (2)A
where I, is the modified Bessel function of the first kind.
Proof. We treat de 2 0, the other case being symmetric. If de = B: - De 2 then B, = De + de.
Mt = min(BB, De) Dj and thus
(B,, D,) (m + de, m)
M, =
By independence of the Poisson processes, the joint pmf at (m + de, m) is
P(B: m de, D: m) - A+ (t) (t) (t) (m + de)!
mn!
Applying Bayes' rule and conditioning on de:
(m d.) = P(B: = m + de, D.= = m)
A+ (t) A (t) ((m de)!
P(d,)
Σ A dd)lr)
Factor (t)dt from numerator and denominator to obtain
(A+ (t) (T))
(A)- ++dd)!!!
(ma (e)
(m + d,) m!!
which exactly equation stated Bessel expression.
The denominator the standard Skellam normalizer, given by the
Li
The slack variable M; represents the "hidden" jumps that cancel out (e.g., birth followed by death), and the Bessel posterior tells us their exact distribution. Notably, such "hidden" jumps are increasingly unlikely with increasing displacement d, due to the factorial growth in the denominator, as observed Figure
We now show that the slack posterior is closed under time-rescaling. In the next section, define the observable bridge as a "mixture" over the slack posterior, and show that it satisfies the bridge consistency equation Jusing the closure property.
20

---

Published as a conference paper at ICLR 2026
Corollary A.7 (Slack posterior is closed under time-rescaling! For any 0 < t,
TT(- da) has the same finctiooal form as Tel |d),
with parameters obtained by replacing (A+ (t), A-)))) by ( (s)). Since A+ Ayw(s),
this closure depends only on the rescaling of w()) and not on any other properties of the process.
Proof. For closure, note that for any & < we have independent thinning
Bg ~ (s){,
D3 ~ Poi(A-(s)),
A+ = ws)
Repeating the same conditioning argument with (B., D.) and d. yields the same functional family equation 11
with parameters (A+(s) (=), (s)) replacing (A, (U)).
A.1.4 COMPOSITION OVER SLACK
We now define the observable bridge as a mixture of the fixed-slack kernels over the slack posterior defined in the previous section.
Definition A.8 (Observable Bridge). Let KaWp) be the fixed-slack bridge kernel,
elo,b (Ta To, v.5 = P(Xa = Ia Xo = To, X, = Tb. M1 = mm),
0 a
and let xe(- de) be the Bessel slack posterior described in A.1.3 Then define the observable bridge in (0,1) at terminal time 1 as
Σ
mi(m di) K k/mn (Ita 21),
di 20.
m==
At intermediate time 0 < t 1, define the observable bridge in (0.t) as:
( zo, Ft)
Te(m de) K
(mm)) (xx I0, zz), de It. 20:
m==
In other words, the observable bridge is the expectation of the fixed-slack bridges over the slack posterior. Using the fixed-slack bridge consistency shown in AA5 and properties of the slack posterior described in A.1.3),
prove the general consistency of the observable bridge.
Theorem A.9 (Consistency of the observable count bridge). Under the Poisson birth-death reference described in section A.1.1 the observable bridge, defined as in A.8 satisfies the discrete form of equation for all 0 < < < 1 and endpoints Eu.I1.
Kopo,1 (x,
Io. x1) = Σ Ks (x Io, xx) Kijo, (xx zo. 20, x1).
(12)
Proof. Fix Tu, IT, d1 = 21 Fo. and an arbitrary bounded test function & on the state space. We prove equation 12 weak form by computing Elp(X,,) Xo Fy X₁ 21] in two ways:
First, we write the expectation as mixture over the slack posterior:
E(((X,)) Xo Ip, X, = x1]
Σ (m di) (i) E(p(X.) | Xo = xo, X1 = Iı, M = m]] (13)
Note that each fixed-slack bridge satisfies bridge-consistency as described in Theorem A.5
e/0.1 (xx 21)
Σ
s/,0 (x., IIo, k/D.) (are 2o, (o)))1).
where K: s/o.1 (I, To. 21) = P(X, = Is
Xo = Iu. X2 = 21, M = m). By taking expectation
over the test function sp, we get:
E(p((,,) Xo To, X1 21, M m)
Σ E(p(X.)) | Xo Zo. X² = .ee M = m] k(() 1 (2z x1)
21

---

Published as a conference paper at ICLR 2026
Plug the above into equation 13 and interchange the sums 2; and m:
(lp((x,) Xo = Jo, X1 = x1]
Σ
Tz (m d1) K K (2re Zo. 01) ())) Xo 20, X² Zt, M m]
(14)
M==
Now consider the first two terms inside the sum over m. Note that they form of the joint likelihood of (M, X)) given (Xo, X1):
T (m d1) K k/ (zt zo, T1) = PMM = m, X2 = Xo = To, X1 = σ1)
Using the definition of conditional probability, we can "flip" this joint density to condition on the intermediate state Xy:
XX, X1) = PPM Xo. X1).
(15)
Once X₁ is fixed, the future terminal state X provides additional information about the slack
jumps that occurred before time t, thereby isolating M, on (0, t). As a result PMM Xy, X is
simply the posterior distribution of M,. following the Bessel form described in Proposition A.6 as
result of the Limerrescaling Corrolary A.7 To summarize:
P(M m X; Ft, Xu 20. X1
x1) = P(M
X1 2t, Xo Ho) = de)
Substituting this back into equation
15 gives the identity:
Ti (m d1), K (m) (x: zo, 21) T (m de), Kojo, (2) 21).
(16)
Plugging the above into equation 14
(Xx( Xo I0, X1 x1
(17)
(m de) ( wo, x (lo((.) Xo X, =x,, 2, mm]
Fe m=0
By the law of total expectation:
(mm de) ((,)(X)) Xo = vo, X, = Fr, M = m] = B(x((x)) Xo = 2o, 20, X₁ = xz).
m
Plugging this into the above equation:
(X,) Xo = zo, X, xi| = Σ
KeualT xo, x1) Ele((.) Xo xo, X₁ = xx.
2t
and the definition of Ksjo,t then gives
[l(((x)) Xo zo, X, xz]= (((s)) Kojo,(10 20).
Putting everything together,
(X,( Xo = To. X1 =x==1]
Σ Σ ((x,) Ksjo,(Zs xos 21) Ka Io x1)
21
Σ ((xx))
zo, K2
K Keo, (0.01))
On the other hand, by the definition K3j0.1.
E(,((x)) Xo Fo. X, 2i p((.)) (x)o, x0,
22

---

Published as a conference paper at ICLR 2026
Since yp is arbitrary, the coefficients of (px,,) must agree for all 2a, that is:
Ka (Its (o. 211 ) = Σ K (T zo, K K2 (1)
which yields equation 12
We have now shown that the observable birth-death bridge satisfies the same bridge consistency property equation I as the Gaussian diffusion bridge, confirming the correctness of the bridge as diffusion kernel.
A.2 LINK WITH SCHRÖDINGER BRIDGES
We now justify the Schrödinger-bridge interpretation stated in Sec. Recall the unconditional birth-death construction
X Xo + B, Dr,
where (Bo)ce(o.1) and (D.) are independent Poisson processes with cumulative intensities (+) w()( (1) = A-w(t), and w0)) = 0, w(1) ) = 1. Let k = VA_A_ denote the (scalar) jump intensity. Then for each zo X. the forward kernel time
has Skellam pmf
21-00
K (21 zo) expl - (1) - (1)] (A (A+(1)
(i)
where L, is the modified Bessel function of the first kind.
Let
Piot (2o, 21) Po (2o) K7jo(21 xo) be the reference joint law of (Xp, X1) induced by the birth-death process and the data prior Po- Over the space couplings
C(po, pI)
{C on X X: C(, - po: C(X,,) -r
consider the entropic projection
Cx E arg min CEC(po.Pi)
KL Pfer))
Fix any C' C(po.Pi) P1) and write expectations with respect C' Ec[.]. By definition,
dC' dy
(T.(C() |ec. log (Xo, X1)
= -C((C) Ec[ log K (X Xo)] + (P),
where H(C) is the Shannon entropy of C and Cx p1) collects all terms depending only the marginals (including (1) and 1c(X1-Xp which are fixed across all couplings with marginals Po.P1).
Using the explicit Skellam form, we can isolate the Bessel contribution: KL (I = (po. p1) Ec[ /x. (x/ HCC)-
To study the limits + oo and K 0+, we use standard asymptotics for I, (NIS) 2025 (10.41
c²
(2/2)² (x) T(u+1) (1 + o(1))
(2)
(1 + o(1)
co,
+ 0+.
2Tz
High-noise limit + For fixed the large- expansion shows that log L, (2x) depends on only through (1/K) terms. Hence, as K + cc,
Ec[log F|x, x. (2m) = (,o, P1) + ox),
where C% (P) does not depend on the choice of coupling C. It follows that
KL Piet)
H(C) (1,, K- OC,
for some constant C' that is independent of C. Maximizing (C)) over C(po. yields the indepen- dent coupling Po 69 Pr, so in the noghenoise limit Cs -> Pn P1-
23

---

Published as a conference paper at ICLR 2026
Low-noise limit + 0*. For small 2, have
log [2(k)) log 2 0(1) as - 0*,
uniformly for integer in any fixed finite range. Applying this with v [X1 gives
Ecl log xx, (2x)) = log 2 - + O(1,, K 0
Substituting into the expression for the KL divergence, we obtain
+ KT.(CII = C + log - +
K k-00
for some constant C' independent of C.
Thus, in the low-noise limit the dominant term in KL(C Pret) is proportional E minimizing the KL over P()).P)) therefore recovers discrete optimal transport with cost [r] -u while the entropy term (C) corresponds to the usual entropic regularization. This proves the characterization stated in Sec.
B LIFTING COUNT BRIDGES TO AGGREGATES
We make central assumptions that enable deconvolution.
Assumption B.1 (Realizability and recoverability). There exists 0 such that for all I [0,
1. Realizability: Qe (ao t,, Po(ao x,,t,2) almost surely
2. Recoverability: The aggregate-to-unit map has local modulus Kite (1): r 0 near 0*,
Dxa(Po(Xo x,, ( x,, x,, 2)) (1) (polao xt, t, 2) Qo(: Xx, t)))
Recoverability means that the aggregate distribution uniquely determines the unit-level distribu- tion-if we know the sum perfectly, we can deduce the summands. This is not always possible. independent, their
Consider the simplest case: if X1, X2 Poisson(A1). Poisson(A2)
is Poisson(Ai + X2). Observing only the sum, we cannot distinguish (Λ₁ = 3, X2 = 2) from (A1 A2 1) both yield Poisson(5). Now if had side information that identified the "component" each Poisson was drawn from we could identify this, but it illustrates the difficulties we will face here.
Recoverability holds when units have sufficient diversity. Formally, it requires that units are condition- ally independent given (x1; 2) and have distinet factorial cumulant signatures eestturrss-sseniallllyy different statistical fingerprints that survive aggregation. For count data, this means units must have different parameters (c.g., different Poisson rates or negative binomial dispersions) that are distinguishable through the covariates z.
Appendix B.2 provides a detailed discussion. practice, this means units should have heterogeneous characteristics captured by foror example, different images associated with the transciptomics is spatial single cell data.
Recoverability faces fundamental limits the number of units G grows. By the central limit theorem, when + oo, the standardized aggregate (Ao - pG)/oG converges to Gaussian regardless of the unit-level distributions. The higher-order cumulants that distinguish different unit configurations vanish at rate (K-G-k-2))/2 for order & > 2, leaving only the mean and variance (Appendix B.3).
This CLT collapse means our method most powerful for moderate G (tens to hundreds of units) where unit heterogeneity is preserved in aggregates. For very large G, additional structure is wicher-either parametric constraints (e.g., unit parameters follow a low-dimensional model), mul- tiple aggregate observations under different conditions, direct observation of some unit-level data.
We illustrate these issues in Fig. in the main text: as the dirichlet concentration increases the group-level mixture weights concentrate. When combined with large group sizes this forces all groups to become identical. This gives clear empirical sense for the limits of deconvolution.
24

---

Published as a conference paper at ICLR 2026
B.1 APPROXIMATELY SAMPLING CONDITIONAL ON THE SUM
The most central part of our deconvolution approach is our projection operation. This appendix formalizes the rescaling operator used Prop. and explains precisely in what sense it approximates the aggregate-conditiona distribution Qa A(X)) ag).
B.1.1
SETUP AND TWO ENSEMBLES
Let X, E 2220 be random count array with (unconstrained) law Po, and define the aggregate statistic
An == A(Xo) E Z
In the main text, we are interested in the microcanonical baardconsstrinll conditional law
Qa( Ao = an) = Pa( A(Xo) = ao), for feasible with Pa ao) > 0,
which is supported on the affine constraint set (xo A(x,) = ao}
key point is that Qa( Ag ag) is not, general, exponential tilt of Po. Exponential tilts instead arise naturally from the canonical (moment-constraineel relaxation, which replaces the hard constraint Au = ag by E[Ao] = ag-
Throughout this subsection, we work with count arrays X. x the aggregate As A(Xo) E 20 is lattice-valued. Hence for any feasible ag with Pa = ao) > O, the microcanonical conditional Po Ao 40) is an ordinary conditional distribution (no measure-zero issue). When defining the deterministic rescaling projection II below, we relax the projection domain to R x to obtain a closed form.
B.1.2 CANONICAL EXPONENTIAL TILT AND FIRST-ORDER EXPANSION
Assumption B.2 (Cramér regularity for the aggregate). The log-mgf
A(A) log Ep. exp Ao} A
is finite open neighborhood 0, twice continuously differentiable there, and has positive- definite covariance
2 A²) Co (Ao) 0. Moreover, V is locally Lipschitz in a neighborhood of
Definition B.3 (Canonical moment-constrainee tilt). For X in the regularity neighborhood, define the tilted law Po. Xo by
d.a.A. (xo) exp(A A(xo) A(A))-
d.
If ao is sufficiently close to == Ep, [Ao] VA()), the canonical parameter X is the unique solution of
VA = 0g:
and Pa.5. is the unique Kl-minimizer of DrL( (- (PP)) over distributions with TE[Ag] ag-
Lemma B.4 (First-order expansion of the canonical parameter). Under Assumption B.2 let 8 a0 Then for [|8]| sufficiently small,
i = + 0(16813).
Proof. Taylor expand VA(A) 0: VA(A) + ΣA R(A) with [[([)]] S / 11122 Solving + EA+ R(A) 5 yields Á = Σ-1(6 36--F R())) and the stated bound.
B.1.3 HOw THE MICROCANONICAL CONDITIONAL RELATES TO CANONICAL TILT We record an identity and standard approximation principle.
25

---

Published as a conference paper at ICLR 2026
Lemma B.5 (Tilting does not change the hard conditional). For any X in the regularity neighborhood and any feasible ag-
Po(- | Ag an) Po.x A an).
Proof. On the event (Ag ) the Radon-Nikodym factor exp(A (A) Ao A(A) is constant, hence cancels in the conditional density ratio.
Lemma B.5 lets us choose the most convenient tilt for analysis of the conditional. The canonical
choice À is particularly useful because under Po.x. the aggregate is centered at the constraint: Ep, [Au]
To connect the hard conditional to the tilt, one needs an additional large-system regime (e.g. many summed contributions, or large aggregate counts) ensuring that conditioning on A dg only induces small correction beyond matching the mean. A classical form is the Gibbs conditioning principle.
Consider sequence of problems indexed by n with independent random arrays X (X(0) Xx. under p(*) and aggregate Ag (m)) = Σg=1 x (elementwise). Assume a uni-
form Cramér condition and local limit theorem for .4"" (m) under the canonical tilt p(*) (where Ån solves (An( = ai and (A) = log [ A (1) H). Then, for any fixed m and any bounded measurable f depending only on (x10) (n)
0,
(m))
E (-44") [1,1 [f]
E plu
[f] -
(n + oo),
where Q- | ag ) denotes the microcanonical conditional.
Essentially, for local observables in a large-sum regime, the conditional behaves like the canonical tilt with parameter X, and Lemma B.4 gives the first-order expansion Ả = --15 + (1 This is the precise meaning of the phrase "admits first-order exponential tilt" in Prop. 4.1 This first-order argument requires large-n justification, which conflicts with our identifiability perspective below, but we provide it as useful perspective on the projection operation here since it may serve blueprint for thinking about reasonable projections in other settings.
B.1.4 THE RESCALING MAP AN ENTROPIC PROJECTION
The operator II in Prop. 4.1 is projection of point onto the hard constraint set. Since points not distributions, we use the standard generalized KL (a Bregman divergence on 12 xxxx
G
D
(d)
DKI [x) =
ΣV 9= d=11
Vg
(d)
log
Ug
(d)
(d) Ug
+ x
with the conventions Oog 0 and log(u/0) 0o for 0.
Lemma B.6 (Closed form of the projection for elementwise sums). Let A(xo) = E, Foo (element-
(d) wise) and assume Eg x g0 > 0 for each d with a > 0. Then the minimizer
TTxx)) arg
min Yo: A(yo)=oo DKL(Yol|xo))
exists, is unique, and is given by the coordinatewise scaling
(d)
(x)
ag
d= 1=.,..,, D.
x (d) Eg
9'0
Proof. The objective is strictly convex in yu on (yo 2 0}, and the constraints are affine, so the
minimizer unique. Form the Lagrangian with multipliers (dd for the D constraints:
(d)
2(y,v) Σ 9.d
yg
(d)
log
y
(d)
_(d)
+ Σ (d)
Σ
Ug 0
(d) (d)
(d)
Tgg
g
26

---

Published as a conference paper at ICLR 2026
(d) //rgg) ) + p(d) 0, hence yo (d) Stationarity gives /L/0yg (d) log(ys /2g E, Ug = =. (d) yields e-uld: 0 " 4g proving the formula.
(1) Dig e-cln Imposing
Since we are then operating count matrices here we additionally introduce randomized rounding algorithm which enables us to exactly sample from a count-valued point on this projected set.
B.2 CONDITIONS FOR RECOVERABILITY
We give conditions under which aggregate supervision can, in principle, identify (unit-level) count distributions from aggregate observations. Throughout
the notation of Sec. 4 Xo
(X10..... Go) -gg Ag == A(Xo), and in the sum case A(xu) Eg Ego. We fix (x,t, and suppress this conditioning the notation when convenient.
B.2.1 FACTORIAL CUMULANT FRAMEWORK
For nonnegative integer- valued random variable define its probability generating function (pgf) Gx (s) = E[sx] for [0, 1). and write the factorial moment generating function (FMGF) as
Mx (t) [[(1 t) Gx tt)
E (-1, 0] (and possibly for small è 0).
Whenever Mx (t) is finite neighborhood of t 0, define the factorial cumulant generating function (FCGF)
K&k
Cx = log (t
tb
where Ka xx(X) (/ (0) are the factorial cumulants. The pgf Gx (hence Cx) determines the law of X uniquely.
Lemma B.7 (Additivity under (conditional) independence). If X and are independent (or condi-
tionally independent given o-field C), then
Cx+r(t)=Cxxle Cx Cy
on any interval where the FCGFs are finite.
B.2.2 ADDITIVE IDENTIHABILITTY OF FACTORIAL-CUMULEE SIGNATURES
Recovering level distributions from the law of sum is not automatic: distinct unit configurations can yield the same aggregate law (e.g., sums of independent Poissons depend only the total rate). To state correct conditions, we separate structural identifiability property from regularity.
Definition B.8 (G-additive identifiability). Let (F(::) v E V be family of real- -al-analytic functions on neighborhood of t = O. We say the family is G-addiivelly identifiable if for any purameters 1h .h,..... V and ,/.,...,'%e
G
F(t: vy) = Σ F(t: v/) (as functionn near 0)
(v....,tG)) and (v!,...,v'' coincide as multisets.
Lemma B.9 (A simple sufficient condition). the family (F(::) th E V is (finitely) linearly independent as analytic functions, i.e., any finite linear combination E Cj (()))) 0 with distinct 3 forces Cj- = 0, then it is G-additively identifiable for every G.
((3))). bring terms to one side and group identical parame- ters. This yields finite linear combination Σ/- ey (s; v5) 0 with integer coefficients ej E Z. Linear independence forces all cy 0, hence matching multiplicities. i.., the multisets coincide. [
Proof. If E ((v))) Eg
For X Poi(A) one has Cx (t) Xt, the family is not additively identifiable: Λit + Aet (X) + A2)t shows only the total rate is recoverable from the sum. Thus recoverability requires families whose FCGEs contain richer (nonlinear) signatures than single linear term.
27

---

Published as a conference paper at ICLR 2026
B.2.3 RECOVERABILITY FROM AGGREGATE LAWS
We now give clean identifiability statement for the sum aggregate under conditional independence. This captures the mathematical core of when deconvolution from aggregates possible.
Theorem B.10 (Recoverability up to permutation from the sum). Fix (x,, t,,t, 2) and suppose that conditional on this information the units X XGo are independent and
Xga (xx, t, 2) Peg. 9 g=1....,G,
where each Pa. is count distribution whose FCGF equals F(; 1b) neighborhood of O. Assume (F(:)): V) G-additively identifiable (Def. B.8). Then the conditional law of the sum
Ao
Σ
X
9=1
uniquely determines the multiset {v1,.... dG} (hence the multiset of unit laws {Po }), up to permutation of units.
Proof. By Lemma B.7
Coo x1,
Cx,o Cx x1, x,, g=1
Σ 1())
Equality in distribution of Ao implies equality of its pgf (hence FCGF) interval, hence equality of the analytic function CA, (-). ). By G-additive identifiability, the multiset of parameters unique.
How covariates make learning possible. Theorem B. the strongest statement one can make from a single sum: because Ag is symmetric in the units, one can in general identify unit distributions only up to permutation. Side information 2 enables learnability once imposes additional structure that breaks this symmetry across groups. Two sufficient structures are:
1. Finite labeled types (composition varies across groups). Suppose each unit has an observed type label 2g E {1,..., K} and Xgo mg share parameters. Then for group with type counts (mı,.... [ nx),
PeL. . so all units of given type
Ca. (t | BK)
Σ nu F(t: 1=1
Across multiple groups whose compositions (m....., ,uk) vary sufficiently (a full-rank design), the shared parameters (vh become identifiable even though within any single group the sum is symmetric.
2. Low-dimensional parameterization by covariates. More generally, assume 3, t (2g) for shared -iwdimmnsionall parameter 7) (e.g., regression model, neural net, or an- other identifiable function class). Then the aggregate FCGF Σg F(;; (xx)) varies with (Z1...., ] if the covariates span the function class and the induced map 7) + Law(A, (Ag =) is injective (locally, in our setting), then 7j is learnable from aggregate supervision.
B.3 LARGE-G G LIMITS: GAUSSIAN COLLAPSE OF AGGREGATE INFORMATION
We explain why learning unit-level structure from aggregates becomes statistically ill-conditioned as the number of units G grows.
Fix (Xe, 1)= Let Xo = (Xio...., Xoo) E & be conditionally independent given (xr, 1, z), with X (x) ~ Pa (- Xx, t, 2),
g g. G,
(e.g., (c( X1, 2,
( x1, 2gg) under model). Let the aggregate random variable be
Ao == A(Xo) E E. and in particular for the sum case A(xo)
Tgo
28

---

Published as a conference paper at ICLR 2026
Write
Ng B[Ag t|,
o = -Var(Ao x,,
Σ Var(X xx,
Assume og > do and Lindeberg (or Lyapunov) condition holds conditional on (x,, 1,2. Then
Ao PG N))
(G + oo).
OG
Moreover, if = -==z
g0 EXgol x,, 2] oo, Benry Esseen bound yields
Ao
PG I x (1,t,2) - (xx)
g= E
El|
go
EX 3
x1,t,
sup Pr #ER
S C
o
for universal constant C.
Let Kg denote cumulants (ordinary cumulants, or factorial cumulants for count data). Under conditional independence, cumulants add:
<o xx,
Σ (x x),
sup, Ni (Xgu x,,t, 2)1 < oo and oz ( then for all k 3.
Kn(Ao xx, k/2) OG--(k- (k-2)/2) of
G+00
Thus, after centering and scaling, the aggregate law is asymptotically Gaussian and is effectively determined by oZ); information about unit heterogeneity enters only through standardized cumulants of size G 2k-k2))
Implication for deconvolution from aggregates. Aggregation is many-to-one map: there typically many distinct collections of unit-level laws {P, 10 { that induce the same (or nearly the same) aggregate law of Ag- The CLT strengthens this ambiguity for large G: once Ao close to Gaussian, distinguishing different unit-level configurations using aggregate observations alone requires resolving the small (k--2) corrections beyond mean and variance. Equivalently, any local "recoverability modulus" that upper-bounds unit-level error aggregate-level error must deteriorate with G unless additional structure is imposed (e.g.. multiple independent aggregates, parametric hing/low-dimensiinnn structure across units, or partial unit-level supervision).
C PROJECTION AND ROUNDING ALGORITHMS
Group rescaling. Algorithm rescales item-level nonnegative values (xx) 1 that each group Gg attains a prescribed aggregate Cy. The procedure linear-time in the number of items and linear in memory the number of groups:
O(B)
Maxtra = O
and is applied in parallel each feature dimension.
29

---

Published as a conference paper at ICLR 2026
Algorithm 5 Group Rescaling to Match Aggregates
Require: item values 24, 2 0 for b I,..., B; groups GI,..., Gai targets Cg 2 0 Ensure: yh 2 with EoEG, Yb = C for all
1: for 1,...,G. G do
2:
5g >-=bbE LoEG Tb
3:
if Sg > then for b E Gg do
4:
5:
yb c rb Cg/Sg
proportional rescaling
6:
end for
7:
else
8:
for E G. do y/ + Gg
9:
P uniform split
10: end for
11: end
12:
end for
13: return {86} }
Randomized rounding. Algorithm 6 independently rounds real value to the nearest integers, preserving the value in expectation. runs in constant time and memory per entry, so applying over B items has
T= OBB)
Mextra 0(1.
Algorithm 6 Randomized Rounding (scalar form)
Require: real value
Ensure: integer E Zzo
1:
[
2: r +-x- x --a
3: sample U Unif 1]
4: U < then
5:
y+a+1
6: else
7: y f-a a
8: end if
9: return
Groupwise exact rounding. Algorithm 7 converts rescaled real values {xt) integers {y6) while exactly preserving each group sum: Dneg, Y/ Cg. Each yb differs from zs by at most 1. The only expensive step is weighted sampling without replacement inside each group. With Gumbel-Top-i sampling the worst-case complexity is
of Σ GGyl log (Go)).
T
Mextra O(B),
and the algorithm is again applied independently coordinates.
30

---

Published as a conference paper at ICLR 2026
Algorithm 7 Groupwise Randomized Rounding with Exact Aggregates (scalar form)
Require: rescaled F3 2 0, groups Gi....: Ga. integer targets C Ensure: integers l/b 2 with ELEG, yb = Cg and l/b. Ihh S
1: for g 1y.1.... , do
work inside group G3
2: for b E G do
3:
06 + [xp]
4:
Tb -- 3b 0b
5:
end for
6:
Ag Ebeg ab
7:
Sg < Cg Ag
b of increments required
8:
if' S 0 then
9:
for E Gg do
10:
3/b < ab
11:
end for
12:
else
13:
sample a subset S C Gg of size Sg without replacement with weights тb for 6 E Gg do if 6 then
14:
15:
17:
t +
else
19:
+
20:
end if end for
21:
22:
end
23:
end for
24:
return {86)
D SYNTHETIC DISTRIBUTIONS
All synthetic tasks use the same base architecture with a 4-layer MLP with 128 dimensional hidden layers. We scale the inputs and outputs in dimension, for example the DFM and CE-CB have x 256 dimensional outputs (since we clip all datasets to use a range of 256 to make for easy tokenization). The energy score models take inputs + noise dim and throughout. We ran all experiments with Adam using both Ir le 3, 2e 4 and present results for the best performing learning rate for each method. We use cosine warmup for the learning rate for 100 steps. For all experiments we use gradient norm clipping to size 1, batch size 256, and train for 500 epochs. For the energy score models, we use exponential model averaging, which is crucial to good performance. Full details are available in the codebase.
noise_din _im 100
For the flow matching we use 0.1 following best practices (we tested larger a but saw large degradations in performance). For the bessel sampler we use VA4A_ 32.
D.1 DISCRETE 8-G B-GAUSSIANS TO 2-MOONS
D.1.1 DATASET
For qualitative evaluation, we adapt the classic continuous "8-Gaussians to 2-Moons" task into a fully discrete, integer-valued setting suitable for count-based flow matching. Each dataset consists of 50,000 paired samples (Io, .1) E z², constructed as follows.
Source distribution (rn). We generate samples from the standard two-moons dataset in IR2 using make_moons with noise level noise 0.1. The moons are shifted to be approximately centered at the origin by subtracting (0.5, 0.25).
Target distribution (x)). We construct an 8-component Gaussian mixture arranged evenly on circle of radius 2.0 in R2_ Each component has isotropic Gaussian noise with variance matching
31

---

Published as a conference paper at ICLR 2026
noise 0.1. A sample is generated by first selecting of the 8 components uniformly random, then drawing from the corresponding Gaussian.
Integerization. Both source and target samples are mapped to the integer lattice by
x + round(clip(2: scale + offset, min va sul, va lue range - 1));
with parameters scale 30.0, offset 80.0, min_value 0, and value.range 196. This procedure ensures that all outputs fall in the discrete vocabulary {0, 1,... 195]², but the scales are chosen so that essentially no values are actually clipped.
D.1.2 RESULTS
We present a visualization the learned trajectories in Fig- and the full details in Table 6 Count bridges achieve uniformly the best performance using the distributional losses, that is the cross entropy or energy scores with the energy score uniformly best.
Table 6: Discrete Moons Results: Noise Two Moons
Method
MMD
W2
Energy
CFM
0.065 + 0.019
0.049 + 0.008
0.874 + 0.246
DFM
0.010 = 0.002
0.010 # 0.002
0.035 = 0.014
Count Bridge (CE)
0,0065 + 0.0023
0.0080 = 0.0009
0.026 = 0.004
Count Bridge (ES)
0.0044 = 0.0018
0.0052 = 0.0007
0.0098 = 0.0029
Count Bridge (MSE)
0.030 = 0.000
0.033 = 0.001
0.366 + 0.015
We also run the Count Bridge across different noise levels, here actually find that our default of A x_ 32 is not optimized, so all results can be considered lower bounds on our performance.
Table 7: Count Bridge (Energy Score) Results Across Different X4 A_ Values
X X_
MMD
W2
Energy
0.0038 0.0012
0.0046 0.0002
0.0075 0.0011
0.0039 0.0015
0.0045
0.0006
0.0080 0.0022
0.0049 0.0003
0.0049
0.0003
0.0095
0.0004
32
0.0052 0.0024
0.0055 # 0.0015
0.011 # 0.006
256
0.0064 0.0020
0.0063 0.0009
0.015 0.004
D.1.3 NOISE TO 2-MOONS FOR DIFFUSION COMPARISONS
Here we compare against a standard Gaussian DDIM model Song et al. (020)) and Discrete Diffusion as in (Sh et al. (2024). Since these models go from noise to target distribution, we cannot do the S-Gaussians to 2-Moons task, so we simply target the 2-Moons. This makes the task substantially casier. For Count Bridges use the same results from the previous table (the more difficult task, but with comparable scores at the endpoint).
Table 8: Discrete Moons Results: Eight Gaussians + Two Moons
Method
MMD
W2
Energy
Count Bridge (ES)
0.0044 + 0.0018
0.0052 + 0.0007
0.0098 + 0.0029
Gaussian Diffusion
0.024 + 0.009
0.017 + 0.006
0.118 + 0.064
Discrete Diffusson
0.017 = 0.010
0.013 = 0.006
0.072 = 0.055
32

---

Published as a conference paper at ICLR 2026
D.2 LOw-RANK GAUSSIAN MIXTURE
D.2.1 DATASET
For synthetic evaluation, we use pre-sampled integer-valued Gaussian mixture dataset that scales with the ambient dimension d while fixing the latent rank at r = Each dataset consists 50,000 paired samples (1) E zd generated according to the following procedure:
Mixture construction. We define a = component Gaussian mixture in latent space R" with 3 (we hold these parameters constant as we scale in d):
Means. Component means drawn from (00, o²1) with scale mean scale/va and shifted to lie near the center the integer range. We set mean sca le = 20.0. Covariances. Each covariance is constructed by sampling eigenvalues from an exponential distribution with scale cov_sca le 10.0, clamped below min_ei genva and conjugating by random orthogonal matrix.
g lue 0.1,
Mixture weights. Weights are drawn from a Dirichlet(I...... prior, yielding a random simplex vector.
Projection to R.. Latent samples E R3 mapped to the ambient space via random projection matrix P E Rddr with entries scaled by projecti on scale/ le/ vr, where projection_sce -nccalle (1 "oisssiisesaalee with noise-scale =
1.0. To avoid degeneracy, isotropic Gaussian noise ~
1.0 is added after projection:
== P
~ MoG,, e ~~ ((0, Ta)-
Integerization. Projected samples E Rd are rounded to the nearest integer and reflected into the bounded range [min va lue, wawueerlllaaaaeee range 1] [0, 255] to ensure validity of DFM.
Scaling in d. The intrinsic latent structure is fixed r = 3, while the output dimension d varied across experiments (e.g. = 5, 16, 32, 128, 256, 512). This construction produces datasets with constant intrinsic complexity but increasing ambient dimension, providing natural test of how models scale in d.
D.2.2
RESULTS
The central scaling results are presented visually in Fig. Here we also present the full experimental details in T
D.3
DECONVOLUTION GAUSSIAN MIXTURE DATASET
We extend the low-rank Gaussian mixture task (Appendix D.2) to evaluate deconvolution capabilities under controlled conditions. Each observation is formed by aggregating group of G unit-level samples into a single count vector.
Group construction. For each group, component proportions drawn from Dirichlet distri- bution with concentration parameter a, yielding group-specific mixture weights. The G unit-level samples are then drawn independently from the corresponding mixture components. Both the ag- gregated sum Xo 2d and the individual unit-level labels 2 E {0, 1 1 are retained, enabling evaluation of methods under both aggregate-only and aggregate+unit supervision.
Experimental variation. We vary two factors that control the difficulty of deconvolution:
Group size: G {4, 8, 32, 128), which determines how many unit-level samples are aggre- gated. Larger groups yield more uniform averages and less information about component heterogeneity.
Dirichlet concentration:
{1, 10, 1000), which controls variability in group-specifis mixture weights. Small a values produce heterogeneous groups (informative for deconvolu- tion), while large values yield nearly uniform group proportions (uninformative).
33

---

Published as a conference paper at ICLR 2026
Table Performance Comparison Across Dimensions and Methods
Dim
Method
NPE
MMD
W2
EMD
0,019 = 0.002
0.716
CFM
0,026 0.17
0.015 + 0.004
0.564 0.460
128
0.026 0018
0.15 0.004
0.565 0.469
0105
128
0.06 a 0.005
0,018 : 0.002
0.759 144
(J053 R.00Y
T,0000 x 0.0INE
0,029 00004
Bridge
0.0054 0.0010
128
0,0988 0.0008
0,058
0.0003
0,055 0.008
UNIT C013
0,23 $ 0.00
1,03 0.18
CFM
0.002
8
0,26 0.007
0.011 = 0.001
0.204 0.067
DFM
52
0.34 0:009
0.111 0.003
0.311 142
128
02222
Count Bridge
0.0038
0.0011
0.0026 + 0.0NNG
0.0077
0J0028
128
0.0050 0.0015
0,0029 0.0003
0,012 0,002
0.11
CFM
0.053 0.011
0017 = 0.001
0,788
128
0.052 0.011
0.015 0.001
0.647 163
0,078 0.001
0.017 0.000
1,20 0.06
128
0.118 0.017
0,025 = 0.004
2,72 0.86
0.0067 0.0014
0,003 : 0,0013
0,025 2 0.007
Count Bridge
0.011 0.001
0,0045 0.0004
0,048 0.007
8
0.145 0.024
0,050 0.001
4.63 1.28
CFM
0.131 0.043
0,066 0.065
3.14 1.72
128
0,22
0079
32
DFM
32
0.089 0.023
0 0.006
1.55 0.85
128
0.100 0.027
0.18 + 0.007
1.99 1.10
(0083 B.OOOS
I
0.021 00002
Bridge
0,031 a 0.0006
128
0.010 0.002
0.034 0.INI7
ILA34 0J807
0.296 0.07
0.42 0.008
15.43 2
128
0.226 0.107
0,019 0.017
17.94 11.55
0.105 0:33
0,02 = 0.007
1.71 1.07
64
DFM
0.126 0.039
0,020 0.005
0,020 0.004
0,0851 0.0005
0,072 0J019
Count Bridge
32
0.827 0.002
0 = 0,0002
0.112 2 0.014
128
0.0065
0.038 0.003
12.89
CFM
32
0.276 0.034
0 0.008
10.50 3.51
128
0.260 0.048
0,32 0.008
10.55 2
DRM
0.236 0.031
0,013 = 0.005
9,06 2.63
128
0.259 0.028
0.050
0.011 0.011
2:92
0.128 B.066
1J014 BJRN
3,50 2.34
Dridge
128
0.151 0.082
0,016 & 0.007
4.55 3.65
0.41 0.007
0.49 0.007
28.45 4.63
CFM
0.006
02116 0.039
B100 0008
0.75 5.07
256
DRM
52
0,228 D.045
0L033 = 0.008
12,33 4.81
128
0.009
Count Bridge
0.092 0.044
0.0099 = 0.0021
1.68 1.30
128
0.105 0.039
0.012 0.001
1,95 : 1.26
CFM
0.471 0.046
0.099 0.005
11.35
128
0.438 0.034
0,488 = 0.005
$3.70 14.23
0.261 0,069
B144 = 0.015
30,49 2 21.88
128
0.319 0,112
0,088 = 0.022
$5,03 35.35
0081 a: 0.028
0,010 0,002
1,46 0.73
Count Bridge
0.085 0.026
0.010 0.001
Dataset parameters. We fix the ambient dimension at d 4. latent rank 3, number of mixture components & = 5, and use the same mixture parameterization as in the low-rank dataset (means scaled by 20.0, covariances scaled by 10.0 with minimum eigenvalue 0.1, projection scale 1.0, isotropic noise 1.0, and integerization into [0, 255)). Each dataset contains 5,000 groups, drawn from base pool of 50,000 pre-sampled mixture samples.
34

---

Published as a conference paper at ICLR 2026
Results. As shown in Fig. deconvolution performance degrades groups become larger and more uniform. This matches the theoretical results in Appendices B.2 and B.3 which establish that deconvolution requires between- group heterogeneity for identification, property that is inherently lost as G grows. We present detailed results across metrics in Tables T0
12
Table 10: Deconvolution Performance: 2 vs Group Size and Dirichlet Concentration
Group Size (n)
a
a = 10
a = 1000
4
0.0091 + 0.0005
0.010 = 0.000
0.011 = 0.000
8
0,011 + 0.001
0.014 #
0.016 0.000
32
0,020 + 0.002
0.023 = 0.002
0,025 = 0.002
128
0.050 + 0.008
0.053 0.006
0.057 0.002
Table 11: Deconvolution Performance: EMD vs Group Size and Dirichlet Concentration
Group Size (n)
a
Q 10
a 1000
0.130 + 0.006
0.195 # 0.025
0.207 = 0.040
8
0.286 + 0.023
0.363 # 0.037
0.483 # 0.010
32
0.530 + 0.051
0.657 + 0.095
0,921 = 0.222
128
2.24 + 0.58
2.22 + 0.54
2.63 0.14
Table 12: Deconvolution Performance: MMD Group Size and Dirichlet Concentration
Group Size (n)
10
1000
0,011 = 0.001
0.011 # 0.002
0.012 0.004
8
0
0.017 # 0.002
0.021 # 0.001
32
0.014
0.020 + 0.003
0.025 0.010
128
0.037 = 0.005
0.045 # 0.007
0.044 + 0.001
To investigate the importance of different rounding approaches our deconvolution implementation we run an ablation study where we substitute our preferred exact rounding approach for two alterna- tives: first a simple deterministic rounding, and second a randomized rounding (where we simply add zero one with probability of the decimal value). Deterministic rounding lead to arbitrarily incorrect results, randomized rounding preserves the expectation, and our exact approach will ensure exactly match the target aggregate value. We find that although our exact approach is superior across three replicates of the n = 128, a = 1 setting the results are very close particularly for the randomized rounding. We believe this could justify substituting the randomized approach since is simpler, although in different regimes this may matter more or less.
Table 13: Performance comparison of different rounding methods with standard errors.
Method
MMD
W₂
EMD
exact
0.037 + 0.005
0.050 + 0.008
2.24 + 0.58 2.33 - 0.43
randomized
0.038 0.000
0.051 + 0.007
round
0.038 + 0.003
0.052 + 0.007
2.32 + 0.44
E NUCLEOTIDE-LEVELL GENE EXPRESSION MODELLING
E DATASET
Nucleotide-level data preprocessing We use the Oneklk peripheral blood mononuclear cells (PBMC) 10X 3' scRNA-seq dataset, originally collected by Yazar et al. (2022). For our analysis, as we are interested in mucleotide-level counts rather than the gene-level counts provided with the initial publication, the preprocessed reads made available by Hingerl et al. (2024). The reads
35

---

Published as a conference paper at ICLR 2026
aligned to the hg38 human reference genome. The resulting BAM files are filtered to include only high quality, UMI-deduplicated reads. The cell type annotations were used as provided in the original dataset.
Gene-level data preprocessing We construct the gene-level count matrices directly from our single- cell coverage matrices rather than following the typical single-cell gene expression preprocessing pipeline. In particular, for each annotated gene and each cell, we take the max count over the nucleotide-level coverage matrix as the count for the gene.
E.2 ARCHITECTURE, TRAINING, AND INFERENCE
E.2.1
INPUTS AND EMBEDDINGS
We model nucleotide level counts on a fixed window of length L=896. For each example we form
CE (I....,C), seq E (,, Sequence context is embedded with a frozen Enformer encoder (EleutherAl checkpoint), yielding per-position embeddings
It
2 E (J, ~~0 Ia,),
E(seq) E R RLXdE
dg = 3072.
We tile the scalars across positions and concatenate
(0)
[((eg)
|| emb(c)) (x
with d, 100 and d.=14. A two-layer SELU MLP projects to the model width d:
x dW2 ((W, E gd
E.2.2 LOCAL ATTENTION BACKBONE
We apply Nau residual self-attention blocks (PyTorch Multthheadhttentiion,, batch-first) with LayerNorm:
( MHA(X(0) X(0 x()), ((+1)) LN(X x
where the residual skip uses the pre-block (0) as in the implementation We use Nam =2 lay- ers,
and h=4 heads. A linear projection followed by softplus produces nonnegative per per-position prediction
to softplus( W X
B
This parameterizes the conditional law q0(- Ft,t, 2, e, seq) used inside the count-bridge reverse kernel (Sec. 3)
E.2.3 LEARNED PROJECTION MODULE TT
When aggregate constraint a-= Σ vo.s is observed, we refine do with a lightweight attention projector that operates across positions. We form
(()
[ æt || ao (Naaa) TR (R
A two-layer SELU MLP lifts to width d, then Npmj=2 self-attention layers (sequence-first API) with residual+LayerNorm are applied:
(mm) = (m)(Y(()), (()) (m))), ((++1) = LN(Y (m)) + y())..
linear head produces an additive correction which we re-softplus for nonnegativity:
İo = softplus W proj (NNaa)) + 2p.
At inference, when ag is present we use To as the endpoint in the reverse step; otherwise we use dg.
?Code uses "global" residual X X (0) within each block. We retained this because it stabilized training with T.=896.
36

---

Published as a conference paper at ICLR 2026
E.2.4 TRAINING OBJECTIVES AND SCHEDULES
Distributional loss (energy score). We train do with the energy score Sp the conditional
Xo X;=z,; (Sec. 3.2p. For each example we draw m i.dd samples a(1) ~ qo( Iz. t, 2, c, seq) via ancestral decoding of the per per-position parameterization and estimate the U-statistic version of Sp with p(r, -')= (0,).. used m=2 in practice.
Aggregate-awand training. With probability Page =0.1 we attach an aggregate 4g and route the forward pass through Hl to obtain žo, then compute the same energy score. This jointly trains Hl to approximate sampling from the mean-conditional Xo A(Xo)=ao X2, while preserving the exact reverse transition of the count bridge.
Cell-type masking. To support both conditional and unconditional generation, we randomly mask the cell-type embedding with probability Pmask (set to zero vector). We used P0..
E.2.5
OPTIMIZATION AND IYPERPARAMETERS
We use Adam, learning rate 12x10-4 cosine warmup for 100 steps, EMA with 0.999, batch size 128, gradient clipping at 1.0. See configs for exact architecture specification.
E.2.6 SAMPLING
At test time we follow Alg. starting from 21 we iterate t. + At each step we sample Xn ~ go Tte tk, 2, seq): if an aggregate is provided replace with =IT (ăo, an, Ita ). We then apply the exact binomial-hyper hyyrreeotppege reverse kernel (Prop.(..I) to obtain It.-,- This guarantees that trajectories remain within the discrete support while leveraging the learned distributional posterior. We use three function evaluations for all results in this application.
E.3
ADDITIONAL RESULTS
In Tab.
E.3
we provide results for gene expression prediction performance, broken down by cell type.
Cell type
Baseline MSE
CB MSE
CD4 ET
3.596
1.402
NK
0.415
0.364
CD4 NC
3.382
1.304
CD8 $100B
2.619
1.002
CD8 ET
1.065
0.540
IN
2.556
1.091
CD8 NC
3.381
1.311
Mem
6.742
3,416
NK R
1.624
0.781
Mono NC
1,485
0.752
Mono C
1,253
0.676
DC
9.302
4,475
Plasma
10763.906
10696.934
CD4 SOX4
3.428
1.323
We also analyze the unit-level profiles of the deconvolved transciptomes. We aggregate the nucleotide- level transcriptomes up to the gene level by computing the maximum count over the gene profile, enabling us to generate a (standard cell by gene) count matrix from our deconvolved profile. We then plot the UMAP of the held-out ground truth alongside the deconvolved transcriptomic profiles, and find that the cells are well- mixed (Fig. 5). suggesting that generated expression distributions match the true distributions.
37

---

Published as a conference paper at ICLR 2026
Gound PVP
Figure 5: UMAP of true single-cell and CB-deconvolved mucleotide-level expression, aggregated the gene level and hued by cell type (as assigned by Yazar et al. 2022)
F SPATIAL TRANSCRIPTOMIC DECONVOLUTION
F.1 DATASET
Preprocessing We used the publicly available mouse brain MERFISH dataset from Vizgen 2021 We subset the data to particular slice (slice I, replicate 2). We used the transcript puncta and nuclear segmentation masks as provided with the dataset. For gene expression, we used the raw transcript counts without applying standard single-cell preprocessing pipelines. For each cell, we resized the DAPI image 256x256 pixels by padding.
Aggregation To simulate a Visium-style spatial transcriptomics dataset, we aggregated the single- cell MERFISH data. A grid of spots defined with center-to-center distance of 100pm and spot radius of 55pm. The gene expression profile for each simulated spot was then generated by summing the transcript counts of all identified cells whose nuclei fell within the circular bounds of that spot.
F.2 ARCHITECTURE, TRAINING, AND SAMPLING
F.2.1 INPUTS AND TOKENIZATION
We model spot-level counts while conditioning on image context and diffusion noise/time tokens. Each training example provides
cml
D.
I, CxCHW (0,1,, EER y E (l,...,C.)) (optional).
Images are patchified by a ViT-style embedder (BatchEmbed) into
Ximg E R Bmxxx Nimg = (H/P)
while counts are converted into small of count patches using a learned projector (CountPatchEnbedding)
X = reshape MLP((a [B, Nemt; d) + Fect:
with Nen learned "pseudo-patches" and Een learnable positional embeddings.
We form auxiliary tokens for time, noise, and (optionally) class:
TimeMLP TTimeMLP(timestee timestep emb(tt) Bxixd
NoiseMLP((We) Bxixd
time
and, labels are used, Emb(y) R BxIxd Concatenating all tokens,
(0) = (; n; X ximg; X + Epus E Bx(Nau Nouteextrs) xd
with single learned positional table Epos covering all tokens.
F.2.2 U-VIT BACKBONE (FUSION AND DECODING)
We process (0) with a U-Net-style VIT:
O ((11)) mid Block
(1/2)
(L)
encoder skips)
decoder (with skips)
38

---

Published as a conference paper at ICLR 2026
where each Block is standard MHA +MLP transformer block with LayerNorm, and decoder blocks attend over skip connections. A final LayerNorm yields Xout E I Bx Na +extrss))) We then drop the auxiliary tokens and split modalities:
X xing = [x img range, :],
X
X x ent range,
Count decoder. Count patches are decoded back to a vector via a small MLP head with nonnega- tivity enforced by Softplus:
to
Softplus (flatttattenxxx
BxD=
This parameterizes q zt It E, y) for the distributional loss and the reverse count-bridge step.
F.2.3 TRAINING OBJECTIVE AND USAGE
train the model to predict the distribution Xo (counts) given multimodal context under the bridge (Xx):
( -x,,(xoxx)) Sp X. I,,t,,,), Xxx.)
using the energy score Sp with p(z, x²) ||x-='1l (5 (0, 2)) and the standard unbiased U-statistic estimator with samples from qo. Time and noise tokens implement the distributional diffusion conditioning: label tokens (if present) enable class-conditional modeling. During reverse sampling
we draw Xu go(- ze2tt I,, t, e, v) and update Ii-A using the exact binomial-hye hhymmiallypeegemettriee
count-bridge kernel (Prop. 3.1
F.2.4 IMPLEMENTATION SPECIFICS
Pabbed uses patch size P on I, (channels C), producing Nimg (H/P)(W/P) tokens of width d. chuntPatchEmmeeeeee counts to Nent tokens width with learned positional embeddings.
Patchification.
projects De-dimensional
Auxiliary tokens. Time token: timestep.embeddee ng followed by linear or ML.P projector (timme, Reime_diin dimm controls concatenated components); noise token: linear to d then 2-layer SiLU MLP; label token: lookup embedding used. All tokens share single learned positional table.
Backbone. Depth L with T./2 encoder and L/2 decoder blocks; each block uses d- dimensional embeddings, h heads, MLP ratio r. LayerNorm, and (optionally) gradient checkpointing. Decoder blocks accept the matching encoder skip.
Heads. Count head: 2-layer GELU MLP over the concatenated count tokens, ending with Softplus. Image head exists but is ignored for the loss.
No weight decay. We exclude token positional tables and count-positional embeddings from weight decay, following ViT practice.
F.2.5
OPTIMIZATION AND HYPERPARAMETERS
We use Adam, learning rate 42x10-4, cosine warmup for 100 steps, EMA with 0.999, batch size 128, gradient clipping 1.0. See configs for exact architecture specification.
F.2.6 SAMPLING
At test time follow Alg. using the aggregates to ensure sampled 2o exactly match the target sum at each intermediate time. We then apply the exact binomial hommial--ypeggommetii reverse kernel (Prop. B.D to obtain Zt,-1"
F.3 ADDITIONAL RESULTS
Comparison with reference-based methods Count Bridges and STDeconvolve are both reference- free methods: that is, they require only aggregate level data, and do not need reference dataset of unit-level measurements. Many deconvolution methods, including RCTD (Cable et al. (022) require single-cell (non-spatial) reference dataset. These methods benefit from unit-level observations, and
39

---

Published as a conference paper at ICLR 2026
Numberat cel Sypes
bpes
Spraman Number of
0.8
0.1
ACT0
Erigge
0.12
Coant Ereg:
PCTD Count Fridge
0.14-
14
0.15
41-
000
23
25
20
25
30
25
35
25
25
20
25
25
(a)
(h)
(c)
Figure 6: Performance of RCTD, STDeconvolve and Count Bridge on MERFISH deconvolution across number of cell types using (a) Jensen Shannon Divergence (JSD), (B) RMSE and (C) Spearman Correlation
Figure 7: UMAP of true deconvolved cell profiles using the Count Bridge.
such solve a more constrained problem but require additional reference data which may not be available in all settings.
Using the MERFISH benchmarking setup also evaluate RCTD, and tabulate the results in Tab. [4 For evaluation, we use Jensen shannon Divergence (ISD) and RMSE metrics as described in Li et al. (2023). We find that Count Bridges perform similarly to RCTD (with higher JSD but lower RMSE), despite the fact that Count Bridges do not have access to reference dataset.
Method
JSD
RMSE
Spearman
STDeconvolve
0.288
0.177
0,255
RCTD
0.161
0.113
0,580
Count Bridge
0.229
0.110
0.332
Table 14: Cell-type deconvolution error for spatial transcriptomic data. Note that RCTD requires single-cell reference dataset for deconvolution.
In Fig. 6 we show the performance of spatial deconvolution methods across varying numbers of cell types. These results evaluate only the recovery of cell type proportions, and do not evaluate full count profiles. Note that RCTD has access to the single-cell level reference data, while STDeconvolve and Count Bridges are fit entirely using aggregate- leeeel data and do not have access to single cell counts.
Inspection of unit-level data generated by Count Bridges In the previous section, we have shown through quantitative metrics that Count Bridges outperform alternatives for reconstructing unit-level gene expression vectors from spot-level aggregates. We next aim to evaluate the extent to which reconstructed gene expression profiles are biologically meaningful. We do this by performing con- ventional single-cell transcriptomic analysis on the synthetic unit-level expression vectors generated by Count Bridges.
First, we plot the UMAP of the expression profiles of the true and generated data in Fig. 7 In this UMAP, we see that the true and deconvolved synthetic data are largely mixed, suggesting that the generated counts realistic and match the distribution of the true counts.
40

---

Published as a conference paper at ICLR 2026
Number of cals annotation
1
reeall
D0P
107
Number of onlll conure bridge groemoed))
Figure 8: Abundance of putative cell types (automatically annotated by celltypist) in synthetic unit- level data generated by Count Bridges through deconvolution, vs. abundance of putative cell types in true unit-level data.
Next, we perform cell type annotation for the generated counts. We preprocess by normalizing counts to 10a per cell (row-normalizing), followed by a log-transform, then annotate cell types using celltypist (Domínguez Conde et al. 2022). From this process, we identify 268 putative cell types in the synthetic unit-level data. The same pipeline, when applied to the real single-cell level measurements, identifies 278 putative cell types. And as shown in Figure 8 the cell type abundances inferred by Count Bridges (without access to unit level data) closely align with the cell type abundances observed in the true unit level data.
Deconvolution of a 10X Visium dataset We next evaluate the deconvolution of spots in a real world 10X Visium dataset profiling the mouse brain. As ground truth is unavailable in this setting, validate model predictions by assessing the extent to which the synthetic deconvolved data reflects known biology.
We use the 10X Visium fluorescence dataset distributed by Palla et al. 2022 which profiles coronal section of mouse brain. To demonstrate the generalization capabilities Count Bridges, we apply the model trained on MERFISH data directly to this Visium dataset without retraining.
In order to correct batch effects and align dimensionality between datasets, we employ moment- matching procedure. For each gene in the MERFISH data, we compute the mean and variance of expression and identify the gene the 10X Visium data that most closely matches these moments. We then map the 10X Visium count matrices to the MERFISH feature space by subsetting to these matched genes.
We deconvolve the 10X Visium data using Count Bridges with a spot-level mean constraint. To evaluate prediction quality, we use a standard single-cell analysis pipeline (as described above) and determine putative cell type annotations using Celltypist (Domínguez Conde et al., 2022). Celltypist identifies 146 putative cell types, suggesting that the synthetic unit-level data recapitulates a significant degree of cell-to-cell variation. Furthermore, the recovered cell types are biologically consistent: the most abundant identified cell type is the oligodendrocyte, which matches the most abundant annotation in the MERFISH mouse brain dataset described above.
41


Reviews:
[{"review_id": "7AxZugG8pD", "reviewer": "Reviewer_K7ei", "summary": "This submission proposes Count Bridges, a novel approach to finding a probability path between end-marginal distributions on integer values. This is partly achieved by modeling what resembles the diffusion term in an SDE as a difference of independent Poisson processes. \n\nUsing an EM algorithm, the method is extended to the setting where the observations obtained from one of the marginals is an aggregate of samples, which then is used to solve important problems in biology. There are multiple real and synthetic data experiments in the paper.", "questions": "* How do you choose $\\Lambda_{-}$,  $\\lambda_{-}$,  $\\Lambda_{+}$ and  $\\lambda_{+}$?\n* Could CBs be used to find a bridge between pairs of bulk data distributions?\n* Could CBs be used to find a bridge between pairs of unit-level data distributions? Where the unit-level counts are inferred using your EM algorithm?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "This is one of the best papers I have read this year. The only reason I did not reward the submission with the highest presentation score is because some concepts were not explained/taken for granted (I discuss this in Weaknesses).\n\nThe originality and quality of the paper is top-class, especially considering the very clever use of the EM algorithm to solve the deconvolution problem. Furthermore, modeling probability bridges directly on the space of count data is of great importance in genomics and transcriptomics. Needless to say, Count Bridges can have a substantial impact on the field of computational biology *and* the method is novel, non-trivial and elegant. \n\nAlthough the results are missing error bars and the quantitative analysis of the results could be more thorough, Count Bridges are evaluated on multiple important biological tasks, outperforming for example a fine-tuned version of the DeepMind produced Enformer. \n\nCount Bridges will surely be an appreciated addition to the class of generative models by the ML community.", "weaknesses": "### Distinction to SBs\nWhy should Count Bridges not be considered as an instance of Schrödinger Bridges? SBs are not constrained to continuous measures, and there is no statement about why CBs are not SBs (there is instead a note on the connection between SBs and CBs in line 172).\n\nI think this requires some clarity to avoid the risk of artificially distancing CBs from SBs in order to promote novelty.\n\n### Clarity\nSurprisingly, I was mainly concerned with the clarity in Sec. 3. I found that key concepts were not defined. \n\n* I found it quite frustrating that $\\Lambda_{-}$,  $\\lambda_{-}$,  $\\Lambda_{+}$ and  $\\lambda_{+}$ were not defined. The definitions/models for these should be in the main text.\n\n* line 200: \"We employ the standard U-statistic estimator (Gneiting & Raftery, 2007; De Bortoli et al., 2025).\" This is not informative: do you mean that you use the weighting scheme as in De Bortoli et al. (2025)? They do not refer to their scheme as standard or a U-statistic **estimator**, nor do they reference Gneiting and Raftery so I could not figure out what the above sentence implied. \n\n*  \"Let $P^{\\kappa}_\\text{ref}$ be the joint induced by the birth-death kernel.\" What is this kernel? It has not been defined at this point?\n\n* The distributions and their parameters pop up in Proposition 3.1 without explanation of where and how they were derived. I found this information in the appendix, but the reader needs to be guided there.\n\n### Related work\nThis is a minor, but I think DestVI might be worth mentioning in the context of ST deconvolution methods. Otherwise I believe the literature review is very comprehensive. Please double check how citations are delivered here, for instance the references between lines 309-310 after (CTMC) should be in parentheses. \n\n### Experiments\nAre the results averaged over multiple runs? Or are these single runs? For a stochastic algorithm like Count Bridges I would expect to see error bars and reported averages.\n\n### Typos\n* line 45: adresss\n* line 170: \"CUDA kernel implementing (Devroye, 2002)\". Either it should be **implementation** or described what in Devroye 2002 has been implemented.\n* line 173: Léonard (2013) should be in parentheses. \n* JDS and RMSE are not defined properly. \n* line 360: The sentence \"STDeconvolve Miller et al. (2022)\" is on the loose.\n* In the Applications section there is inconsistencies in how CBs are referenced: count bridges (line 366), CB, Count bridges and Count Bridges. \n* The introduced abbreviation CB is only used in Sec. 6.1?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wiwo8ahR0e", "reviewer": "Reviewer_S8Cu", "summary": "This paper proposes Count Bridge, an integer-native diffusion framework for count data that replaces Gaussian additive noise with two Poisson processes (birth/death). This yields sequential ±1 jump dynamics that preserve integer structure. The denoiser is distributional and trained with a strictly proper energy score suited to count-space geometry. The authors derive closed-form intermediate conditionals (Binomial / Hypergeometric / Bessel components) and implement principled local-bridge sampling. They extend Count Bridge to deconvolution by treating unit-level X0 as latent and using a generalized EM scheme: a projection-guided diffusion E-step (the projection is applied at each reverse timestep to guide local-bridge sampling under the aggregate constraint) and an M-step that trains the model via an aggregate-level energy score loss. Empirical results on synthetic data and several real deconvolution tasks (single cell -nucelotide expression, bulk RNA data deconvolution, spatial spot deconvolution) show competitive performance versus baselines such as CFM and DFM or Enformer, CIBERSORTx, STDeconvolve.", "questions": "1.\tEnformer-like models are typically valued not just for prediction accuracy but as tools to study sequence-to-expression regulatory mechanisms (e.g., variant effect prediction, attribution maps, motif/TF signal localization). The manuscript reports that Count Bridge outperforms a fine-tuned Enformer on cell-type-specific expression prediction in single-cell data. A key question for me is whether Count Bridge can support the same kinds of downstream regulatory analyses that make it useful in genomics.\n2.\tIn the spatial transcriptomic deconvolution task, I strongly encourage the authors to validate their spatial deconvolution results on genuine spatial transcriptomics data. Specifically, please evaluate the method on at least one sequencing-based platform (e.g., 10x Visium) rather than only on spots synthesized from MERFISH single-cell imaging. \n\nThe reviewer wrote the review. LLM was employed only to correct grammatical errors.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1.\tIntroduces an integer-native diffusion process (birth–death Poisson kernel) that naturally preserves counts.  \n2.\tDerives analytic intermediate conditionals (Binomial / Hypergeometric / Bessel), which support principled local-bridge sampling.\n3.\tPractical deconvolution pipeline: Gives a workable EM-style approach (projection-guided diffusion + aggregate-level loss) to infer unit-level counts from aggregate observations. \n4.\tBroad evaluation: Tests on synthetic and multiple real-world deconvolution tasks, with comparisons to relevant baselines.", "weaknesses": "1. Identifiability and aggregation scale. Pure aggregate supervision is intrinsically ill-posed; performance can degrade as group size increases or between-unit heterogeneity decreases. Please provide quantitative sensitivity analyses showing performance vs. group size and vs. within-group heterogeneity (e.g., varying variance of unit distributions), and explain the practical limits (aggregation scales)  where the proposed EM is reliable. \n2. For the nucleotide-level gene expression modeling task, the authors compare to Count Bridge, which reportedly improves on the fin-tuned Enformer’s direct sequence to cell type-specific expression. I observe that the MSE for the 'plasma' cell type in Appendix E is several orders of magnitude higher than for other cell types. This is concerning and requires clarification. \n3. Please include an ablation over projection and discretization choices (KL rescale + multinomial, rounding/min‑distance, learned Πψ, etc.) and report sensitivity of downstream metrics.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gAMz6AzM4C", "reviewer": "Reviewer_xPQ5", "summary": "This paper introduces Count Bridges, a novel stochastic bridge framework for modeling integer-valued count data, extending diffusion-style generative modeling to discrete domains such as single-cell RNA-seq. The method formulates a birth–death bridge process that connects two count distributions through continuous-time dynamics, allowing for closed-form conditional laws and exact likelihood estimation. The authors further propose an EM-style learning scheme that enables deconvolution of aggregated transcriptomic data—e.g., recovering single-cell distributions from bulk or spatial measurements. Experiments on synthetic and real datasets show that Count Bridges outperform existing discrete flow and diffusion models in both modeling fidelity and biological interpretability.", "questions": "- Many single-cell datasets are zero-inflated. Does Count Bridges handle this naturally via the birth–death process, or is special preprocessing required?\n- How does runtime and memory scale with dimensionality (number of genes) and group size in aggregate deconvolution?\n- Can the learned bridge latent representations be linked to biological factors (e.g., cell states, differentiation trajectories)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Count Bridges constructs a birth–death bridge with a closed-form conditional distribution, enabling precise likelihood estimation for integer data. The EM-style extension for transcriptomics holds practical value: it provides an actionable workflow for deconvoluting single-cell distributions from aggregated observations derived from bulk or spatial sequencing. Experimentally, the method demonstrates robust improvements over relevant discrete baselines across synthetic and real-world data tasks presented in the paper. The paper is thorough in both mathematical derivations and implementation details, ensuring the work combines theoretical rigor with practical applicability.", "weaknesses": "- While being compared against CFM, DFM, and some biological baselines, a direct comparison with other recent count - specific or general discrete diffusion models (beyond Blackout Diffusion) on the proposed tasks could provide a more complete picture.\n- The computational complexity of Count Bridges versus discrete diffusion or flow models is not reported; scalability to large transcriptomic datasets is unclear.\n- While deconvolution results are promising, the paper provides limited biological case studies (e.g., linking recovered cell states to known pathways).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WmpVd2Va3y", "reviewer": "Reviewer_F1ZA", "summary": "This paper introduces Count Bridges, a stochastic bridge framework for modeling and deconvolving integer-valued count data, particularly motivated by RNA-seq and spatial transcriptomics. The proposed model is based on a Poisson birth-death process on the integers, yielding closed-form conditionals that enable exact sampling and tractable training. The framework supports both generative modeling of count distributions and aggregate-to-unit deconvolution via an Expectation Maximization procedure treating unit-level counts as latent variables.\nThe authors demonstrate applications to synthetic benchmarks and two biological tasks: deconvolving bulk RNA-seq data and inferring single-cell profiles from spatial transcriptomic spots. The method achieves strong empirical performance, outperforming baselines such as flow matching, discrete flow matching, CIBERSORTx, and STDeconvolve.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Novelty and conceptual clarity: The introduction of a bridge-based model tailored to integer-valued data is original and addresses an underexplored area between diffusion-style generative models and biological count modeling. The stochastic birth–death formulation provides a principled way to interpolate between discrete distributions.\n\n- Mathematical rigor: The theoretical development is clear and well-supported, including proofs and connections to Schrödinger bridges and optimal transport. The paper demonstrates careful thought in defining the discrete bridge process and the distributional training objective.\n\n- Relevance to biology: Extending generative modeling to integer-valued transcriptomic data and demonstrating utility in deconvolution tasks is timely and important for large-scale single-cell and spatial genomics.\n\n- Empirical results: The experiments convincingly show both generative and deconvolution performance, with clear benchmarks and biological relevance. The model’s ability to outperform reference-free and even reference-based baselines is notable.", "weaknesses": "1. Presentation and flow\n\nWhile mathematically solid, the readability could be improved. Some definitions (e.g., lines 130–140 in the original text) appear abruptly, and notation (such as $A(X_0), Π𝜓$) could be introduced more gently. Reordering to first build intuition before formalism might help readers from outside the generative modeling community.\n\n2. Clarity of deconvolution setup\n\nThe biological deconvolution experiments are appealing but could use more methodological clarity. For example:\n\n- How exactly are unit-level latents initialized or sampled during the EM procedure?\n\n- How are aggregates simulated from real datasets (e.g., in MERFISH) and how sensitive are results to the aggregation scheme?\n\n- What metrics are used to evaluate count-level accuracy beyond JSD and RMSE?\n\n3. Comparative baselines\n\nWhile the model outperforms selected methods, the comparison to continuous diffusion models (e.g., Gaussian-based or Poisson-approximated variants) could be expanded, since this would better position Count Bridges as a discrete alternative rather than a purely new category.\n\n4. Experimental depth\n\nThe evaluation focuses mainly on summary-level results. Including visualizations of reconstructed single-cell profiles, or ablations on the projection module $Π𝜓$, would strengthen the empirical case.\n\n5. Theoretical framing\n\nThe connection between Count Bridges and Schrödinger bridges is very interesting but only briefly touched upon. Expanding this connection, perhaps through an interpretation in terms of entropy-regularized optimal transport, would make the theoretical section more accessible and conceptually richer.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LaamvuTSjt", "reviewer": "Reviewer_TMxD", "summary": "The paper introduces Count Bridges, a generative framework for modeling and deconvolving integer-valued count data. The core of Count Bridges is a stochastic bridge process on the integers, defined by a Poisson birth-death mechanism. The model is trained using a distributional scoring loss, which is necessary to respect the discrete geometry of the data, as opposed to simpler point estimates used in continuous diffusion. For deconvolution, the authors extend this framework via an Expectation-Maximization (EM) algorithm that treats unit-level counts as latent variables: the E-step uses a projection-guided sampling process to impute unit-level data consistent with an observed aggregate sum, while the M-step trains the model by evaluating its performance at the aggregate level. The authors demonstrate the utility of their method on two large-scale biological applications: predicting single-cell gene expression from DNA sequence (outperforming a fine-tuned Enformer model) and deconvolving both bulk RNA-seq and spatial transcriptomic spots into single-cell profiles, benchmarking against methods like CIBERSORTx and STDeconvolve.", "questions": "- What type of metric $\\rho$ do you use in the count space?\n\n- How does the experiment for RNA prediction from DNA work? What I understand is that, for every base pair on the DNA context, you condition noise-based generation of counts with the genomic context obtained via Enformer. Is it true? How do you run Enformer then, and how are the two settings comparable?\n\n- You mention that you are using images to condition the bridge parameterization in the spot experiment. Did you try ablating the conditioning on the images and evaluating the performance after this? I find it unexpected that including image information serves as a performance boost for this task.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Mathematically, I find the paper very sound and compelling. A lot of papers have been released that train diffusion-type models on categorical data, but little research has extended these concepts to counts, which are the main type of modality in transcriptomics. The method appears correct and elegant, and potentially serves as a relevant first step for follow-up research in the field. I also really appreciate the connection between the model and the EM algorithm, which I find very original and effective to guide generation towards a certain aggregated profile.", "weaknesses": "Overall, my main points of criticism are in the readability and experimental evidence. I think the paper could be slightly more rigorous in presenting some concepts, and the biological experiments are not as comprehensive as they could be. Nonetheless, I hold a good opinion about the paper, and I am happy to discuss it during the rebuttal phase. \n\n**Text.**\n\n- L126-136: I recommend being a bit more thorough in defining the terms. For example, define $A$ as a covariance and not *noise*, make the sentence in L132-135 more clear (what is the subject?). \n\n- I would also define the concept of ±. When inspecting the math, it is clear that it represents the addition and removal of counts following the Poisson process, where both adding and removing counts have their own time-resolved rate parameter. However, you can make it explicit from the text. Specifically, the introduction to the model is a bit abrupt and, to me, could benefit from a bit more natural language. I would explain that the counterpart to the continuous formulation of bridges is a birth/death-based Poisson process, with subsequent introduction to the individual terms. This could potentially improve the flow. \n\n- A similar abrupt introduction is the connection with Schrödinger Bridges. The $\\kappa$ parameter and $\\pi$ notation for the process are previously undefined. I believe this breaks the flow of the read. Also, the concept of iteration, does it refer to an iterative proportional fitting? I find this connection quite unclear. \n\n- For the distributional scoring loss, I recommend potentially using another symbol than $w$, as it is used to the success probability of $N_t$ before. A similar concept holds for $A$, it is used both as a covariance for the diffusion term and as an aggregation function. Also $\\Pi$ is interchangeably used as the correcting projection for endpoint $x_0$ and the count bridge joint defined in 175. I personally feel all these aspects hinder the ease of reading a little bit. \n\n- Typo L 159: $(X_s)_{t\\in[0,1]}$, I think $t$ should be replaced by $s$. \n\n- L258: Is the notation $A(\\textbf{X}_{g0})$ correct? $A$ is defined as an aggregation, but here it is applied to a single entry.  \n\n**Content.** \n\n- It's not clear to me how the bulk deconvolution task works. How do you match single cells to bulk profiles? Do you build a synthetic dataset? How do you match endpoints to train the bridge? I feel that all these details should be briefly introduced in the main; otherwise, the results are hard to interpret. Also, deconvolution is a very established approach with other methods like MUSiC [1]. I am not suggesting an exhaustive benchmark, but maybe a more comprehensive comparison would be appropriate. \n\n- In general, I feel the experiment presentation is a bit disorganized. The metrics are not introduced, the captions are very short, and the benchmarks are very restricted. It would have been interesting to see some predictions of important genes performed by the model on the spatial slides, or to plot some 2D embeddings to validate that the generated cells are realistic. \n\n[1] Wang, X., et al. (2019). MuSiC: bulk tissue cell type deconvolution with multi-subject single-cell expression references. Nature Communications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)