---
title: 'CountsDiff: Diffusion for Count Data'
url: https://www.emergentmind.com/topics/countsdiff
type: topic
---

# CountsDiff: Diffusion for Count Data

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 [2604.03779].

## 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 [2604.03779].

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 $i$ extant units dies independently at rate $\mu(t)$, so the total rate of moving from $i$ to $i-1$ is $i\mu(t)$ [2604.03779].

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

$$
p:[0,1]\to[0,1],\qquad p(0)=1,\quad p(1)=0,
$$

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

$$
q(x_t\mid x_0)=\binom{x_0}{x_t}p(t)^{x_t}(1-p(t))^{x_0-x_t}.
$$

The conditional law between two times $0\le s<t\le 1$ is

$$
q(x_t\mid x_s) = \binom{x_s}{x_t} \left(\frac{p(t)}{p(s)}\right)^{x_t} \left(1-\frac{p(t)}{p(s)}\right)^{x_s-x_t}.
$$

The relation between schedule and death rate is

$$
\mu(t)=-\frac{p'(t)}{p(t)}.
$$

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

$$
\mathrm{SNR}(t)=\frac{p(t)}{1-p(t)},
$$

and adopts the cosine schedule

$$
p(t)=\cos^2\!\left(\frac{\pi t}{2}\right)
$$

as its preferred schedule.

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

$$
Q^{(\mathrm{rev})}_{i,j}(t) = (x_0-i)\frac{p'(t)}{1-p(t)}\big(\delta_{i,j}-\delta_{i+1,j}\big).
$$

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

$$
y_t=x_0-x_t.
$$

Learning the reverse dynamics is therefore reduced to predicting $y_t$ from $(x_t,t)$, using

$$
\hat y_t = \big(\mathrm{NN}_\theta(x_t,t)\big)^+,\qquad (a)^+=\log(1+e^a).
$$

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

CountsDiff trains by sampling $t$, drawing $x_t\sim \mathrm{Bin}(x_0,p(t))$ elementwise, setting $y_t=x_0-x_t$, predicting $\hat y_t$, and optimizing the weighted objective

$$
\mathbb E_{t\sim \phi}\left[w(t)\big(\hat y_t-y_t\log \hat y_t\big)\right].
$$

The paper shows that when

$$
w(t)=\frac{-p'(t)}{(1-p(t))\,\phi(t)},
$$

the objective recovers the exact negative log-likelihood of the reverse process. With $t\sim\mathrm{Unif}(0,1)$, this becomes

$$
w(t)= -\frac{p'(t)}{1-p(t)}.
$$

For the cosine schedule, the proposed weighting is

$$
w(t)=\frac{\pi}{2}\sin(\pi t).
$$

A central point in the paper is that explicit weighting functions alter optimization emphasis across times while preserving the same pointwise minimizer, since $\hat y_t=y_t$ minimizes the loss pointwise [2604.03779].

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

$$
\hat y^{(\gamma)} = (\hat y\mid c)^\gamma \,\hat y^{\,1-\gamma}.
$$

CountsDiff also departs from monotone reverse trajectories through attrition, the count-data analogue of churn or remasking. For a reverse step from time $t$ to $s<t$, define

$$
\sigma_{t,s}\in[0,\sigma^{\max}_{t,s}],\qquad \sigma^{\max}_{t,s}=\min\!\left(1,\frac{1-p(s)}{p(t)}\right),
$$

and

$$
\beta_{t,s} = \frac{p(s)-(1-\sigma_{t,s})p(t)}{1-p(t)}.
$$

The reverse-step sampler is

$$
x_s = n_t + b_t, \qquad n_t\sim \mathrm{Bin}(x_t,1-\sigma_{t,s}), \qquad b_t\sim \mathrm{Bin}(x_0-x_t,\beta_{t,s}).
$$

When $\sigma_{t,s}=0$, the step is pure birth; when $\sigma_{t,s}>0$, some present units disappear, enabling non-monotone reverse trajectories. The paper adopts a rescaled strategy,

$$
\sigma_{t,s} = \eta_{\mathrm{rescale}}\,\sigma^{\max}_{t,s},
$$

with $\eta_{\mathrm{rescale}}$ 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:

$$
\hat y_{\mathrm{clipped}} = \lfloor \hat y\rfloor + \xi, \qquad \xi\sim \mathrm{Bernoulli}\big(\hat y-\lfloor \hat y\rfloor\big).
$$

The paper motivates this by a specific low-count pathology: deterministic rounding creates a mode-collapse-at-zero failure mode whenever $\hat y<0.5$, 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 $(0.78,4.71,0.10,0.12,0.28)$ become $(0.55,1.99,0.07,0.08,0.21)$ under CountsDiff, $(0.19,0.46,0.01,0.02,0.05)$ under Gaussian diffusion, and $(3.06,9.22,1.89,2.49,7.27)$ under masked diffusion [2604.03779].

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 $5.726$, IS $9.124\pm0.049$ | baseline |
| CIFAR-10, FI continuous, unconditional, no attrition | FID $5.439$, IS $9.088\pm0.156$ | continuous analog improves FID |
| CIFAR-10, FI continuous, $\gamma=1.0$, $\eta_{\mathrm{rescale}}=0.01$ | FID $5.198$, IS $9.640\pm0.167$ | moderate guidance and small attrition help |
| CIFAR-10, cosine continuous, $\gamma=1.0$, $\eta_{\mathrm{rescale}}=0.01$ | FID $5.261$, IS $9.853\pm0.083$ | slightly better IS, slightly worse FID |
| CelebA, 30M model, cosine, $\gamma=1.0$, $\eta_{\mathrm{rescale}}=0.01$ | FID $7.637$ | viable high-dimensional generation |
| CelebA, 60M model, cosine, $\gamma=1.5$, $\eta_{\mathrm{rescale}}=0.005$ | FID $4.948$ | improved large-model result |

The paper further reports that increasing guidance to $\gamma=2.0$ improves IS further but sharply worsens FID; for cosine on CIFAR-10 this gives FID $11.546$ and IS $9.929\pm0.128$. It also reports a clear qualitative effect of attrition: increasing $\eta_{\mathrm{rescale}}$ 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 [2604.03779].

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 $0.094$, RMSE $1.401$, bias $0.004$, log(scFID) $-9.253$ | better RMSE and scFID than ReMDM single |
| Fetal atlas, 50% MCAR, 5-sample imputation | RMSE $1.195$ | essentially tied with ReMDM’s $1.196$ |
| Fetal atlas, 25% low-biased MNAR, single imputation | Spearman $0.418$, RMSE $0.871$, bias $0.299$, log(scFID) $-6.409$ | strong under realistic low-count missingness |
| Fetal atlas, 25% low-biased MNAR, 5-sample imputation | RMSE $0.580$, bias $0.300$ | best RMSE and best bias |
| Heart atlas, 50% MCAR, single imputation | RMSE $7.035$ | better than ReMDM’s $8.877$ |
| Heart atlas, 50% MCAR, 5-sample imputation | RMSE $6.209$ | close to ReMDM’s $6.109$ |

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 [2604.03779].

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 $p$-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 $\mathbb N_0$ 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.

Source: https://www.emergentmind.com/topics/countsdiff