Count Bridges enable Modeling and Deconvolving Transcriptomic Data

Imagine you are trying to understand a massive, bustling city. You have two ways of looking at it:

The High-Resolution View: You have a drone that can see every single person, their exact location, and what they are doing. This is like single-cell data in biology.
The Blurry View: You only have a satellite photo that shows "neighborhoods." In each neighborhood, you see a single number representing the total population, but you can't tell who is there, how many people are in each house, or what they are doing. This is like bulk data or spatial transcriptomics (where a "spot" on a slide contains 10–50 cells mixed together).

For a long time, scientists had a hard time going from the Blurry View back to the High-Resolution View. They could guess the average number of people in a neighborhood, but they couldn't reconstruct the specific, messy, integer-based reality of individual lives (e.g., "House A has 3 people, House B has 0, House C has 7").

Enter Count Bridges.

The Problem: The "Integer" Puzzle

Most computer models for generating data (like AI that creates images) are built for smooth, continuous numbers (like 3.14159...). But biology is made of discrete counts. You can have 3 molecules of a gene, or 4, but never 3.5.

Existing AI tools tried to force these "whole number" problems into smooth models, which was like trying to fit a square peg in a round hole. They either ignored the "whole number" nature of the data or used methods that were too rigid to handle the complexity of biological mixtures.

The Solution: The "Count Bridge"

The authors built a new type of AI called Count Bridges. Think of it as a magical, reversible elevator specifically designed for whole numbers.

The Analogy of the "Birth-Death Elevator":
Imagine you have a stack of blocks (representing molecules).

The Forward Trip (Adding Noise): The elevator starts at the bottom (your real data) and goes up. As it rises, it randomly adds blocks (births) or removes blocks (deaths). By the time it reaches the top, the stack is a chaotic, random mess.
The Reverse Trip (Denoising): Now, you want to go back down to the bottom to see the original stack. The AI learns the "rules" of how blocks were added and removed. It doesn't just guess; it calculates the exact probability of how many blocks were there at the start, step-by-step, ensuring it always lands on a whole number.

This is the "Bridge": a path that connects a messy, noisy state back to a clean, specific integer state, respecting the fact that you can't have half a molecule.

The Superpower: Deconvolution (The "Un-Mixing" Trick)

The real magic happens when you only have the Blurry View (the aggregate).

Imagine you are handed a bucket of mixed LEGO bricks from 10 different kids. You know the total number of red, blue, and green bricks in the bucket, but you don't know which kid built what.

Old methods could tell you the average number of bricks per kid.
Count Bridges can reconstruct the exact set of bricks for each individual kid.

How?
The paper uses a clever trick called an Expectation-Maximization (EM) loop:

Guess: The AI guesses what the individual kids' LEGO sets might look like.
Check: It adds them all up. Does the total match the bucket you were given?
Adjust: If the total is too high or too low, it gently reshuffles the guesses (like a "projection") to fit the bucket's total perfectly.
Learn: It uses this "corrected" guess to get smarter for next time.

By repeating this, the AI learns to "un-mix" the bucket, revealing the unique, high-resolution LEGO sets of every single kid (cell) inside.

Why This Matters in the Real World

The authors tested this on two huge biological problems:

Bulk RNA-Seq (The "Smoothie" Problem): Scientists often take a tissue sample, blend it all together, and measure the average gene activity. It's like drinking a smoothie and trying to guess exactly which fruits were in it. Count Bridges can taste the smoothie and tell you, "Ah, this specific cell had 5 strawberries, and that one had 2," effectively reconstructing the whole fruit salad from the juice.
Spatial Transcriptomics (The "Pixel" Problem): New microscopes take pictures of tissue, but each "pixel" on the image actually contains a cluster of 20–50 cells. It's like looking at a low-res photo of a crowd. Count Bridges can take that low-res pixel and generate a high-res list of exactly who is standing where, revealing the hidden architecture of the tissue.

The Bottom Line

