Worldwide Federated Training of Language Models

## Abstract

Abstract The reliance of language model training on massive amounts of computation and vast datasets scraped from potentially low-quality, copyrighted, or sensitive data has come into question practically, legally, and ethically. Federated learning provides a plausible alternative by enabling previously untapped data to be voluntarily gathered from collaborating organizations.
However, when scaled globally, federated learning requires collaboration across heterogeneous legal, security, and privacy regimes while accounting for the inherent locality of language data; this further exacerbates the established challenge of federated statistical heterogeneity. We propose a Worldwide Federated Language Model Training ( WorldLM ) system based on federations of federations , where each federation has the autonomy to account for factors such as its industry, operating jurisdiction, or competitive environment. WorldLM enables such autonomy in the presence of statistical heterogeneity via partial model localization by allowing sub-federations to attentively aggregate key layers from their constituents.
Furthermore, it can adaptively share information across federations via residual layer embeddings.
Evaluations of language modeling on naturally heterogeneous datasets show that WorldLM outperforms standard federations by up to 1.91 × 1.91\times 1.91 × , approaches the personalized performance of fully local models, and maintains these advantages under privacy-enhancing techniques.

## 1 Introduction

Contemporary Language Models (LMs) are trained in data centers with numerous hardware accelerators [, ] on large scraped corpora [, ]. This paradigm has restricted training to well-funded actors and brought concerns regarding centralization [, sec. 4.7.1], copyright [], and the limited supply of data []. While small organizations may own valuable data, the per-batch synchronization constraints of stochastic gradient descent (SGD) [, ] make collaboration impractical. Federated Learning (FL) has been proposed as an alternative [, , ] due to its lower synchronization constraints, limited data movement, and ability to integrate privacy-enhancing techniques [, ].

This work addresses the novel challenges that global-scale FL applications face. The first challenge is related to federated governance []: how to create federated systems including participants with different legal, privacy, and security concerns or combine existing systems seamlessly without sacrificing performance. The second challenge is the statistical heterogeneity of naturally distributed data, as organizations may have datasets that differ in terms of language (e.g., for federations spanning multiple countries) and genres (e.g.,
publishers operating in different domains). Significant heterogeneity has been shown to slow model convergence or even cause divergence in standard single-federation settings [, , ]. This issue is worsened at a global level, as external factors [] can implicitly create different clusters of data in the federation, making it necessary to optimize a participant’s parameters with respect to both the global distribution and its peers.

Our proposed Worldwide Federated Language Model (WorldLM) training system addresses both challenges by building federations of federations, as seen in Fig. 1. This brings several advantages:

- 1.
Arbitrary Federation Structure: WorldLM allows each federation to adapt to its unique legal, privacy, and security requirements and also permits easy integration between previously isolated organizations, with differential privacy [] protocols.
- 2.
Partially-personalized Aggregation: WorldLM takes advantage of cluster relationships within federations by segmenting models into a backbone serving as a shared feature extractor trained using standard FL algorithms [] and partially-personalized *key* layers relevant to a node. These key layers are aggregated using attention over the keys of its peers. Results show that the aggregation procedure of WorldLM allows it to achieve a perplexity lower by $11.5$-$45.3\%$ on a novel hierarchical partitioning of The Pile [] dataset (e.g, Fig. 2) and on a language-partitioned mC4 dataset []. Furthermore, WorldLM converges to similar performance when using differential privacy, whilst standard FL diverges completely.
- 3.
Cross-federation Information Sharing: To address intra-federation heterogeneity, where a participant may contain data that is more similar to another sub-federation than to its peers, WorldLM can route *residual* parameters to the most relevant sub-federation in a communication-efficient fashion based on key similarity. Including such residuals allows WorldLM to partially maintain performance even when the underlying sub-federations contain dissimilar participants while requiring the transfer of a limited number of additional transformer blocks. We find
the last $1-3$ blocks to be sufficient.

The novel optimization algorithm and residual information routing of WorldLM tackle the challenges of global LM training without diverging far from the infrastructure of production systems []. Thus, it provides an effective solution for training LMs across legal, socioeconomic, and cultural boundaries, diversifying the pool of available data away from its current geographic [] concentration.

Figure: Figure 1: WorldLM federations exchange information in the form of models containing a backbone, personalized layers ($\mathcal{B,K,V}$), and lower-dimensional residual embeddings serving as keys and values ($\mathcal{K,V}$). While full models are exchanged between parents and children, residuals are dynamically routed to the most appropriate sub-federation to be used in attention-based aggregation.
Refer to caption: x1.png

## 2 Background

Well-known performance scaling laws [, ] drive LM training. The resulting models [, , , ] have led to numerous decentralized downstream applications [, , , , ].
However, since the Ring AllReduce algorithm [] used to synchronize the model state for every batch [, ] requires high cross-GPU bandwidth, pre-training still requires large centralized datacenters [].

Douillard et al. [] relax the communication requirements of training via Local SGD [], a data-parallel technique synchronizing independent model replicas. Based on the observation that publically available data may not keep up with the demands of LM training [], future work [] extended generative pre-training to Federated Learning (FL) settings and showed performance comparable to centralized training even for billion-scale models. This shift enables collaboration across geographically distributed actors with valuable data currently serving as passive data providers [, ]. Given the potential of FL to reach untapped data reservoirs, we are motivated by these results to tackle the challenges that emerge when applied at a *global* scale.

