在AWS上重新定义云原生联邦学习架构

将云原生联邦学习架构重新定义于AWS平台之上

机器学习(ML),尤其是深度学习,需要大量的数据来提高模型的性能。客户经常需要用来自不同地区、组织或AWS账户的数据来训练模型。由于隐私要求、数据传输成本高或操作复杂性等原因,要将这些数据集中起来进行ML是一项具有挑战性的任务。

联邦学习(FL)是一种分布式ML方法,它在分布式数据集上训练ML模型。FL的目标是通过使用更多的数据来提高ML模型的准确性,同时保护分布式数据集的隐私和局部性。FL增加了用于训练ML模型的数据量,特别是与稀有和新事件相关的数据,从而产生更通用的ML模型。在AWS上,现有的合作伙伴开源FL解决方案包括FedMLNVIDIA FLARE。这些开源软件包通过在虚拟机上运行而不使用AWS上可用的云原生服务来部署在云上。

在本博客中,您将学习如何在AWS上构建一个云原生FL架构。通过在AWS上使用基础设施即代码(IaC)工具,您可以轻松部署FL架构。另外,云原生架构充分利用了各种具备安全性和运营卓越性的AWS服务,从而简化了FL的开发过程。

我们首先讨论FL的不同方法和挑战。然后我们演示如何在AWS上构建一个云原生FL架构。构建这个架构的样例代码在GitHub上可供参考。我们使用AWS Cloud Development Kit(AWS CDK)通过一键部署来部署这个架构。样例代码演示了这样一个场景,即服务器和所有客户端都属于同一组织(同一AWS账户),但由于数据本地化的要求,它们的数据集无法集中起来。样例代码支持用于训练神经网络模型的水平和同步FL。FL客户端使用的ML框架是TensorFlow。

联邦学习概述

FL通常涉及一个中央FL服务器和一组客户端。客户端是执行本地训练的计算节点。在FL训练循环中,中央服务器首先向一组客户端发送一个通用的全局模型。客户端使用本地数据训练全局模型,然后将本地模型提供给服务器。服务器将本地模型聚合成一个新的全局模型,然后开始新的训练循环。可能会有几十个训练循环,直到全局模型收敛或达到训练轮次的阈值为止。因此,在中央FL服务器和客户端之间进行ML模型交换时,并不需要将训练数据移动到中央位置。

根据客户端类型,FL可以分为两大类:跨设备和跨数据中心。跨设备FL通过在大量设备上保留训练数据本地进行训练,如手机或物联网设备,这些设备的网络连接有限且不稳定,训练一个共同的全局模型。因此,跨设备FL的设计需要考虑频繁加入和退出FL客户端的情况。

跨数据中心FL在分布在不同组织和地理分布式数据中心的数据集上训练一个全局模型。由于数据保护规定、操作挑战(如数据复制和同步)或高成本等原因,这些数据集被禁止移出组织和数据中心地区。与跨设备FL相比,跨数据中心FL假设组织或数据中心具有可靠的网络连接、强大的计算资源和可访问的数据集。

FL已经应用于各个行业,如金融医疗保健医学电信等,其中隐私保护至关重要或需要数据本地化。FL已经用于在多家金融机构之间训练全局模型以进行金融犯罪检测。全局模型的表现比只使用本地数据集训练的模型提高了20%。在医疗保健领域,FL已被用于根据多家医院的电子医疗记录预测住院患者的死亡率。全局模型在所有参与医院中的预测死亡率优于本地模型。FL还被用于脑肿瘤分割。与在中央位置收集分布式数据集进行训练的模型相比,脑肿瘤分割的全局模型表现相似。在电信领域,FL可应用于边缘计算、无线频谱管理和5G核心网络。

有许多其他的方法来分类联邦学习:

  • 水平或垂直 – 根据分布式数据集中的特征划分,联邦学习可以分为水平或垂直。在水平联邦学习中,所有分布式数据集具有相同的特征集。在垂直联邦学习中,数据集具有不同的特征组,需要额外的通信模式来根据重叠的特征对样本进行对齐。
  • 同步或异步 – 根据联邦学习服务器上的聚合策略,联邦学习可以分为同步或异步。同步联邦学习服务器将一组选定的客户端的本地模型汇总到全局模型中。异步联邦学习服务器在接收到客户端的本地模型后立即更新全局模型,从而减少等待时间并提高训练效率。
  • 中心式或点对点 – 典型的联邦学习拓扑是中心式的,其中一个中央联邦学习服务器协调一组客户端。另一种联邦学习拓扑是点对点的,没有任何中央联邦学习服务器,客户端从相邻的客户端汇总信息以学习模型。

联邦学习中的挑战

