Discrete Diffusion for Single-Cell Gene Expression Modeling

该论文提出了离散细胞模型（DCM），一种直接在离散域学习细胞表征的扩散框架，旨在克服现有连续潜在表示方法的局限，并在单细胞基因表达建模及扰动响应预测任务中显著超越了包括 scVI 和 scLDM 在内的最先进方法。

要点

这篇论文介绍了一种名为**“离散细胞模型”（Discrete Cell Models, DCM）**的新方法，它就像是为单细胞基因数据量身定做的“超级生成器”。

为了让你轻松理解，我们可以把细胞想象成一个巨大的乐高城堡，而基因就是构成城堡的乐高积木。

1. 以前的做法：把乐高融化成泥（连续模型）

在 DCM 出现之前，科学家们研究细胞时，习惯把乐高积木（基因表达量，即整数计数）先融化成泥巴（连续空间）。

• 为什么这么做？ 就像用泥巴捏形状比用一块块积木拼更容易控制，数学上处理“泥巴”（连续数值）比处理“积木”（离散的整数 0, 1, 2...）要顺手得多。
• 有什么问题？ 细胞里的基因表达是非黑即白的：要么有 0 个分子（没表达），要么有 1 个分子（开始表达）。这就像乐高积木，你要么有一块，要么没有。
  • 把积木融化成泥巴后，模型会生成"0.5 块积木”这种现实中不存在的东西。
  • 从 0 到 1 的变化（从无到有）和从 100 到 101 的变化（只是多了一点点噪音），在生物学上意义完全不同，但泥巴模型很难区分这种本质区别。

2. 新做法：直接玩积木（DCM 离散扩散）

这篇论文提出的 DCM 就像是一个直接玩乐高积木的机器人。它不再把数据融化，而是直接在“积木”的世界里工作。

• 核心概念：离散扩散（Discrete Diffusion）
  想象一下，你有一堆拼好的乐高城堡（真实的细胞数据）。

  1. 破坏过程（前向扩散）： 机器人随机把城堡里的积木一块块拿走，换成“空白”的标记（MASK），直到城堡变成一堆乱码。
  2. 重建过程（反向扩散）： 机器人看着这堆乱码，根据它学到的规律，一步步把“空白”替换回正确的积木，直到城堡恢复原样。

  以前的模型是在“泥巴”里修补，而 DCM 是在“积木”堆里修补，每一步都严格遵循“必须是整数块”的规则。

3. 它有多厉害？（实验结果）

作者把 DCM 和目前最顶尖的几种方法（比如 scLDM、scGPT 等）进行了比赛，就像让不同的建筑师重建同一个乐高城堡。

• 比赛一：无条件的“盲拼”（Dentate Gyrus 数据集）

  • 结果： DCM 拼出来的城堡，在形状相似度（W2 距离）上比第二名好了一倍，在细节纹理（MMD 指标）上更是领先了 5 倍！
  • 比喻： 如果其他模型拼出来的城堡看起来像“大概像那么回事”，DCM 拼出来的则连窗户上的花纹都一模一样。而且它用的积木块（模型参数）还更少，更聪明。

• 比赛二：带指令的“定制拼”（Replogle 数据集）

  • 场景： 科学家给指令：“在肝脏细胞里，把基因 A 关掉，看看城堡会变成什么样？”
  • 结果： DCM 在预测这种“基因改造”后的细胞状态时，整体结构（W2 距离）达到了世界最好水平。
  • 小瑕疵： 在极其细微的局部连接（MMD 指标）上，它比最强的对手稍微弱了一点点。
  • 原因猜测： 可能是因为 DCM 太专注于“积木块”本身的准确性，而在处理“积木块之间复杂的互动关系”时，还需要更精细的“胶水”（条件机制）。但这就像是一个刚学会玩积木的天才，大局观已经无敌了，细节还在打磨中。

4. 为什么这很重要？