### 2.1 Global Federated Systems

Standard federated training proceeds in rounds that involve broadcasting a shared set of parameters to all clients, local optimization of said parameters on private data, and then aggregation []. Its roots in the principle of data minimization [, ] make FL an attractive paradigm across international regulatory frameworks [, ] (Section A.1).
However, while it is communication-efficient and scalable, FL suffers from statistical and system heterogeneity, defined as clients having different data distributions, training speeds, or communication bandwidth. For example, Charles et al. [] and Zhou et al. [] observe that training on more clients within a round brings diminishing returns due to near-orthogonal updates. These behaviors are analogous to the limitations of large-batch training  [], indicating the need for more sophisticated aggregation strategies, unlike the Local SGD literature [], which assumes homogeneous data. Similarly, systems heterogeneity can result in varying client completion times [, ], causing stragglers and slow convergence.

Since high-quality language data is liable to run out within the decade [], the quality of data that FL makes available is also paramount. For example, obtaining data from regions with a reduced internet presence may alleviate the challenges associated with low-resource languages [, ]. We argue that the impact of data and systems heterogeneity increases as the reach of the federated system grows, which causes governance to emerge as a new challenge.

Federated governance, hitherto considered primarily in the medical domain [, ], requires determining how to create systems spanning actors with different legal, privacy, and security concerns and how to combine already existing federations. For example, unlike US-based organizations, training an LM on the private data of EU-based organizations requires abiding by the GPDR [], which restricts cross-border data transfer, while mitigating bias  [, Art. 10.2f,fa] (see Section A.1). Although differential privacy (DP) [] or secure aggregation [] may enable collaboration across legally heterogeneous regions, not using them within a region when possible avoids computational and performance costs [, , , ]. Similar hierarchical structures enable collaboration in platform ecosystems [, ], OSS [], and vertical FL on medical data[, sec. 2.4.1].

Privacy-enhancement techniques (PETs), such as DP [] or secure aggregation, have become popular in FL since client updates may leak private information [].
Researchers and practitioners consider DP the de facto method for alleviating privacy concerns in FL [, ], as it provides statistical bounds on the leakage of private information.

### 2.2 Related Work

Previous works in standard federated learning have tackled aspects of data and systems heterogeneity relevant to our work [, , , ]. However, they are unable to account for the complexities of building federations-of-federations.
Personalized Federated Learning (PFL) [] attempts to tackle the performance disparity between clients induced by the heterogeneity of client data by creating personalized models. The most common means of achieving this is a local adaptation (fine-tuning) of the federated model after training [, , , ]. Split approaches such as those of Arivazhagan et al. [] and Li et al. [] train models locally but only average and update the layers up to a cut layer $q$, with the intuition that earlier layers represent feature extractors.

Another means of tackling data heterogeneity is clustering clients based on a similarity metric between the flattened model embeddings [, , , ]. However, such approaches reduce the compatibility of client models to a single scalar and do not consider the structure of the model. Clusters may also exist naturally based on characteristics like location or language, which become relevant if the clients and servers controlling them are geographically correlated. However, most previous works in this area have focused solely on communication efficiency rather than data heterogeneity [, , , ]. As such, they broadcast a global model rather than server personalization.

Mhaisen et al. [] do consider scenarios where the data distribution of edge servers is taken into account and propose optimal user-edge assignment. While promising, their work assumes that edge servers have sufficiently low communication latency to efficiently train with FedSGD despite the original work of McMahan et al. [] showing FedSGD to be up to two orders of magnitude slower than FedAvg in terms of convergence speed. Another work that addresses data heterogeneity within hierarchical FL is Wu et al. []; however, their framework focuses on well-known client-level fine-tuning strategies and never explicitly utilizes the hierarchical structure to enhance personalization.

## 3 WorldLM

Given the interlocking nature of legal, privacy, and security concerns, we assume that organizations find it easier to collaborate with enterprises operating in the same geographic region, legal jurisdiction, or industry. Therefore, a key design objective for our system is facilitating these types of collaborations. In edge-device FL, the large number of potentially ephemeral participants with fine-grained preferences leads to a fully decentralized collaborative learning paradigm [, ] where the server-client federated network is replaced by a graph encoding client compatibility. For training LMs, where the participating entities are stable organizations, we argue that a federation of federations approach, as portrayed in Fig. 1, provides an attractive compromise between standard federated learning and decentralized learning. To optimize in this setting, WorldLM takes advantage of a custom aggregation procedure and information-sharing mechanism.

### 3.1 Partially-personalized Aggregation

Figure: Algorithm 1 $\mathrm{Fit}$: execution procedure for a given sub-federation.

Federated learning can train powerful feature extractors beneficial to all WorldLM participants due to its meta-learning properties[, , ]. However, actors in the federated training process may hold heterogeneous data; for example, it may be in a different language or come from a separate domain (e.g., one organization may publish news while another focuses on scientific publications). Thus, the feature extractor needs to be adapted to the data of each actor and sub-federation, which requires a departure from the standard FL objective. Thus, WorldLM takes inspiration from split-learning techniques [, ] and partitions the model into a backbone $\mathcal{B}$, comprising the majority of the model’s parameters, and a set of partially personalized key layers $\mathcal{K}$ specific to each node. In practice, we expect the key layers to include the final $1$-$3$ transformer blocks.

