使用 TFF 进行联合学习研究

概述

TFF 是一个可扩展的强大框架，通过在实际代理数据集上模拟联合计算来进行联合学习 (FL) 研究。本页面描述了与研究模拟相关的主要概念和组件，以及在 TFF 中进行各种研究的详细指南。

TFF 中研究代码的典型结构

在 TFF 中实现的研究 FL 模拟通常包括三种主要的逻辑类型。

  1. 单个 TensorFlow 代码段（通常为 tf.function），它会封装在单个位置（例如客户端或服务器）上运行的逻辑。此代码在编写或测试时通常没有任何 tff.* 引用，且可以在 TFF 之外重用。例如，联合平均中的客户端训练循环就是在此级别上实现的。

  2. TensorFlow Federated 编排逻辑，它会通过将第 1 点中的各个 tf.function 封装成 tff.tf_computation 从而将其绑定在一起，然后使用抽象（如 tff.federated_computation 中的 tff.federated_broadcast 和 tff.federated_mean ）对其进行编排。相关示例请参阅联合平均编排。

  3. 外部驱动程序脚本，它能模拟生产 FL 系统的控制逻辑，从数据集中选择模拟客户端，然后在这些客户端上执行第 2 点中定义的联合计算。例如，Federated EMNIST 实验驱动程序。

联合学习数据集

TensorFlow Federated 托管了多个数据集，这些数据集代表可以通过联合学习解决的实际问题的特征。

注：这些数据集也可以被任意基于 Python 的 ML 框架（如 Numpy 数组）使用，如 ClientData API 中所述。

数据集包括：

  • StackOverflow。一个用于语言建模或监督学习任务的真实文本数据集，训练集中有 342,477 个唯一用户和 135,818,730 个样本（句子）。

  • Federated EMNIST。EMNIST 字符和数字数据集的联合预处理，其中每个客户端对应一个不同的编写器。完整的训练集包含 3400 个用户和来自 62 个标签的 671,585 个样本。

  • Shakespeare。基于威廉·莎士比亚全集的较小的字符级文本数据集。该数据集由 715 个用户（莎士比亚戏剧中的角色）组成，其中每个样本对应给定戏剧中的角色所说的一组连续台词。

  • CIFAR-100。CIFAR-100 数据集跨 500 个训练客户端和 100 个测试客户端的联合分区。每个客户都有 100 个独特的样本。分区以在客户端之间创建更现实异构性的方式完成。有关详情，请参阅 API。

  • Google Landmark v2 数据集。该数据集由各种世界地标的照片组成，图像按摄影师分组以实现数据的联合分区。提供两种形式的数据集：较小的数据集包括 233 个客户端和 23080 个图像，较大的数据集包括 1262 个客户端和 164172 个图像。

  • CelebA。名人面部样本（图像和面部属性）数据集。该联合数据集将每个名人的样本组合在一起形成一个客户端。共有 9343 个客户端，每个客户端至少包含 5 个样本。该数据集可以按客户端或按样本分为训练组和测试组。

  • iNaturalist。一个由不种物种的照片组成的数据集。该数据集包含 1203 个物种的 120300 个图像。提供七种形式的数据集。其中一种按摄影师分组，包含 9257 个客户端。其余数据集按拍摄照片的地理位置分组。这六种数据集包含 11 - 3606 个客户端。

高性能模拟

虽然 FL 模拟的时钟时间不是评估算法的相关指标（因为模拟硬件不代表真实的 FL 部署环境），但是快速运行 FL 模拟的能力对于提高研究效率至关重要。因此，TFF 投入了大量资源来提供高性能的单机和多机运行时。相关文档正在编写中，但现在您可以参阅使用 Kubernetes 进行高性能模拟教程、有关使用加速器进行 TFF 模拟的说明，以及有关设置 GCP 上的 TFF 模拟的说明。默认情况下，高性能 TFF 运行时处于启用状态。

针对不同研究领域的 TFF

