CountsDiff: Diffusion for Count Data

• CountsDiff is a framework that models count data as sums of independent Bernoulli survivors, preserving ordinality and unbounded support explicitly.
• It parameterizes the forward process with a survival schedule p(t) and utilizes explicit loss weighting, continuous-time training, and classifier-free guidance.
• Empirical evaluations on synthetic vectors, images, and scRNA-seq imputation demonstrate its performance advantages and stability over traditional diffusion models.

CountsDiff is a diffusion framework for data that live natively on the natural numbers, $\mathbb N_0=\{0,1,2,\dots\}$. It is designed for count-valued, discrete, ordinal data without forcing them into either a continuous space or a finite categorical space. The framework extends Blackout diffusion through a direct parameterization by a survival probability schedule $p(t)$ and an explicit loss weighting, and it incorporates continuous-time training, classifier-free guidance, attrition or churn or remasking-style reverse dynamics, and randomized rounding. The paper presents CountsDiff as a count-native alternative to Gaussian diffusion, which treats counts as real-valued and discretizes afterward, and to token-style discrete diffusion, which treats each integer as an unrelated class and ignores ordinality and unbounded support (Soatto et al., 4 Apr 2026).

1. Count-native diffusion as a modeling problem

CountsDiff is motivated by the claim that standard diffusion frameworks are structurally misaligned with count data. In the paper’s formulation, Gaussian diffusion naturally handles ordered values but assumes a continuous ambient space, while categorical diffusion handles discreteness but not ordinality or unbounded support. CountsDiff instead treats counts as sums of independent Bernoulli survivors under a birth-death process, preserving ordering and extending naturally to arbitrarily large counts (Soatto et al., 4 Apr 2026).

The framework is introduced for domains in which measurements are literally counts. The paper emphasizes RNA-seq, scRNA-seq, ATAC-seq, metagenomic reads, and related abundance measurements as natural targets. It also uses image pixels as a stress test, treating pixel intensities as count-valued variables in order to study scaling behavior, qualitative artifacts, and the effect of diffusion design choices in a familiar domain. This suggests that the paper is interested not only in biological count assays, but also in the broader question of whether diffusion can be reformulated directly on $\mathbb N_0$ without reverting to continuous relaxations.

CountsDiff is explicitly positioned as an extension and clarification of Blackout diffusion. Blackout is treated as a special case with no guidance, fixed schedule and weighting, and no attrition. The novelty is therefore not the pure-death/binomial idea alone, but the reparameterization by $p(t)$, the explicit weighted objective, continuous-time implementation, guidance, attrition-based non-monotone reverse dynamics, and randomized rounding.

2. Forward and reverse processes on $\mathbb N_0$

CountsDiff starts from an inhomogeneous pure-death process on the natural numbers. If $x_t\in\mathbb N_0$ denotes the surviving count at time $t\in[0,1]$, the forward generator is

$$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$

where $\mu(t)\ge 0$ is a time-dependent death rate. At state $i$, each of the $p(t)$0 extant units dies independently at rate $p(t)$1, so the total rate of moving from $p(t)$2 to $p(t)$3 is $p(t)$4 (Soatto et al., 4 Apr 2026).

The paper’s central simplification is to parameterize the process by a differentiable monotone decreasing survival probability schedule

$p(t)$5

rather than primarily by $p(t)$6. For any such $p(t)$7, the forward marginals have the binomial form

$p(t)$8

The conditional law between two times $p(t)$9 is

$\mathbb N_0$0

The relation between schedule and death rate is

$\mathbb N_0$1

The paper also introduces the Bernoulli or binomial analogue of diffusion signal-to-noise ratio,

$\mathbb N_0$2

and adopts the cosine schedule

$\mathbb N_0$3

as its preferred schedule.

For the pure-death forward model, the reverse process is a pure-birth process. Given the latent clean count $\mathbb N_0$4, the reverse generator is

$\mathbb N_0$5

The unknown quantity is the number of units lost in the forward process, written as

$\mathbb N_0$6

Learning the reverse dynamics is therefore reduced to predicting $\mathbb N_0$7 from $\mathbb N_0$8, using

$\mathbb N_0$9

3. Objective, weighting, guidance, and non-monotone reverse dynamics

CountsDiff trains by sampling $p(t)$0, drawing $p(t)$1 elementwise, setting $p(t)$2, predicting $p(t)$3, and optimizing the weighted objective

