Arc 虛擬細胞挑戰賽：入門指南

Arc 研究所 (Arc Institute) 最近推出了虛擬細胞挑戰賽。參賽者需要訓練一個模型，該模型能夠預測在（部分）未見過的細胞型別中沉默某個基因的效果，他們將這項任務稱為情境泛化。對於幾乎沒有生物學背景的機器學習工程師來說，其中的術語和所需的背景知識可能相當令人生畏。為了鼓勵參與，我們以更適合其他學科工程師理解的形式重新闡述了這項挑戰。

目標
訓練一個模型，以預測使用 CRISPR 技術沉默一個基因對細胞產生的影響。

在原子世界中進行實驗成本高昂、耗時費力且容易出錯。如果我們能夠在不接觸培養皿的情況下測試數千種候選藥物，那會怎麼樣？這就是虛擬細胞挑戰賽的目標——建立一個模型（很可能是一個神經網路），能夠精確模擬當我們改變某個引數時細胞會發生什麼。鑑於縮短反饋迴圈通常是加速進展的最佳方式，一個能夠準確完成這項任務的模型將產生重大影響。

要訓練這個神經網路，我們需要資料。針對這項挑戰，Arc 整理了一個包含約 30 萬個單細胞 RNA 測序圖譜的資料集。在繼續之前，重溫中心法則可能會有所幫助。本文將以此為基礎，提供你參加挑戰賽所需的最少生物學知識。

訓練資料

訓練集由一個稀疏矩陣和一些相關的元資料組成。具體來說，我們有 22 萬個細胞，每個細胞都有一個轉錄組。這個轉錄組是一個稀疏行向量，其中每個條目是相應基因（我們的列）編碼的 RNA 分子（轉錄本）的原始計數。在這 22 萬個細胞中，約有 3.8 萬個是未擾動的，意味著沒有基因被 CRISPR 技術沉默。這些對照細胞至關重要，我們很快就會看到原因。

為了更具體地理解資料集，讓我們選擇一個基因 TMSB4X（資料集中被沉默最頻繁的基因），並比較一個對照細胞和一個受擾動細胞中檢測到的 RNA 分子數量。

我們可以看到，與對照細胞相比，TMSB4X 基因被沉默的細胞的轉錄本數量大大減少了。

為挑戰賽建模

敏銳的讀者可能會想，為什麼不直接測量基因沉默前後的 RNA 分子數量——為什麼我們還需要對照細胞呢？不幸的是，讀取轉錄組會破壞細胞，這個問題讓人想起了觀察者效應。

這種無法在擾動前後測量細胞狀態的情況帶來了許多問題，因為我們被迫使用一群基線（也稱為對照、未擾動）細胞作為參考點。即使在擾動之前，對照細胞和受擾動細胞也並非完全同質。這意味著我們現在必須從異質性引起的噪聲中分離出我們的真實訊號，即擾動的影響。

更正式地，我們可以將被擾動細胞中觀察到的基因表達建模為

$$\hat{X}_p \sim \hat{T}_p(\mathcal{D}_{\text{basal}}) + H(\mathcal{D}_{\text{basal}}) + \varepsilon, \quad \varepsilon \sim P_\varepsilon$$

其中

• $\hat{X}_p$: 在受到擾動 $p$ 的細胞中觀察到的基因表達測量值。
• $\mathcal{D}_{\text{basal}}$: 未受擾動的基線細胞群體的分佈。
• $\hat{T}_p(\mathcal{D}_{\text{basal}})$: 擾動 $p$ 對該群體造成的真實效應。
• $H(\mathcal{D}_{\text{basal}})$: 基線群體的生物學異質性。
• $\varepsilon$: 特定於實驗的技術噪聲，假定其獨立於未受擾動的細胞狀態和 $\mathcal{D}_{\text{basal}}$。

STATE：Arc 提供的基線模型

在虛擬細胞挑戰賽之前，Arc 釋出了 STATE，這是他們自己嘗試用一對基於 Transformer 的模型來解決這個挑戰。這為參賽者提供了一個強大的起點，因此我們將詳細探討它。

