單細胞測量技術(shù)發(fā)展迅速，徹底改變了生命科學。我們已經(jīng)從測量幾十個細胞擴展到數(shù)百萬個細胞，從一種模式擴展到多個高維模式。單個細胞水平上的大量信息為訓練機器學習模型提供了一個很好的機會，幫助我們更好地理解 intrinsic link of cell modalities ，這可能會對合成生物學和 drug target discovery 產(chǎn)生變革。

這篇文章介紹了模態(tài)預測，并解釋了我們?nèi)绾斡没?NVIDIA GPU 的 RAPIDS cuML 實現(xiàn)取代基于 CPU 的 TSVD 和內(nèi)核嶺回歸（ KRR ），從而加速了 NeurIPS Single-Cell Multi-Modality Prediction Challenge 的獲勝解決方案。

使用 cuML ，只修改了六行代碼，我們加速了基于 scikit 學習的獲勝解決方案，將訓練時間從 69 分鐘縮短到 40 秒：速度提高了 103.5 倍！即使與 PyTorch 中開發(fā)的復雜深度學習模型相比，我們觀察到 cuML 解決方案對于這種預測挑戰(zhàn)來說更快更準確。

多虧了單細胞技術(shù)，我們可以測量同一單細胞內(nèi)的多種形態(tài)，如 DNA 可達性（ ATAC ）、 mRNA 基因表達（ GEX ）和蛋白質(zhì)豐度（ ADT ）。圖 1 顯示了這些模式之間的內(nèi)在聯(lián)系。只有可獲得的 DNA 才能產(chǎn)生 mRNA ，而 mRNA 又被用作生產(chǎn)蛋白質(zhì)的模板。

當期望從一個模態(tài)預測另一個模態(tài)時，模態(tài)預測的問題自然出現(xiàn)。在 2021 NeurIPS challenge 中，我們被要求預測從 ATAC 到 GEX 以及從 GEX 到 ADT 的信息流。

若一個機器學習模型能夠做出好的預測，那個么它一定已經(jīng)了解了細胞的復雜狀態(tài)，它可以為細胞生物學提供更深入的見解。擴展我們對這些調(diào)控過程的理解，對于藥物靶點的發(fā)現(xiàn)也具有革命性意義。

模態(tài)預測是一個多輸出回歸問題，它提出了獨特的挑戰(zhàn)：

High cardinality.例如， GEX 和 ADT 信息分別以長度為 13953 和 134 的矢量描述。

Strong bias.數(shù)據(jù)收集自 10 個不同的捐贈者和 4 個地點。培訓和測試數(shù)據(jù)來自不同的站點。捐贈者和站點都強烈影響數(shù)據(jù)的分布。

稀疏、冗余和非線性。模態(tài)數(shù)據(jù)稀疏，列高度相關(guān)。

在這篇文章中，我們專注于 GEX 到 ADT 預測的任務，以證明單個 – GPU 解決方案的效率。我們的方法可以擴展到使用多節(jié)點多 GPU 架構(gòu)的具有更大數(shù)據(jù)量和更高基數(shù)的其他單細胞模態(tài)預測任務。

使用 TSVD 和 KRR 算法進行多目標回歸

作為基線，我們使用了密歇根大學鄧凱文教授的 NeurIPS 模態(tài)預測挑戰(zhàn)“ GEX 到 ADT ”的 first-place solution 。核心模型的工作流程如圖 2 所示。訓練數(shù)據(jù)包括 GEX 和 ADT 信息，而測試數(shù)據(jù)只有 GEX 信息。

任務是預測給定 GEX 的測試數(shù)據(jù)的 ADT 。為了解決數(shù)據(jù)的稀疏性和冗余性，我們應用截斷奇異值分解（ TSVD ）來降低 GEX 和 ADT 的維數(shù)。

特別是，兩種 TSVD 模型分別適用于 GEX 和 ADT ：

對于 GEX ， TSVD 適合訓練和測試的級聯(lián)數(shù)據(jù)。

對于 ADT ， TSVD 僅適合訓練數(shù)據(jù)。

在鄧的解決方案中， GEX 的維度從 13953 大幅降低到 300 ， ADT 從 134 大幅降低到 70 。