$p(t)$4

The paper shows that when

$p(t)$5

the objective recovers the exact negative log-likelihood of the reverse process. With $p(t)$6, this becomes

$p(t)$7

For the cosine schedule, the proposed weighting is

$p(t)$8

A central point in the paper is that explicit weighting functions alter optimization emphasis across times while preserving the same pointwise minimizer, since $p(t)$9 minimizes the loss pointwise (Soatto et al., 4 Apr 2026).

The framework imports classifier-free guidance into count diffusion. If $\mathbb N_0$0 denotes class conditioning and $\mathbb N_0$1 is the guidance strength, the guided prediction is

$\mathbb N_0$2

CountsDiff also departs from monotone reverse trajectories through attrition, the count-data analogue of churn or remasking. For a reverse step from time $\mathbb N_0$3 to $\mathbb N_0$4, define

$\mathbb N_0$5

and

$\mathbb N_0$6

The reverse-step sampler is

$\mathbb N_0$7

When $\mathbb N_0$8, the step is pure birth; when $\mathbb N_0$9, some present units disappear, enabling non-monotone reverse trajectories. The paper adopts a rescaled strategy,

$x_t\in\mathbb N_0$0

with $x_t\in\mathbb N_0$1 as a tunable sampling hyperparameter.

Because the network output is real-valued while the reverse binomial process requires a natural number, CountsDiff uses randomized rounding:

$x_t\in\mathbb N_0$2

The paper motivates this by a specific low-count pathology: deterministic rounding creates a mode-collapse-at-zero failure mode whenever $x_t\in\mathbb N_0$3, which occurs frequently in sparse count data.

4. Empirical evaluation on synthetic count vectors and images

The paper first evaluates CountsDiff on synthetic sparse count vectors in order to isolate the consequences of respecting count structure. These data are 10-dimensional negative-binomial samples with multiplicative size factors, about 50% zeros, and maxima around 50. CountsDiff is compared with a Gaussian diffusion model in log-space and a masked categorical diffusion model. The main reported pattern is that CountsDiff and masked diffusion learn marginals reasonably well, masked diffusion has much worse sliced Wasserstein distance and inflated variance, and Gaussian diffusion fails outright with severe mode collapse. In one subset of dimensions, true variances $x_t\in\mathbb N_0$4 become $x_t\in\mathbb N_0$5 under CountsDiff, $x_t\in\mathbb N_0$6 under Gaussian diffusion, and $x_t\in\mathbb N_0$7 under masked diffusion (Soatto et al., 4 Apr 2026).

The image experiments on CIFAR-10 and CelebA are used as a high-dimensional stress test rather than as the primary target domain. They are nonetheless important because they isolate the effect of schedule, guidance, and attrition in a well-studied setting.

Setting	Reported result	Interpretation in the paper
CIFAR-10, Blackout-equivalent FI discrete, unconditional, no attrition	FID $x_t\in\mathbb N_0$8, IS $x_t\in\mathbb N_0$9	baseline
CIFAR-10, FI continuous, unconditional, no attrition	FID $t\in[0,1]$0, IS $t\in[0,1]$1	continuous analog improves FID
CIFAR-10, FI continuous, $t\in[0,1]$2, $t\in[0,1]$3	FID $t\in[0,1]$4, IS $t\in[0,1]$5	moderate guidance and small attrition help
CIFAR-10, cosine continuous, $t\in[0,1]$6, $t\in[0,1]$7	FID $t\in[0,1]$8, IS $t\in[0,1]$9	slightly better IS, slightly worse FID
CelebA, 30M model, cosine, $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$0, $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$1	FID $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$2	viable high-dimensional generation
CelebA, 60M model, cosine, $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$3, $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$4	FID $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$5	improved large-model result

The paper further reports that increasing guidance to $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$6 improves IS further but sharply worsens FID; for cosine on CIFAR-10 this gives FID $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$7 and IS $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$8. It also reports a clear qualitative effect of attrition: increasing $$Q^{(\mathrm{fw})}_{i,j}(t)= i\,\mu(t)\big(\delta_{i-1,j}-\delta_{i,j}\big),$$9 smooths images, while large values oversmooth and erase texture and even perspective. This is presented as the count-domain counterpart of churn or remasking behavior in token diffusion.