The backbone parameters $\mathcal{B}$ are trained using FL aggregation algorithms such as FedAvg [] or FedOPT []. As shown in Algorithm 1 (L.$8-12$), for each round $k$, the server broadcasts $\mathcal{B}$ to its children $c\in C_{q}$. When a child begins execution, it replaces its previous backbone with the one received from the parent (L.$5$), and then executes its own sub-federation (L.$10-12$) and/or local training (L.$9$) and computes pseudo-gradients $\Delta_{c}^{k}$. These pseudo-gradients are then aggregated and applied to the backbone of the server via $\mathrm{ServerOpt}$ (L.$15$). It is important to note that local training (L.$9$) can be executed either in-parallel with the node model treated as a client or sequentially. Similarly, when observing the performance of a hierarchical FL approach we distinguish between sequential and parallel steps as parallel steps are averaged across clients.
For this work, we choose the sequential approach, meaning that for the first level of the hierarchy, the root executes by itself, whilst for the final level, all leaves are executing simultaneously.

Unlike traditional personalization approaches, WorldLM optimizes the key layers $\mathcal{K}$ while considering the models of its parent and children. It achieves this by aggregating each layer within $\mathcal{K}$ using an attention mechanism [, ] where the parameters of that given layer across the nodes of the federation represent the queries, keys, and values. To exploit the locality of data and the similarity of nodes within the same sub-federation, the attention mechanism is applied within a local context, either with respect to the parent of a given node or its children. This procedure allows WorldLM to handle statistical heterogeneity when the nodes within a sub-federation have a similar data distribution.

### 3.2 Cross-Federation Information Sharing

The attention-based aggregation of a node’s children (L.$17$) will reduce to simple unweighted averaging when a cluster relationship does not exist in a sub-federation. In contrast, the aggregation between a node and its parent (L.$6$) will focus almost exclusively on the keys of the node and thus approach a personalization-layer strategy. Similarly, if a given node has a data distribution that differs significantly from its peers, its key layers will be highly dissimilar and ignored in the aggregation.

However, being dissimilar to its peers does not preclude a node from containing valuable data for another sub-federation. Inspired by residual connections [], we introduce the concept of residual key layers, which are dynamically chosen based on their dissimilarity to peers and routed to a more appropriate sub-federation. For each layer of its post-aggregation key $\mathcal{K}$, a node selects at most $\nu_{\mathcal{K}}$ of its children’s layers with the lowest similarity score (L.18), e.g., by the dot product. It then sends these residual layer embeddings to the highest permissible level of the federation as dictated by the participants’ legal, security, and privacy concerns. This server then routes the residuals to its highest-similarity children based on a cached set of keys from the previous round (L.8) by placing them in $R^{k}$. When the residuals reach an edge node, they are placed in $A^{k}$ and later aggregated. Aggregating in the edge nodes (L.$4-6$) has the benefit of diffusing information into the entire sub-federation right after local training.

## 4 Experimental Design

Figure: Figure 2: Data-perspective upon a hierarchical dataset constructed from The Pile []. The LHS contains two naturally heterogeneous and quantity-skewed groupings of data sources, corresponding to organizations accessing data from the internet or the medical domain. We construct such groupings using the internet-based Common Craw (CC) and Wikipedia (WK) versus the medial data of PubMed Abstracts (PBA) and PubMed Central (PBC). To test the effectiveness of WorldLM when such a cluster relationship is absent, we swap the position of the two smaller datasets.
Refer to caption: x2.png

Given the recent emergence of federated generative pre-training and the lack of benchmark datasets, we evaluate WorldLM on tasks that approximate realistic scenarios for its application: (a) organizations in different industries collaborating to train an LM despite holding different genres of text, (b) organizations trying to train an LM despite holding data in different languages.

To simulate organizations holding different genres of text, we partition *The Pile* [] into its constituent heterogeneous datasets and construct federations of federations by bottom-up building different mixtures of datasets. For example, as seen in Fig. 2, if the children of a node hold data from the Pile Common Crawl, Wikipedia, PubMed Central, and PubMed Abstracts datasets, then the node itself will hold data from them, proportionate to their size. For this dataset, we use the common *gpt-neox-20b* English tokenizer also adopted by Sani et al. []

To simulate geographically distributed systems, we use a subset of Multilingual Colossal Common Crawl (*mC4*)[], covering high and low-resource languages []. We then build sub-federations based on language families similarly to *The Pile*. Given the larger vocabulary size [] and consequent model size, we use a single federated structure partitioned with the high-resource French (FR) and Italian (IT) on one side and the lower-resource Bulgarian (BG) and Ukrainian (UK) on the other. For IID experiments, we use the standard Cleaned Colossal Common Crawl (*C4*) English dataset, partitioned into equal-sized shards with the same tokenizer as *The Pile*.

### 4.1 Tasks