这就好比以前我们试图用油画（连续模型）来描绘像素画（离散基因数据），虽然也能画得很像，但总归有点“失真”。

DCM 证明了：直接用最原始的数据格式（离散计数）来建模，不仅能更真实地还原细胞的“本来面目”，还能更准确地预测基因被修改后会发生什么。

总结一句话：
这篇论文告诉我们，与其把细胞数据“软化”成好处理的泥巴，不如直接尊重它“积木”的本质。DCM 就是那个直接玩积木的超级机器人，它拼出的细胞模型更真实、更精准，为未来设计“虚拟细胞”和开发新药打下了更坚实的基础。

技术摘要

这是一篇发表于 ICLR 2026 的会议论文《DISCRETE DIFFUSION FOR SINGLE-CELL GENE EXPRESSION MODELING》（用于单细胞基因表达建模的离散扩散模型）的详细技术总结。

1. 研究背景与问题 (Problem)

• 核心问题：现有的单细胞转录组（scRNA-seq）生成模型主要依赖连续潜在表示（continuous latent representations）。这些模型将本质上离散且稀疏的基因计数（整数）数据映射到连续向量空间中进行建模，生成后再通过取整或分布采样恢复为离散值。
• 现有方法的局限性：
  • 表示浪费：连续模型将概率质量分配给非整数值（即不可能存在的测量状态），浪费了模型容量。
  • 度量不匹配：计数空间的自然度量并非欧几里得距离。例如，0 到 1 个转录本（表达有无）的生物学差异，与 100 到 101 个转录本（采样噪声）的差异截然不同。连续嵌入难以自然捕捉这种不对称性。
  • 生物学特性丢失：基因调控网络中的关系往往依赖于基因产物的功能存在（对于低表达基因是二态的、随机的），离散模型能天然捕捉这种特性，而连续模型需要额外学习。
  • 信息论视角：当真实数据生成过程是离散的，任何连续松弛都会引入表示间隙，迫使模型学习离散化边界而非离散空间内的结构。
• 目标：开发一种直接在离散域（discrete domain）中学习的生成框架，以保留单细胞数据的离散和稀疏特性，从而更准确地建模复杂的生物学场景（如细胞类型特异性转录响应和遗传扰动）。

2. 方法论 (Methodology)

作者提出了离散细胞模型（Discrete Cell Models, DCM），这是一个基于**分数熵离散扩散（Score Entropy Discrete Diffusion, SEDD）**的框架。

• 数据表示：

  • 将单细胞基因表达谱表示为离散序列 $x \in \mathcal{X}^M$，其中 $\mathcal{X} = \{0, 1, ..., K\}$ 是分箱或原始的表达计数词汇表。
  • 每个元素 $x_i$ 代表基因 $i$ 在单个细胞中的表达水平（整数计数）。

• 扩散过程：

  • 前向过程：定义了一个连续时间的离散状态马尔可夫过程，通过独立的 Token 级转换逐步破坏干净的表达谱。采用吸收扩散结构（absorbing diffusion structure），所有 Token 最终都向特殊的'MASK'状态转移。
  • 反向过程：通过**具体分数（Concrete Scores）**来参数化反向过程。具体分数定义为相邻离散状态之间数据分布的比率（$p_t(y)/p_t(x)$），这是连续扩散中 $\nabla \log p_t(x)$ 的离散模拟。

• 模型架构：

  • 骨干网络：使用基于 Transformer 的DiT (Diffusion Transformer) 作为分数网络（Score Network）。
  • 条件机制：支持无条件生成和多条件生成（如细胞类型、遗传扰动）。条件信息（如细胞类型 One-hot 编码、扰动标签的蛋白质语言模型嵌入）通过AdaLN (Adaptive LayerNorm) 机制注入到网络中。
  • 上下文处理：使用全基因表达谱作为 Transformer 的上下文长度（约 17k 基因），利用 Flash Attention 提高效率，并使用特殊'PAD' Token 处理未表达或未选中的基因。

