[论文解读] Towards Federated Foundation Models: Scalable Dataset Pipelines for Group-Structured Learning

Zachary Charles, Nicole Mitchell+3 more|arXiv (Cornell University)|Jul 18, 2023

Privacy-Preserving Technologies in DataComputer Science被引用 3

一句话总结

本文介绍了 Dataset Grouper，一个可扩展的、与框架无关的库，用于通过用户定义的分区方式，从现有数据集创建大规模分组结构数据集——尤其适用于联邦学习。该库支持包含数十亿个分组的数据集，并首次实现了对拥有数亿至数十亿参数的基座模型进行联邦训练，揭示了在大规模下 FedAvg 的行为更类似于元学习方法，而非经验风险最小化。

ABSTRACT

We introduce Dataset Grouper, a library to create large-scale group-structured (e.g., federated) datasets, enabling federated learning simulation at the scale of foundation models. This library facilitates the creation of group-structured versions of existing datasets based on user-specified partitions and directly leads to a variety of useful heterogeneous datasets that can be plugged into existing software frameworks. Dataset Grouper offers three key advantages. First, it scales to settings where even a single group's dataset is too large to fit in memory. Second, it provides flexibility, both in choosing the base (non-partitioned) dataset and in defining partitions. Finally, it is framework-agnostic. We empirically demonstrate that Dataset Grouper enables large-scale federated language modeling simulations on datasets that are orders of magnitude larger than in previous work, allowing for federated training of language models with hundreds of millions, and even billions, of parameters. Our experimental results show that algorithms like FedAvg operate more as meta-learning methods than as empirical risk minimization methods at this scale, suggesting their utility in downstream personalization and task-specific adaptation. Dataset Grouper is available at https://github.com/google-research/dataset_grouper.

研究动机与目标

• 解决联邦学习与基座模型研究中缺乏大规模分组结构数据集的问题。
• 从 TensorFlow 和 HuggingFace Datasets 等现有数据集出发，实现高效、内存优化的异构联邦数据集创建。
• 支持模拟现代基座模型规模的联邦训练，包括参数量达数亿至数十亿的模型。
• 提供与框架无关的解决方案，可扩展至无法完全加载到内存中的数据集，采用流式处理与分层数据格式。
• 在前所未有的规模下，实证研究联邦优化算法（如 FedAvg 和 FedSGD）的行为，尤其关注其与元学习和个性化的关系。

提出的方法

• Dataset Grouper 使用用户定义的分区函数，将基础数据集中的样本分组为客户端专属的组，实现灵活且可定制的联邦数据生成。
• 该库支持三种数据格式——内存中、分层式与流式处理，通过避免将完整数据集加载到 RAM 中，实现内存效率优化。
• 系统设计用于扩展至包含数百万甚至数十亿个分组的数据集，通过流式迭代最小化峰值内存使用量。
• 该系统与 TensorFlow 和 HuggingFace 等现有机器学习框架集成，可直接用于训练流水线。
• 通过生成 C4 等大规模文本数据集的分组结构版本，支持大规模联邦语言建模，其数据分布呈长尾分布且序列长度较长。
• 该框架支持预训练与微调场景，支持个性化与差分隐私感知的训练。

Figure 1: Per-group statistics of the new group-structured (i.e. federated) language modeling datasets.

实验结果

研究问题

• RQ1 如何高效生成可扩展至超过标准系统内存容量的大规模分组结构联邦学习数据集？
• RQ2 在联邦设置下训练大规模基座模型时，FedAvg 在多大程度上表现为元学习方法，而非经验风险最小化方法？
• RQ3 每个客户端的本地小批量数量等超参数如何影响大规模联邦训练中预个性化与后个性化阶段的性能？
• RQ4 能否通过现有优化算法与可扩展的数据流水线，实际实现十亿参数模型的联邦训练？
• RQ5 数据分区策略与通信效率对大规模联邦学习中模型收敛性与个性化性能的影响如何？

主要发现

• Dataset Grouper 实现了首次在 C4 数据集的分组结构版本上对参数量为 1 亿至 10 亿的解码器仅用 Transformer 模型进行联邦训练，证明了大规模可行性的实现。
• 当每个客户端使用 64 个批次时，FedAvg 的后个性化损失为 0.008，而 FedSGD 在相同条件下保持 3.3 的后个性化损失，表明 FedAvg 在此设置下具有更优的泛化能力。
• 在处理的总 token 数相等的情况下，FedAvg 在 τ=4 时达到最低的预个性化损失（3.8）与后个性化损失（0.006），表明其在本地更新频率上实现了最优权衡。
• 当通信非瓶颈时，后个性化损失在不同 τ 值下保持稳定，表明在使用足够训练 token 的前提下，对本地更新调度具有鲁棒性。
• 即使在包含数十亿个分组的数据集上，流式与分层格式的峰值内存使用量也低于 2 MB，而内存格式需要数百 GB，证明了系统的可扩展性。
• 结果表明，FedAvg 在大规模下更像是一种元学习方法，而非经验风险最小化，且在下游个性化任务中表现强劲。

(a) In-Memory.