**Table 1: Hyperparameters for WorldLM, the federated learning rate $\boldsymbol{\eta_{s}}$ and momentum $\boldsymbol{\mu_{s}}$ [] are applied by a standard FL server or each server of WorldLM, $\boldsymbol{|\mathcal{K}|}$ is the number of blocks used for the key in WorldLM. In contrast, $\boldsymbol{\nu_{\mathcal{K}}}$ is the number of layers we select for each residual across all clients. Finally, $\mathbf{S_{C}}$ are the parameters of the learning rate scheduler synchronized across sequential steps.**
| Model (size) | #Rounds | $\boldsymbol{\eta_{s}}$ | $\boldsymbol{\mu_{s}}$ | $\boldsymbol{|\mathcal{K}|}$ | $\boldsymbol{\nu_{\mathcal{K}}}$ | $\mathbf{S_{C}(\boldsymbol{\alpha},~{}\boldsymbol{\eta_{max}},~{}T)}$ |
| --- | --- | --- | --- | --- | --- | --- |
| English ($\mathbf{75}$M) | $12$ | $0.2$ | $0.9$ | 1 | 1 | $(10^{-2},~{}8\times 10^{-4},~{}3\times 10^{3})$ |
| English ($\mathbf{125}$M) | $21$ | $0.2$ | $0.9$ | 3 | 1 | $(10^{-2},~{}6\times 10^{-4},~{}5\times 10^{3})$ |
| Multi ($\mathbf{250}$M) | $12$ | $0.2$ | $0.9$ | 1 | 1 | $(10^{-2},~{}8\times 10^{-4},~{}5\times 10^{3})$ |

For all experiments, we use decoder-only transformers [] for the language modeling task. Given that we are concerned with partially personalized FL, we compare the personalized local perplexity of WorldLM against three relevant baselines representing alternative scenarios. First, we compare against our closest alternative, standard FL with momentum [, , ], which has no hierarchical structure and would be difficult to integrate across legally and organizationally heterogeneous participants given the aforementioned governance issues. Second, we compare against independently trained fully personalized local models which do not share information, thus being heavily constrained by the local dataset size. This may result in heavy memorization due to overfitting to the limited amount of data of one organization. Third, we compare against centralized training of a global model on all the node’s data pooled together, irrespective of privacy concerns. The centralized baseline approximates the upper bound performance for a global model.

Language Modelling: We focus on an underrepresented segment: groups having access to a meaningful but not sufficient number of GPUs. For example, organizations holding $8$ NVIDIA A40s (as we use in this work). In such scenarios, while billion-scale LMs may be out of reach due to sheer VRAM and bandwidth requirements, medium models may be practically useful (e.g., by providing highly relevant embeddings) without possessing the externalities of pre-trained models. Thus, given the hardware, we use three scales of models ($75$M, $125$M, and $250$M). Our federated training uses the parameters shown in Table 1 while the local parameters are available in Table 4. We compare models that have executed for a given number of sequential steps, the step-wise execution we chose for WorldLM (Algorithm 1) implies that each level of a hierarchy (Fig. 1) trains in parallel; however, the levels execute sequentially relative to each other in three stages.

Privacy and Security: To validate the effectiveness of WorldLM for enhancing privacy and security, we simulate differentially private training, in which the leaves of a hierarchy contain potentially sensitive information and thus use *DPFedAvg* [, ] instead of standard averaging. This requires gradient clipping and the injection of Gaussian noise. To this end, we assumed that the Pile-CC and WK leaf clients of The Pile require DP and inject noise with $\sigma=0.5$ while clipping their gradients to the median $l_{2}$ norm of the previous round (for the first round we use a bound of $1.0$ as done for centralized ML []). Given the hierarchical training of WorldLM and the personalization involved, we postulate that it can more effectively account for the DP noise.

## 5 Evaluation

**Table 2: Language modeling personalized performance (over local client test sets) of WorldLM in terms of perplexity (the lower, the better). We compare against standard FL, fully local, and centralized models. The latter is trained on the union of all local training sets. WorldLM outperforms standard FL across all Non-IID dataset partitions, reaching a perplexity almost twice lower for *mC4*. While fully local models may seem to perform well, they are unlikely to generalize out of distribution and are prone towards divergence due to small dataset sizes, as seen in Fig. 3.**
|  |  | Collaborative | Non-Collaborative |  |  |
| --- | --- | --- | --- | --- | --- |
| Dataset | Model | WorldLM | FL | Local | Centralized |
| Pile | 75M | $\phantom{1}\mathbf{73.82\pm 44.18}$ | $107.31\pm 52.50$ | $\mathbf{40.66\pm 25.28}$ | $85.81\pm 24.42$ |
| Pile | 125M | $\phantom{1}\mathbf{48.34\pm 32.41}$ | $\phantom{1}53.92\pm 24.24$ | $\mathbf{24.83\pm 12.47}$ | $29.61\pm 13.17$ |
| MC4 | 250M | $\phantom{1}\mathbf{80.47\pm 68.53}$ | $153.27\pm 95.47$ | $\mathbf{45.47\pm 31.13}$ | $72.21\pm 49.78$ |
| C4 | 75M | $167.31\pm\phantom{1}2.92$ | $\mathbf{145.32\pm\phantom{1}3.53}$ | N/A | $67.01\pm\phantom{1}1.67$ |

Our results for the language modeling task (Table 2) show that WorldLM offers the desired compromise between global performance and local personalization. WorldLM is capable of outperforming standard FL in terms of local personalized performance for datasets that exhibit statistical heterogeneity, while obtaining similar performance on the IID C4 dataset. Furthermore, it approaches a fully-centralized baseline for the $75$M model trained on The Pile and the $250$M model trained on mC4. As can be seen in Fig. 3, the partially-personalized aggregation of WorldLM allows it to approach the performance of a centralized model despite data never leaving each organization’s premises. This holds the promise of international collaboration towards the federated training of LLMs. Since WorldLM executes in three stages (Algorithm 1) for a $3$-level hierarchy, Fig. 3 exhibits temporary perplexity spikes in the convergence curve (e.g, the triangle formed by rounds $18-21$ in Fig. 3). We chose to show all three stages since they correspond to different sequential steps. Thus, the root executes by itself while all nodes on a level execute together. For comparisons not dependent on the stage, consult Table 2.