联合优化算法

在 TFF 中，根据所需自定义程度的不同，可以采用不同的方法对联合优化算法进行研究。

此处提供了联合平均算法的最小独立实现。举例来说，代码包括用于本地计算的 TF 函数、用于编排的 TFF 计算，以及 EMNIST 数据集上的驱动程序脚本。这些文件可按照 README 中的详细说明轻松适应自定义的应用和算法更改。

您可以在此处查看联合平均算法更为普遍的实现。此实现支持更复杂的优化技术，包括在服务器和客户端上使用不同的优化器。您可以在此处查看包括联合 k-means 聚类算法在内的其他联合学习算法。

模型更新压缩

模型更新的有损压缩可以降低通信成本，进而减少总体训练时间。

要复制最近的论文，请参阅本研究项目。要实现自定义压缩算法，请参阅基线项目中的 comparison_methods 作为示例，如果您尚不熟悉，请参阅 TFF 聚合器教程。

差分隐私

TFF 可与 TensorFlow 隐私库互操作，以研究新的算法，从而对使用差分隐私的模型进行联合训练。有关使用基本 DP-FedAvg 算法和扩展程序进行 DP 训练的示例，请参阅此实验驱动程序。

如果要实现自定义 DP 算法并将其应用于联合平均算法的聚合更新，可以实现一个新的 DP 均值算法作为 tensorflow_privacy.DPQuery 的子类，然后使用查询实例创建 tff.aggregators.DifferentiallyPrivateFactory。实现 DP-FTRL 算法的示例可以在此处找到。

联合 GAN（如下所述）是 TFF 项目的另一个示例，它实现了用户级别的差分隐私（例如此处的代码）。

鲁棒性和攻击

TFF 还可以用于模拟联合学习系统上的针对性攻击以及 Can You Really Back door Federated Learning? 中所考虑的基于差分隐私的防御。这是通过使用潜在的恶意客户端构建迭代过程来实现的（请参阅 build_federated_averaging_process_attacked）。targeted_attack 目录中包含更多详细信息。

  • 新的攻击算法可以通过编写客户端更新函数来实现，该函数是 Tensorflow 函数。有关示例请参阅 ClientProjectBoost。
  • 新的防御可通过自定义 'tff.utils.StatefulAggregateFn'（聚合客户端输出以获得全局更新）来实现。

有关模拟的示例脚本，请参阅 emnist_with_targeted_attack.py。

生成对抗网络

GAN 提供了一种有趣的联合编排模式，这种模式看上去和标准的联合平均略有不同。它们涉及两种不同的网络（生成器和判别器），每种网络使用自己的优化步骤进行训练。

TFF 可用于研究 GAN 的联合训练。例如，最近研究工作中展示的 DP-FedAvg-GAN 算法就是在 TFF 中实现的。此研究工作演示了将联合学习、生成模型和差分隐私相结合的有效性。

个性化

联合学习设置中的个性化是一个活跃的研究领域。个性化的目的是为不同的用户提供不同的推理模型。此问题可能有不同的解决方法。

其中一种方法是让每个客户端使用自己的本地数据微调单个全局模型（使用联合学习进行训练）。这种方法与元学习有关，请参阅此论文。emnist_p13n_main.py 中给出了此方法的示例。要探索和比较不同的个性化策略，您可以进行如下操作：

  • 通过实现 tf.function 来定义个性化策略，该函数从初始模型开始，使用每个客户端的本地数据集训练和评估个性化模型。build_personalize_fn 中给出了一个示例。

  • 定义一个 OrderedDict，将策略名称映射到相应的个性化策略，并将其用作 tff.learning.build_personalization_eval 中的 personalize_fn_dict 参数。

另一种方法是完全以本地方式训练部分模型，从而避免训练完全全局的模型。本博文中介绍了这种方法的实例化。这种方法也与元学习有关，请参阅本文。要探索部分本地联合学习，您可以：