您可以使用在联邦学习架构中在联邦学习服务器和客户端上运行的算法解决以下挑战:

  • 数据异构性 – 联邦学习客户端的本地数据可能因特定地理位置、组织或时间窗口而异(即数据异构性)。数据异构性影响全局模型的准确性,导致更多的训练迭代和更长的训练时间。已经提出了许多解决数据异构性影响的方法,例如优化算法,客户端之间的部分数据共享,以及领域适应。
  • 隐私保护 – 本地和全局模型可能通过对抗性攻击泄露私人信息。已经提出了许多用于联邦学习的隐私保护方法。可以使用安全聚合方法来保护联邦学习服务器和客户端之间交换的本地模型的隐私。本地和全局差分隐私方法通过向本地或全局模型添加噪声来限制隐私损失,从而提供隐私和模型准确性之间的可控权衡。根据隐私需求,可以使用不同的隐私保护方法的组合。
  • 联邦分析 – 联邦分析在不违反隐私要求的情况下对分布式数据集进行统计测量。联邦分析不仅对训练之前的分布式数据集进行数据分析,还对推理中的模型监控起到重要作用。

尽管联邦学习算法存在这些挑战,但构建一个提供端到端联邦学习操作的安全架构至关重要。构建这样的架构的一个重要挑战是实现部署的便捷性。该架构必须协调联邦学习服务器和客户端进行模型构建、训练和部署,包括客户端之间的持续集成和持续开发(CI/CD),追溯性,以及联邦学习服务器和客户端的身份验证和访问控制。这些特性与集中式机器学习操作(ML Ops)类似,但由于涉及更多的参与者,实施起来更具挑战性。该架构还需要灵活性以实现不同的联邦学习拓扑和同步或异步聚合。

解决方案概述

我们在AWS上提出了一个基于云原生的联邦学习架构,如下图所示。该架构包括一个中央联邦学习服务器和两个联邦学习客户端。实际上,跨区域客户端的数量可以达到数百个。联邦学习服务器必须在AWS云上,因为它由云上提供的一套微服务组成。联邦学习客户端可以在AWS上或客户现场。联邦学习客户端托管自己的本地数据集,并拥有自己的IT和ML系统用于训练ML模型。

在联邦学习模型训练期间,联邦学习服务器和一组客户端交换ML模型。也就是说,客户端从服务器下载一个全局ML模型,进行本地训练,并将本地模型上传到服务器。服务器下载本地模型,将本地模型汇总成一个新的全局模型。这个模型交换过程是一个联邦学习训练循环。联邦学习训练循环重复进行,直到全局模型达到给定的准确性或训练轮次达到阈值。

FL架构

图1-云原生FL架构,用于FL服务器和FL客户端之间的模型训练。

先决条件

要实施此解决方案,您需要一个AWS帐户,以启动用于中央FL服务器和两个客户端的服务。本地客户端需要安装AWS命令行接口(AWS CLI),该接口允许访问FL服务器上的AWS服务,包括Amazon Simple Queue Service(Amazon SQS)、Amazon Simple Storage Service(Amazon S3)和Amazon DynamoDB

联邦学习步骤

在本节中,我们将在图1中介绍的建议架构中进行详细介绍。在FL服务器上,AWS步骤函数状态机按照图2所示的工作流运行,执行图1中的步骤0、1和5。状态机初始化服务器上的AWS服务(步骤0),并迭代FL训练的轮数。对于每个训练轮次,状态机发送一个Amazon Simple Notification Service(Amazon SNS)通知到主题global_model_ready,并附带一个任务令牌(步骤1)。然后状态机暂停并等待包含任务令牌的回调。有订阅global_model_ready主题的SQS队列。每个SQS队列对应一个FL客户端,并排队将服务器发送到客户端的通知。

图2-步骤函数状态机的工作流程。

每个客户端从其分配的SQS队列中持续获取消息。当接收到global_model_ready通知时,客户端从Amazon S3下载全局模型(步骤2),并开始本地训练(步骤3)。本地训练生成本地模型。然后客户端将本地模型上传到Amazon S3,并将本地模型信息与接收到的任务令牌一起写入DynamoDB表(步骤4)。

我们使用Amazon S3和DynamoDB来实现FL模型注册表。我们使用Amazon S3存储全局和本地模型。我们使用DynamoDB表存储本地模型信息,因为本地模型信息可能在FL算法之间不同,需要DynamoDB表支持的灵活架构。

我们还启用了DynamoDB流以触发Lambda函数,这样每当记录写入DynamoDB表(接收到新的本地模型)时,将触发Lambda函数来检查是否已经收集到所需的本地模型(步骤5)。如果是,则Lambda函数运行聚合函数将本地模型聚合成全局模型。生成的全局模型被写入Amazon S3。该函数还向步骤函数状态机发送回调和从DynamoDB表检索到的任务令牌。状态机根据条件(例如,达到训练轮次的数量阈值)确定是否应继续进行FL训练或停止。