Similar results can be observed on our heterogeneous partition of The Pile in Fig. 4. However, the smaller heterogeneity between English texts causes a smaller gap between WorldLM and standard FL. We use the $75$M model and The Pile to evaluate the resilience of WorldLM to differential privacy and changes in the data hierarchy (Fig. 2) due to being compute-limited.

Figure: (a)
Refer to caption: x3.png

Figure: (a)
Refer to caption: x5.png

We evaluate the robustness of WorldLM to differential privacy and using an alternative arrangement of The Pile which does not contain an inherent cluster relationship. The main limitation of WorldLM is an imperfect ability to reconcile hierarchies with a high degree of intra-federation heterogeneity. As can be seen Table 3 and Fig. 5 swapping the data of WK with PBA (Fig. 2) results in worse performance compared to standard FL due to the $\mathcal{K}$ layers of the root being unable to agree during attention-based aggregation. Crucially, as can be observed from Fig. 5, the personalized layers of the other nodes, together with the residual mechanism of WorldLM allow them to maintain performance despite this decrease for the root. Consequently, for a majority of the participating organizations WorldLM serves as a superior alternative to FL from a personalization perspective.

Table 3 shows that WorldLM is generally more robust than standard FL to the gradient clipping and noise that DP injects into the models of two leaf nodes. Standard FL diverges immediately due to its inability to suppress the impact of DP on the global model. The personalized keys of WorldLM, on the other hand, can ignore the impact of the noise entirely, as seen in Fig. 6. Furthermore, the additional per-level momentum mechanism of WorldLM allows it to stabilize the backbone training. This indicates that additional means of accounting for noise may be highly beneficial for standard FL approaches as well. For example, first pre-training on non-DP clients.

**Table 3: Results evaluating privacy and robustness. WorldLM is highly resilient to DP with $\sigma=0.5$ being applied over two of its leaf participants due to its aggregation procedure and because they start from a better initialization due to parent training. By contrast, the high amount of noise causes the entire global model of standard FL to diverge, without any personalized portions to maintain the performance of the non-DP members. Regarding robustness, WorldLM suffers a decrease in performance once the cluster relationship in Fig. 2 is broken in Pile (A) but remains competitive.**
| Method | Pile | $\mathbf{\mathrm{DP}_{\mathrm{CC},\mathrm{WK}}}$ | $\mathbf{\mathrm{DP}_{\mathrm{PBC},\mathrm{PBA}}}$ | Pile (A) |
| --- | --- | --- | --- | --- |
| WorldLM | $\phantom{1}\mathbf{73.82\pm 44.18}$ | $\mathbf{101.78\pm\phantom{1}88.48}$ | $\phantom{1}\mathbf{103.68\pm\phantom{1}90.53}$ | $140.05\pm 100.52$ |
| FL | $107.31\pm 52.50$ | $724.56\pm 251.89$ | $724.24\pm 250.98$ | $\mathbf{107.31\pm\phantom{1}52.50}$ |

Figure: (a)
Refer to caption: x9.png

Figure: (a)
Refer to caption: x11.png

## 6 Conclusion

WorldLM supports the extension of federated learning (FL) to the challenging setting of worldwide optimization of language models (LMs). Our results indicate that systems based on federations-of-federations can compete with standard FL and centralized optimization for the medium-sized LMs affordable to small organizations and groups, given their hardware. Our results show that WorldLM can outperform standard FL under realistic federated topologies and data distributions constructed using naturally heterogeneous datasets. Furthermore, they also indicate our method to be robust under the constraints of differential privacy, unlike standard FL. Thus, WorldLM is an effective approach for addressing the nascent sub-field of worldwide LM pre-training. We open several new research opportunities such as: (a) bringing the benefits of WorldLM to the larger models that wealthy organizations can afford, (b) tackling broader forms of statistical heterogeneity, and (c) applying it to parameter-efficient fine-tuning. We hope this will help democratize LM training across national boundaries and address the societal concerns regarding its governance.

## Acknowledgments and Disclosure of Funding

This work was conducted with the support of the Ministry of Education of Romania through the Credit and Scholarship Agency.

## Appendix A Appendix

**Table 4: Architecture details and local training parameters for our $75$M and $250$M models. They represent the number of transformer blocks, hidden model dimension, number of attention heads, the linear layer expansion ratio and the parameters of Adam.**
| Model (size) | #Blocks | $\boldsymbol{d}$ | #Heads | Exp. Ratio | $\mathbf{(\boldsymbol{\beta_{1}},~{}\boldsymbol{\beta_{2}})}$ | $\boldsymbol{|\textbf{Vocab}|}$ |
| --- | --- | --- | --- | --- | --- | --- |
| English ($\mathbf{75}$M) | 3 | 896 | 16 | 4 | $(0.9,~{}0.95)$ | $50$K |
| English ($\mathbf{125}$M) | 12 | 768 | 12 | 4 | $(0.9,~{}0.95)$ | $50$K |
| Multi ($\mathbf{250}$M) | 3 | 896 | 16 | 4 | $(0.9,~{}0.95)$ | $250$K |

