SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Here, we detail various statistics of the federated language-modeling datasets we propose and discuss in Section 4. In Table 6, we present the per-client statistics of the dataset, and compare these to per-client statistics of existing federated language modeling datasets supplied by TensorFlow Federated [11], Leaf [12], FedNLP [13], and FedScale [15]. We see that at larger percentiles, FedC4, FedWiki, FedBookCO, and FedCCnews contain many more words per client. FedBookCO also contains dramatically more words per client at lower percentiles than previous datasets. In general, datasets like FedC4 and FedCCnews exhibit a dramatic variance in statistics across clients. This can make the datasets more challenging for federated algorithms, and potentially more representative of heavy-tailed settings in practice.
In Table 7, we compare the per-example (i.e. per-sequence) statistics of the same datasets. We note that this information is important for large-scale language modeling tasks, which may require long sequences of tokens for training. We see that for nearly every percentile, the datasets we introduce contain sequences of significantly larger lengths. This is especially true of larger percentiles, at which point FedBookCO and FedWiki contains examples with thousands if not millions of words.
To better visualize and compare the distribution of words per client, and given the large number of clients in each dataset, we present a letter value plot [92] of these distributions. The result is in Figure 9. This plot gives various quantiles of the distribution of the number of words per client. We see that many of these datasets exhibit large amounts of variance between quantiles. FedC4 has an especially heavy tail, with some clients having fewer than 10 words, while others have tens or even hundreds of millions of words.
Figure 9: A letter value plot of the number of words held by each client, in the FedC4, FedWiki, FedBookCO, and FedCCnews datasets.
Example application scenarios. FedBookCO (each client is a book) and FedWiki (each client is a Wikipedia article) map to applications where each client is an "expert" in a certain topic. In the context of modern LLM pipelines where each sequence is broken up into multiple examples of a fixed length, the fact that each client has a single sequence is much less important than the total number of words. This is especially true for FedBookCO where the  $10^{\text{th}}$  percentile data sequence length is 24K words, while the maximum sequence length for LLMs today is  $O(1\mathrm{K})$  words.
FedC4 is typical of a group-structured pre-training (or second stage pre-training) dataset, which might be practically encountered in practice with documents (corporate, medical, legal, etc.) or emails. FedCCnews is a subset of FedC4 with similar long-tailed characteristics and is suitable for faster experimentation of second stage pre-training or fine-tuning. It can also be used to study the effects of pretraining data contamination.
C Full Experimental Details
We describe the full details of the experiments in Section 5. We discuss the datasets and preprocessing in Appendix C.1, the model in Appendix C.2, the federated algorithms in Appendix C.3, the hyperparameter choices in Appendix C.4, the reported metrics in Appendix C.5, and the hardware configuration in Appendix C.6.
C.1 Datasets and Preprocessing
We use FedC4 with domain-level partitioning for all experiments, which we discuss in Section 4. We use a WordPiece tokenizer [79] with a pre-trained BERT vocabulary [80] of size 30523. For each client, we concatenate all of the text in its examples into sequences of tokens of length 129, padding the last sequence as needed. In each sequence  $x_{1:129}$ , we predict  $x_{2}$  given  $x_{1}$ ,  $x_{3}$  given  $x_{1:2}$  and so on until  $x_{129}$  given  $x_{1:128}$ . This leads to a total of 128 predictions per sequence. We batch the sequences with a batch size of 16 and apply "take" and "repeat" operations to ensure that each client has exactly 64 batches. For evaluation, we follow the same procedure. However, we use the "validation" split of the C4 dataset, rather than the "train" split.
C.2 Model
We use a decoder-only transformer model whose size is roughly on the order of the BERT base or GPT-2 small. It has 12 layers, 12 attention heads, and hidden layers of dimension 768. As discussed above, it makes 128 predictions in each sequence using the causal language modeling loss (i.e. next-token prediction with cross-entropy loss).
C.3 Federated Algorithms
We use two federated learning algorithms, FedAvg and FedSGD [81]. In both algorithms, at each round  $t$  the server model  $x^{t}$  is sent to all the clients participating in round  $t$  (i.e. the cohort at round  $t$ ). Throughout, 16 clients participate in each round (i.e. the cohort size is 16). As discussed in Section 3, we shuffle the clients globally once and iterate successively through the stream of shuffled clients in windows of size 16. In both algorithms, each client uses this model to compute a gradient at each of its 64 batches. However, the algorithms differ in where these gradients are computed.
For FedAvg, we actually use the slightly generalized FedOpt framework [30] in which FedAvg uses both a client optimizer and a server optimizer. Throughout, we use SGD as the client optimizer and Adam as the server optimizer. After a client  $c$  computes a gradient, it updates its model locally using the client optimizer (i.e. SGD), starting at the broadcast model  $x^t$ . After  $K = 64$  updates, this results in some updated model  $x_c^t$ . The client then sends  $\Delta_c^t \coloneqq x^t - x_c^t$  to the server. For FedSGD, the clients do not locally update their model. Each gradient is computed at the broadcast model  $x^t$ , and each client  $c$  sends the average  $\Delta_c^t$  of these gradients to the server.
Once the server has received  $\Delta_c^t$  for each participating client, it averages them uniformly (as weighted and uniform are the same in our setting) to produce some quantity  $\Delta^t$ . The server treats this as an estimate of the gradient of the empirical risk function at the model  $x^t$  and applies the server optimizer (i.e. Adam) to update the model accordingly.
For both algorithms, we perform 3125 rounds of federated training. This means that the server model is updated 3125 times and that throughout the course of training,  $3125 \times 16 \times 64 = 3.2\mathrm{M}$  batched gradients are computed (as each client has 64 batches, and 16 clients participate in each round). If we consider each client to be computing gradients simultaneously, and form "meta-batches" of gradients of size 16 over the batched gradients, this means that we are doing roughly  $3125 \times 64 = 200\mathrm{K}$  "meta-batched" gradient computations in total.
Note that when doing personalization for the purposes of evaluation, as in Table 5 and Figure 5, we personalize the model (trained via FedAvg or FedSGD) using the same client training scheme as in FedAvg: Clients perform 64 steps of SGD on the model broadcast to them. The batches of data are formed in the exact same way as in training.
C.4 Optimizer Hyperparameters
As discussed in Appendix C.3, for both FedAvg and FedSGD we use a server optimizer of Adam. We tune only the learning rate of this optimizer (the server learning rate, denoted  $\eta_s$ ) and fix the Adam hyperparameters of  $\beta_1 = 0.9$ ,  $\beta_2 = 0.999$  and  $\varepsilon = 10^{-8}$ . For FedAvg we also use a client optimizer of SGD. We tune the learning rate of this optimizer (the client learning rate, denoted  $\eta_c$ ). We tune both learning rates  $\eta_s$ ,  $\eta_c$  over the range  $\{10^{-4}, 10^{-3}, \dots, 10^0\}$ . We select the learning rates that minimize average training loss across rounds. See Table 8 for a summary.
Placement	Optimizer	Hyperparameter	Value	Tuning Range
Server	Adam	Learning Rate (ηs)	N/A	{10-4, 10-3, ..., 100}
First Moment Decay (β1)	0.9	N/A
Second Moment Decay (β2)	0.999	N/A
Numerical Stability Term (ε)	10-8	N/A
Clients	SGD	Learning Rate (ηc)	N/A	{10-4, 10-3, ..., 100}
Table 8: Optimizer hyperparameters.
Algorithm	Server LR Schedule	Server LR (ηs)	Client LR (ηc)
FedAvg	Constant	10-3	10-1
Warmup + Exponential Decay	10-3	10-1
Warmup + Cosine Decay	10-3	10-1
FedSGD	Constant	10-4	N/A
Warmup + Exponential Decay	10-3	N/A
Warmup + Cosine Decay	10-3	N/A
Table 9: Tuned learning rates for FedAvg and FedSGD, with varying server learning rate schedules.
In Figure 4 we also experiment with learning rate scheduling. Note that these are only applied to the server optimizer. The client learning rate (for FedAvg) is held constant throughout an experiment. The learning rate schedule is applied across the 3125 training rounds. We compare constant learning rates to learning rates with (1) exponential decay and (2) cosine decay. For both of these decay schedules, we perform linear warmup (starting at 0) for the first 312 rounds ( $10\%$  of the total number of rounds). We then decay for the remaining rounds, with a final server learning rate of 0. In such cases, the server learning rate parameter  $\eta_{s}$  refers to the maximum learning rate attained (i.e. at round 312) and is tuned just as above.
The best performing (tuned) learning rates for each algorithm and schedule are given in Table 9. Generally, we found that a client learning rate of  $\eta_c = 10^{-1}$  worked well throughout. For FedAvg, a server learning rate of  $\eta_s = 10^{-3}$  worked well, though we see little to no difference between server learning rate schedules Figure 4a. For FedSGD, we find that we could only use  $\eta_s = 10^{-4}$  for constant learning rates, but learning rate schedules allowed us to use  $\eta_s = 10^{-3}$  and led to improved convergence Figure 4b.
When performing personalization evaluation (Table 5 and Figure 5), we use the same client optimizer of SGD, and use the client learning rate that led to the best training performance for FedAvg (i.e.  $\eta_c = 10^{-1}$ , as in Table 9).
C.5 Reported Metrics
In Figure 4, we report a causal language modeling loss (i.e. the logarithm of the perplexity) at each training round of FedAvg and FedSGD. How we compute these averages depends on the federated algorithm used. For FedSGD, each client computes gradients and loss values across 64 batches, at the same model. The client averages the loss across these 64 batches and sends the result to the server. The server averages these quantities across the 16 clients participating in a round.
We do the same thing for FedAvg, except we must keep in mind that for each client, the 64 different loss values are computed at different models. This is because in FedAvg, each client locally updates its model as it computes gradients. The average of these 64 loss values is still sent to the server and averaged across the 16 clients participating in a round. However, because of the local training, this loss represents a different quantity. In short, it accounts for both how good the broadcast model is, and how well the model adapts to a client's data. By contrast, the loss reported for FedSGD only represents how good the broadcast model is.
In Table 5 and Figure 5, there is no such ambiguity. The pre-personalization loss is the average causal language modeling loss before a client personalizes the model to its own data. The post-personalization loss is the average causal language modeling loss after a client personalizes to its own data.
C.6 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.
Here, we present detailed personalization evaluation results on FedCCnews and FedWiki in Appendix D.1, followed by detailed ablation results in Appendix D.2.
D.1 Personalization Results
Here we present the results of performing pre-personalization and post-personalization evaluation on the FedC4-trained models, but using other datasets. In addition to the evaluation on FedBookCO presented in Section 5, we present results on FedCCnews (Figures 10 and 11), and FedWiki (Figures 12 and 13).
For FedBookCO and FedCCnews, we report the personalization results on the entire dataset. Due to the large number of clients in FedWiki, we randomly sample 20K clients in the dataset for personalization evaluation.
Figure 10: Median pre-personalization (left) and post-personalization (right) loss over FedCCnews clients while training on FedC4. Error bars indicate the 10th and 90th percentiles.
Figure 11: Histograms of pre-personalization (left) and post-personalization (right) loss on FedCCnews after FedC4 training.
Figure 12: Median pre-personalization (left) and post-personalization (right) loss over FedWiki clients while training on FedC4. Error bars indicate the 10th and 90th percentiles.
Figure 13: Histograms of pre-personalization (left) and post-personalization (right) loss on FedWiki after FedC4 training.
D.2 Ablation: Batches per Client
To better explore the phenomena discussed in Section 5, we perform a modified experiment where we repeat the training from Section 5, but vary the number of batches  $\tau$  each training client yields. As above, we repeat and truncate clients' datasets so that each client has exactly 1024 examples. For example, when  $\tau = 64$ , and we use a batch size of 16, this means that for FedAvg and FedSGD, each client computes  $\tau = 64$  mini-batch gradients. For FedAvg, this means that the client does  $\tau = 64$  steps of training, while for FedSGD, this means that the client sends the average of these  $\tau = 64$  batches back to the server. We vary the number of batches  $\tau$  over  $\{1, 4, 16, 64\}$ . We do so by repeating and truncating clients' data so that they have 16, 64, 256, and 1024 examples, respectively. We then perform different amounts of training, one in which we equalize the number of communication rounds, and one in which we equalize the total number of tokens seen across all clients.
Note that fixing the number of examples per client at 1024 and simply changing the batch size would be a more elegant way to do this kind of ablation, as it would allow normalizing the number of communication rounds and tokens simultaneously. Unfortunately, the batch size is often dictated by compute constraints (e.g. what fits in the memory of a given hardware device), and cannot be increased in an unbounded fashion in realistic machine learning settings, especially on-device settings which are common in FL. Thus, we instead focus on varying  $\tau$ .
Equalizing communication rounds. We do 5000 rounds of training for each number of batches per client  $\tau$ , using the same warmup and cosine decay schedule discussed above. Throughout training, we perform personalization evaluation on the models. We give the pre-personalization results in Figure 14, and the post-personalization results in Figure 15. For more detailed numerical results, see Appendix D.2.
Figure 14: Median pre-personalization loss across FedC4 validation clients, with different numbers of batches used in each client's local computation. Error bars indicate the 10th and 90th percentiles. All runs are equalized to perform the same number of communication rounds in total.
Figure 15: Median post-personalization loss across FedC4 validation clients, with different numbers of batches used in each client's local computation. Error bars indicate the 10th and 90th percentiles. All runs are equalized to perform the same number of communication rounds in total.
We see that for FedSGD, the pre- and post-personalization losses do not change much with the number of batches per client  $\tau$ . For FedAvg, on the other hand, lower values of  $\tau$  attain better
pre-personalization loss, while higher values achieve better post-personalization loss. Note that when the number of batches per client is  $\tau = 1$ , FedAvg and FedSGD are effectively the same algorithm (up to differences in normalization), and perform only a single update step per client each round. Notably, FedAvg attains a better trade-off between the pre- and post-personalization metrics. Specifically, FedAvg with  $\tau = 4$  can match the best pre-personalization performance of FedSGD (FedAvg with  $\tau = 4$  and FedSGD with  $\tau = 64$  both attain a median pre-personalization loss of 4.2), while performing significantly better on the post-personalization metrics (a median loss of 1.9 for FedAvg with  $\tau = 4$  versus 3.4 for FedSGD with  $\tau = 64$ ). In short, these results seem to suggest that "client drift" [28] is less of an impediment to federated learning, and more indicative of a trade-off between minimizing pre- and post-personalization loss functions.
Algorithm	Loss	Batches per Client (τ)
1	4	16	64
FedAvg	Pre-Personalization	4.4	4.2	4.8	5.2
Post-Personalization	3.5	1.9	0.009	0.008
FedSGD	Pre-Personalization	4.5	4.4	4.4	4.2
Post-Personalization	3.6	3.4	3.4	3.3
Table 10: Median pre-personalization and post-personalization loss after training with FedAvg and FedSGD, with different numbers of batches per client, keeping the total number of communication rounds constant.
Equalizing tokens. Next, we perform an analogous experiment, but where each setting of  $\tau$  (the number of batches per client), is trained for a different number of rounds, so that in each setting the same number of tokens is processed (over all clients). For example,  $\tau = 1$  trains for  $64 \times$  more communication rounds than for  $\tau = 64$ . Each training run processes roughly 10.5 billion tokens in total. We give the pre-personalization results in Figure 16, and the post-personalization results in Figure 17. For more detailed numerical results, see Appendix D.2. Notably, the post-personalization loss diverged at intermediate stages for  $\tau = 1$ , but recovered (albeit with high variance across clients) afterwards.
Figure 16: Median pre-personalization loss across FedC4 validation clients, with different numbers of batches used in each client's local computation. Error bars indicate the 10th and 90th percentiles. All runs are equalized to process the same number of tokens in total.
Here, we see a notably different story than in Figures 14 and 15. In particular, we see that for both FedAvg and FedSGD, lower values of  $\tau$  lead to lower pre-personalization loss. However, for both algorithms as long as  $\tau$  is sufficiently large (at least 4 in this case), the post-personalization loss essentially does not change with  $\tau$ . This suggests that when communication is not a bottleneck, we can attain good pre- and post-personalization loss by using FedAvg with a small (but not too small) value of  $\tau$ . In other words, by performing enough rounds of FedAvg with a moderate  $\tau$ , we can attain a model that does well before and after personalization.
Figure 17: Median post-personalization loss across FedC4 validation clients, with different numbers of batches used in each client's local computation. Error bars indicate the 10th and 90th percentiles. All runs are equalized to process the same number of tokens in total.
Algorithm	Loss	Batches per Client (τ)
1	4	16	64
FedAvg	Pre-Personalization	3.6	3.8	4.3	5.2
Post-Personalization	3.8	0.006	0.007	0.007
FedSGD	Pre-Personalization	3.6	3.7	3.9	4.2
Post-Personalization	3.9	3.5	3.3	3.3
E Memory Usage
In this section, we provide details of how much memory are used by the various dataset formats in Section 3.1. Recall that we consider three formats: in-memory, hierarchical, and streaming. We compare the amount of time required to iterate over various federated datasets in these formats in Table 3. Here, we instead detail the peak memory usage (in megabytes) when iterating over the same datasets. We note that these were collected on a single CPU, and do not include the time to do operations like shuffling or batching.
The results are in Table 12. We see that in-memory formats use much larger amounts of memory, as they load the entire dataset into memory. By contrast, hierarchical and streaming formats use significantly less peak memory. Moreover, their peak memory usage does not scale with the total size of the dataset, unlike in-memory formats. We note that streaming can use slightly more memory, though only at most  $2\mathrm{MB}$  more in all experiments.
Table 11: Median pre-personalization and post-personalization loss after training with FedAvg and FedSGD, with different numbers of batches per client, keeping the total number of tokens processed constant.
Dataset Format	In-Memory	Hierarchical	Streaming
CIFAR-100	156	0.40	0.74
FedCCnews	1996	0.08	1.16
FedBookCO	6643	0.001	0.10
Table 12: Peak memory usage, in megabytes, when iterating over federated datasets. We iterate over all examples in all group datasets, in serial, on a single CPU. We compare a federated CIFAR-100 dataset (partitioned across 100 groups, each with 100 examples), FedCCnews (in which examples are split across users at a domain level), and FedBookCO (in which examples are split across users at a title level). See Section 4 for more details on the latter two datasets.