# 单细胞模态预测：使用 RAPIDS cuML 将训练速度提升 103 倍 单细胞测量技术飞速发展，彻底改变了生命科学研究的格局。我们从仅能测量几十个细胞，扩展到如今可同时分析数百万个细胞；从单一模式跃升至多个高维数据模式。单个细胞层面蕴含的海量信息，为训练机器学习模型提供了绝佳契机，帮助我们更深刻地理解细胞模态之间的内在关联——这对合成生物学和药物靶点发现具有变革性意义。 本文将深入解析模态预测的概念，并展示如何用基于 NVIDIA GPU 的 RAPIDS cuML 替代 CPU 上的 TSVD 和内核岭回归（KRR），从而加速 NeurIPS Single-Cell Multi-Modality Prediction Challenge 的获胜方案。仅需修改六行代码，训练时间便从 69 分钟锐减至 40 秒——速度提升了 103.5 倍！即便与 PyTorch 开发的复杂深度学习模型相比，cuML 方案同样更快、更精准。 --- ## 什么是单细胞模态预测？ 通过单细胞技术，我们能够同时测量同一细胞内的多种模态，例如 DNA 可达性（ATAC）、mRNA 基因表达（GEX）和蛋白质丰度（ADT）。图 1 展示了这些模态之间的内在逻辑：只有开启的 DNA 区域才能转录出 mRNA，而 mRNA 随后作为模板合成蛋白质。 **模态预测问题**由此自然浮现：期望从一个模态预测另一个模态。在 2021 NeurIPS 挑战赛中，任务要求预测从 ATAC 到 GEX 以及从 GEX 到 ADT 的信息流。 如果机器学习模型能够做出准确预测，那么它必然已经捕捉到了细胞的复杂状态，从而为细胞生物学提供更深入的洞察。扩展对这些调控机制的理解，同样对药物靶点发现具有革命性价值。 ### 模态预测的独特难点 模态预测是一个多输出回归问题，面临以下挑战： - 高基数（High cardinality）：例如，GEX 和 ADT 数据分别以长度为 13953 和 134 的向量表示。 - 强偏差（Strong bias）：数据采集自 10 位不同捐赠者和 4 个地点，训练集与测试集来自不同站点，捐赠者和站点均显著影响数据分布。 - 稀疏、冗余且非线性（Sparse, redundant, and nonlinear）：模态数据高度稀疏，列之间相关性极强，且存在非线性关系。 本文聚焦于 GEX 到 ADT 预测任务，以证明单 GPU 解决方案的高效性。该方法可扩展至多节点多 GPU 架构，适用于更大规模、更高基数的单细胞模态预测任务。 --- ## 获胜方案：TSVD + KRR 多目标回归 作为基线，我们采用了密歇根大学邓凯文教授在 NeurIPS 模态预测挑战“GEX 到 ADT”中的 第一名解决方案。核心模型工作流程如图 2 所示： **图 2. 模型概览。方块表示输入与输出数据，方块旁的数字表示维度。** 训练数据包含 GEX 和 ADT 信息，测试数据仅包含 GEX。任务是根据给定的 GEX 预测测试数据的 ADT。 ### 降维：截断奇异值分解（TSVD） 为解决数据的稀疏性与冗余性，应用 TSVD 降低 GEX 和 ADT 的维度： - 对于 GEX：TSVD 应用于训练集与测试集的级联数据。 - 对于 ADT：TSVD 仅应用于训练数据。 在邓的方案中，GEX 维度从 13953 大幅缩减至 300，ADT 从 134 降至 70。主成分数量（300 和 70）通过交叉验证和调优确定，是模型的超参数。 ### 回归：内核岭回归（KRR） 将训练数据的 GEX 和 ADT 简化版本与 RBF 核一同输入 KRR。推理时，使用训练好的 KRR 模型执行以下步骤： 1. 预测测试数据 ADT 的简化版本。 2. 应用 TSVD 的逆变换。 3. 恢复测试数据的 ADT 预测。 > **小提示**：TSVD 是稀疏数据降维的常用技术，广泛应用于特征工程。此处 TSVD 同时缩减特征（GEX）和目标（ADT）的维度，使下游多输出回归模型更加简洁——因为 TSVD 输出在列之间更独立。 选择 KRR 作为多输出回归模型的原因：与 SVM 相比，KRR 同时计算所有输出列，而 SVM 一次仅预测一列，因此 KRR 能高效学习非线性关系且速度更快。 --- ## 使用 cuML 实现 GPU 加速方案 cuML 是 RAPIDS 库中的一员，提供一套 GPU 加速的机器学习算法，包含高度优化的 TSVD 和 KRR 模型。仅需少量代码调整，即可将基线模型从 scikit-learn 实现快速迁移至 cuML 实现。 ### 代码对比：仅需更改六行 在下面的代码示例中，只需要更改 六行代码，其中三行是导入。为简化，省略了预处理和工具代码。 #### 基线 sklearn 实现： ``` from sklearn.decomposition import TruncatedSVD from sklearn.gaussian_process.kernels import RBF from sklearn.kernel_ridge import KernelRidge tsvd_gex = TruncatedSVD(n_components=300) tsvd_adt = TruncatedSVD(n_components=70) gex_train_test = tsvd_gex.fit_transform(gex_train_test) gex_train, gex_test = split(get_train_test) adt_train = tsvd_adt.fit_transform(adt_train) adt_comp = tsvd_adt.components_ y_pred = 0 for seed in seeds: gex_tr,_,adt_tr,_=train_test_split(gex_train, adt_train, train_size=0.5, random_state=seed) kernel = RBF(length_scale = scale) krr = KernelRidge(alpha=alpha, kernel=kernel) krr.fit(gex_tr, adt_tr) y_pred += (krr.predict(gex_test) @ adt_comp) y_pred /= len(seeds) ``` #### RAPIDS cuML 实现： ``` from cuml.decomposition import TruncatedSVD from cuml.kernel_ridge import KernelRidge import gc tsvd_gex = TruncatedSVD(n_components=300) tsvd_adt = TruncatedSVD(n_components=70) gex_train_test = tsvd_gex.fit_transform(gex_train_test) gex_train, gex_test = split(get_train_test) adt_train = tsvd_adt.fit_transform(adt_train) adt_comp = tsvd_adt.components_.to_output('cupy') y_pred = 0 for seed in seeds: gex_tr,_,adt_tr,_=train_test_split(gex_train, adt_train, train_size=0.5, random_state=seed) krr = KernelRidge(alpha=alpha, kernel='rbf') krr.fit(gex_tr, adt_tr) gc.collect() y_pred += (krr.predict(gex_test) @ adt_comp) y_pred /= len(seeds) ``` ### 关键差异说明 - **内核指定方式不同**：cuML 的 KernelRidge 不支持独立内核对象，需在构造函数中指定 `kernel='rbf'`（因为 cuML 尚未支持高斯过程）。 - **显式垃圾收集**：当前版本的 cuML 实现需要 `gc.collect()`，因为循环中创建了某种引用循环，不手动回收则对象不会自动释放。详见 [GitHub 仓库](https://github.com/daxiongshu/rapids_nips_blog) 中的完整笔记本。 --- ## 性能对比结果 我们将 cuML 实现的 TSVD + KRR 与 CPU 基线以及挑战中的其他顶级方案进行比较： - **GPU 方案**：在单个 V100 GPU 上运行 - **CPU 方案**：在双 20 核 Intel Xeon CPU 上运行 - **评估指标**：均方根误差（RMSE） 结果显示：cuML 实现比 CPU 基线快 103 倍，得分（RMSE）因随机性略有下降，但仍优于比赛中任何其他模型。 我们还与两种深度学习模型进行了比较： | 排名 | 模型 | 框架 | |------|--------------------------|--------| | 第4名 | 多层感知机 (MLP) | PyTorch | | 第2名 | 图神经网络 (GNN) | PyTorch | 两个深度学习模型都在单个 V100 GPU 上运行，包含大量层和数百万参数，容易过拟合。相比之下，TSVD + KRR 只需训练 少于 30K 参数。图 4 显示，cuML TSVD + KRR 不仅速度更快，而且准确度更高。 **图 4. 性能与训练时间对比（横轴为对数刻度）** --- ## 常见问题 ### Q1: 为什么 cuML 的 TSVD + KRR 比深度学习模型更快更准？ **A**: 深度学习模型（如 MLP 和 GNN）需要训练数百万参数，容易在小规模数据集上过拟合。而 TSVD 先降维至 300/70 维度，KRR 仅有不到 30K 参数，更简单、更正则化，因此泛化能力更强，且 GPU 加速使训练极快。 ### Q2: 代码中 `adt_comp = tsvd_adt.components_.to_output('cupy')` 的作用是什么？ **A**: cuML 的 `components_` 属性默认在 GPU 上以 cuDF 或 cupy 格式存储。调用 `.to_output('cupy')` 将其转换为 cupy 数组，以便后续与预测结果进行矩阵乘法（`@` 操作），确保所有计算都在 GPU 上完成。 ### Q3: 为什么 cuML 的 KernelRidge 不能使用独立内核对象？ **A**: 当前 cuML 版本（截至本文撰写）尚未支持 `scikit-learn` 式的内核对象（如 RBF 内核）。因此需要在 `KernelRidge` 构造函数中直接指定 `kernel='rbf'` 或其他字符串。未来版本可能添加支持。 ### Q4: 如果我想用多 GPU 处理更大数据，该如何扩展？ **A**: 本文演示了单 GPU 方案。对于更大数据量和更高基数，可使用 RAPIDS 的 `dask-cuML` 结合多节点多 GPU 架构，将数据分片并在多个 GPU 上并行训练多个 TSVD 和 KRR 模型，然后聚合结果。 --- ## 实用小提示 - **调整超参数时注意**：TSVD 的 `n_components` 和 KRR 的 `alpha`、`kernel` 参数需要通过交叉验证调优。可借助 `cuml.model_selection.GridSearchCV`（支持 GPU）加速搜索。 - **显式垃圾收集**：在循环中每次拟合后调用 `gc.collect()` 可避免 GPU 内存泄漏，尤其当使用多个种子并行训练时很关键。 - **数据准备**：确保 GEX 和 ADT 数据已归一化（例如对数归一化），因为 TSVD 对尺度敏感。 --- ## 总结 RAPIDS cuML 凭借闪电般的速度和用户友好的 API，成为加速单细胞数据分析的利器。通过少量代码更改（仅六行），即可将现有的 scikit-learn 工作流提升至 GPU 级别，实现 103 倍训练加速。 在处理单细胞模态预测时，强烈建议从 cuML TSVD 开始降低数据维度，再配合 KRR 进行多输出回归，以获取最佳性能与精度平衡。 立即尝试完整代码：访问 [GitHub 仓库](https://github.com/daxiongshu/rapids_nips_blog) 获取 RAPIDS cuML 实现。