Figure: (a)
Refer to caption: x13.png

Figure: (a)
Refer to caption: x15.png

### A.1 The Legal Context of LLM Training

The surge in popularity of language models, notably exemplified by the release and widespread adoption of ChatGPT, has accelerated the integration of AI into various sectors. This expansion has subsequently encouraged the development of regulatory frameworks to govern AI technologies. A pioneering effort in this domain is the European Union’s Artificial Intelligence Act (the EU AI Act), which represents the first comprehensive legal framework of its kind, anticipated to set a precedent for global AI regulation []. The Act encompasses, among various provisions, rigorous data governance guidelines (Art. 10), including adherence to the General Data Protection Regulation (GDPR). This introduces significant challenges for AI developers, particularly concerning the international transfer of data and data de-biasing processes.

Restrictions in cross-border transfer of data: The EU’s GDPR imposes stringent conditions on the international transfer of personal data, particularly to third countries deemed to lack adequate protection for personal data. The criterion for ’adequacy’, as established in Schrems v DPC (C-262/14)) and Recital 104, requires a level of protection ’essentially equivalent’ to that of the EU, a high bar for international data transfers, especially to developing countries. Additionally, the requirement for periodic reviews of adequacy for jurisdictions considered equivalent, alongside mandated safeguard measures (Art. 49) for transfers to non-equivalent third countries, introduces a layer of uncertainty and financial burden for businesses engaged in data transfer. Moreover, the EU is not the only jurisdiction tightening controls over data transfer. China, for instance, has enacted laws and supplementary provisions mandating a security assessment by regulator for the transfer of ’important data’ or personal data exceeding specific thresholds, barring certain exemptions. The varied landscape of regulations poses challenges for accessing diverse local datasets while adhering to disparate, and sometimes significantly different, regulations across jurisdictions. Hierarchical FL offers an efficient solution to maintain compliance by storing data and model within its jurisdiction of origin, avoiding cross-border transfers.

Mitigation of data bias: The EU AI Act mandates rigorous oversight throughout the entire lifecycle of the data used in AI models and obligates the implementation of ’appropriate measures to detect, prevent, and mitigate potential biases’ (Art. 10.2f,fa). It underscores the growing importance of accessing diversified data sources, particularly those from jurisdictions that are underrepresented. To accommodate the acquisition of such data while adhering to individual privacy and local regulations, a novel training paradigm, comprising multiple layers of federation both among clients and jurisdictions, is necessitated to address the limitations of the traditional single-layer FL framework.

The right to information: Transparency regarding the data collected and utilized in training, along with the rights of AI developers to access such information, has garnered considerable attention from both regulators and content creators. The GDPR grants individuals the right to obtain all information stored by a service provider (Art. 15, Rec. 63 & 64), including details on the application of this data in training models. Furthermore, the EU AI Act mandates that model providers compile and disclose a comprehensive summary of the training content publicly (Art. 52c). In the US, legal disputes such as "Times v OpenAI" have underscored the debate over the extent to which the fair use doctrine under US copyright law protects the utilization of copyrighted materials in the training of AI models. This case also ignites broader discussions about the adequacy of the current legal framework in safeguarding content creators against the opaque practices of LLM training. These challenges have led to legislative proposals, including the Generative AI Copyright Disclosure Bill in the US House of Representatives, aimed at enhancing transparency and accountability. Hierarchical FL, by its design, offers inherent advantages over centralized training models by delineating clear data provenance — identifying the sources of data and their contributors. This attribute of Hierarchical FL positions it favorably in addressing concerns related to informational rights and data privacy, presenting a more transparent framework for data utilization in AI development.

Energy efficiency: The EU AI Act advocates for the environmentally sustainable development of AI systems by proposing the formulation of a Code of Conduct. This Code is intended to establish explicit objectives and key performance indicators (Art. 69), mirroring the core values of the EU but also responding to growing concerns within the industry and broader society regarding the energy consumption associated with the training and use of AI.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 3
presentation: 2
questions: See **Weaknesses**.
soundness: 1
strengths: 1. This work conducts extensive experiments to show its effectiveness.
summary: This paper proposes that WorldLM addresses hierarchical federated learning and creates federations of federations, where different federations enforce various regulations in terms of data protection and other competitive constraints. The work conducts extensive experiments to show the effectiveness of the proposed work.
weaknesses: 1. Section 3.1 mentions that the children under root $q$ has much more critical data heterogeneity than those under parent $p$. I think the authors should justify that this really happens. In my opinion, the setting in clustered FL usually assumes that the clients within a cluster have very similar data distribution while quite distinct from those not in the same cluster [1]. However, in this case, clients may not have such information, i.e., they don't know the data distribution with each other. Besides, I believe the assumption is very strong in hierarchical FL [2]. For example, in terms of data distribution, the heterogeneity between a local orthopedics hospital and a local ophthalmology hospital may be much more severe than the heterogeneity between two orthopedics hospitals in two different cities. 
2. The presentation of the algorithm can be improved. One example is the residual layers. After going through the details of the paper, do you mean the personalized layers (decoders/encoders if the model is a transformer)? The authors should put more information on the design of a model. 
3. Table 1 and Table 2 show that "local" under "Non-collaborate" performs much better than World LM. I cannot find the details about this baseline. In common practice, it should use local data to train a client's exclusive model. If that is how you implement the baseline, I am curious about the effectiveness of the proposed work. 
4. Figure 1 should explicitly mention the meaning of different notions. I had no idea of the figure until I dived into the details of Section 3. 

