🧬 biology

Better Models, Faster Training: Sigmoid Attention for single-cell Foundation Models

本文证明，在单细胞基础模型中用 sigmoid 注意力机制替代 softmax，可借助理论上有界的导数实现更优的表征质量、更快的训练速度和更强的稳定性，并得到了高度优化的 Triton 内核实现的支持。

原作者： Vijay Sadashivaiah, Georgios Dasoulas, Judith Mueller, Soumya Ghosh

发布于 2026-05-01

📖 1 分钟阅读☕ 轻松阅读

原作者： Vijay Sadashivaiah, Georgios Dasoulas, Judith Mueller, Soumya Ghosh

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ⚕️ 这是一篇未经同行评审的预印本的AI生成解释。这不是医疗建议。请勿根据此内容做出健康决定。阅读完整免责声明

以下是用简单语言和创意类比对该论文的解释。

宏观图景：解读“生命语言”的更佳方式

想象你正在尝试教计算机理解人类细胞的“语言”。在这种语言中，每个基因都是一个单词，而整个细胞则是一个句子。科学家们一直使用一种名为Transformer的强大工具（即那种能写文章或与你聊天的同类人工智能）来阅读这些细胞句子。

然而，这些计算机用来理解单词之间关系的标准工具被称为Softmax 注意力机制。本文作者认为，对于生物学而言，这种标准工具就像试图用一本僵化、刻板的规则书去应对活细胞中混乱的现实。他们提议将其替换为一种名为Sigmoid 注意力机制的新工具，称其速度更快、更稳定，且实际上能帮助人工智能更好地学习。

以下是他们的具体拆解：

1. 问题所在：“零和博弈”

旧方法（Softmax）：
想象你是一位老师，手中只有有限数量的“注意力代币”（比如 100 张贴纸）要分发给教室里的学生。如果你给了学生 A 50 张贴纸，你就必须从其他人那里拿走这些贴纸。这就是Softmax的工作原理。它迫使人工智能做出决定：“我将重点关注基因 X，因此我必须忽略基因 Y。”

这对生物学为何不利：
在真实的细胞中，基因往往协同工作。单个基因可能同时受到多个不同调节因子的控制。这就像一名学生需要同时得到三位不同老师的帮助。Softmax 的“零和”规则迫使人工智能只选择一位老师，这与生物学实际运作方式不符。

新方法（Sigmoid）：
Sigmoid 注意力机制就像如果学生值得，就给他们每人一张贴纸，且没有数量限制。如果学生 A、学生 B 和学生 C 都需要帮助，老师可以给所有人贴纸。这允许人工智能说：“这个基因很重要，而且那个基因也很重要”，而无需为了关注其中一个而忽略另一个。

2. 稳定性问题：“走钢丝者”

旧方法（Softmax）：
Softmax 就像一位对风非常敏感的走钢丝者。如果数值变得太大（当基因序列很长时就会发生这种情况），数学计算可能会失控。“风”会将走钢丝者吹下绳索，导致训练崩溃。在论文中，他们通过移除安全网（梯度裁剪）并使用非常长的序列来测试这一点。Softmax 模型从悬崖上跌落，其数学计算爆炸了10,000 倍。

新方法（Sigmoid）：
Sigmoid 就像一位在宽阔平坦桥梁上行走的人。它拥有内置的护栏。无论数值变得多大，数学计算都保持在安全、可预测的范围内。在同样的“无安全网”测试中，Sigmoid 模型保持完美行走，从未崩溃。

3. 速度问题：“参差不齐的拼图”

挑战：
生物数据是混乱的。一个细胞可能只有 500 个基因（短句），而另一个可能有 10,000 个基因（长篇小说）。为了将它们一起处理，计算机通常必须用“空白空间”（零）填充较短的序列，使它们看起来长度相同。这就像试图通过将短诗塞满填充泡沫来将其与长篇小说装入同一个盒子。

解决方案：
作者构建了一种名为TritonSigmoid的全新超快速引擎。