Count Bridges is a new mathematical tool that finally allows AI to handle "whole number" data correctly. It acts as a bridge between the messy, aggregated data we often have and the precise, individual-level data we desperately need. It doesn't just guess averages; it reconstructs the full, detailed story of individual cells, helping scientists understand diseases, tissues, and life itself with much sharper clarity.

1. Problem Statement

Modern biological assays, particularly RNA sequencing (RNA-seq) and spatial transcriptomics, generate integer-valued count data. However, a critical limitation exists: many technologies produce aggregated measurements rather than single-cell resolution.

Bulk RNA-seq: Aggregates thousands to millions of cells, yielding average expression profiles that obscure cellular heterogeneity.
Spatial Transcriptomics (e.g., Visium): Captures 10–50 cells per "spot," resulting in mixed signals.
The Challenge: Existing generative models (diffusion, flow matching) are primarily designed for continuous Euclidean data or treat discrete data as unordered categories (masking). They fail to respect the ordinal structure of counts (where $n$ is distinct from $n+1$ ) and struggle to systematically deconvolve aggregated observations back into unit-level (single-cell) profiles without external reference atlases.

2. Methodology: Count Bridges

The authors introduce Count Bridges, a stochastic bridge process defined on the integer lattice $\mathbb{Z}^d$ using Poisson birth-death dynamics.

A. The Stochastic Bridge Process

Unlike Gaussian diffusion models that add continuous noise, Count Bridges model the transition between a source distribution $X_0$ and a target $X_1$ using independent Poisson birth ( $B_t$ ) and death ( $D_t$ ) processes:
$X_t = X_0 + B_t - D_t$

Jump Intensities: The processes are governed by time-dependent intensities $\Lambda_\pm(t) = \lambda_\pm w(t)$ , where $w(t)$ is a shape function.
Slack Variables: The process tracks "slack" ( $M_t = \min(B_t, D_t)$ ), representing jumps that cancel each other out (e.g., a birth followed immediately by a death).
Closed-Form Conditionals: The authors derive exact, tractable conditionals for the bridge. Given endpoints $X_0$ $X_{0}$ and $X_1$ $X_{1}$ , the intermediate state $X_s$ $X_{s}$ is sampled by:
1. Sampling the slack $M_t$ from a Bessel distribution (derived from the Skellam distribution of the difference of Poissons).
2. Sampling the total jumps $N_t$ and birth jumps $B_t$ using Binomial and Hypergeometric distributions, respectively.
Theoretical Foundation: This setup is shown to be an instance of the Schrödinger bridge problem for integer-valued data, solving an entropy-regularized optimal transport problem where the cost function is the $L_1$ distance (Manhattan distance) between counts.

B. Distributional Scoring Loss

Because the space is discrete, standard regression losses (MSE) or cross-entropy are insufficient. Cross-entropy fails to capture the lattice geometry and becomes intractable for joint modeling in high dimensions.