**References:**
[1] Structured Federated Learning through Clustered Additive Modeling
[2] Hierarchical Federated Learning with Multi-Timescale Gradient Correction

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see above Weaknesses.
soundness: 2
strengths: 1. To address the data heterogeneity, authors propose attention-based aggregation and residual embeddings, which are very suitable for language model training.

2. The experiments are well-designed, different aspects (such as FL v.s. centralized, privacy, local task performance) are covered. These results are promising.
summary: In this submission, the authors present WorldLM, a method that can enable collaborations among different organizations/data owners to train language model.
weaknesses: 1) Scale can be large.
It is great that the authors proposed specific-designed FL method for language model (LM). As we know, LM becomes useful when the scale is very large, i.e., LLM. Thus whether the proposed method can be scaled up is really important. Currently, the authors conduct experiments with a largest size of 400M parameters, which is still small. As the experiments with large size might be difficult, could the authors discuss more details when the size of LM reaches billion scale, given their experiments with 75M ~ 400M? For example, discuss potential challenges in communication, memory requirements, or convergence behavior as the model size increases to billion-scale. Additionally,  how the attention-based aggregation and residual embedding mechanisms might behave at larger scales.

2) Communication should be discussed and experimentally tested.
As the presented WorldLM aims collaborations among different data owners, it is necessary to analyze and further experimentally test the communication costs, at lease for the designed simulations, i.e., (a) and (b) in the first paragraph of section 4. It would be better to have quantitative comparisons of communication between WorldLM and other federated learning approaches. Additionally, it would be helpful to discuss how communication costs scale with model size and number of participants in the federation.


3) Technique novelty can be clearly stated.
In section 2.2, related work about PFL, client clustering and hierarchical systems are simply discussed. It is better to clearly state the technique novelty of the proposed method compared to existing methods. In this way, readers can better understand the technique contribution of the proposed method in the field of FL. 
For example, authors an provide a concise summary table or paragraph that explicitly compares WorldLM's key features (e.g., attention-based aggregation, residual embeddings) with those of existing methods like PFL, client clustering approaches, and hierarchical systems. This would help readers quickly grasp the unique contributions of WorldLM.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: I found the comparison against FL and Centralized a little bit confusing and it 
would be great if the authors can elaborate more:

1. I don't get why 400M WorldLM can be better than Centralized, and would
appreciate some explanation here. Is it because WorldLM personalized further 
more against the global data? 

If this is the case, this paper should compare with more personalization-based
method? Even training a lora adapter for each person seems to be a stronger
method?

2. " For the 400M, we only had the resources for comparison against centralized, 
as standard FL is much less computationally efficient" -- I am not sure I get this --
simulate FL should be cheaper than simulating the centralized setting?
soundness: 2
strengths: - The idea seems to make sense, although its practical application could be better motivated
summary: In this paper, the authors focused on a hierarchical way of training federated models.
The key idea is to use attention-based aggregation and residual embedding sharing
to enable learning over "federations of federations" efficently.
weaknesses: - I would appreciate if the authors can motivate the problem more concretely, rather than 
in a high-level way
- I am a little bit confused by the comparison and the key insights we can get from these
results. (see Questions)

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 2
questions: - Why is simple averaging sufficient for the backbone while key layers require attention-based aggregation?
- Why was this specific proportion of backbone vs. key layers chosen (I noticed the 30% explanation of transfer learning)
- What are the "meta-learning properties of FL" that motivate these choices?
- Why is sequential training across levels (root → regional → edge) necessary rather than parallel training?
- Why did the authors choose this attention mechanism and do the results indicate that it’s not resolving interference in the weight updates?
soundness: 3
strengths: * The paper addresses an important practical problem in distributed LM training.
* The attention-based aggregation mechanism is an interesting approach to handling heterogeneous data (though it does not seem to be a complete solution, it suggests an interesting direction).
* Experimental design and evaluation strategy:
  * Evaluation across multiple model sizes (75M-400M parameters). Scaling experiments show competitive performance with standard FL. 
  * Comprehensive testing combining perplexity metrics with LLM benchmarks (though these can be reported more clearly)
  * Demonstrates robust performance under differential privacy constraints, outperforming standard FL approaches in these settings.
summary: This paper presents WorldLM, a hierarchical federated learning approach for training language models across organizations, each with its own domain data. The organizations are assumed to have a hierarchy of similar datasets, and it is an appealing idea to exploit this structure in the parameter updates. The key contributions are: 
1. a "federation of federations" architecture that allows organizations to collaborate while managing data heterogeneity; 
 2. a model decomposition into shared backbone and personalized key layers with attention-based aggregation mechanism to address distribution mismatch across federations; 
3. empirical demonstration that this approach outperforms standard federated learning on multilingual and domain-specific datasets, while maintaining strong performance under differential privacy constraints. 

