基于对比学习的冷冻电镜单颗粒图像聚类算法

颜 阳，郑清炳，张东旭，李少伟，葛胜祥，张 军，夏宁邵

(厦门大学公共卫生学院，福建厦门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]之类的方法分析颗粒的三维异构性.