使用 TFF 进行联邦学习研究

概览

TFF 是一个可扩展的强大框架，通过在现实的代理数据集上模拟联邦计算来开展联邦学习 (FL) 研究。本页面介绍了与研究模拟相关的主要概念和组件，并提供了在 TFF 中开展各类研究的详细指南。

TFF 中研究代码的典型结构

TFF 中实现的联邦学习模拟通常包含三种主要逻辑。

  1. 个别的 TensorFlow 代码片段，通常为 tf.function，用于封装在单个位置（例如在客户端或服务器上）运行的逻辑。这些代码通常是在没有任何 tff.* 引用的情况下编写和测试的，并且可以在 TFF 之外重复使用。例如，联邦平均算法中的客户端训练循环就是在这一层实现的。

  2. TensorFlow Federated 编排逻辑，它通过将上述 1 中的各个 tf.function 封装为 tff.tensorflow.computation，然后使用诸如 federated_language.federated_broadcast 和 federated_language.federated_mean 等抽象在 federated_language.federated_computation 内部对它们进行编排。例如，请参阅此联邦平均算法的编排。

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

联邦学习数据集

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

数据集包括

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

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

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

  • CIFAR-100。 CIFAR-100 数据集的联邦划分版本，分布在 500 个训练客户端和 100 个测试客户端上。每个客户端拥有 100 个唯一样本。这种划分方式旨在在客户端之间创建更真实的异构性。有关更多详细信息，请参阅 API。

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

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

  • iNaturalist 一个由各类物种照片组成的数据集。该数据集包含 120,300 张图像，涵盖 1,203 个物种。该数据集提供七种版本。其中一种按摄影师分组，包含 9257 个客户端。其余数据集按拍摄照片的地理位置分组。这六种版本的数据集包含 11 - 3,606 个客户端。

高性能模拟

虽然联邦学习模拟的挂钟时间并不是评估算法的相关指标（因为模拟硬件并不代表真实的联邦学习部署环境），但能够快速运行联邦学习模拟对于研究效率至关重要。因此，TFF 在提供高性能的单机和多机运行时方面投入了大量精力。相关文档正在开发中，但目前请参阅关于 TFF 加速器模拟 的说明，以及关于 在 GCP 上使用 TFF 设置模拟 的说明。高性能 TFF 运行时默认处于启用状态。

针对不同研究领域的 TFF

联邦优化算法

根据所需的自定义程度，可以在 TFF 中以不同方式开展联邦优化算法研究。

联邦平均算法 的最小独立实现版本可以在 此处 找到。该代码包括用于本地计算的 TF 函数、用于编排的 TFF 计算，以及一个以 EMNIST 数据集为例的 驱动脚本。按照 README 中的详细说明，这些文件可以轻松调整以用于自定义应用和算法更改。

可以在 此处 找到更通用的联邦平均算法实现。此实现允许使用更复杂的优化技术，包括在服务器和客户端上使用不同的优化器。其他联邦学习算法，包括联邦 k-means 聚类，可以在 此处 找到。

模型更新压缩

模型更新的有损压缩可以降低通信成本，从而缩短整体训练时间。

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

差分隐私

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

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

联邦 GAN（在下文中描述）是实现用户级差分隐私的另一个 TFF 项目示例（例如，此处代码所示）。

稳健性和攻击

TFF 也可用于模拟针对联邦学习系统的定向攻击，以及你真的能给联邦学习留后门吗？中考虑的基于差分隐私的防御措施。这是通过构建一个包含潜在恶意客户端的迭代过程来完成的（请参阅 build_federated_averaging_process_attacked）。targeted_attack 目录包含更多详细信息。

  • 可以通过编写一个 TensorFlow 函数作为客户端更新函数来实现新的攻击算法，请参阅 ClientProjectBoost 获取示例。
  • 新的防御措施可以通过自定义 'tff.utils.StatefulAggregateFn' 来实现，该函数对客户端输出进行聚合以获取全局更新。

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

生成对抗网络 (GAN)

GAN 构成了一种有趣的 联邦编排模式，它看起来与标准的联邦平均算法略有不同。它们涉及两个不同的网络（生成器和判别器），每个网络都使用各自的优化步骤进行训练。

TFF 可用于联邦训练 GAN 的研究。例如，近期工作 中提出的 DP-FedAvg-GAN 算法就是 在 TFF 中实现 的。这项工作证明了结合联邦学习、生成模型和 差分隐私 的有效性。

个性化

联邦学习环境下的个性化是一个活跃的研究领域。个性化的目标是为不同的用户提供不同的推理模型。解决这个问题的方法可能有很多。

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

另一种方法是通过完全在本地训练模型的一部分，从而避免训练完整的全局模型。此方法的实例在 这篇博客文章 中有描述。这种方法也与元学习有关，请参阅 这篇论文。