基于對比學習的冷凍電鏡單顆粒圖像聚類算法

顏 陽，鄭清炳，張東旭，李少偉，葛勝祥，張 軍，夏寧邵

(廈門大學公共衛生學院，福建廈門361102)

冷凍電鏡技術是一種著名的結構生物學分析方法.相比于同樣被廣泛使用的兩種結構生物學的研究手段：X射線和核磁共振，冷凍電鏡技術具有不需要結晶、解析范圍廣、能夠捕捉生物分子動態變化等優點[1].冷凍電鏡技術在近些年來發展迅速，廣泛應用于生物大分子結構重構[2]，其中包括非洲豬瘟病毒[3]以及部分新型冠狀病毒相關蛋白質[4]的重構.盡管冷凍技術能夠在一定程度上保護樣品，但在高計量的電子輻射下，樣本仍然無法承受[5]，導致冷凍電鏡拍攝得到的圖像信噪比極低.Frank等[6]通過單顆粒分析(single particle analysis，SPA)方法很大程度上克服了冷凍電鏡低信噪比的問題[7].之后伴隨著拍攝技術[8]、制樣方法、重構軟件的進步[9]，SPA的重構分辨率不斷進步，甚至達到了原子級別的分辨率[10].

SPA的一個重要步驟是單顆粒圖像聚類：通過將挖取后的目標顆粒圖像根據投影角的不同分成不同的簇實現.早期算法直接利用聚類平均圖生成三維初始模型，聚類效果會直接影響三維初始模型的質量[11].當前聚類的主要意義在于幫助挑選出完整的、有價值的顆粒，拋棄誤選的雜質.同時通過聚類平均圖可以在二維初步觀測顆粒基本形態，得到對稱性等形態信息，為接下來進行的步驟提供參考.

由于噪聲干擾過大、顆粒異構多樣、數據量過大等原因，設計用于冷凍電鏡圖像的聚類算法是一個很大的挑戰.早期由于計算資源的限制，電鏡顆粒聚類算法首先使用對應分析(correspondence analysis,CA)、主成分分析(principal component analysis,PCA)等方法對圖像進行降維處理，之后使用層次聚類法聚類[12]，此類方法需要對圖像先進行全局校準(使二維旋轉角相同)，該過程占用計算資源較大，會限制參與聚類顆粒數量.之后Sigworth[13]提出了基于最大似然法的聚類算法應用于電鏡圖像，該方法被Scheres等[14]進一步擴展后稱之為ML2D算法.ML2D是最廣泛使用的電鏡圖像聚類算法之一，已經被整合到電鏡圖像處理軟件Xmipp[15]和Relion[16]上.

隨著人工智能的迅猛發展，許多研究者將深度學習技術應用到電鏡顆粒的聚類任務上來，但需要預訓練.例如，自編碼器[17](autoencoders，AE)具有能夠利用低層次的特征形成更加高級的抽象特征的優勢[18]，因此，迭代式基于變分AE的多參考對齊模型(IterVM)[19]和級聯降噪AE(CDAE)[20]算法均先通過預訓練AE提取圖像特征，之后再利用這些特征聚類.另外，IterVM使用迭代過程中產生的類平均圖訓練聚類模型，若過程中聚類精確度較差則會很大程度影響模型的訓練；CDAE需要事先準備的加噪仿真數據集進行模型預訓練，在實際應用中難以滿足該條件.

為了提高冷凍電鏡單顆粒圖像聚類精度，簡化流程，本研究提出了一種基于對比學習的無監督電鏡圖像聚類算法：CL-Clustering.該算法不需要帶標簽的數據集或者人工合成數據集對模型預訓練，且不需要聚類迭代過程中的二維校準處理，聚類精度高并能夠高效處理大規模冷凍電鏡圖像數據.

1 CL-Clustering算法流程

圖1 編碼器網絡架構Fig.1Architecture of encoder network

CL-Clustering算法流程可分為三步：

1) 將數據集進行數據增強處理.根據自然環境同一投影角拍攝到的顆粒圖像(即應聚為同一類的圖像)的差異性，有針對性的對數據集進行數據增強處理.該步驟對同一圖像分別進行相互獨立的兩次數據增強，此時同一圖像兩次數據增強得到的圖像類似于兩張真實環境應聚為同一類的圖像；

2) 利用對比學習訓練基于殘差網絡的電鏡特征編碼器，該過程中同一圖像的數據增強在特征空間被拉近，訓練完成后將原始數據集輸入到編碼器得到圖像特征；

