PaperlayerView exact revision on arXiv

Source-linked AI summary

CRNDiff: Count-Native Diffusion Framework via Chemical Reaction Networks

Yuxuan Qiu, Praful Gagrani, Tetsuya J Kobayashi

arXiv:2609.31149v1cs.LG

TL;DR

Count-valued data motivate diffusion models that operate directly on the nonnegative integer lattice, while existing count-native approaches are tied to particular transition mechanisms. CRNDiff uses chemical reaction networks and tilted Feynman–Kac steering to condition a frozen generator. On heart-atlas targets, it improves conditional fidelity across three populations and supports marker preservation and downstream classification.

  • Problem

    Count-valued data lie on an unbounded nonnegative integer lattice, while existing count-native diffusion models are formulated around specific transition mechanisms rather than a general structured framework.

    Source: §1 · p. 1

  • Method

    CRNDiff specifies count-space diffusion with chemical reaction networks and uses marginal tilting followed by Feynman–Kac correction to condition a frozen generator without retraining.

    Source: §1 · p. 2, §3.2 · p. 6

  • Results

    Across three heart-atlas target cell types, CRNDiff improves conditional fidelity, preserves marker-level differential-expression structure, supports downstream classification, and reduces genealogical concentration relative to FK steering.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The heart-atlas results suggest that separating generator learning from sampler-level conditioning is useful for rare target populations.

    Source: §6 · p. 9

    Terminal-time selection uses second-order excess covariance, and finite-particle, density-ratio, or target-discriminator errors can shift rare-target purity–fidelity trade-offs.

    Source: §6 · p. 9

Abstract

from arXiv · show

Scientific measurements such as single-cell RNA (scRNA) sequencing often take the form of nonnegative integer counts, whereas continuous-state diffusion models approximate this discrete structure using continuous coordinates. Building on stochastic chemical reaction networks (CRNs), a class of count-native Markov jump processes, we introduce CRNDiff, a structured framework that combines count-space diffusion with inference-time conditioning on rare subpopulations. An independent birth--death instantiation yields a closed-form transition kernel for forward noising. This kernel enables reverse sampling via forward-filtering backward-sampling (FFBS) and supports data-driven selection of the terminal noising time, eliminating the need for a validation sweep. This tractability also lets us introduce tilted Feynman--Kac (FK) steering, a method for sampling target subpopulations from a frozen generator without retraining. By tilting posterior marginals before FK particle correction, steering mitigates importance-weight concentration when the target population is rare. Using scRNA-seq data from the human heart cell atlas, we test the ability of CRNDiff to generate cell-type-specific distributions. Across the three evaluated target populations, CRNDiff achieves the highest conditional fidelity among the evaluated generative models, with larger mean purity margins for rarer target populations. Generated cells preserve marker-level differential-expression structure. Replacing real training cells for the target classes with generated cells yields downstream classification performance approaching that of the real-data reference.

1 INTRODUCTION

CRNDiff uses chemical reaction networks to structure count-space diffusion and support inference-time conditioning. On heart-atlas data, it achieves the highest conditional fidelity among evaluated models across three target populations.

1 INTRODUCTION

  • Existing count-native diffusion models use specific transition mechanisms, motivating a structured framework for diverse count-state transitions and state-dependent rates.CRNs specify admissible count changes through reaction channels and their rates through propensity functions.Source: §1 · p. 1
  • An independent birth–death CRN yields closed-form transition kernels and bridge-based reverse sampling, while factorial-cumulant dynamics select terminal time from excess-covariance decay without a validation sweep.The forward process preserves empirical coordinate means and converges to a product-Poisson stationary distribution.Source: §1 · p. 2
  • Tilted Feynman–Kac steering conditions a frozen generator by adapting reverse proposals toward target marginals, then correcting residual joint mismatch without retraining.The marginal-first tilt is intended to mitigate importance-weight concentration for rare targets.Source: §1 · p. 2
  • Across three heart-atlas target populations, CRNDiff outperforms evaluated baselines on conditional-fidelity metrics, with larger mean purity margins for rarer targets.Generated cells preserve differential-expression rankings, and replacement training data yield downstream classification performance approaching the real-cell reference.Source: §1 · p. 2

2 BACKGROUND

The background frames diffusion as a continuous-time Markov jump process on count space, with reverse sampling based on tractable forward kernels. It then introduces endpoint conditioning, Feynman–Kac steering, and chemical reaction networks as a structured way to specify count-valued dynamics.