While the work addresses an important problem and presents some interesting approaches, there are several significant limitations in the presentation, justification of design choices, and evaluation methodology.
weaknesses: 1. Presentation and Motivation:
The paper's introduction and related work sections attempt to cover both technical and policy aspects of federated learning, but in doing so, fail to provide a clear technical foundation. While the data regulation context is interesting to learn about, it comes at the expense of a precise technical exposition. At times, the paper mentions low-level technical concepts (e.g., RingAllReduce, local SGD) without proper explanation. The presentation would benefit from a clearer explanation of key technical concepts, such as a precise definition of a federation or “federation of federations”, and a discussion of FL aggregation methods.

The transition from presenting challenges to proposing a "federation of federations" solution lacks sufficient justification. While the abstract and introduction emphasize federated governance and legal/privacy challenges across organizations, it’s unclear how the proposed hierarchical setup addresses these challenges. The gaps between the motivation, formal problem setup (Section 3.1), and algorithm details (Section 3.2) is not fully justified. The term "Worldwide" seems unnecessarily broad for a technical ML paper.

One of the questions addressed by this paper is how we can update a model’s parameters from gradients computed in different tasks, without much interference. This question is heavily studied in multi-task learning, where a single neural network is trained on a collection of different tasks (e.g., see these two influential papers https://arxiv.org/abs/2001.06782, and https://arxiv.org/abs/1705.07115). Could the authors discuss how the proposed approach relates to/differs from these multi-task learning techniques? It seems like the attention mechanism proposed by the authors attempts to achieve similar objectives, but, as I write below, it is unclear whether it leads to consistent improvements because interference seems to be high (Fig 3).

2. Technical Framework:
The mathematical presentation of hierarchical data distributions, while formal, adds complexity without clear benefit. The use of terms like "LDA" with values 0.0001 and 1000 (lines 189-190) is unclear (is this Latent Dirichlet Allocation?). Could the authors provide concrete examples of how their abstract mathematical formulations (Section 3.1) translate to real-world data scenarios?

3. Algorithm Design:
Several critical design choices lack proper justification:
- Why is simple averaging sufficient for the backbone while key layers require attention-based aggregation?
- Why was this specific proportion of backbone vs. key layers chosen (I noticed the 30% explanation of transfer learning, but I wasn’t convinced by its importance).
- What are the "meta-learning properties of FL" that motivate these choices?
- Why is sequential training across levels (root → regional → edge) necessary rather than parallel training?
- Why did the authors choose this attention mechanism and do the results indicate that it’s not resolving interference in the weight updates?

4. Evaluation and Results:
The experimental results raise several concerns:
- Figure 3 shows large perplexity spikes between levels, suggesting potential instability, but the implications aren't discussed. Could the authors provide an explanation for these perplexity spikes and discuss how they affect the stability and performance of the model?
- Table 3 / Figure 4 seem to suggest the proposed method is in fact not very robust to statistical heterogeneity. Restructuring the federation hierarchy to combine non-heterogeneous data results in almost 2x increase in perplexity scores. More detailed experiments regarding the method’s robustness could be insightful. For example, instead of entirely swapping the WK and PBA datasets in the robustness experiment (Fig 2), could the authors show the results as X% of WK and PBA are swapped, varying X from 0 to 100?
- It's unclear whether perplexity is measured on consistent validation sets across levels?
- I understand that the models here are not on the multi-billion parameter scale, yet, the high perplexity values (in the hundreds) and large fluctuations seem somewhat too high for a modern LM. Could the authors provide some context for these perplexity values (e.g. comparisons to models of similar size) or discuss why these values are reasonable given their experimental setup?
- Lack of ablations for robustness to DP results: it’s unclear which aspect of the proposed method (hierarchical modeling, algorithm details, attention aggregation method?) is responsible for the robustness compared to standard DP.
- The MosaicML gauntlet results are presented as percentage improvements without raw scores, making it difficult to assess their significance given the known variance in these benchmarks (unless I am missing something).
- The sequential training approach introduces practical deployment challenges that aren't addressed

5. Practical Limitations:
The sequential training across levels appears to be a major limitation that isn't thoroughly discussed. Questions about deployment readiness, training efficiency, and stability during level transitions remain unanswered.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. How scalable is WorldLM for much larger models or significantly larger numbers of federations?
2. How does WorldLM handle emerging data distributions in dynamic environments?
3. What specific legal and privacy frameworks were considered in the evaluation?
soundness: 2
strengths: 1. The federations-of-federations approach allows for adaptable collaboration across various jurisdictions, making it feasible to integrate global, region-specific, or industry-specific data in a way that respects privacy constraints.

2. The backbone with personalized key layers effectively captures and adapts to local variations in data, enhancing performance in heterogeneous settings.

3. WorldLM is robust in applying differential privacy, even where traditional federated learning might struggle, which is a critical advantage for handling sensitive information.
summary: The paper proposes WorldLM, a federated language model training system that leverages a "federation of federations" structure. This approach enables collaboration among organizations across different legal, security, and privacy jurisdictions to train language models on heterogeneous data without sharing sensitive information. WorldLM uses a partially personalized aggregation technique and cross-federation information sharing via residual layer embeddings, which addresses statistical heterogeneity in the data. Evaluation on diverse datasets shows that WorldLM outperforms traditional federated learning, closely matches the performance of localized models, and remains effective under differential privacy conditions.
weaknesses: 1. The method shows diminished effectiveness when data within a federation lacks inherent similarity, suggesting a need for improved aggregation techniques for highly diverse datasets.

2. While WorldLM works well on medium-sized language models, scaling to larger models could be resource-intensive, especially for smaller organizations with limited computational resources.