STATE 由兩個模型組成：狀態轉換模型 (State Transition Model, ST) 和狀態嵌入模型 (State Embedding Model, SE)。SE 旨在為細胞生成豐富的語義嵌入，以提高跨細胞型別的泛化能力。ST 是“細胞模擬器”，它接收對照細胞的轉錄組或由 SE 生成的細胞嵌入，以及一個表示目標擾動的獨熱編碼向量，然後輸出受擾動後的轉錄組。

狀態轉換模型 (ST)

狀態轉換模型是一個相對簡單的 Transformer，它使用 Llama 骨幹網路，處理以下輸入：

1. 一組協變數匹配的基線細胞的轉錄組（或 SE 嵌入）。
2. 一組表示每個細胞基因擾動的獨熱向量。

使用一組協變數匹配的對照細胞與配對的目標細胞，應有助於模型識別出我們預期擾動的實際效果。對照組張量和擾動張量都透過獨立的編碼器（即帶有 GELU 啟用函式的 4 層 MLP）進行處理。如果直接在基因表達空間中工作（即生成完整的轉錄組），它們會將輸出透過一個學習到的解碼器。

ST 使用最大均值差異 (Maximum Mean Discrepancy) 進行訓練。簡而言之，模型學習最小化兩個機率分佈之間的差異。

狀態嵌入模型 (SE)

狀態嵌入模型是一個類似 BERT 的自編碼器。要更深入地理解這一點，我們首先需要補充一些生物學基礎知識。

一點生物學知識補充

一個基因由外顯子（蛋白質編碼區）和內含子（非蛋白質編碼區）組成。DNA 首先被轉錄成前體 mRNA，如上圖所示。然後細胞進行選擇性剪接。這基本上是“挑選外顯子”，並剪掉所有內含子。你可以把基因想象成一份製作桌子的宜家說明書。透過省略一些部件，人們也可以造出一張三條腿的桌子，或者稍加努力還能造出一個奇怪的書架。這些不同的物體就類似於蛋白質異構體，即由同一個基因編碼的不同蛋白質。

回到模型

有了這個基本理解，我們就可以繼續討論 SE 模型的工作原理了。記住，我們 SE 的核心目標是建立有意義的細胞嵌入。要做到這一點，我們必須首先建立有意義的基因嵌入。

為了生成單個基因嵌入，我們首先獲取該基因編碼的所有不同蛋白質異構體的氨基酸序列（例如，TMSB4X 的序列為 $\texttt{SDKPDMAEI}$...）。然後，我們將這些序列輸入到 ESM2，這是一個來自 FAIR 的 150 億引數的蛋白質語言模型。ESM 為每個氨基酸生成一個嵌入，我們將它們進行均值池化，以獲得一個“轉錄本”（即蛋白質異構體）的嵌入。

現在我們有了所有這些蛋白質異構體的嵌入，我們只需對它們進行均值池化，就能得到基因嵌入。接下來，我們使用一個學習到的編碼器將這些基因嵌入投影到我們的模型維度，如下所示：

$$\tilde{g}_j = \text{SiLU}(\text{LayerNorm}(g_j \mathbf{W}_g + \mathbf{b}_g))$$

我們現在已經得到了一個基因嵌入，但我們真正想要的是一個細胞嵌入。為此，Arc 將每個細胞表示為按對數倍數表達水平排名的前 2048 個基因。

然後，我們從我們的 2048 個基因嵌入中構建一個“細胞句子”，如下所示： $$\tilde{\mathbf{c}}^{(i)} = \left[\mathbf{z}_{\text{cls}}, \tilde{\mathbf{g}}_1^{(i)}, \tilde{\mathbf{g}}_2^{(i)}, \ldots, \tilde{\mathbf{g}}_L^{(i)}, \mathbf{z}_{\text{ds}}\right] \in \mathbb{R}^{(L+2) \times h}$$