2.1 CONTINUOUS-TIME DIFFUSION MODELS

  • The forward process is chosen to approach a tractable stationary distribution, and the reverse process evolves from qT back to the data distribution q0 = D.In practice, π replaces qT after sufficient forward relaxation.Source: §2.1 · p. 2
  • A forward transition kernel and clean-state posterior define Markov bridges; backward sampling gives exact FFBS when the posterior and terminal distribution are exact.In practice, the posterior is approximated by a learned model trained on clean and forward-noised states.Source: §2.1 · pp. 2–3

2.1.1 CONDITIONAL GENERATION

  • Endpoint tilting reweights the posterior for a changed endpoint distribution while preserving the forward kernel and bridge, allowing adaptation without retraining.Exact sampling also requires the target terminal marginal; initializing from the common stationary distribution is the usual large-T approximation.Source: §2.1.1 · p. 3
  • Feynman–Kac steering distributes endpoint correction along the reverse trajectory when the target occupies a region with little probability under the original sampler.Intermediate potentials affect sampling efficiency, while the endpoint distribution retains the desired correction.Source: §2.1.1 · p. 3

2.2 CHEMICAL REACTION NETWORKS

  • Chemical reaction networks specify count-space Markov jump processes through species, reaction changes, and state-dependent propensities; unavailable reactants have zero propensity.This keeps the dynamics in the nonnegative integer orthant.Source: §2.2 · pp. 3–4
  • Probability-generating functions encode count distributions, and coefficient extraction supplies transition kernels; first-order PGF equations include births, deaths, conversions, and branching reactions.The PGF representation implements count-unit addition and removal through multiplication and differentiation.Source: §2.2 · p. 4

3 CRN DIFFUSION MODEL

The section instantiates CRNDiff with an exactly solvable birth–death noising process, then uses its tractability for data-calibrated stopping and inference-time conditioning of a frozen generator.

3.1 NOISING WITH A BIRTH–DEATH MODEL

  • Birth–death reactions keep counts on the nonnegative integer lattice; the exact kernel combines surviving initial counts with independently created Poisson counts for reverse bridge sampling.Conditional on the initial count, the result is a binomial number of survivors plus an independent Poisson number of new counts.Source: §3.1 · pp. 4–5, §3 · p. 4
  • The process preserves coordinate means while each factorial cumulant of order m ≥2 decays as e^-mt.These dynamics track how correlations and deviations from independent Poisson statistics are erased by noising.Source: §3.1 · p. 5
  • The terminal time is when the strongest positive excess-covariance mode reaches the finite-sample product-Poisson noise floor, a second-order criterion rather than a guarantee of full convergence.The threshold is determined in closed form by the number of genes and cells, without Monte Carlo sampling.Source: §3.1 · p. 5

3.2 CONDITIONAL DENOISING WITH TILTED FK STEERING

  • Tilted FK steering reweights reverse-posterior marginals toward a target, then corrects residual joint mismatch while the generator remains frozen.Target-specific preparation estimates generator and target marginals and fits a residual discriminator.Source: §3.2 · p. 6
  • For rare targets, less reliable marginal ratios motivate stronger tempering; at fixed τ, marginal tilting cannot distinguish targets with identical one-dimensional marginals.The marginal tilt adapts the proposal using marginal information, while FK correction addresses the remaining joint discrepancy.Source: §3.2 · p. 6

4 RELATED WORK

CRNDiff is situated among discrete and count-native generative models, stochastic-kinetics foundations, inference-time steering methods, and single-cell generative models.

Discrete and count-native generative models

  • Discrete diffusion and count-native methods use varied transitions, while CRNDiff uses birth–death noising as one instance of a broader reaction-network formulation.The cited count methods include death or thinning, thinning and thickening, and birth–death processes.Source: §4 · p. 7

Stochastic kinetics and generating functions

  • CRNDiff’s kernels, bridges, and factorial-cumulant dynamics instantiate established exact solutions for monomolecular and first-order systems.The operator representation draws on Doi–Peliti theory and related work on CRNs and factorial-moment hierarchies.Source: §4 · p. 7

Conditional generation and inference-time steering

  • Unlike selection and value-guided resampling, CRNDiff adapts a frozen generator’s reverse proposal toward target marginals before joint density-ratio correction.The cited work notes that rare targets can concentrate importance weights and cause genealogical degeneracy in SMC steering.Source: §4 · p. 7

Single-cell generative models

  • Where scVI, scANVI, and CFGen use latent representations and count-valued decoders, CRNDiff models count-valued dynamics directly and conditions populations at sampling time.Source: §4 · p. 7