类比： 想象一辆通常必须驶向社区中每一栋房子的送货卡车，即使是空房子（填充部分）也不例外。TritonSigmoid 是一辆智能卡车，它确切知道哪些房子是空的，并完全跳过它们。
结果： 这个新引擎速度极快。在最新的超级计算机芯片（NVIDIA H100）上，其运行速度达到515 TFLOPS（每秒万亿次计算）。这比当前最佳的 Softmax 引擎更快，甚至比之前的 Sigmoid 尝试更快。它使这些生物模型的训练速度提高了高达10%。

4. 结果：更聪明的头脑，更快的训练

该团队训练了四种不同的人工智能模型，以观察哪种模型学习得更好。他们使用了一个包含 1.31 亿个细胞的庞大数据集。

更好的学习： 使用 Sigmoid 的模型不仅学得更快，而且学得更好。它们区分不同细胞类型（例如区分心脏细胞和脑细胞）的能力比 Softmax 模型提高了25%。
更低的错误率： Sigmoid 模型在预测基因模式时犯的错误更少。
无崩溃： 在压力测试中，Softmax 模型崩溃并被迫重启，浪费了时间和金钱。Sigmoid 模型则毫无故障地完成了任务。

总结

该论文声称，对于理解单细胞生物学中复杂、混乱且多变的数据，标准的"Softmax"注意力机制过于僵化且不稳定。通过切换到Sigmoid 注意力机制，他们创建了一个系统，能够：

允许多重连接（基因可同时受多种因素影响）。
保持稳定，即使数学计算变得激烈（不再崩溃）。
运行更快，通过智能跳过空白数据（填充部分）。

他们已将这种新的、更快的引擎（TritonSigmoid）免费开放，以便其他科学家利用它构建更好的生物模型。

以下是论文“更好的模型，更快的训练：用于单细胞基础模型的 Sigmoid 注意力”的详细技术总结。

1. 问题陈述

本文解决了训练生物基础模型（特别是针对单细胞 RNA 测序，即 scRNA-seq）时的两个关键挑战：

Softmax 的训练不稳定性： 标准 Transformer 模型使用Softmax 注意力，它将注意力权重在 token 之间进行归一化，使其总和为 1。这创造了一种“零和”竞争机制，即增加对某个基因的注意力会迫使对其他基因的注意力减少。在生物学背景下，基因通常受多个独立因素的共同调控，这种竞争假设在生物学上是不现实的。此外，随着序列长度的增加，Softmax 容易出现数值不稳定性（梯度爆炸），特别是在注意力熵坍缩到少数 token 上时。
生物数据变异性的低效性： 生物序列（细胞）的长度具有极大的变异性（从数百到数千个表达基因不等）。与可以打包成固定长度的文本不同，生物样本必须单独处理，并通过**填充（padding）**至批次中的最大长度。现有的高效注意力内核（如 FlashAttention-2 或 FlashSigmoid）要么缺乏原生填充支持，要么无法高效处理生物数据的“参差不齐（jagged）”特性，导致显著的计算浪费。

2. 方法论

作者提出用Sigmoid 注意力替代 Softmax，并开发了一种专用的、高性能的 GPU 内核来支持它。

A. Sigmoid 注意力机制

Sigmoid 注意力不对行进行归一化（如 Softmax 那样），而是对点积分数应用逐元素的 Sigmoid 函数（ $\sigma(x) = \frac{1}{1+e^{-x}}$ ）：
$\text{SigmoidAttn}(X) = \sigma\left(\frac{QK^\top}{\sqrt{d}} + b\right)V$

解耦： 每个注意力权重独立计算。一个查询可以同时强烈地关注多个键，更好地模拟共调控。
偏置项： 添加偏置 $b = -\log(n)$ 以近似 Softmax 的缩放行为，同时保持逐元素的独立性。

B. 理论稳定性分析

本文提供了 Sigmoid 稳定性的理论依据：

有界导数： Sigmoid 函数的导数全局有界于 $0.25 $（$ \sigma'(x) = \sigma(x)(1-\sigma(x))$）。相比之下，Softmax 的 Lipschitz 常数随分数幅度的增加呈指数级增长。
对角雅可比矩阵： Sigmoid 注意力的雅可比矩阵是对角的，意味着不同 token 的梯度是解耦的。Softmax 具有稠密的雅可比矩阵，耦合了所有 token 的梯度，导致分数增大时出现不稳定性。

