SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Federated algorithms. We use two prototypical FL algorithms: FedAvg and FedSGD [81]. In each federated round, we select the next cohort of 16 clients. Local training is done on client data batched to 16 examples with a sequence length of 128 tokens. We repeat client data as necessary to ensure that all clients have 1024 examples. For FedAvg, we run 3125 rounds of federated training. During each round, each client in that round's cohort takes 64 gradient steps. Thus, the federated training will involve roughly 200K batched gradient computations in total. For FedSGD, we use the same setup, except that clients do not locally update their own model when computing local gradients. Instead, these 64 minibatch gradients are averaged into a single large-batch gradient and sent to the server.
Optimizer hyperparameters. For FedAvg, we use the client/server-optimizer framework proposed by Reddi et al. [30]. We use SGD for the client optimizer and Adam for the server optimizer. FedSGD only has a server optimizer, which we also set to Adam. We only tune the learning rate(s), tuning over  $\{10^{-4}, 10^{-3}, \dots, 10^{0}\}$ , and selecting the learning rate(s) that minimize average training loss. For details and a full list of optimizer hyperparameters, see Appendix C.
Hardware configuration. We run our experiments using a TPU Pod slice consisting of 16 TPU v3 chips in a 4x4 topology, configured to use a multi-machine inter-chip interconnect mesh. Each TPU v3 chip contains two TensorCores, 32 GiB of high-bandwidth memory, 87.25 GiB RAM, 900 GBps bandwidth, and 123 teraflops peak compute.
Iteration efficiency. We test the efficiency of Dataset Grouper in practical large-scale simulations. Specifically, we measure the time it takes for each round of federated training and what portion of that time is spent iterating over data, including preprocessing. We perform 100 rounds of FedAvg for varying cohort sizes (the number of clients per round) and present the results in Table 4. We see that dataset iteration takes under  $10\%$  of the total runtime, even for larger cohort sizes. This is despite the fact that dataset iteration is done entirely on the host, while the training time is parallelized between multiple TPU slices. Further improvements in the data pipeline can only lead to a marginal speedup, highlighting the efficiency and scalability of the streaming dataset format in Section 3.1.
Federated learning rate schedules. Large-scale training on non-partitioned data generally involves a variety of important techniques, such as learning rate scheduling, to attain good performance. In order to determine how best to scale federated training to larger-scale settings, we investigate the use of various learning rate schedules for FedAvg and FedSGD. In both cases, we apply the learning rate schedule at the server (see [30] for a discussion of client versus server optimizers). We use constant
(a) FedAvg.
(b) FedSGD.
Figure 4: Training loss of FedAvg and FedSGD on FedC4 with different learning rate schedules. The per-round training loss is computed by (a) averaging over all batches seen by a given client within the round, and (b) averaging over all clients that participate in the round.
(c) Comparison w/ cosine decay.
learning rates, warmup with exponential decay, and warmup with cosine decay. Whenever we use warmup, we warm up for  $10\%$  of the total number of training rounds, and decay for the remainder.
We compare the resulting training loss for FedAvg and FedSGD in Figure 4. Notably, we see that learning rate scheduling leads to significant improvements in the behavior of FedSGD, while FedAvg is robust to different choices. This reflects the fact that these learning rate schedules were developed in the context of SGD, which involves applying many unbiased gradient estimates. FedSGD operates similarly, computing an unbiased gradient estimate at each round. By contrast, FedAvg involves biased gradients, often called "pseudo-gradients" [30], which may not be the gradient of any loss function [82]. Our results suggest that developing effective learning rate scheduling techniques for FedAvg is an open question, and may involve coordinating client and server learning rates.
We also see that FedAvg appears to attain a significantly lower train loss than FedSGD. We stress that this is due to how the training loss is computed. For both algorithms, it is computed by averaging the loss of all batches seen by a client and then averaging that quantity across all clients. However, the client trains as it sees data batches in FedAvg. Therefore, the client's local model adapts to its own distribution (leading to a lower loss), while in FedSGD the client does not adapt its local model. We explore this difference, which is connected to meta-learning, below.
Federated evaluation and personalization. Partitioned datasets enable group-structured learning, as well as group-level (or federated) evaluation, which may be particularly informative for measuring downstream performance across heterogeneous data splits. To demonstrate this, we use Dataset Grouper to generate an evaluation dataset from FedC4 by using its held-out validation split. We use the same partition structure as before, grouping examples according to their base domain. Because of this group structure, we can compute histograms of metrics across all groups, rather than just an average metric across all examples.
We take the resulting models trained by FedAvg and FedSGD (with constant learning rates, though we see similar results for all learning rate schedules we considered above), and compute two separate metrics for each validation client. First, we compute the average loss of the model on all examples held by the client. We refer to this as the pre-personalization loss. We then fine-tune the model for a single epoch on the client's dataset (using a client optimizer of SGD with a tuned learning rate). After personalization, we compute the average loss again, resulting in the post-personalization loss.
Algorithm	Pre-Personalization Loss	Post-Personalization Loss
10thperc.	Median	90thperc.	10thperc.	Median	90thperc.
FedAvg	5.13	5.64	6.27	0.002	0.012	0.934
FedSGD	4.38	4.93	5.40	1.25	3.38	4.53
Table 5: Validation loss of FedAvg and FedSGD, before and after personalizing on a client's dataset. Percentiles are computed across all clients in the FedC4 validation dataset.
Figure 5: Histograms of pre- and post-personalization loss across all FedC4 validation clients.
We present quantiles of these metrics in Table 5. Intriguingly, they show that the FedSGD-trained model works better for pre-personalization, but the FedAvg-trained model is much more effective at personalizing to the client's data. To further illustrate this, we consider histograms of the two distributions (across all clients) in Figure 5. This suggests a more dramatic shift. While the FedAvg and FedSGD-trained models are close in pre-personalization performance (though FedSGD does better), the post-personalization distribution for FedAvg is extremely light-tailed.
Task-specific personalization. Pre-trained foundation models are typically employed on a range of downstream tasks. In this spirit, we use the models trained on FedC4 to perform pre- and post-personalization evaluation on FedBookCO. The results, Figures 6 and 7, are similar to but less drastic than those of Figure 5. The pre-personalization loss of FedAvg is slightly larger than FedSGD (5.0 vs. 4.3 in the last checkpoint) while its post-personalization loss is smaller (2.9 vs. 4.0). Similar trends hold for FedCCnews and FedWiki datasets; cf. Appendix D. Overall, these results show that FedAvg's superior personalization performance is robust to shifts in the distribution over clients.
This phenomenon highlights connections between federated learning and meta-learning previously noted in the literature [46, 83-88]. In short, we see that FedAvg acts as a meta-learning algorithm (specifically, the Reptile algorithm [89]) where it quickly minimizes the loss of a client after a few local gradient steps (i.e., after personalization). It does not behave like an algorithm designed to minimize the empirical risk. By contrast, FedSGD operates much like SGD in the centralized setting, attempting to minimize the average loss across all examples. To the best of our knowledge, Figure 5 constitutes some of the strongest empirical evidence of the connection between federated and meta-learning to date. The scale of the FedC4 dataset (enabled by Dataset Grouper) is critical here, as clients have sufficiently large amounts of data to exacerbate client drift [28] and cause tension between loss minimization and meta-learning.
Scaling to larger models. To further demonstrate the scalability of Dataset Grouper, we train a transformer model with 1 billion parameters on the FedC4 dataset. In contrast to the results above, we train with 4 batches per client (rather than 64). Despite this, we see in Figure 8 that FedSGD still sees
Figure 6: Median pre-personalization (left) and post-personalization (left) loss over FedBookCO clients while training on FedC4. The shaded region indicates the 10th and 90th percentiles.
Figure 7: Histograms of pre- and post-personalization loss on FedBookCO after FedC4 training.
Figure 8: Pre-personalization loss of FedAvg and FedSGD across all FedC4 validation clients on a 1 billion parameter transformer model.
improved pre-personalization loss compared to FedAvg. Moreover, both algorithms see improved pre-personalization compared to Table 5, highlighting the effect of increasing the larger size.
The intersection of foundation models and group-structured data is a fertile area for research. We provide tooling for creating group-partitioned datasets for use in large-scale research simulation. We acknowledge that there are inherent risks in this endeavor. As detailed by Koch et al. [90], the typical dynamics of dataset use in machine learning research tend to enshrine certain datasets as "sole benchmarks" in the field, even on tasks for which they were not designed. Tooling aimed at allowing for flexible and reproducible dataset creation risks further entrenchment of these sole benchmarks by expanding the scope of tasks to which they are applied. However, we posit that Dataset Grouper's pipeline approach will prove to be a sustainable mechanism for ensuring availability of datasets whose intended use cases match their application in research, and can potentially reduce the enshrinement of benchmarks in areas such as federated learning.
There are a wide array of other research benefits enabled by Dataset Grouper, especially as it delivers scalable and efficient partitioned dataset pipelines compatible with foundation model training and fine-tuning. This crucially enables the exploration of phenomena that only emerge at the scale of foundation models. Our experiments are intended as a demonstration of the scaling capabilities unlocked by Dataset Grouper to the billion parameter regime. Our empirical findings indicate several interesting future directions. Most excitingly (and speculatively), the tendency of FedAvg to metalearn tasks suggests that it could provide a better "base" model for the personalization of foundation models or adaptation to downstream tasks. Moreover, there is a need to design tailored learning rates and default optimization recipes for the wider applicability of federated training algorithms. We hope that Dataset Grouper will spur further research in training, evaluation, finer grained analyses, and diagnoses of foundation models with group-structured data.
The authors thank Keith Rush, H. Brendan McMahan, Sean Augenstein, and Liam Collins for fruitful discussions and helpful comments. The authors would also like to thank Keith Rush for help with training code and multi-TPU support. Finally, the authors would like to thank Zheng Xu, Shanshan Wu, Keith Rush, and Arun Ganesh for helpful conversations on group-structured datasets.
[1] Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista A. Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafael G. L. D'Oliveira, Hubert Eichner, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascon, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konečný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrede Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasmus Pagh, Hang Qi, Daniel Ramage, Ramesh Raskar, Mariana Raykova, Dawn Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Praneeth Vepakomma, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and Open Problems in Federated Learning. Found. Trends Mach. Learn., 14(1-2):1-210, 2021.
[2] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks. In ICML, pages 1126-1135. PMLR, 2017.
[3] Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam D. Smith. Calibrating Noise to Sensitivity in Private Data Analysis. In Theory of Cryptography Conference, volume 3876 of Lecture Notes in Computer Science, pages 265-284, 2006.
[4] Martín Abadi, Andy Chu, Ian J. Goodfellow, H. Brendan McMahan, Ilya Mironov, Kunal Talwar, and Li Zhang. Deep Learning with Differential Privacy. In CCS, pages 308–318, 2016.
[5] Natalia Ponomareva, Hussein Hazimeh, Alex Kurakin, Zheng Xu, Carson Denison, H. Brendan McMahan, Sergei Vassilvitskii, Steve Chien, and Abhradeep Guha Thakurta. How to DP-fy ML: A Practical Guide to Machine Learning with Differential Privacy. J. Artif. Intell. Res., 77: 1113-1201, 2023.
[6] Brian Lester, Rami Al-Rfou, and Noah Constant. The Power of Scale for Parameter-Efficient Prompt Tuning. In EMNLP, pages 3045-3059, 2021.
[7] Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-Rank Adaptation of Large Language Models. In ICLR, 2022.
[8] Xiao Liu, Kaixuan Ji, Yicheng Fu, Zhengxiao Du, Zhilin Yang, and Jie Tang. P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks. arXiv Preprint, 2021.
[9] Nicholas Carlini, Florian Tramer, Eric Wallace, Matthew Jagielski, Ariel Herbert-Voss, Katherine Lee, Adam Roberts, Tom B Brown, Dawn Song, Ulfar Erlingsson, et al. Extracting training data from large language models. In USENIX Security Symposium, volume 6, 2021.
[10] Hannah Brown, Katherine Lee, Fatemehsadat Mireshghallah, Reza Shokri, and Florian Tramèr. What Does it Mean for a Language Model to Preserve Privacy? In FAccT, pages 2280-2292, 2022.
[11] TensorFlow Federated: Machine Learning on Decentralized Data. https://www.tensorflow.org/federated.
[12] Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H. Brendan McMahan, Virginia Smith, and Ameet Talwalkar. LEAF: A benchmark for federated settings. arXiv Preprint, 2018.
[91] Derek G Murray, Jiri Simsa, Ana Klimovic, and Ihor Indyk. tf.data: A machine learning data processing framework. arXiv Preprint, 2021.
[92] Heike Hofmann, Hadley Wickham, and Karen Kafadar. Letter value plots: Boxplots for large data. Journal of Computational and Graphical Statistics, 26(3):469-477, 2017.
A Software 19
A.1 Installation, Usage and Examples 19
A.2 Dataset Hosting 20
A.3 Dataset Licenses 20
A.4 Software Development and Maintenance 21
B Dataset Statistics 21
C Full Experimental Details 22
C.1 Datasets and Preprocessing 23
C.2 Model 23
C.3 Federated Algorithms 23
C.4 Optimizer Hyperparameters 23
C.5 Reported Metrics 24
C.6 Hardware Configuration 25
D Additional Experimental Results 25
D.1 Personalization Results 25
D.2 Ablation: Batches per Client 27
E Memory Usage 29
A Software
All code for Dataset Grouper can be found on GitHub, including all applicable licenses, disclaimers, usage instructions, and examples. We showcase some functionalities below.
A.1 Installation, Usage and Examples
Installation. Dataset Grouper can be installed as a Python package using the pip command pip install dataset-grouper.
Basic usage. At its core, Dataset Grouper can be used to partition datasets (eg. from TensorFlow Datasets [16]) across groups. It can be used with any embarrassingly parallel user-defined partitioning function of the signature get_key_fn example) -> group_id.
Below we give a simple example where we partition the MNIST dataset according to label (so that we form one group per label).
Listing 1: Using Dataset Grouper to partition MNIST
1 import apache_beam as beam   
2 import dataset_grouper as dsgp   
3 import tensorflow_datasets as tfds   
4   
5 # First we download the MNIST dataset from TFDS.   
6 datasetbuilder  $=$  tfds.Builder("mnist")   
7 dataset.Builder.download_andprepare(..)   
8   
9 # Next, we define a function that maps MNIST examples to their group.   
10 def get_label_fn(x):   
11 label  $=$  x["label"].numpy()   
12 return str.label).encode("utf-8")   
13   
14 # Finally, we build the partitioning pipeline, and run it using Apache Beam.   
15 mnist=Pipeline  $\equiv$  dsgp.tfds_to_trecords(   
16 dataset_buildersdaset.Builder,   
17 split  $=$  "train",   
18 get_key_fn  $\equiv$  get_label_fn,   
19 file_path_prefix  $= \dots$    
20 ）   
21 with beam.Pipeline() as root:   
22 mnistpipeline(root)
Dataset partitioning examples. We give additional examples of using Dataset Grouper to partition datasets into groups:
- Feature-based partitioning: This script allows partitioning of a dataset by one of its features, allowing for natural heterogeneous partitions. The language modeling datasets in Section 4 use this approach. For example, to create FedCCnews and FedC4 we partition based on the URL.
- Random partitioning: This script assigns each example to client at random.
- Heterogeneous partitioning using a Dirichlet process: This script assigns each example to a client based on a Dirichlet process. This is an embarrassingly parallel version of the popular LDA-based method that is popular in the federated learning literature [e.g. 71]. Together with the random partitioning, this can be used, for instance, to study the effect of the level of data heterogeneity on FL algorithms.
The specific commands used to create the datasets introduced in Section 4 can be found in the GitHub repository at https://github.com/google-research/dataset_grouper/tree/main/dataset_grouper/examples/datasets.
$^{6}$ https://github.com/google-research/dataset_grouper
<sup>7</sup>https://pypi.org/project/dataset-grouper/
Usage in training loops. Dataset Grouper saves partitioned datasets to the widely used TFRecord format and builds grouped dataset pipelines via tf.data [91]. It can be used as shown below.
Listing 2: Iterating over group-partitioned datasets via Dataset Grouper
1 import dataset_grouper as dsgp   
2   
3 # Load the dataset by passing the file. paths and the TFDS dataset name   
4 partitioned_dataset  $=$  dsgp.PartitionedDataset( file_pattern  $= \dots$  .   
6 tfds_features  $= 1$  c4" # Or any other TFDS dataset name.   
7 ）   
8 #Obtain an iterator of clients/groups.   
9 client_STREAM  $=$  partitioned_dataset.build_group_STREAM()   
10 # Iterate over clients.   
11 for client_dataset in client_STREAM: # client_dataset is an iterable of examples.   
12 for example in client_dataset.as_numpy_iterator(): # Process this example.
Typically, FL processes cohorts of clients. This can easily be achieved by applying a batch operation on the client_STREAM object above; see the training code below for details.
Training code examples. The GitHub repository contains JAX code to reproduce the experiments discussed in Section 5 and Appendix C. It also gives sample code to demonstrate the integration of Dataset Grouper with TensorFlow Federated [11].
A.2 Dataset Hosting
Dataset Grouper does not directly host any datasets and instead provides tools for downloading, preparing, and partitioning publicly available datasets. As with most design decisions, this is a trade-off. It enables Dataset Grouper to focus on utilities that apply to a broad range of datasets (notably, all datasets hosted by TensorFlow Datasets [16] and HuggingFace Datasets [17]). However, it also means that we cannot guarantee the hosting of various datasets in perpetuity. That being said, the tools underlying Dataset Grouper generalize to a wide array of settings (namely, any setting in which the base dataset can be accessed via an Apache Beam pipeline), and as such we believe Dataset Grouper provides a concrete benefit to machine learning researchers and practitioners working with group-structured data.
A.3 Dataset Licenses
Before using Dataset Grouper to partition a "base" dataset, users should also ensure that their use case falls under the license of that base dataset and that they abide by the associated terms and services. For example, below, we discuss the end license of the four (non-partitioned) datasets we use to create the partitioned datasets in Section 4.
BookCorpusOpen. We access this dataset through HuggingFace datasets. This dataset was originally derived from smashwords.com, and therefore should be used in compliance with the associated terms of service. We encourage readers to read the datasheet for BookCorpus created by Bandy and Vincent [78]. This datasheet identifies potential deficiencies of the dataset, including problematic and skewed content, possible copyright issues, and book duplication.
CC-News. We access this dataset through HuggingFace datasets. This dataset is publicly available and hosted by https://commoncrawl.org/, and therefore should be used in compliance with the associated terms of use.
C4. We access this dataset through TensorFlow Datasets. This dataset is based on a publicly available dataset hosted by https://commoncrawl.org/ and therefore should be used in compliance with the associated terms of use.
Wikipedia. We access this dataset through TensorFlow Datasets. This dataset is created through Wikipedia downloads (https://dumps.wikimedia.org/). Each example is derived from single Wikipedia articles and subjected to post-processing and cleaning to strip markdown and unwanted sections (references, etc.). This dataset should be used in compliance with Wikipedia's terms of use.
A.4 Software Development and Maintenance
Dataset Grouper is actively developed and maintained using standard open-source workflows. In particular, we accept responsibility for reviewing and acting on issues and contributions from the community, as long as they abide by a contributor license agreement discussed in the repository.
Source	Dataset	Partition on	#Clients	#Words	#Words per client
10thperc.	25thperc.	Median	75thperc.	90thperc.
Ours	FedC4	Domain	15.6M	132B	82	220	815	3.3K	0.5B
FedWiki	Article	6.5M	3B	39	75	198	486	70K
FedBookCO	Book	18K	1.2B	24K	32K	52K	81K	4M
FedCCnews	Domain	8.8K	0.3B	303	1.1K	5K	20K	8.4M
Existing	Amazon Reviews	Account	1.5M	4.3B	278	565	1.1K	2.3K	5K
Stack Overflow	Account	0.3M	2B	1.2K	1.7K	2.7K	5.1K	11K
Reddit	Account	1.7M	1.2B	58	111	257	675	1720
Blog Corpus	Account	17K	0.1B	551	908	2K	5.3K	13K
Shakespeare	Role/play	715	0.4M	14	45	175	0.6K	1.6K
Gigaword	Synthetic	100	0.3M	3.0K	3.1K	3.1K	3.2K	3.2K
Table 6: Detailed version of Table 1: A summary of the per-client statistics of the new language modeling datasets we introduce using Dataset Grouper, and of previous benchmark datasets supplied by TFF, Leaf, FedNLP, and FedScale.
Source	Dataset	Partition on	#Examples	#Words	#Words per Exmaple
10thperc.	25thperc.	Median	75thperc.	90thperc.
Ours	FedC4	Domain	0.36B	132B	49	88	191	417	783
FedWiki	Article	6.5M	3B	39	75	198	486	70K
FedBookCO	Book	18K	1.2B	24K	32K	52K	81K	4M
FedCCnews	Domain	0.7M	0.3B	78	151	316	548	842
Existing	Amazon Reviews	Account	68M	4.3B	3	10	28	67	155
Stack Overflow	Account	0.1B	2B	3	7	13	20	29
Reddit	Account	33M	1.2B	7	11	21	42	81
Blog Corpus	Account	0.5M	0.1B	6	28	105	248	460
Shakespeare	Role/play	16K	0.4M	4	8	12	29	63
Gigaword	Synthetic	10K	0.3M	21	26	31	36	41
Table 7: Detailed version of Table 1: A summary of the per-example (i.e., per-sequence) statistics of the new datasets we introduce using Dataset Grouper, and of previous benchmark datasets supplied by TFF, Leaf, FedNLP, and FedScale.