3) 使用K-means++將提取后的特征進行聚類，由于數據增強中包括圖像隨機旋轉與翻折，提取后特征具有旋轉不變性，因此可以避免在聚類時考慮對比過程中圖像二維旋轉差異，只需要在生成類平均圖時統一校準即可.

上述步驟中，數據增強和模型訓練時數據集被等分為多個批次(batch)處理，模型訓練完成后再統一將所有圖像編碼得到特征并聚類.可以看出，CL-Clustering算法直接使用需要聚類的數據集訓練模型，不需要使用任何帶人工標簽的數據集或人工合成數據集對模型進行預訓練，同時利用對比學習訓練得到的編碼器提取的特征維度低且在特征空間區分度強，可以同時對大量數據進行聚類且聚類結果穩定.

1.1 基于殘差網絡的電鏡特征編碼器

在冷凍電鏡單顆粒圖像的聚類任務中，原圖像受噪聲污染非常嚴重，因此直接對原圖像進行聚類得到的精度非常低，同時直接使用原圖像運算量過大，因此本研究先提取圖像特征后進行聚類計算.

本文使用的電鏡特征編碼器如圖1所示.主干網絡基于殘差網絡[21](ResNet18)搭建.傳統的卷積神經網絡(CNN)通過卷積層和下采樣層的不斷堆疊搭建，殘差網絡通過引入捷徑分支(shortcut)有效解決了傳統CNN梯度消失/爆炸問題以及退化問題.本文使用的殘差網絡結構如圖1所示，主干網絡包括17個卷積層(conv)和一個全連接層(FC)，在經過預處理(圖中前兩層)后，每兩個卷積層之間使用捷徑分支連接，為了維度匹配，其中三個捷徑分支添加一個卷積層對數據降維.模型使用Relu激活函數，同時在每個卷積層后使用批標準化處理，能夠加速網絡的收斂并提升準確率.經過主干網絡處理后，使用兩層全連接層得到128維特征zi，該特征將進一步用于對比學習.

1.2 數據增強策略

聚類任務中，需要將同一投影角的冷凍電鏡圖像分為一類，除了信噪比極低這個因素外，這個過程中還存在以下難點：1) 同一投影角下冷凍電鏡顆粒圖像會在二維空間隨機旋轉，如果在聚類過程中考慮該旋轉，在數據量較大時將會占用大量計算資源，因此殘差網絡編碼器提取的特征需要具有旋轉不變性，這樣可以在聚類時避免二維旋轉角的校準，只需直接對顆粒進行聚類；2) 由于顆粒挑選和挖取過程中存在一定誤差，難以保證顆粒正好在圖像居中位置，可能存在顆粒在圖像中朝某方向的整體偏移，聚類時要使投影角相同但存在偏移的顆粒聚為一類；3) 同一投影角的顆粒之間可能存在微小形變，同樣要使這類顆粒能順利聚為一類.為了解決以上難點，通過數據增強配合對比學習訓練讓深度學習模型在編碼中忽略二維旋轉翻折、像素分布差異、顆粒(同一投影)微小形變、噪聲這些與顆粒投影角無關的圖像信息，提取反映投影角的顆粒形態特征這一有用信息，達到編碼器對同一投影角得到的不同圖像編碼盡量相似的目的.

對比學習過程中，模型將學習同一圖像的不同數據增強的相似點，數據增強的形式將會決定對比學習得到的信息的質量[22].根據電鏡圖像特點，選擇對冷凍電鏡單顆粒圖像應用以下數據增強形式：1) 二維空間隨機旋轉0°～180°；2) 以50%的概率隨機左右翻轉整幅圖像；3) 標準化處理，具體方式為

(1)

4) 對圖像進行隨機大小和長寬比的裁剪，相對于原圖，裁剪后圖像大小比例范圍為0.3～1.0，長寬比范圍為0.9～1.1.具體步驟如圖2所示，在原圖像范圍內隨機選定一塊矩形區域，矩形的大小和長寬比隨機在預設范圍內選擇，裁剪下該區域，再將該區域縮放到和原圖一樣的大小，該過程中顆粒形變的幅度在隨機范圍內變化，其中有些時候明顯大于現實可能存在的情況，這是為了幫助深度學習模型提取更豐富尺度的特征，并提升模型的泛化性.以上步驟中，標準化是為了加快模型收斂速度、減少圖像像素值分布差異對模型的干擾，隨機地旋轉、翻折主要是為了模擬同一投影角的圖像二維旋轉、翻折，隨機裁剪主要是為了減弱背景(或噪聲)因子的權重且使模型面對缺失信息不敏感.以上所有數據強化步驟都有助于提高模型的穩定性和魯棒性.

