如何用基于NVIDIA GPU的RAPIDS cuML實(shí)現(xiàn)取代KRR?

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

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

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

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

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

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

模態(tài)預(yù)測(cè)是一個(gè)多輸出回歸問(wèn)題，它提出了獨(dú)特的挑戰(zhàn)：

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

Strong bias.數(shù)據(jù)收集自 10 個(gè)不同的捐贈(zèng)者和 4 個(gè)地點(diǎn)。培訓(xùn)和測(cè)試數(shù)據(jù)來(lái)自不同的站點(diǎn)。捐贈(zèng)者和站點(diǎn)都強(qiáng)烈影響數(shù)據(jù)的分布。

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

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

使用 TSVD 和 KRR 算法進(jìn)行多目標(biāo)回歸

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

任務(wù)是預(yù)測(cè)給定 GEX 的測(cè)試數(shù)據(jù)的 ADT 。為了解決數(shù)據(jù)的稀疏性和冗余性，我們應(yīng)用截?cái)嗥娈愔捣纸猓?TSVD ）來(lái)降低 GEX 和 ADT 的維數(shù)。

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

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

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

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

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

預(yù)測(cè)測(cè)試數(shù)據(jù)的 ADT 的簡(jiǎn)化版本。

應(yīng)用 TSVD 的逆變換。

恢復(fù)測(cè)試數(shù)據(jù)的 ADT 預(yù)測(cè)。

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

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

選擇 KRR 作為多輸出回歸模型。與 SVM 相比， KRR 同時(shí)計(jì)算輸出的所有列，而 SVM 一次預(yù)測(cè)一列，因此 KRR 可以像 SVM 一樣學(xué)習(xí)非線性，但速度更快。

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

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

I在下面的代碼示例中，我們只需要更改六行代碼，其中三行是導(dǎo)入。為簡(jiǎn)單起見(jiàn)，省略了許多預(yù)處理和實(shí)用程序代碼。

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)核的語(yǔ)法與 scikit learn 略有不同。我們沒(méi)有創(chuàng)建獨(dú)立的內(nèi)核對(duì)象，而是在 KernelRidge 的構(gòu)造函數(shù)中指定了內(nèi)核類(lèi)型。這是因?yàn)?cuML 還不支持高斯過(guò)程。

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

后果

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

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

我們還將我們的解決方案與兩種深度學(xué)習(xí)模型進(jìn)行了比較：

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

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

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

圖 4.績(jī)效和培訓(xùn)時(shí)間對(duì)比。橫軸為對(duì)數(shù)刻度。

結(jié)論

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

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

使用 /daxiongshu/rapids_nips_blog GitHub 存儲(chǔ)庫(kù)上的代碼嘗試這個(gè) RAPIDS cuML 實(shí)現(xiàn)。