我們在句子中新增一個 $\texttt{[CLS]}$ 標記和 $\texttt{[DS]}$ 標記。 $\texttt{[CLS]}$ 標記最終被用作我們的“細胞嵌入”（非常像 BERT），而 $\texttt{[DS]}$ 標記用於“解耦特定於資料集的影響”。雖然基因是按對數倍數表達水平排序的，但 Arc 透過一種類似於位置嵌入的方式將轉錄組資訊融入，從而進一步強調每個基因的表達強度。透過一個奇特的“軟分箱”演算法和兩個 MLP，他們建立了一些“表達編碼”，然後將這些編碼新增到每個基因嵌入中。這應該會根據基因在轉錄組中的表達強度來調節每個基因嵌入的量級。

為了訓練模型，他們對每個細胞遮蔽 1280 個基因，模型的任務是預測它們。這 1280 個基因的選擇標準是使其具有廣泛的表達強度範圍。對於喜歡圖形化解釋的讀者，下圖展示了細胞句子的構建過程。

評估

瞭解你的提交將如何被評估是成功的關鍵。Arc 選擇的 3 個評估指標是擾動判別、差異表達和平均絕對誤差。鑑於平均絕對誤差很簡單，且顧名思義，我們將在分析中省略它。

擾動判別

擾動判別旨在評估你的模型在揭示不同擾動之間的相對差異方面的表現。為此，我們計算測試集中所有測得的受擾動轉錄組（我們試圖預測的真實值，$y_t$）和所有其他受擾動轉錄組（$y_p^n$）與我們預測的轉錄組 $\hat{y}_t$ 之間的曼哈頓距離。然後，我們按照以下方式對真實值相對於所有轉錄組的位置進行排名：

$$r_t = \sum_{p \neq t} \mathbf{1}\{d(\hat{y}_t, y_p) < d(\hat{y}_t, y_t)\}$$

之後，我們用轉錄組的總數進行歸一化：

$$\text{PDisc}_t = \frac{r_t}{T}$$

其中 $0$ 將是一個完美的匹配。你的預測的總體得分是所有 $$\text{PDisc}_t$$ 的平均值。然後將其歸一化為：

$$\text{PDiscNorm} = 1 - 2\text{PDisc}$$

我們乘以 2 是因為對於隨機預測，大約一半的結果會更近，一半會更遠。

差異表達

差異表達旨在評估你正確識別出真正受影響的基因中有多大比例被顯著影響。首先，對每個基因使用帶並列校正的 Wilcoxon 秩和檢驗來計算一個 $p$ 值。我們對我們預測的擾動分佈和真實擾動分佈都進行此操作。

接下來，我們應用 Benjamini-Hochberg 程式，基本上是一些統計方法來調整 $p$ 值，因為對於 $20,000$ 個基因和一個 $p$ 值閾值為 $0.05$ 的情況下，你預計會有 $1,000$ 個假陽性。我們將我們預測的差異表達基因集表示為 $G_{p,pred}$，真實差異表達基因集表示為 $G_{p,true}$。

如果我們的集合大小小於真實集合大小，則取集合的交集，然後除以真實差異表達基因的數量，如下所示：

$$DE_p = \frac{G_{p,pred} \cap G_{p,true}}{n_{p,true}}$$

如果我們的預測集大小大於真實集大小，我們選擇預測中差異表達最顯著的子集（我們“最確信”的預測，記作 $\tilde{G}_{p,pred}$），取其與真實集的交集，然後除以真實集的數量。

$$DE_p = \frac{\tilde{G}_{p,pred} \cap G_{p,true}}{n_{p,true}}$$

對所有預測的擾動重複此操作，並取平均值以獲得最終分數。

結論

如果這項挑戰激起了您的興趣，該如何開始呢？幸運的是，Arc 提供了一個 Colab 筆記，其中詳細介紹了訓練其 STATE 模型的全過程。此外，STATE 很快將登陸 transformers 庫，因此，使用他們的預訓練模型將會非常簡單，只需

import torch
from transformers import StateEmbeddingModel

model_name = "arcinstitute/SE-600M"
model = StateEmbeddingModel.from_pretrained(model_name)

input_ids = torch.randn((1, 1, 5120), dtype=torch.float32)
mask = torch.ones((1, 1, 5120), dtype=torch.bool)
mask[:, :, 2560:] = False 
outputs = model(input_ids, mask)

祝所有參賽者好運！

本文最初發佈於此。