圖2 圖像的隨機裁剪Fig.2Random crop of image

對所有圖像進行以上所有數據強化步驟，強化過程用函數g表示，則有：

(2)

圖3 單顆粒冷凍電鏡圖像數據增強效果圖Fig.3Augmentation of single particle cryo-electron microscopy images

1.3 模型的訓練策略

本研究使用對比學習[23]訓練特征編碼器，其中，訓練目標為：同類圖像經過神經網絡輸出的特征盡可能的相似，不同類圖像神經網絡輸出的特征盡可能的不同.在該設定下，特征空間中同一圖像的不同增強圖像的距離將會被拉近，不同圖像的增強圖像的距離將會被拉遠.值得注意的是，非同一圖像數據增強的圖像也可能是同類圖像，但概率很低，對整體訓練效果影響很小.神經網絡對特征的提取表示函數為：

(3)

其中，fEC為電鏡特征編碼器函數，zi為通過編碼器提取的特征.使用余弦相似度計算特征之間的距離：

(4)

(5)

(6)

1.4 對提取特征的聚類

1) 在Z中隨機挑選zi作為第一個類中心C1.

2) 利用輪盤賭選擇法依次隨機選出盡可能分散的k個類中心.具體地，若當前已選出k個類中心，則根據式(7)計算特征zj被隨機選為下一個類中心的概率：

p(zj)=

(7)

其中,Di(zj)表示特征zj與被選為第i個類中心的特征zi的矩陣歐式距離.式(7)的分子表示求特征zj與已選中類中心的矩陣歐式距離，并找出最小值，使離已有類中心越近的特征被選為下一個類中心的概率越小.分母表示所有特征zj與已選中類中心的最小歐式距離之和，起到歸一化的作用.

4) 對每個特征zi，計算其與第k個類中心的距離Dk(zi)，如式(8)所示，將該特征劃分到距離最近的類中心所屬類中：

(8)

其中，L(zi)表示特征zi的標簽.

5) 在所有特征得到類標簽后如式(9)所示重新確定每一類的類中心：

(9)

其中Nk為被劃分到第k類的特征數目.

6) 重復步驟4)以及步驟5)直到類標簽保持不變或者達到最大迭代次數.聚類完成后，根據所有特征的類標簽可以得到所有圖像的類歸屬.

2 實驗和結果分析

2.1 數據集的生成與介紹

為了測試CL-Clustering算法的性能，本研究利用歐洲電子顯微鏡數據庫[25]公開的3個高分辨率電鏡三維重構顆粒模型生成具有標簽的仿真單顆粒冷凍電鏡圖像數據集，分別為0.42 nm分辨率的熱休克蛋白(GroEL，EMD-5001)[26]、0.26 nm分辨率的β半乳糖苷酶(EMD-6840)[27]以及0.36 nm的間隙連接蛋白(INX-6，EMD-9973)[28].這3種顆粒中EMD-5001對稱性為D7，EMD-6840的對稱性為D2，EMD-9973對稱性為C8，三者基礎形狀、分辨率以及對稱性有較大差別，可以較全面評價目標算法的性能.3種顆粒三維展示圖如圖4所示.

圖4 用于數據生成的三種單顆粒結構Fig.4Three single-particle structures for datasets generation

對于每個結構，隨機選取10個不同投影角的投影(其中會考慮結構的對稱性，避免因對稱性出現投影角不同而投影相同的情況)；對于每個投影使其在二維空間旋轉，每旋轉1度生成一個圖像，圖像大小為128 × 128.每個投影生成的360張圖像標記為同一類的圖像.對每個生成圖像添加離焦值(defocus)為1.5～2.0 μm的襯底轉換函數CTF(加速電壓300 keV，球差2.7 mm，相位襯度比7%)以及高斯噪聲，得到信噪比分別為0.10和0.05的兩套加噪圖像，其中信噪比通過圖像方差除以添加噪聲的方差獲得.通過以上步驟，每個顆粒得到2×3×10×360=21 600張仿真圖像，總共獲得了3×21 600=64 800張帶類別標簽的仿真圖像用于算法評價.圖5展示了仿真圖像的示意圖.在仿真數據集實驗中，一次實驗使用一套相同顆粒且信噪比相同的3 600張圖像進行編碼器訓練、特征提取、特征聚類.