每个FL客户端使用以下示例代码与FL服务器交互。如果要自定义FL客户端的本地训练,可以修改localTraining()函数,只要返回的值是上传到FL服务器的local_model_namelocal_model_info即可。您可以选择任何ML框架来训练FL客户端的本地模型,只要所有客户端使用相同的ML框架即可。

# 第2步:从SQS队列接收通知和模型文件名client.receiveNotificationsFromServer(sqs_region, client_queue_name)# 第3步:下载全局模型并在本地训练local_model_name, local_model_info = client.localTraining(global_model_name, s3_fl_model_registry)# 第4步:将本地模型和本地模型信息上传到FL服务器client.uploadToFLServer(s3_fl_model_registry, local_model_name, dynamodb_table_model_info, local_model_info)

用于在服务器上运行聚合函数的Lambda函数具有以下示例代码。如果要自定义聚合算法,需要修改fedAvg()函数和输出。

# 第5步:在Lambda函数中聚合本地模型def lambda_handler(event, context):  # 从DynamoDB Stream触发的事件中获取task_name  task_name = event['Records'][0]['dynamodb']['Keys']['taskName']['S']    # 从DynamoDB表中读取交易信息transactions = readFromFLServerTaskTable(os.environ['TASKS_TABLE_NAME'], task_name) # 从指定的客户端中读取本地模型信息  # token是来自Step Functions状态机的回调令牌  local_model_info, round_id, token = receiveUpdatedModelsFromClients(transactions, task_name)    # fedAvg函数将本地模型聚合成全局模型,并将全局模型存储在S3中 global_model_name, avg_train_acc, avg_test_acc, avg_train_loss, avg_test_loss = fedAvg(local_model_info, round_id)  # 输出发送到Step Function状态机    output = {'taskName': task_name, 'roundId': str(round_id), 'trainAcc': str(avg_train_acc), 'testAcc': str(avg_test_acc), 'trainLoss': str(avg_train_loss), 'testLoss': str(avg_test_loss), 'weightsFile': str(global_model_name)}   # 发送回调到Step Functions状态机,报告该令牌标识的任务已成功完成 step_client = boto3.client('stepfunctions') out_str = json.dumps(output)    step_client.send_task_success(taskToken=token, output=out_str)

这种架构有两个创新设计。首先,FL服务器使用无服务器服务,如Step Functions和Lambda。因此,不会保留用于FL服务器的计算实例,从而最大程度地减少了计算成本。其次,FL客户端从其分配的SQS队列中拉取消息,并将模型和信息上传或下载到FL服务器上的服务中。这种设计避免了FL服务器直接访问客户端的资源,这对于为FL客户端提供私密和灵活的IT和ML环境(本地或AWS云上)非常重要。

云原生的优点

这种架构是云原生的,并通过使用具有成熟的安全性和运营卓越的AWS服务来提供端到端的透明性。例如,您可以使用跨账户客户端来扮演角色,以访问FL服务器上的资源。对于本地客户端,客户端上的AWS CLI和AWS SDK for Python (Boto3)会自动提供FL服务器和客户端之间的安全网络连接。对于在AWS云上的客户端,您可以使用AWS PrivateLink和具有在传输和静止时的数据加密的AWS服务来保护数据。可以使用Amazon CognitoAWS Identity and Access Management (IAM) 来对FL服务器和客户端进行身份验证和访问控制。对于部署训练后的全局模型,您可以使用Amazon SageMaker中的ML Ops功能。

云原生架构还支持与自定义的ML框架和联邦学习算法和协议集成。例如,您可以选择一个ML框架来训练FL客户端的本地模型,并在服务器上作为Lambda函数运行的脚本中定制不同的聚合算法。此外,您还可以修改Step Functions中的工作流程,以适应服务器和客户端之间的不同通信协议。

云原生架构的另一个优势是利用云提供的基础设施即代码工具轻松部署。您可以使用AWS 云开发工具包 (AWS CDK) 和 AWS 云形成 进行一键部署。

结论

全球范围内不断实施新的隐私法律,技术基础设施迅速扩展到多个地区并延伸到网络边缘。联邦学习帮助云客户以隐私保护的方式使用分布式数据集训练精确的机器学习模型。联邦学习还支持数据本地化,并且由于不需要移动或共享大量原始数据,可能节省成本。

您可以开始为您的用例试验和构建云原生联邦学习架构。您可以自定义架构以支持各种机器学习框架,如TensorFlowPyTorch。您还可以自定义以支持不同的联邦学习算法,包括异步联邦学习聚合算法差分隐私算法。您可以利用Amazon SageMaker中的 ML Ops 功能来启用此架构。