5 EXPERIMENTS

The experiments test count-mixture steering and conditional generation on adult human heart atlas data. CRNDiff shows strong conditional fidelity across three target cell types, alongside marker preservation and downstream classification results.

5 EXPERIMENTS

  • The evaluation covers a synthetic count mixture and heart-atlas endothelial, myeloid, and neuronal cells, using an unconditionally trained CRNDiff generator for target conditioning.Conditional fidelity uses held-out test cells as the real-data reference.Source: §5 · p. 7, §5.2 · p. 8

5.1 A COUNT MIXTURE WITH MATCHED-MARGINAL TARGETS

  • Tilted FK steering reduces genealogical concentration and approaches the i.i.d. target reference on energy distance, mean marginal total variation, and ancestor diversity.The comparison uses the same frozen generator; the target comprises rare components c1 and c5, together 1% of the training mixture.Source: §5.1 · p. 8, Table 1 · p. 8

5.2 GENERATION AND EVALUATION ON SINGLE-CELL DATA

  • CRNDiff has the highest mean purity and PCC and the lowest mean W1 and MMD2 for all three target cell types.Its mean purity margin over the strongest baseline increases as training abundance decreases; neuronal mean purity approaches the real-cell reference.Source: §5.2 · p. 8, Table 2 · p. 9
  • On the frozen CRNDiff generator, larger best-of-N pools and stronger value-guided resampling can raise purity, sometimes matching or exceeding tilted FK steering, often with higher W1 and MMD2.Ablations also find that tilted FK steering preserves more distinct initial ancestors, with metric-dependent distributional trade-offs.Source: §5.2 · p. 8
  • Unconditionally, CRNDiff has the lowest sliced-W1 and near-reference detected-gene variability, while its lower library-size variability and median Fano factor indicate dispersion limits.These results assess generators without conditional steering.Source: §5.2 · pp. 8–9
  • Generated cells achieve the highest mean macro F1, myeloid F1, and neuronal F1 when used as replacement training data, while marker-ranking results vary by cell type and metric.CRNDiff has the highest mean top-100 marker overlap and Spearman correlation for endothelial and myeloid cells; neuronal marker results are metric-dependent.Source: §5.2 · p. 9

6 DISCUSSION

The discussion presents CRNDiff as count-space diffusion structured by reaction channels and propensity functions. It reports heart-atlas results across three target cell types and identifies limits in terminal-time calibration and rare-target steering.

6 DISCUSSION

  • In the birth–death instance, closed-form kernels, Markov bridges, and factorial-cumulant dynamics support reverse sampling, terminal-time selection, and conditioning a frozen generator.The framework specifies count-space diffusion through reaction channels and propensity functions.Source: §6 · p. 9
  • Across three target cell types, CRNDiff improves conditional fidelity, preserves marker-level differential-expression structure, supports downstream classification, and reduces genealogical concentration relative to FK steering.The authors describe the heart-atlas results as suggesting that separating generator learning from sampler-level conditioning is useful for rare target populations.Source: §6 · p. 9
  • Terminal-time calibration uses second-order excess covariance, while finite-particle, density-ratio, or target-discriminator errors can shift rare-target purity–fidelity trade-offs.Higher-order dependence may require additional diagnostics.Source: §6 · p. 9

LARGE LANGUAGE MODEL STATEMENT

The appendix gives exact-sampling and birth–death results, and describes a spectral rule for terminal-time selection. The paper also reports how AI tools were used in writing and coding, while excluding them from theory, experiments, and interpretation.

LARGE LANGUAGE MODEL STATEMENT

  • LLM-based tools assisted with related-work retrieval and manuscript polishing, and AI coding assistants helped with implementation and debugging; authors reviewed the cited references, text, and code.The authors state that LLMs were not used to develop the framework, formulate or prove mathematical claims, design experiments, generate data, or interpret results.Source: LARGE LANGUAGE MODEL STATEMENT · p. 10

A.1 REVERSE PROCESS, FFBS, AND ENDPOINT TILTING

  • The appendix states that FFBS is exact on any grid when initialized from qT, and that endpoint tilting leaves the bridge unchanged because the forward transition process is unchanged.Reverse propensities need not be mass action even when the forward propensities are.Source: §A.1 · p. 13

A.2 GENERATING FUNCTIONS AND BIRTH–DEATH SOLUTION

  • For independent birth–death coordinates, the generating-function solution needs no factorization of the initial generating function; setting g0(z)=z^n0 and extracting coefficients gives the transition kernel.Factorial-cumulant decay follows by differentiation of the corresponding expression.Source: §A.2 · p. 14