(i)行為添加CTF后的投影圖像; (ii)行為信噪比為0.10的仿真電鏡圖像; (iii)行為信噪比為0.05的加噪仿真電鏡圖像.圖5 生成單顆粒冷凍電鏡圖像Fig.5Generated cryo-electron microscopy images of single-particle

評價目標算法在仿真數據集上的性能有兩個指標，分別是聚類準確度和歸一化互信息量，下文中二者用A與I表示.聚類準確度用于計算正確的預測標簽占整個數據集的比例，計算如式(10)所示：

(10)

其中，Ltrue為真實標簽，Lpred(xi)為算法預測xi的標簽，T表示預測標簽到真實標簽的最佳映射.

歸一化互信息I用來計算聚類預測標簽與真實標簽的相似度[29].其計算方法如式(11)所示：

(11)

其中，X、Y為標簽集合，X={Lpred}，Y={Ltrue}，P(i,j) 表示標簽為i和j的數據的交集出現的概率，H(X)和H(Y)分別為X、Y的熵.I的值在0到1之間，越接近1表示算法預測標簽越接近真實標簽.

現實世界拍攝的冷凍電鏡圖像噪聲比仿真圖像更加復雜，除了信噪比極低外，還存在雜質、破損顆粒干擾聚類.為了進一步評價算法性能，本研究使用了T20S蛋白酶體真實拍攝的冷凍電鏡圖像數據集進行實驗.選擇EMPIAR-10025數據集[30]中包含的20張原始拍攝圖像的子集，這些圖像經過了運動補償以及對比度傳遞函數校準處理，經過顆粒挑選以及顆粒挖取后總共得到15 552張冷凍電鏡單顆粒圖像用于真實數據集聚類實驗.

以上實驗中使用的模型基于PyTorch 1.9.0框架編寫完成，模型在配有NVIDIA RTX 2080Ti GPU以及Intel Xeon Bronze 3204 CPU的服務器上訓練，并開展后續對照實驗.訓練過程中，采用了圖像隨機旋轉(0°～180°)、隨機翻折(50%概率)、標準化、隨機裁剪(大小比例范圍為0.3～1.0，長寬比范圍為0.9～1.1)的數據增強方式，模型訓練中使用的批大小為60，利用1.1節中的編碼器提取圖像特征并根據式(6)計算損失函數，優化算法為隨機梯度下降(stochastic gradient descent,SGD)，循環100輪.仿真數據集和真實數據集選用相同的訓練方式和參數.

2.2 仿真數據集實驗結果

本研究使用了PCA+K-means算法以及ML2D算法作為對照.PCA+K-means算法[31]是一個被廣泛應用到各種領域的聚類算法，并已經被整合到冷凍電鏡圖像處理軟件Spider[32].該算法使用PCA對原圖像進行降維，之后使用K-means對降維后的特征進行聚類.ML2D[14]算法是基于最大似然法的軟分類算法，同時包括對噪聲的建模，是使用最廣泛的冷凍電鏡單顆粒圖像聚類算法之一，實驗中使用Xmipp調用ML2D實現聚類.以A和I為指標，CL-Clustering算法以及兩種對照算法聚類效果評價如表1所示.整體上CL-Clustering算法在兩項指標上優于PCA+K-means算法以及ML2D算法.

得到聚類標簽后，為了便于直觀評價聚類效果，根據聚類標簽將分為同一類的圖像進行二維校準處理，之后對校準后的圖像取平均圖.對于效果好的聚類，生成平均圖時其疊加的噪聲將會相互抵消，平均圖信噪比將會大大高于原始圖像并顯示出原顆粒的形態信息.圖6展示了利用PCA+K-means算法(對應(i)行)、ML2D算法(對應(ii)行)以及CL-Clustering算法(對應(ii)行)聚類標簽生成的類平均圖.通過觀察可以發現利用PCA+K-means算法聚類標簽生成的類平均圖相比原顆粒結構失真較為嚴重；利用ML2D算法聚類標簽生成的類平均圖個別類信

表1 聚類效果對比

噪比很低；利用CL-Clustering算法得到的標簽生成的類平均圖擁有更多的顆粒細節信息，并且顆粒形態與圖4所示的用來生成仿真數據集的高精度三維顆粒結構高度相似.

