FAX: Scalable and Differentiable Federated Primitives in JAX

我们介绍了 FAX，这是一个基于 JAX 构建的库，旨在支持数据中心和跨设备应用中的大规模分布式和联邦计算。FAX 利用 JAX 的分片机制，实现了对 TPU 和最先进的 JAX 运行时（包括 Pathways）的本地化支持。FAX 将联邦计算的构建块作为 JAX 中的基元进行嵌入，从而实现了三个关键优势。首先，FAX 计算可以转换为 XLA HLO。其次，FAX 提供了联邦自动微分的完整实现，大大简化了联邦计算的表达。最后，FAX 计算可以解释为现有的生产跨设备联邦计算系统。我们展示了 FAX 在数据中心提供了易于编程、高效和可扩展的联邦计算框架。FAX 可在 https://github.com/google-research/google-research/tree/master/fax 上获得。

FAX旨在解决大规模分布式和联邦计算中的问题，包括数据中心和跨设备应用程序。它试图通过将联邦计算构建块嵌入JAX原语中来简化联邦计算的表达。

FAX利用JAX的分片机制，使得可以原生地针对TPU和最先进的JAX运行时进行定位。FAX计算可以被翻译成XLA HLO，提供了联邦自动微分的完整实现，可以将计算解释为现有的生产跨设备联邦计算系统。

FAX提供了易于编程、高效、可扩展的联邦计算框架，可以用于数据中心。FAX嵌入了联邦计算构建块，可以将计算翻译成XLA HLO。FAX提供了联邦自动微分的完整实现，大大简化了联邦计算的表达。FAX可用于现有的生产跨设备联邦计算系统。

最近的相关研究包括《Federated Learning: Strategies for Improving Communication Efficiency》、《Federated Learning with Non-IID Data》等。