更优的模型、更快的训练：用于单细胞基础模型的 Sigmoid 注意力机制 Better Models, Faster Training: Sigmoid Attention for single-cell Foundation Models
用 sigmoid 替代 softmax，单细胞基础模型训练更快更稳、表征更优
前置知识
自注意力机制 (Self-Attention)
Transformer 的核心组件，通过 $QK^\top/\sqrt{d}$ 计算序列中每对 token 的相关性，再用 softmax 归一化得到权重，最后对 $V$ 加权求和。
本文的全部讨论都围绕这一机制的两种实现差异展开，必须先理解其计算流程。
Softmax 归一化
将一组实数通过指数化和归一化映射为概率分布，$\text{softmax}(a)_i = e^{a_i}/\sum_j e^{a_j}$。它会让所有权重互相竞争总和为 1。
softmax 的竞争式归一化正是本文要替换掉的关键设计点，理解它才能理解 sigmoid 的优势。
单细胞 RNA 测序 (scRNA-seq) 与基础模型
对单个细胞的全转录组测序，得到每个细胞表达哪些基因、表达量多少。基础模型以每个基因作为 token，在数百万细胞上做掩码语言建模。
本文的所有实验都建立在 scRNA-seq 数据上，读者需要知道这种数据为什么需要长上下文。
FlashAttention 风格的 IO-aware 内核
通过分块 (tiling)、不实例化完整 $n\times n$ 注意力矩阵、利用片上 SRAM 而非 HBM，显著降低显存与加速 attention 的 GPU 实现思路。
TritonSigmoid 的工程基础就是这套设计，不了解 FlashAttention 就难理解它为何高效。
Jacobian 矩阵与 Lipschitz 常数
Jacobian 描述函数对各输入维度的偏导，Lipschitz 常数则刻画函数对输入扰动的最大放大倍数，$\|f(x)-f(y)\| \le L\|x-y\|$。
论文用这两个工具证明 sigmoid 的 Jacobian 是对角且全局 1/4-Lipschitz，是其稳定性的理论核心。
研究动机
现有单细胞基础模型几乎全部沿用 NLP 领域的 softmax 注意力，但转录组数据给注意力机制带来三重压力。第一是数据形态：单细胞基因表达序列高度不均匀（每个细胞从几百到上万基因不等），必须 padding 到统一长度，浪费大量算力；如图 1 所示，CellxGene 数据集中 43% 的细胞表达基因超过 2,048，迫使模型必须支持 4K–16K 长上下文。第二是训练稳定性：softmax 在长序列+高学习率下经常崩溃，文献记录其局部 Lipschitz 常数随分数大小指数增长，注意力熵容易坍缩到一两个 token，导致梯度爆炸。第三是生物学先验不匹配：基因常被多个转录因子协同调控 (co-regulation)，而 softmax 的零和归一化天然强迫注意力在 token 间互相竞争，无法同时给多基因强权重，限制了表征质量。
本文的目标是本文目标有三：(1) 找到一种即插即用的注意力替代机制，能在六类单细胞数据集上稳定获得更低的验证损失、更优的细胞类型聚类结构；(2) 提供一个真正支持 padding 的高效 GPU 内核，使长上下文单细胞模型在 H100 等现代架构上训练可行，在 16K 上下文下达到 515 TFLOPS；(3) 在理论上解释该机制为何更稳定，并在压力测试中实证其在 8K 上下文、关闭梯度裁剪的 160M 模型训练中不崩溃。
与已有工作不同的是，Ramapuram 等人 (2025) 已经在 NLP/视觉上验证了 sigmoid 注意力的潜力，但既未系统研究生物学场景，也未提供兼容 padding 的内核；FlashSigmoid 虽然快但不支持 padding，标准 PyTorch sigmoid 又只有 41 TFLOPS。本文正好填补了这一交叉空白：把 sigmoid 注意力系统化地搬到单细胞基础模型上，并补齐了原生 padding 支持与 NVIDIA Blackwell 兼容性这两个工程拼图。
核心方法
方法的核心思路非常直观：把注意力打分里的 softmax 换成逐元素的 sigmoid，再补一个工程级的 GPU 内核让这件事在生物数据上真正跑得快。先说直觉：softmax 像在一桌菜里只准挑一道最想吃的，sigmoid 则是给每道菜独立打分，可以同时给多道菜高分，这正契合基因被多个调控因子共同作用的生物学直觉。技术路线上分两层：理论层证明 sigmoid 的导数全局不超过 $1/4$、Jacobian 是对角的，因此天然没有跨 token 梯度耦合；工程层基于 Triton 实现 TritonSigmoid，通过块稀疏跳过 padding、tanh 近似 $\sigma$、反向拆分为两个 kernel 等手段，在 H100 上冲到 515 TFLOPS 并原生支持任意长度序列混合。
最本质的创新在于把注意力机制从「竞争式归一化」转为「独立打分」。softmax 的权重 $\text{softmax}(a)_i$ 中每个 $a_i$ 都耦合进分母，导致一个权重上升必让其他下降，且局部 Lipschitz 常数随分数大小指数爆炸；而 sigmoid 权重 $\sigma(q_i^\top k_j/\sqrt{d} + b)$ 只依赖单个点积，导数 $\sigma'(x) = \sigma(x)(1-\sigma(x)) \le 1/4$ 恒成立，Jacobian 对角化让每个位置的梯度互不耦合，从根本上消除 softmax 训练崩溃的根源。这一区别不仅是工程优化，而是注意力范式层面的转变——它在生物学中天然对应基因的协同调控。
方法步骤详情
TritonSigmoid 实现分六步：(1) 输入 $Q,K,V\in\mathbb{R}^{Z\times L\times H\times D}$ 与每样本真实长度 $n_q,n_k$，按 FlashAttention-2 思路做块划分；(2) 前向先判断 query 块是否完全在 padding 区域，若是则写零跳过 (块稀疏)；否则加载 $K,V$ 块算 $S = QK^\top\alpha + b$，padding 位置加 $-\infty$；(3) 用硬件 $\tanh$ 近似 $\sigma(S)\approx 0.5(\tanh(S/2)+1)$ 累加到输出；(4) 反向拆两个 kernel：dQ 重算 $P=\sigma(S)$ 再算 $\partial S = P\odot(1-P)\odot(\partial O V^\top)$；dK/dV 公式对称但转置视角；(5) 反向过程重算 sigmoid 激活省内存，$K$ 转置读取最大化吞吐；(6) 通过 Triton JIT 自动调优选最优 tile 配置，并兼容 torch.compile。
技术新颖性
技术新颖性体现在三点：(1) 第一个原生支持 padding 的高效 sigmoid 注意力内核，相对 FlashSigmoid 解决了生物序列必须 padding 的痛点，相对 PyTorch 原生 sigmoid 提升了 14.58× 速度；(2) 利用 sigmoid 的元素级特性简化计算，去掉了 softmax 的行归一化与指数运算，在硬件层面用更少指令获得更高 FLOPS；(3) 给出完整的 Jacobian-Lipschitz 理论分析，把 sigmoid 的全局 $1/4$ 有界导数与对角 Jacobian 形式化为谱范数上界，并把这一性质直接关联到「无梯度裁剪 8K 训练不崩溃」这一可观测现象，这是已有文献未做的端到端联系。
实验结果
实验分三层证据。第一层是内核效率：TritonSigmoid 在 16K 上下文、head_dim=128、无 padding 时达 515.6 TFLOPS 前向，相对 FlashAttention-2 提升 43%、相对 FlashSigmoid 提升 17%；25% padding 下相对 PyTorch 原生 sigmoid 取得 14.58× 前向加速，padding 仅带来 9.3% TFLOPS 损耗。第二层是端到端训练速度：sigmoid 在 16×H100 上所有配置都快于 softmax，1.4B 模型 8K 上下文节省 645 GPU 小时 (7.5%)。第三层是表征质量：sigmoid 在 6 个 scRNA-seq 数据集、2K 与 4K 上下文下均取得更低验证损失，scIB 框架下 silhouette 6/6 数据集胜出，Heart OFT 28 对细胞类型 MMD 全部胜出、平均提升 25.0%。压力测试中 softmax 在 55.6K 步发散 (梯度爆炸 4 个数量级)，而 sigmoid 顺利完成 80K 步。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Attention kernel (16K context, head_dim=128, no padding) | Forward TFLOPS | 515.6 (TritonSigmoid) | 360.6 (FlashAttention-2) / 439.7 (FlashSigmoid) | +43% vs FA2 / +17% vs FlashSigmoid |
| Attention kernel (25% padding, head_dim=128) | Forward TFLOPS | 397.5 (TritonSigmoid) | 约 41 (PyTorch 原生 sigmoid) | 约 9.7× 加速 vs PyTorch |
| End-to-end training (1.4B 模型, 8K 上下文) | GPU hours / 131.6M samples | 7958 (sigmoid) | 8603 (softmax) | -7.5%（节省 645 GPU 小时） |
| End-to-end training (400M 模型, 4K 上下文) | GPU hours / 131.6M samples | 1739 (sigmoid) | 1832 (softmax) | -5.1% |
| Cell-type separation (Heart OFT, 28 对) | Average pairwise MMD | 全部 28 对胜出，平均提升 25% | Softmax 4K 模型 | +25.0% 平均 MMD |
| Validation loss (Adolescent Brain, 4K) | Masked LM loss | 2.526 ± 0.000 (sigmoid) | 2.537 ± 0.000 (softmax) | -0.011 (更低更好) |
| Validation loss (Heart OFT, 4K) | Masked LM loss | 2.615 ± 0.000 (sigmoid) | 2.631 ± 0.000 (softmax) | -0.016 (更低更好) |
| Stress test (160M, 8K, no grad clip) | Training steps to divergence | 80,000 步全部稳定 (sigmoid) | 约 55,600 步灾难性发散 (softmax) | 训练不崩溃 |
局限与改进
作者明确指出四点限制：(1) 全部实验在 scRNA-seq 数据上进行，是否能推广到多组学 (如 ATAC-seq、空间转录组、蛋白质表达) 尚待验证；(2) 评估只在 6 个公开 CellxGene 数据集上，未涵盖所有组织与疾病状态；(3) TritonSigmoid 暂不支持 NVIDIA Blackwell (B200) 架构，需要后续 JIT 编译适配；(4) 主实验模型规模上限为 1.4B 参数，更大的基础模型 (如 7B+) 上 sigmoid 相对 softmax 的优势是否会保持还需要探索。我的额外观察还包括：bias 项 $b=-\log(n)$ 是按 Ramapuram 等人的设置硬编码，没有扫参验证其最优性；MMD 指标只在 Heart OFT 一个数据集上统计，结论的稳定性有待复现；压力测试只跑了一次随机种子，缺乏对失败概率的统计估计。
独立分析的弱点
独立分析三个值得改进的弱点。第一个是评估协议的偏差：所有 6 个评测数据集都来自 CellxGene，与训练数据 (同一平台的 131.6M 细胞) 同分布，外部数据集上的 zero-shot 性能未充分验证，这会让结果偏乐观；改进方向是引入 GEO 上的独立数据集或人工设计的跨平台 benchmark。第二个是 ablation 不足：论文没有单独验证 $b=-\log(n)$ 这个 bias 设计对最终性能的贡献，也没有比较 sigmoid 与 gating、ReLU、cosine 等其他元素级注意力的差异，读者难以判断收益是来自 sigmoid 本身还是「去掉 softmax 归一化」这个更大的范畴；改进方向是补一组对照实验，比如对比 $\sigma(x/\sqrt{d})$ 与 $\sigma(\text{score}+b)$ 的消融。第三个是压力测试的工程性偏强：只展示了「softmax 在某次随机种子下崩溃」，缺乏崩溃概率统计与崩溃早期预警信号，建议加一组在 50–100 个随机种子下的发散率分布，并把梯度范数阈值与崩溃时间做成可监控指标。
未来方向
未来可延伸方向至少有三类。第一类是机制延伸：把 sigmoid 注意力扩展到多组学 (CITE-seq、ATAC+RNA 联合、空间转录组)，以及扩展到 protein language model 与 DNA 语言模型，验证其跨域通用性。第二类是工程延伸：补齐 TritonSigmoid 对 NVIDIA Blackwell 的支持，实现 FP8/INT8 量化与张量并行，使其能服务 10B+ 参数规模；并把 padding-aware 思路反向迁移到 softmax 内核，让所有长上下文模型都受益。第三类是理论延伸：进一步分析 sigmoid 注意力与线性注意力、gated attention (如 Qiu et al. 2025) 的关系，给出在何种模型规模、何种数据稀疏度下应选用何种机制的判据，把「稳定性 + 表达力」的 Pareto 前沿画清楚。
复现评估
复现评估总体友好。代码层面，作者开源了 TritonSigmoid 的 PyTorch 集成 (github.com/MSDLLCpapers/triton-sigmoid)，兼容 torch.compile，工程上易上手。数据层面，预训练与评测全部使用 CellxGene 公开快照 (2024-06-11 版, 131.6M 细胞) 与 6 个公开评估数据集，无需对接 Merck 内部资源。算力门槛较高：主实验使用 16×H100 80GB SXM5，1.4B 模型全量训练需 4,000+ GPU 小时；但 160M 模型与全部压力测试可压缩到 1–2 块 H100 上完成。附录 A.2 给出完整超参数 (12 层、hidden=768、12 heads、AdamW、$\text{lr}=10^{-4}$、gradient clip 1.0、BF16)。复现难度中等偏高——核心 GPU 内核对 Triton 编程能力有一定要求，建议先跑 160M 模型验证管线再扩展到更大规模。
论文图表
左侧直方图展示每个细胞表达基因数从 200 到 16,000 的高频分布；右侧 CDF 显示累积比例，标注 2,048/4,096/8,192/16,384 四个上下文阈值对应的覆盖率 (57.7%/82.0%/96.6%/100%)。
直观说明为什么单细胞数据必须 padding 到长上下文、且 padding-aware 内核不可少，是 motivation 章节的核心证据。
表格给出 Sigmoid 2K、Softmax 2K、Sigmoid 4K、Softmax 4K 在 6 个评测数据集上的均值验证损失与 95% 置信区间。Sigmoid 4K 在全部 6 个数据集上都取得最佳 (粗体)。
提供验证损失的全部原始数字，让读者可以独立验证「sigmoid 全面优于 softmax」这一结论的统计显著性。