5. Biological count assays and scRNA-seq imputation

The biologically most natural application in the paper is scRNA-seq imputation. CountsDiff is adapted using RePaint: after each reverse step, observed entries are reset to their noised ground-truth values, while missing entries are resampled, and this requires no retraining (Soatto et al., 4 Apr 2026).

The datasets are the human fetal cell atlas and human heart cell atlas. The paper evaluates fetal 50% MCAR, fetal 25% low-biased MNAR, and heart 50% MCAR, using sample-level Spearman correlation, RMSE, bias, and log(scFID), where lower is better.

Task	CountsDiff	Comparison highlighted in the paper
Fetal atlas, 50% MCAR, single imputation	Spearman $\mu(t)\ge 0$0, RMSE $\mu(t)\ge 0$1, bias $\mu(t)\ge 0$2, log(scFID) $\mu(t)\ge 0$3	better RMSE and scFID than ReMDM single
Fetal atlas, 50% MCAR, 5-sample imputation	RMSE $\mu(t)\ge 0$4	essentially tied with ReMDM’s $\mu(t)\ge 0$5
Fetal atlas, 25% low-biased MNAR, single imputation	Spearman $\mu(t)\ge 0$6, RMSE $\mu(t)\ge 0$7, bias $\mu(t)\ge 0$8, log(scFID) $\mu(t)\ge 0$9	strong under realistic low-count missingness
Fetal atlas, 25% low-biased MNAR, 5-sample imputation	RMSE $i$0, bias $i$1	best RMSE and best bias
Heart atlas, 50% MCAR, single imputation	RMSE $i$2	better than ReMDM’s $i$3
Heart atlas, 50% MCAR, 5-sample imputation	RMSE $i$4	close to ReMDM’s $i$5

The paper emphasizes that CountsDiff is the least biased method in the fetal low-biased MNAR scenario, which it treats as scientifically important because low-count missingness is a realistic scRNA-seq failure mode. It also reports a qualitative stability claim: ReMDM tends to oversample outliers, while CountsDiff is more stable. Parameter counts are also compared directly: CountsDiff has about half as many parameters as ReMDM on fetus and one-fourth as many on heart, because ReMDM’s output layer scales with maximum count whereas CountsDiff predicts a count residual scalar per feature rather than a full categorical distribution over values.

Ablations reinforce the role of the new design parameters. Moderate attrition improves both scFID and energy distance for cosine and continuous Blackout schedules, with larger gains for cosine. Moderate guidance helps, while extreme values hurt. Discrete Blackout-style training is substantially worse than continuous schedule versions. Reverse-step ablations show graceful degradation around 10–30 steps, noticeable deterioration below 5, and collapse at 1 step.

6. Strengths, limitations, and place in the diffusion landscape

The principal strengths attributed to CountsDiff are structural. It is native to count data, respects discreteness, ordinality, and unbounded support, has exact binomial marginals, and admits an exact negative log-likelihood objective for suitable weighting. It also imports into the count setting the design space that became central in modern diffusion models: schedule choice, explicit weighting, classifier-free guidance, and churn or remasking-like sampling. The paper’s broader claim is that this formulation avoids the quantization mismatch of Gaussian diffusion and the exploding output-space problem of tokenizing counts (Soatto et al., 4 Apr 2026).

The paper is equally explicit about limitations. The presented instantiation is intentionally simple and likely not optimal. It repeatedly notes substantial room for improvement in $i$6-schedule design, loss weighting, and attrition schedules. Attrition scheduling is described as underexplored, and the paper suggests future work on value-dependent attrition rates. The image experiments are used to show viability rather than to establish state-of-the-art image generation. For imputation, evaluation remains difficult because scFID, Spearman, RMSE, and bias can disagree. More broadly, the method is specialized to count-valued data; if the data are genuinely continuous nonnegative values, the paper suggests combining CountsDiff with JUMP’s Poisson randomization trick rather than applying it directly.

A common misconception would be to treat CountsDiff as merely another discrete diffusion model. The paper argues for a narrower and more specific interpretation. CountsDiff is not only discrete; it is discrete on $i$7 with binomial thinning as the corruption mechanism, direct survival-schedule parameterization, and a reverse process that predicts how many counts should be restored. In that sense, its central contribution is less the use of diffusion in the abstract than the reformulation of diffusion so that count structure is the primitive object rather than an afterthought.