C. TritonSigmoid 内核实现

为了使 Sigmoid 注意力适用于生物数据，作者开发了TritonSigmoid，这是一个使用 Triton 语言定制的 GPU 内核。关键创新包括：

原生填充支持： 它实现了块稀疏计算，完全跳过查询和键两侧的填充块，以避免浪费 FLOPs。
融合操作： 它将注意力计算融合到单个内核中，使用硬件优化的 tanh 近似来替代 Sigmoid 函数。
反向传播分解： 它将反向传播拆分为两个内核（一个用于 $Q$ 梯度，一个用于 $K/V$ 梯度），以消除原子操作并即时重新计算激活值，从而提高内存效率。
硬件优化： 专为 NVIDIA H100/Blackwell 架构设计，支持 BF16 精度。

3. 主要贡献

TritonSigmoid 内核： 一个高性能实现，在 NVIDIA H100 GPU 上实现了515 TFLOPS的算力。它优于 FlashAttention-2（361 TFLOPS）和 FlashSigmoid（440 TFLOPS），并且是首个支持变长生物序列高效填充的内核。
实证性能提升： 证明了在六个不同的单细胞数据集上，Sigmoid 注意力产生的学习表示优于 Softmax。
训练稳定性： 证明了 Sigmoid 注意力在没有梯度裁剪的情况下，能够防止长上下文（8K token）场景中的灾难性训练发散，而这是 Softmax 模型中观察到的失效模式。

4. 结果

A. 计算效率

吞吐量： 与标准 PyTorch 注意力相比，TritonSigmoid 在前向传播中实现了高达7.15 倍的加速，在反向传播中实现了2.44 倍的加速。
填充效率： 在 25% 填充的情况下，TritonSigmoid 的前向传播速度比 FlashAttention-2 快29%。由于填充导致的 TFLOPS 下降极小（仅约 9.3%），证明该内核有效地跳过了掩蔽区域。
端到端训练： 使用 Sigmoid 注意力训练的模型比 Softmax 对应模型快3–10%，且随着上下文长度的增加，差距会扩大（例如，对于 14 亿参数模型，在 8K 上下文下快 7.5%）。

B. 模型质量（表示学习）

在六个保留的单细胞数据集（大脑、血液、结肠、肺、心脏）上进行了评估：

验证损失： Sigmoid 模型在所有数据集和上下文长度（2K 和 4K）上均一致实现了更低的验证损失。
细胞类型分离： Sigmoid 模型在细胞类型之间显示出高 25% 的最大均值差异（MMD），表明在嵌入空间中具有更好的分离度。
生物学保守性： Sigmoid 在所有 6 个数据集上实现了更优越的细胞类型凝聚力（轮廓系数），并在 4/6 个数据集上获得了更好的聚合生物学保守性分数。
可视化： UMAP 可视化证实，与 Softmax 相比，Sigmoid 嵌入将细胞类型聚类得更紧密、更清晰。

C. 稳定性压力测试

设置： 在 8K token 序列上训练 1.6 亿参数模型，不使用梯度裁剪。
结果：
- Softmax： 在约 55k 步时发生灾难性发散。全局梯度范数爆炸了四个数量级（从 $100 $增至$ 1.6 \times 10^6$），注意力分数呈指数级增长。
- Sigmoid： 在 80k 步内训练稳定，梯度有界（10–100 范围），注意力分数受控，证实了理论上的稳定性保证。

5. 意义

本文确立了Sigmoid 注意力作为生物基础模型中 Softmax 的理论基础扎实且实证更优的替代方案。

生物学相关性： 通过去除 Softmax 的人工竞争约束，它更好地与基因调控（共激活）的生物学特性相一致。
实际部署： 通过TritonSigmoid解决了变长序列的效率问题，使得训练此前不稳定或计算成本过高的大规模、长上下文生物模型成为可能。
开源： 作者发布了该实现，为单细胞生物学社区提供了一个可即插即用的 Softmax 替代品，能够提供更快的训练速度、更好的稳定性和更高质量的表示。

代码可用性： https://github.com/MSDLLCpapers/triton-sigmoid