• 训练目标：

  • 采用**去噪交叉熵（Denoising Cross-Entropy）**损失函数。
  • 对于吸收情况，SEDD 的目标简化为加权交叉熵：$L = \mathbb{E} [\sum \text{CrossEntropy}(p_\theta(x_0 | x_t, t, c), x_0)]$。这使得模型能够直接预测原始计数 $x_0$，同时自然地处理离散、高维的基因表达数据。

3. 关键贡献 (Key Contributions)

1. 范式转变：首次将离散扩散模型（SEDD）直接应用于原始单细胞转录组计数数据，摒弃了传统的连续潜在空间松弛方法。
2. 统一架构：提出了一个端到端的 DCM 框架，支持无条件和多条件（细胞类型 + 扰动）生成，无需像 scLDM 那样复杂的“编码器 - 扩散器 - 解码器”两阶段架构。
3. 理论论证：从信息论、度量空间差异和生物学二态性角度，有力论证了直接在离散空间建模优于连续松弛。
4. 性能突破：在多个基准测试中实现了新的最先进（SOTA）性能，特别是在全局分布对齐（W2 距离）上表现卓越。

4. 实验结果 (Results)

作者在两个主要基准数据集上评估了 DCM，并与 scVI, CPA, STATE, scGPT, scLDM (当前 SOTA) 等进行了对比。评估指标包括 MMD$^2_{RBF}$（衡量高阶统计和分布相似性）和 2-Wasserstein 距离 (W2)（衡量全局几何对齐和均值/协方差）。

• 无条件生成 (Dentate Gyrus 数据集)：

  • W2 距离：DCM (5.913) 比领先的连续扩散基线 scLDM (10.615) 降低了近 2 倍 的误差。
  • MMD$^2_{RBF}$：DCM (0.019) 比 CFGen (0.075) 提升了 5 倍，比 scLDM (0.102) 也有显著提升。
  • 效率：DCM 仅使用 500 万参数（5M），远小于 scLDM 的两阶段架构，且训练更简单。

• 条件生成 (Replogle 扰动数据集)：

  • W2 距离：DCM 在所有基线中取得了最佳成绩。在 Parse 1M 基准上，DCM (5.428) 比 scLDM (12.457) 提升了约 56%；在完整 Replogle 数据集上，DCM (10.03) 比 scLDM (11.292) 提升了 13%。
  • MMD$^2_{RBF}$：DCM 在 Parse 1M 上表现优异（0.020 vs 0.027），但在完整 Replogle 数据集上略逊于 scLDM (0.688 vs 0.200)。
  • 分析：作者指出，W2 的显著提升表明 DCM 能准确恢复每种扰动条件下的平均表达谱和基因级方差。MMD 的差距可能源于离散扩散在捕捉高阶基因间相关性（特别是涉及细胞类型与扰动的复杂交互）方面不如连续流匹配（Flow Matching）有效，或者是因为条件机制（简单的加法拼接）未能完全捕捉复杂的交互作用。

5. 意义与结论 (Significance & Conclusion)

• 生物学建模的忠实度：该工作证明了当生成模型的状态空间与生物测量数据的离散、稀疏结构相匹配时，能显著提升表示能力。这为构建更真实的“虚拟细胞”奠定了基础。
• 可扩展性：离散扩散框架不仅适用于单细胞转录组，其原则也可扩展到其他基于计数的分子检测技术。
• 未来方向：虽然 DCM 在全局分布对齐（W2）上取得了 SOTA，但在高阶依赖结构（MMD）的建模上仍有提升空间，特别是针对复杂的条件交互。未来的工作将集中在改进条件机制（如引入交叉注意力）以更好地捕捉基因间的相关性。

总结：DCM 通过直接在离散计数空间进行扩散建模，克服了传统连续模型在单细胞数据上的表示局限，在生成质量和分布对齐上均取得了显著突破，确立了离散扩散作为细胞生物学基础模型的一个重要发展方向。