A.3 SPECTRAL TERMINAL-TIME CRITERION

  • The terminal-time criterion equates the strongest positive second-order signal with the empirical Poisson threshold, whose population excess covariance is zero under the product-Poisson null.The threshold is the largest normalized eigenvalue expected from sample fluctuation alone, using the Marchenko–Pastur edge.Source: §A.3 · p. 14
  • The criterion calibrates decay of detectable positive excess-covariance structure, but does not itself bound the discrepancy between qTO and π.For the independent birth–death model, factorial cumulants of order m ≥2 decay as e−mt, while thresholding uses the strongest positive second-order mode.Source: §A.3 · p. 14

A.4 CONDITIONAL SAMPLER DETAILS

The appendix details limits of marginal tilting, the sampler's reverse-chain weighting, and how the count-native process relates to continuous diffusion. A terminal-time example illustrates the covariance-based calibration rule.

A.4 CONDITIONAL SAMPLER DETAILS

  • Targets with identical coordinate marginals yield identical marginal ratios and tilted posteriors at every reverse step, so marginally tilted reverse processes cannot distinguish their joint structure.Intermediate approximations affect weight variance and resampling efficiency, while the terminal density-ratio correction determines the formal endpoint target.Source: §A.4 · p. 15
  • The single-cell sampler uses 32 reverse steps and 2nout particles to generate nout cells, with systematic resampling triggered when ESS/M falls below 0.5, subject to a two-step minimum gap.Final output selection uses weighted systematic resampling and is separate from chain-internal resampling.Source: §A.4 · p. 15

A.5 CONNECTION TO DIFFUSION MODELS

  • In the large-system-size limit, the reverse CRN has the familiar reverse-diffusion structure; at finite system size, the sampler instead uses exact count kernels at V=1.Neither forward nor reverse diffusion approximation is required by that sampler.Source: §A.5 · p. 16

B.1 TERMINAL-TIME CALIBRATION

  • For the illustrated example, the excess-covariance signal reaches the empirical Poisson floor at TO=3.09; the 0.25TO reconstruction differs visibly, while TO and 1.5TO reconstructions are similar to the data and each other.The example illustrates that further noising beyond the selected time produces little visible change in the reconstructed marginal.Source: §B.1 · p. 17, Figure 2 · p. 17

B.2 TWO-DIMENSIONAL COUNT MIXTURE

The synthetic count mixture tests whether steering recovers a rare target joint distribution when an alternative has the same coordinate marginals. It compares marginal tilt, FK steering, and tilted FK steering on distributional agreement and ancestry.

B.2 TWO-DIMENSIONAL COUNT MIXTURE

  • The target components c1 and c5 comprise 1% of training-mixture mass, while c3 and c7 have matching coordinate marginals but a different joint distribution.The eight-component mixture is defined on the count lattice {0, . . . , 63}^2.Source: §B.2 · p. 17
  • Energy distance assesses joint-distribution discrepancy, whereas mean marginal total variation cannot distinguish the target pair from the alternative pair.The energy-distance reference is an equally sized i.i.d. sample from the exact target distribution.Source: §B.2 · p. 18
  • Tilted FK steering concentrates on the target pair while retaining more distinct initial ancestors than FK steering, which leaves many output descendants from relatively few ancestors.Marginal tilt alone generates samples near all four rare components rather than distinguishing the target pair.Source: §B.2 · p. 18, Figure 4 · p. 19
  • Before output sampling, tilted FK endpoint particles are closer to the target first-coordinate marginal, while FK steering requires a larger redistribution of probability mass.The target and alternative component pairs have identical one-dimensional marginals.Source: §B.2 · p. 19, Figure 5 · p. 19

B.3 SINGLE-CELL EXPERIMENTAL DETAILS

The single-cell experiments use shared data splits and preprocessing to assess conditional generation, biological structure, and downstream utility. They report low-dimensional localization, differential-expression agreement, and classification using generated cells in place of real target-class training cells.

