Science China｜用增强的联邦学习应对药物发现数据小和偏的困境
2021年7月27日，中科院上海药物所蒋华良和郑明月课题组/上海科技大学联合华为云的医疗智能体 (EIHealth) 平台，在SCIENCE CHINA Life Sciences发表了题为Facing small and biased data dilemma in drug discovery with enhanced federated learning approaches的文章。
该工作证明了联邦学习在预测药物相关性质方面的适用性，也证明了联邦学习在应对药物发现数据“小”和“偏”的问题中的重要作用。
AI应用于药物研发，存在数据困境
众所周知，人工智能在药物研发中的应用需要大量的高质量数据，来实现良好的预测性能。
然而，AI制药的发展往往面临数据获取的挑战。药物研发数据 (尤其是药物开发后期阶段的数据) 生成过程耗时、昂贵且保密性极强，严重影响了制药公司及相关数据方分享数据的意愿。
此外，科学研究中的数据偏差也困扰着人工智能在药物发现中的应用。比如，一个特定的药物分子的性质在不同来源的记录中存在很大的差异，即使是在相同的科学测试中。
联邦学习，从“数据孤岛”走向“共同富裕”
许多药物研发数据被封藏在一个个数据孤岛之中，这对于药物开发来说是巨大的损失。
那么，有没有方法可以改善药物研发的数据困境呢？这就要说到联邦学习。
在机器学习领域，联邦学习的独特性就在于可以实现“共同富裕”。在联邦学习中，参与者是独立自治的个体，组织通过合理的激励机制和利益分配机制来激励参与者积极加入联邦学习。这是传统的机器学习所没有的。
联邦学习代表了一种场景，即多个参与者可以集体训练一个模型而不共享原始数据，在充分利用来自不同来源的分布式数据的同时，不泄露数据的敏感信息 (如机密性和知识产权的数据) 。这种新兴的分散式机器学习范式有望极大地提高人工智能驱动的药物发现的成功率。
除了在参与者的数据集上集体训练一个全局模型外，在联邦学习中，每个客户可以拥有一个个性化的模型。个性化联邦学习是将连成一片的孤岛，再次一一分开 (孤岛形成是有原因的，每个孤岛都有自己的特色，如同针对老人的模型和针对小孩的模型是不一样的) 。个性化模型的建立考虑到了本地数据的特色，有助于提高模型预测的准确率。
为药物发现建立一个通用的联邦学习框架
在这篇文章中，上海药物所蒋华良和郑明月课题组/上海科技大学联合华为云团队，搭建了一个通用的联邦学习框架 (图1) ，并在FATE (Federated AI Technology Enabler) (https://github.com/FederatedAI/ FATE) 上进行了测试。FATE是一个开源项目，旨在为联邦学习提供一个安全的计算框架。
与Gboard这种跨设备的联邦学习应用不同，药物发现的联邦学习是跨数据孤岛训练的 (跨孤岛联邦学习, cross-silo federated learning) 。这样的学习方式包括一个协调服务器和几个合作者，并配备了联邦学习客户端程序。合作者可以是大型制药公司、生物技术初创公司，甚至是拥有自己的数据的学术实验室。
图1 一个用于药物发现的联邦学习系统的工作流程。在联合训练中：(i) 协调服务器向每个客户端传送 (Broadcast) 最新的共享全局模型；(ii) 客户端在本地进行模型更新的计算，(iii) 对模型更新进行加密和上传；(iv) 最后，协调服务器安全地汇总所有加密的模型更新内容，并使用它们来更新共享全局模型以进行下一轮训练。训练结束后，选择最佳模型进行推广，并为拥有自己标注数据的用户进行模型的个性化定制 (rollout & Customizing) 。
联邦学习的流程如下。
在每一轮跨孤岛的联合训练中， (i) 协调服务器向每个客户端传送最新的共享全局模型， (ii) 每个客户端通过执行训练程序在本地计算模型更新， (iii) 根据安全聚合协议 (Secure Aggregation protocol) 加密并上传模型更新， (iv) 最后，协调者服务器安全地聚合所有加密的模型更新并使用它们来更新共享全局模型。这个过程需要进行多轮训练，直到模型收敛或达到停止的标准。然后选择最佳模型进行推广。
用不同任务模拟跨数据孤岛的联邦学习过程
在本文的研究中，联合团队使用三个任务来模拟跨数据孤岛的联合学习过程：基于化学结构进行药物溶解度、激酶抑制活性和hERG心脏毒性的预测。
这些案例中的数据集涵盖化合物的化学空间、测量方法、实验条件、非标准表示方法和数据规模方面的差异。这些来自不同来源的真实世界的药物性质数据集代表了不同客户的数据分布，可以用来研究药物发现如何从联邦学习中获益。
7个水溶性数据集 (数据集F1-F4，数据集C1-C3) 从7个不同的来源收集。在模拟中，数据集F1-4的所有者是想要参与联邦学习的合作方，而数据集C1-3的所有者是想要使用联邦学习训练的模型的用户。数据集F1-4的1/10的分子被作为验证集，以防止过度拟合。另有1/10的分子作为测试集，与不同的模型进行比较。
激酶抑制数据集包含四个数据集。所有这四个数据集被用于联合训练。
hERG数据集从不同的来源收集。来自Cai等人、Pubchem NCATS、Pubchem JHICC和ChEMBL的数据集被模拟为参加联合训练的客户，其余数据集被模拟为测试集。
结果1：联邦学习应对非IID数据非常有用
在传统的集中式机器学习药物发现应用中，为了包含更多的数据，研究人员从不同的来源收集数据，并假设数据是独立同分布 (IID，independent identically distributed，随机过程中，如果随机变量服从同一分布，并且互相独立，那么这些随机变量是独立同分布) 。然而，由于科学测试的测量中引入了很高的数据偏差，而科学测试是由不同的人在不同的实验环境中进行的，所以通常会违反这一假设。
结果2：5个epochs产生最佳预测性能
如图2B所示，研究构建了一个六层的多层感知器 (MLP，Multilayer Perceptron) 模型。与传统的集中式模型相比，联邦式模型具有相同的网络结构，只是在既定的本地训练epoch (训练集中的全部样本都在训练模型中走了一遍，并返回一次，为一个epoch) 之后，每个参与者的模型更新将被加密并上传到协调服务器，随后协调服务器将新模型返回给每个参与者。
其中，模型更新频率，即客户端在上传加密的模型更新前应在本地运行多少个epoch，是一个有影响的超参数。如表2所示，在5个独立的运行中，每5个epochs的本地训练性能中的模型更新，平均产生了最佳的预测性能。
结果3：联合训练优于单独训练，Federated模型优于Union模型
在这项研究中，数据集F1、F2、F3和F4被用来模拟参与联邦学习模型训练过程的参与者，C1、C2和C3被用来模拟没有参与训练但希望用自己的数据对联邦模型进行个性化的客户。我们将联合建模与个性化和集中建模进行了比较，以每个参与者的测试集上的平均绝对误差 (MAE) 值来衡量模型的性能。一般来说，在单个数据集上训练的子模型在他们自己的内部测试集上取得了较高的性能，但在其他测试集上的性能要低得多，表明这些子模型不能很好地泛化。
相比之下，Federated模型 (联邦学习模型) 和Union模型 (先集中数据再训练) 在跨数据孤岛的数据集上的预测性能有很大的提高。对于数据集F1-4，联邦学习模型的MAE值普遍低于本地训练的相应子模型，并且在外部客户C1-3的测试中保持了预测能力。
在Union模型中，来自不同来源的数据被简单地汇集在一起，以一种非隐私的方式进行训练。研究结果显示，Federated模型的表现比Union模型更好。
为了研究它们在学习方面的差异，我们比较了Union模型和共享的Federated模型中全连接层的权重分布 (图2C-E) 。与初始化的权重分布相比，Union模型的权重分布在集中训练后基本没有变化，而federated模型的权重分布变化明显，更多的权重集中在0上。在相同的网络架构和相同的参数群中，更多的权重为0意味着模型的规则化和简单化，这可能会使泛化效果更好。
这表明，联邦学习具有内在的正则化 (regularization) 效应，在面对高偏差数据时，可能有助于提高泛化性能。
结果4：改进网络架构--从MLP到RFCN，改进聚合算法--从FedAvg到FedAMP
为了研究不同的网络架构和联邦学习聚合算法对联邦学习性能的影响，除了之前的MLP架构和FedAvg聚合算法，研究人员还测试了一个名为RFCN (residual fully connected neural network) 的神经网络架构和一个名为FedAMP (federated learning approach named federated attentive message passing) 的个性化联邦学习方法。
最初发现RFCN善于对非欧氏空间数据 (底层数据结构没有特殊或时间关系) 进行建模，如基因组数据。在此，我们还比较了RFCN架构与MLP在建模高维扩展连接性指纹 (ECFP) 数据方面的表现。
关于客户端更新的聚合，我们比较了FedAMP和最常用的FedAvg算法。FedAMP是一种个性化的联邦学习聚合算法，它鼓励模型参数相似的客户端进行更强的协作，因此该算法自适应地发现了客户端之间隐藏的协作关系，并通过给不同的客户端分配不同的模型来增强其协作效果。
研究结果显示，RFCN模型在集中式学习或联邦式学习中的表现均优于MLP模型，这表明RFCN在为ECFP数据建模时与MLP相比是一个更好的神经元网络架构。
结果5：抑制激酶性活性预测的案例
结果6：预测hERG影响的案例
药物诱导的hERG阻滞是心脏毒性的主要原因之一。然而，许多实验检测都可用于评估hERG的影响，这在hERG作用的记录值中会引起很大的偏差。以前的研究侧重于合并不同来源的数据，并构建一个集中的模型来拟合数据，这可能导致有偏见和过度拟合的模型。
在我们的研究案例中，使用来自不同来源的hERG抑制性数据构建了一个联合hERG分类模型。研究结果显示，FedAVG+MLP模型和FedAMP+MLP模型的表现优于单个数据集的建模，但始终不如Union+MLP模型。
联邦学习在制药中的未来
在讨论和结论中，作者对联邦学习的未来进行了展望。
一，联邦迁移学习将进一步扩大特征空间和样本量。
从更大的联邦学习范围来看，本研究建立的框架只关注模拟有相同特征空间 (分子ECFP指纹) 作为输入的参与者，这是横向联邦学习 (horizontal federated learning) 。也有纵向联邦学习 (vertical federated learning) 方案，可以应对有不同特征类型作为输入的参与者。
此外，水平和纵向联邦学习的组合，被称为联邦迁移学习 (federated transfer learning) ，可以有效地应对共享一些特征类型和样本的参与者，但也有自己的专有特征。联邦迁移学习将进一步扩大我们的特征空间和样本量，同时考虑到多个参与者的特征空间和样本空间的联合。例如，为了预测候选药物的临床结果，我们需要整合来自多方的具有共享和专有特征的数据，包括相关的制药公司、医院和病人。联邦迁移学习可能为每一方产生大的附加值。
二，联邦学习虽然还存在一些问题，但对药物发现来说很有吸引力。
虽然联邦学习仍可能存在一些安全问题和恶意或非恶意的失败，如模型更新中毒、数据中毒和数据管线失败，但它已经引起了大量的关注，并一直在快速改进和发展。这种模式为通过安全的分布式训练整合机密数据集提供了可能，这在以前被认为是不切实际的，但对药物发现来说绝对是有吸引力的。鉴于药物发现中的预测模型经常在非常有限的领域内工作，如果有机会利用来自多个机构的更大的和不同的数据孤岛，将会提高药物发现中预测模型的通用性。
联邦学习代表了一种新的机器学习范式，保护隐私的特点将鼓励越来越多的机构充分利用他们的数据，将来会有越来越多的数据被提供给最新的机器学习模型，从而解决药物发现中的 "小数据"困境。
联邦学习的设置也使得为不同的用户/客户提供个性化的模型成为现实，从而缓解了数据偏差的问题，在实际应用场景中实现更好的预测性能并形成更明智的建议。