主成分 300 和 70 的數(shù)量是通過交叉驗證和調(diào)整獲得的模型的超參數(shù)。然后將訓練數(shù)據(jù)的 GEX 和 ADT 的簡化版本與 RBF 核一起饋送到 KRR 中。根據(jù)鄧的方法，在推理時，我們使用經(jīng)過訓練的 KRR 模型執(zhí)行以下任務：

預測測試數(shù)據(jù)的 ADT 的簡化版本。

應用 TSVD 的逆變換。

恢復測試數(shù)據(jù)的 ADT 預測。

圖 2.模型概述。塊表示輸入和輸出數(shù)據(jù)，塊旁邊的數(shù)字表示尺寸。

通常， TSVD 是對稀疏數(shù)據(jù)執(zhí)行降維的最常用選擇，通常在特征工程期間使用。在這種情況下， TSVD 用于減少特征（ GEX ）和目標（ ADT ）的尺寸。目標的降維使下游多輸出回歸模型更加容易，因為 TSVD 輸出在列之間更加獨立。

選擇 KRR 作為多輸出回歸模型。與 SVM 相比， KRR 同時計算輸出的所有列，而 SVM 一次預測一列，因此 KRR 可以像 SVM 一樣學習非線性，但速度更快。

使用 cuML 實現(xiàn) GPU 加速解決方案

cuML 是 RAPIDS 庫之一。它包含一套 GPU 加速機器學習算法，可提供許多高度優(yōu)化的模型，包括 TSVD 和 KRR 。您可以將基線模型從 scikit 學習實現(xiàn)快速調(diào)整為 cuML 實現(xiàn)。

I在下面的代碼示例中，我們只需要更改六行代碼，其中三行是導入。為簡單起見，省略了許多預處理和實用程序代碼。

Baseline sklearn implementation:

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 implementation:

from cuml.decomposition import TruncatedSVD from cuml.kernel_ridge import KernelRidge import gctsvd_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 內(nèi)核的語法與 scikit learn 略有不同。我們沒有創(chuàng)建獨立的內(nèi)核對象，而是在 KernelRidge 的構(gòu)造函數(shù)中指定了內(nèi)核類型。這是因為 cuML 還不支持高斯過程。

另一個區(qū)別是當前版本的 cuML 實現(xiàn)需要顯式垃圾收集。在這個特定的循環(huán)中創(chuàng)建了某種形式的引用循環(huán)，并且在沒有垃圾收集的情況下不會自動釋放對象。有關(guān)更多信息，請參閱 /daxiongshu/rapids_nips_blog GitHub 存儲庫中的完整筆記本。

后果

我們將 TSVD + KRR 的 cuML 實施與 CPU 基線和挑戰(zhàn)中的其他頂級解決方案進行了比較。 GPU 解決方案在單個 V100 GPU 上運行， CPU 解決方案在雙 20 核 Intel Xeon CPU 上運行。競爭的度量是均方根誤差（ RMSE ）。

我們發(fā)現(xiàn)， TSVD + KRR 的 cuML 實現(xiàn)比 CPU 基線快 103 倍，由于管道中的隨機性，分數(shù)略有下降。然而，比分仍然比比賽中的任何其他車型都好。

我們還將我們的解決方案與兩種深度學習模型進行了比較：

第四名解決方案：Multilayer Perceptron (MLP)

第二名解決方案：Graph Neural Network (GNN)

這兩個深度學習模型都在 PyTorch 中實現(xiàn)，并在單個 V100 GPU 上運行。這兩個深度學習模型都有許多層，需要訓練數(shù)百萬個參數(shù)，因此容易對該數(shù)據(jù)集進行過度擬合。相比之下， TSVD + KRR 只需訓練少于 30K 的參數(shù)。圖 4 顯示，由于其簡單性， cuML TSVD + KRR 模型比深度學習模型更快、更準確。

圖 4.績效和培訓時間對比。橫軸為對數(shù)刻度。

結(jié)論

由于其閃電般的速度和用戶友好的 API ， RAPIDS cuML 對于加速單細胞數(shù)據(jù)的分析非常有用。通過少量的代碼更改，您可以提升現(xiàn)有的 scikit 學習工作流。

此外，在處理單細胞模態(tài)預測時，我們建議從 cuML TSVD 開始，以減少下游任務的數(shù)據(jù)維度和 KRR ，從而實現(xiàn)最佳加速。

使用 /daxiongshu/rapids_nips_blog GitHub 存儲庫上的代碼嘗試這個 RAPIDS cuML 實現(xiàn)。