The authors employ a Distributional Scoring Rule, specifically the Energy Score with a metric $\rho(x, x') = \|x - x'\|_2^\beta$ (typically $\beta=1$ ).
This allows the model to learn the full conditional distribution $q_\theta(X_0 | X_t, t)$ rather than just the mean, preserving the discrete nature of the data.

C. Deconvolution via Expectation-Maximization (EM)

To handle aggregated data (where only the sum $A_0 = \sum X_{g,0}$ is observed, not individual $X_{g,0}$ ), the authors propose an EM-style framework:

E-Step (Guided Sampling): The model generates latent unit-level samples $\hat{X}_0$ $\hat{X}_{0}$ conditioned on the aggregate constraint. This is achieved via a projection-guided diffusion process where, at each reverse step, the predicted sample is projected onto the constraint set $\{x : \sum x_g = a_0\}$ ${x : \sum x_{g} = a_{0}}$ .
- The projection $\Pi$ is derived as a first-order approximation to the conditional distribution, effectively performing a rescaling of the predicted counts to match the observed aggregate sum.
M-Step (Training): The model is trained using the aggregated data. The loss is computed on the aggregates of the generated samples, using a lifted version of the energy score: $S_\rho(A(X), a)$ .

3. Key Contributions

Count Bridges Framework: A novel stochastic bridge process on integers using Poisson birth-death dynamics, providing exact, closed-form conditionals for sampling and training.
Distributional Loss for Counts: The application of strictly proper scoring rules (Energy Score) to discrete generative modeling, overcoming the limitations of cross-entropy and mean-squared error.
Principled Deconvolution: An EM-based approach that treats unit-level counts as latent variables, enabling the inference of single-cell profiles from aggregated bulk or spatial data without requiring external reference atlases (reference-free).
Theoretical Connection: Establishing the link between Count Bridges and discrete Optimal Transport, showing that the bridge parameters control the entropy regularization strength.

4. Results

The paper evaluates Count Bridges on synthetic benchmarks and two major biological applications.

A. Synthetic Benchmarks

Discrete 8-Gaussians to 2-Moons: Count Bridges outperformed Continuous Flow Matching (CFM) and Discrete Flow Matching (DFM) across Wasserstein-2 ( $W_2$ ), Maximum Mean Discrepancy (MMD), and Energy Score metrics. CB produced trajectories that respected the underlying geometry (OT-like), whereas DFM trajectories were decoupled from geometry.
Scalability: In low-rank Gaussian mixtures with increasing ambient dimensions (up to 512), Count Bridges demonstrated superior scaling compared to baselines.

B. Biological Applications

Nucleotide-Resolution Bulk RNA-seq Deconvolution:
- Task: Predicting single-cell gene expression profiles from bulk RNA-seq data, conditioned on DNA sequence context.
- Performance: Count Bridges significantly outperformed fine-tuned Enformer (a state-of-the-art sequence-to-expression model) in predicting nucleotide-level counts (MSE reduced from ~2.59 to ~0.60).
- Deconvolution: When deconvolving bulk profiles into single-cell profiles, CB achieved state-of-the-art performance in JSD, RMSE, and Spearman correlation compared to CIBERSORTx and MuSiC, while providing full count profiles rather than just cell-type proportions.
Reference-Free Spatial Transcriptomic Deconvolution:
- Task: Resolving multicellular Visium spots into single-cell count profiles using only spot aggregates and cell images (no external scRNA-seq reference).
- Performance: CB outperformed STDeconvolve (a leading reference-free method) in JSD, RMSE, and Spearman correlation for cell-type proportions.
- Biological Validity: The deconvolved single-cell profiles were analyzed via UMAP and cell-type annotation (Celltypist). The generated data closely matched the distribution and cell-type abundance of true single-cell data, demonstrating the model's ability to learn realistic biological heterogeneity.

5. Significance

Native Discrete Modeling: Count Bridges offer a principled alternative to continuous approximations for biological count data, respecting the integer and ordinal nature of molecular counts.
Reference-Free Deconvolution: The ability to deconvolve aggregates without external reference atlases is a major advancement, reducing reliance on potentially mismatched or unavailable single-cell datasets.
Multi-Modal Integration: The framework successfully integrates sequence context (DNA), spatial context (images), and aggregate constraints, enabling high-resolution reconstruction of biological systems from coarse measurements.
Foundation for Future Work: The paper lays the groundwork for rigorous discrete generative modeling, with potential applications extending beyond transcriptomics to mass cytometry, fluorescence imaging, and other count-based scientific measurements.

Limitations Noted by Authors

Continuous Approximation: When counts are very large, continuous Euclidean models may still perform comparably.
Identifiability: Deconvolution becomes statistically ill-conditioned as group sizes grow or between-group heterogeneity decreases (the "Gaussian collapse" of aggregate information).
Projection Approximation: The projection step used in the E-step is a first-order surrogate; while effective, it lacks deep theoretical guarantees for all regimes.