圖6 聚類平均圖Fig.6Clustering average images

圖7 T20S蛋白酶體聚類平均圖對比Fig.7Comparison of clustering averages of T20S proteasome

表2為三種算法占用的計算時間.其中CL-Clustering算法統計時長時包括了模型訓練時間(100輪).ML2D占用時間最長，CL-Clustering次之，PCA+K-means占用時間最短，但其性能表現較差.

表2 三種算法計算占用時間

2.3 真實數據集實驗結果

本部分使用T20S蛋白酶體真實拍攝的冷凍電鏡圖像數據進行實驗，模型的參數設置和訓練步驟與仿真數據集的實驗相同.分別使用CL-Clustering以及ML2D算法對15 552張冷凍電鏡單顆粒圖像進行聚類實驗，實驗中所有圖像都參與CL-Clustering編碼器的訓練，類數都設定為50.兩種算法得到的聚類平均圖如圖7所示，CL-Clustering得到了更豐富投影角的類平均圖.根據CL-Clustering類平均圖(圖7(a))的清晰度挑選33個類(按由左到右、由上到下的順序，包含顆粒數目為：296, 232, 335, 367, 334, 254, 322, 377, 301, 310, 340, 292, 353, 325, 301, 322, 278, 258, 279, 295, 320, 336, 245, 315, 328, 316, 283, 350, 224, 294, 317, 304, 320)共計10 123個顆粒進行三維重構；根據ML2D聚類平均圖(圖7(b))的清晰度挑選15個類(按由左到右、由上到下的順序，包含顆粒數目為：1 129, 513, 670, 153, 1 588, 956, 280, 128, 852, 451, 187, 1 255, 729, 262, 257)共計9 410個顆粒進行三維重構.

由于真實數據集投影角未知，即真實數據集沒有標簽，無法使用A、I等指標定量評價聚類效果，因此在該實驗中模擬現實世界電鏡聚類算法應用場景，在聚類完成后根據聚類平均圖挑選顆粒，使用挑選后的顆粒進行三維重構，根據重構效果評價聚類算法效果.

圖8為根據兩種算法聚類平均圖(圖7)挑選的顆粒進行三維重構得到的結構.三維重構分為兩個步驟，分別為初始模型的構建以及三維精修，結果對應圖8的第一行和第二行.精修后得到結構的分辨率根據傅里葉殼相關函數[33]計算得到.實驗中使用的初始模型的構建以及三維精修算法為Relion集成的方法[34]，其中初始模型構建時設置的對稱性為C1，三維精修時設置的對稱性為D7.如圖8所示，CL-Clustering挑選顆粒在經過兩步三維重構處理后生成了高分辨率(0.352 nm)的三維生物分子結構，該結構與圖7得到的類平均圖相匹配.作為對照，ML2D挑選顆粒重構得到了0.357 nm分辨率的三維生物分子結構，二者結構高度相似.

圖8 T20S蛋白酶體三維重構結果對比Fig.8Comparison of reconstruction result of T20S proteasome

3 結 論

本研究針對單顆粒冷凍電鏡圖像的特點，提出了基于對比學習的深度學習聚類算法CL-Clustering.該算法根據電鏡圖像特點選擇圖像增強方式，同時使用了對比學習訓練編碼器，使訓練得到的編碼器能夠提取有利于聚類的圖像特征，同時讓聚類過程免于二維校準.為了評價目標算法，本研究構建了帶有標簽的仿真冷凍電鏡單顆粒圖像數據集，同時使用真實拍攝的冷凍電鏡圖像測試目標方法.在仿真數據集以及真實數據集上，CL-Clustering都展現出了優秀的性能.

未來該研究仍有許多可以改進的地方：1) 嘗試使用性能更加優越的主干網絡；2) 對于對稱性較高的蛋白質分子，非同類圖像的數據增強屬于同類圖像的概率會更大，會影響算法精度，可以嘗試在訓練過程中引入聚類，迭代進行模型訓練與特征聚類，根據聚類結果動態修正損失函數的計算，使編碼器的訓練過程盡量規避該現象；3) 使用更多類型的真實電鏡圖像數據集評價目標聚類算法；4)相比較聚類，三維分類對重構的影響更大.鑒于CL-Clustering的類平均圖擁有更好的類別覆蓋性，可以嘗試從中提取輔助三維分類的信息，比如由平均圖采樣重構出子結構再結合三維PCA[35]之類的方法分析顆粒的三維異構性.