B.3 SINGLE-CELL EXPERIMENTAL DETAILS

  • The atlas data use an 8:1:1 train-validation-test split, shared across methods, with 2,000 highly variable genes and test cells as the real-data reference.Endothelial, myeloid, and neuronal training populations contain 80,463, 18,422, and 3,168 cells, respectively.Source: §B.3 · p. 20
  • Marginal ratios are estimated from target training cells and unconditional samples, then added to posterior logits with τ = 0.4 for endothelial cells and τ = 0.2 for myeloid and neuronal cells.Residual discriminators are gradient-boosted classifiers trained to distinguish real target cells from marginally tilted samples.Source: §B.3 · p. 20
  • The single-cell sampler uses 32 reverse steps, 16 endpoint draws per intermediate potential, and resampling when ESS/M < 0.5, subject to a two-step minimum gap.Final output cells are selected separately by weighted systematic resampling from terminal particles.Source: §B.3 · p. 21

B.4 BIOLOGICAL FIDELITY AND DOWNSTREAM UTILITY

  • CRNDiff-generated neuronal cells concentrate near the real neuronal population in fixed PCA projections, while baseline projections are more dispersed.These embeddings describe low-dimensional geometry and do not establish full count-space fidelity.Source: Figure 6 · p. 22, §B.4 · p. 23
  • Differential-expression agreement is evaluated by top-100 signed log-fold-change overlap and Spearman correlation across all 2,000 gene rankings.Evaluation resamplings reuse fixed generated-cell pools rather than regenerating cells.Source: §B.4 · p. 23, Table 4 · p. 23
  • Generated-cell training yields macro F1 and target-class F1 values close to the real-data reference on the evaluated eleven-class classification task.Generated cells replace real training cells for three target classes; the other eight classes retain real training cells.Source: §B.4 · pp. 23–24, Table 5 · p. 24

B.5 SAMPLER DIAGNOSTICS AND ABLATIONS

Sampler diagnostics compare marginal tilt, FK steering, and tilted FK steering on one frozen generator, including their distributional outcomes, ancestry, and sensitivity to tilt strength.

B.5 SAMPLER DIAGNOSTICS AND ABLATIONS

  • Marginal tilt alone yields dispersed projections and low purity, particularly for myeloid and neuronal targets, indicating that coordinate-marginal adaptation is insufficient in these tested settings.The ablation contrasts marginal tilt with FK correction and their combination on the frozen generator.Source: §B.5 · p. 24, Table 6 · p. 24
  • Tilted FK steering improves neuronal distributional metrics and preserves greater genealogical diversity than FK steering for every evaluated target, but does not uniformly improve all metrics.FK steering has better endothelial distributional metrics, while myeloid results are mixed; the two methods also use separately fitted discriminators.Source: §B.5 · pp. 24–25, Table 6 · p. 24
  • Despite more intermediate resampling events for myeloid and neuronal targets, tilted FK steering preserves greater final genealogical diversity and has higher median pre-resampling ESS at typical events.The diagnostic resampling counts exclude terminal weighted output selection.Source: §B.5 · p. 25, Table 7 · p. 26, Figure 9 · p. 26
  • Tempered tilt settings improve all three distributional metrics over τ = 1 across the three targets, but endothelial purity is highest at τ = 1 and no single exponent optimizes every outcome.The tested τ = 0.2 versus τ = 0.4 trade-off depends on target population and metric.Source: Table 8 · p. 26, §B.5 · p. 27

B.6 SEPARATING GENERATOR QUALITY FROM CONDITIONING

This section first assesses unconditional generator quality, then compares inference-time conditioning strategies on a frozen CRNDiff generator. The strategy trade-offs differ by target population, especially for rare neuronal cells.

B.6 SEPARATING GENERATOR QUALITY FROM CONDITIONING

  • CRNDiff has the lowest unconditional sliced-W1 and the closest detected-gene coefficient of variation to the test sample, but lower median Fano factor and library-size coefficient of variation.scANVI most closely matches median Fano, while CFGen most closely matches library-size variability.Source: §B.6 · p. 27, Table 9 · p. 27
  • CRNDiff most closely matches test-sample endothelial and neuronal fractions, while MDLM is closest for myeloid; neuronal cells remain below 6% of generated pools.The number available for conditional selection depends on that fraction and the candidate budget.Source: §B.6 · p. 27, Table 9 · p. 27
  • For myeloid cells, best-of-N at 8× improves all three distributional metrics at similar purity, whereas 16× increases purity but worsens fidelity.Source: §B.6 · p. 28, Table 10 · p. 28
  • For rare neuronal cells, tilted FK steering achieves higher purity and lower W1 than every tested best-of-N and value-guided setting, while fine-tuning improves neuronal distributional fidelity.Best-of-N retains a large purity gap at the largest tested candidate-pool budget; tilted FK steering retains higher purity without retraining.Source: §B.6 · p. 28, Table 10 · p. 28
Loading 2609.31149v1…