Discrete Diffusion for Single-Cell Gene Expression Modeling

This paper introduces Discrete Cell Models (DCM), a novel diffusion-based framework that operates directly on discrete gene count data to outperform state-of-the-art continuous latent methods in both unconditional and conditional single-cell gene expression modeling.

The Gist

Imagine you are trying to teach a computer to understand how a living cell works. A cell is like a tiny, bustling factory where thousands of different machines (genes) are running, turning on and off, and producing parts (RNA molecules).

To "see" inside this factory, scientists use a technique called single-cell RNA sequencing. Instead of giving you a smooth video of the factory, this technique gives you a count sheet. It's a list that says: "Gene A made 0 parts, Gene B made 5 parts, Gene C made 12 parts."

The Old Way: Trying to Smooth Out the Rough Edges

For a long time, computer models trying to learn from these count sheets had a weird habit. They would take these whole numbers (0, 5, 12) and force them into smooth, continuous numbers (like 0.04, 5.23, 11.99).

Think of it like this: Imagine you are trying to teach a child to count apples.

• The Reality: You have 0 apples, 1 apple, or 2 apples. You can't have 1.5 apples.
• The Old Models: They tried to teach the child that 1.5 apples is a valid thing. They would say, "Maybe you have 1.2 apples?" The model would spend a lot of brainpower trying to figure out what "1.2 apples" looks like, even though that's impossible in the real world.

The authors of this paper argue that this is a waste of time and creates confusion. If the data is made of whole numbers (discrete), the model should think in whole numbers, not smooth decimals.

The New Way: Discrete Cell Models (DCM)

The team introduces a new framework called Discrete Cell Models (DCM). Instead of smoothing out the data, they keep it "chunky" and discrete, just like the real world.

They use a technique called Discrete Diffusion. Here is a simple analogy for how it works:

1. The "Noise" Game: Imagine you have a perfect, clear sentence written on a piece of paper.
2. Adding Noise: You slowly start scribbling over the words, replacing them with random symbols until the sentence is just gibberish. This is the "forward" process.
3. The Learning: The computer watches this happen millions of times. It learns the rules of how the words get scrambled.
4. The Magic: Then, the computer tries to do it in reverse. It starts with the gibberish and slowly removes the noise, word by word, to reconstruct the original perfect sentence.

In the world of cells, the "sentence" is the list of gene counts. The "gibberish" is a cell with random gene counts. The computer learns how to turn the gibberish back into a realistic, healthy cell.

Why is this better?

The paper shows that by sticking to whole numbers, the new model (DCM) is much better at two things:

1. Creating Fake Cells (Unconditional Generation): When asked to invent a new cell from scratch, DCM creates cells that look and act much more like real ones than the old models. It's like a master chef who knows exactly how many eggs to use, rather than a chef who guesses "maybe 2.3 eggs."
2. Predicting Changes (Conditional Generation): Scientists often want to know: "What happens if we turn off Gene X?" or "What happens in a liver cell vs. a skin cell?"
   • The old models sometimes got the average right but messed up the details.
   • DCM is incredibly good at predicting the average outcome (the overall shape of the cell's reaction). It's like a weather forecaster who is spot-on about the temperature but maybe slightly off on the exact wind speed.

The Results

The researchers tested their new model against the current "champions" of the field (like scVI and scLDM).

• On a test called the Dentate Gyrus (a specific type of brain cell), their model was 5 times better at matching the real data distribution than the previous best model.
• On a test involving genetic perturbations (changing genes), their model was the best at predicting the overall outcome, beating all other competitors.

The Bottom Line

This paper is a reminder that sometimes, the simplest approach is the best. If the data is made of whole numbers (like counting molecules), don't try to force it into a smooth, continuous world. By respecting the "discrete" nature of biology, the new model builds a more accurate "virtual cell" that helps scientists understand how life works at a fundamental level.

It's like realizing that to build a perfect Lego castle, you shouldn't try to melt the bricks into a smooth clay shape; you should just use the bricks exactly as they are.

Technical Summary

1. Problem Statement

Single-cell RNA sequencing (scRNA-seq) data consists of discrete, sparse, integer-valued count matrices representing mRNA molecules. However, the dominant generative modeling approaches (e.g., scVI, scLDM, scGPT) rely on continuous latent representations. These methods typically:

1. Embed discrete counts into continuous vector spaces.
2. Model the data using continuous diffusion or flow matching.
3. Recover discrete counts only at the sampling stage via rounding or distribution fitting.

The Authors' Critique:
The paper argues that this "continuous relaxation" introduces fundamental representational limitations:

• Wasted Capacity: Continuous models assign probability mass to non-integer values (impossible states).
• Metric Mismatch: The natural metric for count data is not Euclidean. The biological difference between 0 and 1 transcript (presence/absence) is distinct from the difference between 100 and 101 (sampling noise), a nuance Euclidean metrics struggle to capture without learning it from data.
• Bimodality: Lowly expressed genes exhibit bimodal behavior (on/off states) that is naturally discrete but requires complex learning in continuous spaces.
• Information Gap: Continuous relaxation forces the model to learn discretization boundaries rather than the intrinsic structure of the discrete space.

2. Methodology: Discrete Cell Models (DCM)

The authors propose Discrete Cell Models (DCM), a framework that applies Score Entropy Discrete Diffusion (SEDD) directly to raw transcript counts, eliminating continuous relaxation.

Core Framework

• Representation: Gene expression profiles are treated as discrete sequences $x \in \mathcal{X}^M$, where $\mathcal{X} = \{0, 1, \dots, K\}$ represents binned or raw counts.
• Forward Process: A continuous-time Markov process progressively corrupts clean data $x_0$ by transitioning tokens toward a special 'MASK' state (absorbing diffusion). The transition is governed by a diffusion matrix $Q_t$.
• Reverse Process: The model learns to reverse this corruption using Concrete Scores. Instead of predicting gradients in continuous space, the model estimates the ratio of probabilities between neighboring discrete states (Hamming distance 1):
  $$s_\theta(x_t, t, c)_{i,v} \approx \frac{p_t(x_t \text{ with } i \to v)}{p_t(x_t)}$$
• Architecture:
  • Backbone: A Diffusion Transformer (DiT) with Adaptive LayerNorm (AdaLN).
  • Context Length: Handles full gene expression profiles ($\approx 17k$ genes) using Flash Attention.
  • Conditioning: Supports unconditional and conditional generation. Conditioning variables (cell type, perturbation identity) are embedded (using one-hot or protein language models) and concatenated with diffusion-time embeddings via AdaLN.
• Training Objective: The objective simplifies to a weighted Denoising Cross-Entropy (DWDSE) loss, which is tractable and directly optimizes the likelihood of the discrete data distribution.

3. Key Contributions

1. First Discrete Diffusion for scRNA-seq: DCM is the first framework to apply discrete diffusion directly to single-cell transcriptomics, avoiding the intermediate continuous latent space used by SOTA methods like scLDM and scDiffusion.
2. End-to-End Architecture: Unlike two-stage approaches (VAE + Diffusion), DCM operates end-to-end on discrete tokens, simplifying the design and reducing the number of hyperparameters.
3. Conditional Generation: The model supports precise modeling of complex biological scenarios, including cell-type-specific responses to genetic perturbations (e.g., gene knockouts).
4. Theoretical Justification: The paper provides a strong argument for why discrete modeling is superior for count data, citing information-theoretic gaps and the biological nature of gene expression (zero-inflation, bimodality).

4. Experimental Results

The authors evaluated DCM on two benchmarks: Dentate Gyrus (unconditional generation) and Replogle (conditional perturbation prediction). Metrics used were W2 Distance (global geometric alignment) and MMD$^2_{RBF}$ (fine-grained statistical similarity).

A. Unconditional Generation (Dentate Gyrus)

DCM significantly outperformed all baselines (scVI, scDiffusion, CFGen, scLDM).

• W2 Distance: DCM achieved 5.913, a nearly 2-fold improvement over the best continuous baseline, scLDM (10.615).
• MMD$^2_{RBF}$: DCM achieved 0.019, a 5-fold improvement over the closest baseline, CFGen (0.075).
• Efficiency: These results were achieved with a 5M parameter model, which is substantially smaller than scLDM's two-stage architecture.

B. Conditional Generation (Replogle Perturbation)

DCM was tested on predicting gene expression under specific genetic perturbations (gene knockouts) across different cell lines.

• W2 Distance: DCM set a new state-of-the-art, achieving 10.03 on the full Replogle dataset (vs. 11.292 for scLDM), representing a 13% improvement in global distributional alignment.
• MMD$^2_{RBF}$: DCM was competitive on the Parse 1M benchmark (0.020 vs. 0.027 for scLDM) but showed higher error on the full Replogle dataset (0.688 vs. 0.200 for scLDM).
  • Analysis: The authors attribute the MMD gap to the conditioning mechanism. The current additive embedding concatenation may fail to capture complex interactions between cell types and perturbations (e.g., a knockdown having different effects in different cell lines), a capability scLDM's cross-attention might handle better.

5. Significance and Conclusion

• Paradigm Shift: The paper establishes that discrete diffusion is a promising direction for foundational models in cellular biology. It challenges the field's default reliance on continuous relaxation for count data.
• Biological Fidelity: By respecting the discrete, sparse, and zero-inflated nature of scRNA-seq data, DCM captures global transcriptomic structures and mean/variance alignments more effectively than continuous models.
• Future Directions: While DCM excels at global alignment (W2), the gap in higher-order dependency modeling (MMD) suggests that future work should focus on improving how discrete models handle complex conditional interactions (e.g., via multiplicative conditioning or attention mechanisms).
• Broader Impact: The principle of matching the generative model's state space to the discrete nature of biological measurements (count-based assays) can be extended beyond transcriptomics to other molecular data types, enabling more faithful "virtual cells."

In summary, DCM demonstrates that operating directly in the discrete domain yields superior generative performance for single-cell data, offering a more natural and efficient alternative to the continuous latent space approaches that have dominated the field.