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.
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.
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.
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?