CountsDiff: A Diffusion Model on the Natural Numbers for Generation and Imputation of Count-Based Data

• URL: http://arxiv.org/abs/2604.03779v1
• Date: Sat, 04 Apr 2026 16:06:14 GMT
• Title: CountsDiff: A Diffusion Model on the Natural Numbers for Generation and Imputation of Count-Based Data
• Abstract summary: CountsDiff is a diffusion framework designed to model distributions on the natural numbers.<n>We propose an initial instantiation of CountsDiff and validate it on natural image datasets.<n>We then highlight biological count assays as a natural use case, evaluating CountsDiff on single-cell RNA-seq imputation.
• Score: 13.773137893785984
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Diffusion models have excelled at generative tasks for both continuous and token-based domains, but their application to discrete ordinal data remains underdeveloped. We present CountsDiff, a diffusion framework designed to natively model distributions on the natural numbers. CountsDiff extends the Blackout diffusion framework by simplifying its formulation through a direct parameterization in terms of a survival probability schedule and an explicit loss weighting. This introduces flexibility through design parameters with direct analogues in existing diffusion modeling frameworks. Beyond this reparameterization, CountsDiff introduces features from modern diffusion models, previously absent in counts-based domains, including continuous-time training, classifier-free guidance, and churn/remasking reverse dynamics that allow non-monotone reverse trajectories. We propose an initial instantiation of CountsDiff and validate it on natural image datasets (CIFAR-10, CelebA), exploring the effects of varying the introduced design parameters in a complex, well-studied, and interpretable data domain. We then highlight biological count assays as a natural use case, evaluating CountsDiff on single-cell RNA-seq imputation in a fetal cell and heart cell atlas. Remarkably, we find that even this simple instantiation matches or surpasses the performance of a state-of-the-art discrete generative model and leading RNA-seq imputation methods, while leaving substantial headroom for further gains through optimized design choices in future work.

Related papers

• Impute-EM: Native Mixed-State Diffusion Models for Heterogeneous Data Imputation [71.74203645336385]
  Impute-EM alternates between imputing missing entries with the current model and refitting a diffusion backbone on completed data.<n>We characterize the update and show that the observed mask-indexed marginals match the targets at the limit.
  arXiv  Detail & Related papers  (2026-09-14T09:39:13Z)
• From Scores to Gibbs Correctors: Accelerating Uniform-Rate Discrete Diffusion Models [58.02657021435213]
  We propose a novel Gibbs-based corrector for discrete diffusion models, termed Gibbs-Accelerated Discrete Diffusion (GADD)<n>We show GADD achieves an overall sampling complexity of $mathcalO(mathrmpolylog (varepsilon-1)$, yielding the first such rate for diffusion-based samplers for uniform-rate discrete diffusion models.
  arXiv  Detail & Related papers  (2026-05-26T17:52:28Z)
• Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design [56.957070405026194]
  We propose an algorithm that enables direct backpropagation of rewards through entire trajectories generated by diffusion models.<n>DRAKES can generate sequences that are both natural-like and yield high rewards.
  arXiv  Detail & Related papers  (2024-10-17T15:10:13Z)
• Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding [84.3224556294803]
  Diffusion models excel at capturing the natural design spaces of images, molecules, DNA, RNA, and protein sequences. We aim to optimize downstream reward functions while preserving the naturalness of these design spaces. Our algorithm integrates soft value functions, which looks ahead to how intermediate noisy states lead to high rewards in the future.
  arXiv  Detail & Related papers  (2024-08-15T16:47:59Z)
• Navigating the Structured What-If Spaces: Counterfactual Generation via Structured Diffusion [20.20945739504847]
  We introduce Structured Counterfactual diffuser or SCD, the first plug-and-play framework leveraging diffusion for generating counterfactual explanations in structured data. Our experiments show that our counterfactuals not only exhibit high plausibility compared to the existing state-of-the-art but also show significantly better proximity and diversity.
  arXiv  Detail & Related papers  (2023-12-21T07:05:21Z)
• Fast Sampling via Discrete Non-Markov Diffusion Models with Predetermined Transition Time [49.598085130313514]
  We propose discrete non-Markov diffusion models (DNDM), which naturally induce the predetermined transition time set.<n>This enables a training-free sampling algorithm that significantly reduces the number of function evaluations.<n>We study the transition from finite to infinite step sampling, offering new insights into bridging the gap between discrete and continuous-time processes.
  arXiv  Detail & Related papers  (2023-12-14T18:14:11Z)
• Bayesian Flow Networks [4.197165999892042]
  This paper introduces Bayesian Flow Networks (BFNs), a new class of generative model in which the parameters of a set of independent distributions are modified with Bayesian inference.<n>Starting from a simple prior and iteratively updating the two distributions yields a generative procedure similar to the reverse process of diffusion models.<n>BFNs achieve competitive log-likelihoods for image modelling on dynamically binarized MNIST and CIFAR-10, and outperform all known discrete diffusion models on the text8 character-level language modelling task.
  arXiv  Detail & Related papers  (2023-08-14T09:56:35Z)
• Score-based Continuous-time Discrete Diffusion Models [102.65769839899315]
  We extend diffusion models to discrete variables by introducing a Markov jump process where the reverse process denoises via a continuous-time Markov chain. We show that an unbiased estimator can be obtained via simple matching the conditional marginal distributions. We demonstrate the effectiveness of the proposed method on a set of synthetic and real-world music and image benchmarks.
  arXiv  Detail & Related papers  (2022-11-30T05:33:29Z)
• Diffusion Causal Models for Counterfactual Estimation [18.438307666925425]
  We consider the task of counterfactual estimation from observational imaging data given a known causal structure. We propose Diff-SCM, a deep structural causal model that builds on recent advances of generative energy-based models. We find that Diff-SCM produces more realistic and minimal counterfactuals than baselines on MNIST data and can also be applied to ImageNet data.
  arXiv  Detail & Related papers  (2022-02-21T12:23:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.