Personalized Federated Learning via Tailored Lorentz Space

Authors: Jiahong Liu, Xinyu Fu, Menglin Yang, Weixi Zhang, Rex Ying, Irwin King
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Personalized Federated Learning, Hyperbolic Geometry

Abstract:
Personalized Federated Learning (PFL) has gained attention for privacy-preserving training on heterogeneous data. However, existing methods fail to capture the unique inherent geometric properties across diverse datasets by assuming a unified Euclidean space for all data distributions. Drawing on hyperbolic geometry's ability to fit complex data properties, we present FlatLand, a novel personalized federated learning method that embeds different clients' data in tailored Lorentz space. FlatLand can directly tackle the challenge of heterogeneity through the personalized curvatures of their respective Lorentz model of hyperbolic geometry, which is manifested by the time-like dimension. Leveraging the Lorentz model properties, we further design a parameter decoupling strategy that enables direct server aggregation of common client information, with reduced heterogeneity interference and without the need for client-wise similarity estimation. To the best of our knowledge, this is the first attempt to incorporate Lorentz geometry into personalized federated learning. Empirical results on various federated graph learning tasks demonstrate that FlatLand achieves superior performance, particularly in low-dimensional settings.

Decision: reject
Reviewer scores: 6, 6, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1729677973222

thread_type:
official_review

<thread_object>
creator: reviewer_ka7l

<content>
summary:
The paper proposes a personalized FL method called Flatland. Flatland follows a standard Federated Averaging framework. Each client $c$ first embeds their data into a Lorentz space with a personalized (learned) curvature coefficient $K_c$. They then run local training of a Lorentz model (neural network consisting of so called Lorentz layers that are designed so that if the inputs are in a Lorentz space then so are the outputs). The parameters of this network that correspond to the first dimension of the embedded data are personalized, per client, parameters, while the remaining parameters are shared. The shared parameters are updated using FedAvg, the personalized parameters are kept on the client. 

The paper shows that the specific decomposition proposed still maintained the Lorentz property of the network. The paper then carries out an empirical evaluation of the method on several different graph datasets, for both node and graph classification tasks. It compares to a range of baselines.

strengths:
* The proposed method is easy to integrate into standard FL training pipelines. It trains using federated averaging and aggregates client updates, meaning it is compatible with privacy preserving techniques such as secure aggregation and differential privacy.

* The originality of the paper is that it is the first to propose using a different Lorentz embedding for each client in order to better capture the differing client data distributions.

* The experiments demonstrate good performance on node classification in real world datasets with small numbers of clients

* The authors provide ablation studies examining the impact of different components of the method

weaknesses:
I do not believe the paper makes a significant enough contribution, either theoretically or empirically:

Regarding theory, there is a lack of any substantial theoretical contributions 

* As far as I can tell Proposition 1 and Corollary 1 are both instant consequences of the definition of a Lorentz network (eq 2), giving a name to the first row and column of the matrix does not require reproving anything. Even more confusing is Corollary 1 and it’s proof, Proposition 1 states that $\forall M, LT(x;M) \in \mathcal{L}$ then of course it is also true that $LT(x;\Phi(M, N)) \in \mathcal{L}$ because it holds for all matrices in the second entry, this does not require a proof as given in the appendix. 
* The remainder of section 6 makes some hand wavy and confusing claims and does not prove any concrete properties or statements about the method. 
* There are many potentially interesting directions here the authors could have taken, e.g. what is the convergence rate, how does the lorentz embedding impact convergence and/or performance, what types of distribution heterogeneity are well modeled by the Lorentz space etc.

Regarding empirical contribution. My overall concern here centers around the empirical evaluation being too narrow. Essentially the evaluation is limited to federated cross-silo, graph datasets, which to me makes it a specialized method rather than a general one, as the authors claim.

* The empirical evaluation is only on graph datasets. The authors present the method as a general personalized federated algorithm and claim that the method is applicable to other data modalities; they present no evidence of this. I am also not convinced that this would be the case, state of the art methods that deal with the most common data modalites (text, image, video, tabular) do not make use of Lorentz spaces. 

* The number of clients in each dataset is very low. Therefore, the experiments do not provide evidence that the method works in a cross-device setting. Crucially, the authors do not mention the cohort size they used in federated averaging, and given the low number of clients in their experiments I am assuming that all clients participate in each round of training. I have concerns about if the method would even work without full client participation in each round (if client A has participated before and client B has not, then A will have already updated their personalized parameters often, while B is starting from initialization, this is going to lead to very heterogeneous shared parameter updates)

* I believe the authors need to provide experiments for other data modalities, with larger numbers of clients as well as partial client participation 

There are some minor typos, e.g. I think in section 5.2 $m^{(l)}$ should be in $\mathbb{R}^m$, not $\mathbb{R}^{m+1}$ and $M^{(l)}$ should be in $\mathbb{R}^{m\times n}$. Also references to tables in section 7 are off and need to be checked e.g. on line 452, Table 3 should be Table 1.

Overall, I would need to see substantial improvements to the theoretical and/or empirical contributions, as suggested in the final bullet points of each section above.

questions:
* Why do the $v$ parameters not appear anywhere in Equation (4)?
* How does Lorentz space actually model heterogeneity? The paper seems to claim that $x_t$ captures the heterogeneity between client distributions, while the client distributions in Flatland i.e. $x_S$ should be the same. This is an interesting idea but no evidence is provided. I would like to see either theoretical or empirical evidence of this fact.
* Why do you decouple the parameters the way you do? This is related to the above point. It seems that the parameters are decoupled this way so that the heterogeneous feature $x_t$ is fed into the personalized portion of the network $m$ in Equation (4), while the homogeneous features $x_S$ are fed into the shared portion of the network, $M$ in Equation (4). However, as far as I can tell this only works for the first layer since the output of layer 1 $m x_t + M^{(1)} x_S$ will be fed into the shared part of layer 2 $M^{(2)}$ meaning that the updates to shared parameters from layer 2 onwards are still impacted by heterogeneous parts of the client data.
* Does the method work for larger numbers of clients, without full client participation each round?
* What do you mean by statistically significant for the results in Table 1? How can you say that FlatLand outperforms GCFL in graph classification when the mean accuracies are very close with very large overlap of the standard deviation intervals. This looks to me like noise rather than signal.
* It is not clearly defined what the embedding dimension is in Section 7.3. Is this the hidden dimension of the model?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730614830321

thread_type:
official_review

<thread_object>
creator: reviewer_dc6w

<content>
summary:
This paper introduces FlatLand, a novel personalized federated learning (PFL) approach that addresses data heterogeneity across clients by leveraging hyperbolic geometry, specifically Lorentz space. Unlike traditional PFL methods that operate in Euclidean space, FlatLand embeds each client's data in a tailored Lorentz space with different curvatures to better capture non-Euclidean properties of the data. Instead of explicitly calculating client similarities or using additional modules for personalization, FlatLand leverages the geometric properties of Lorentz space to handle heterogeneity naturally through different curvatures. Also, they use parameter decoupling for shared and local models, similar to some other personalized FL models.

strengths:
This paper presents an innovative approach to personalized federated learning by leveraging hyperbolic geometry, specifically through the use of Lorentz spaces with tailored curvatures for different clients. The theoretical foundation is well-developed, with mathematical proofs and insights connecting geometric properties to data heterogeneity. The experimental results show promising improvements over existing methods, particularly in low-dimensional settings, which could be valuable for resource-constrained environments. The parameter decoupling strategy based on geometric properties is elegant and eliminates the need for explicit client similarity calculations, which is a common computational burden in some of the other personalized federated learning approaches using clustering approaches.

weaknesses:
The paper's evaluation primarily focuses on graph-based tasks, leaving questions about its generalizability to other types of data and models. While they mention the potential applicability to other domains, this claim remains unverified even in the experimental sense. The method introduces additional complexity through hyperbolic geometry and requires curvature estimation, which could make implementation and tuning more challenging in practice. The computational overhead of working in Lorentz space compared to Euclidean space is not thoroughly discussed, and the paper doesn't address potential numerical stability issues that often arise when working with hyperbolic spaces. Additionally, while the method shows improvements over baselines, some of the gains are modest, particularly in the graph classification tasks, suggesting that the added complexity might not always justify the performance benefits.
The main concerns of the paper are as follows:
 - The paper presents its parameter decoupling strategy as a novel contribution without comparing it with similar approaches in existing literature. The absence of comparisons with previous decoupling methods, particularly [A], makes it difficult to assess the true novelty and advantages of their approach.
- The paper lacks clarity in explaining crucial algorithmic details, particularly regarding curvature estimation. While curvature is central to the method's performance, the paper doesn't adequately explain how it's estimated, updated, or integrated into the training process. The absence of analysis on the impact of different curvature values and their stability leaves important implementation questions unanswered.
- The evaluation is heavily focused on graph-based tasks, with no experimental validation on other common federated learning applications like computer vision or natural language processing (such as FEMNIST). This narrow focus raises questions about the method's applicability and effectiveness in broader contexts, especially given the additional complexity introduced by hyperbolic geometry.
- The paper lacks a thorough analysis of computational overhead compared to Euclidean-based methods. There's no discussion of memory requirements for storing different Lorentz spaces or how the method scales with an increasing number of clients. The communication costs and practical implications in real-world deployments remain unexplored.


[A] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In ICML, volume 139 of Proceedings of Machine
Learning Research, pp. 2089–2099. PMLR, 2021.

questions:
In addition to the questions and concerns above, I have some other questions:

- Could you elaborate on how the curvature parameters are practically updated during the federated learning process? Specifically, what happens if the estimated curvature is unstable or changes significantly between rounds, and how do you ensure this doesn't negatively impact the model's convergence?
- Could you discuss any potential challenges or modifications needed when applying your method to these domains, particularly regarding the relationship between data heterogeneity and geometric properties in these different contexts?
- Could you provide insights into the scalability of your approach as the number of clients increases, especially considering the additional overhead of maintaining different Lorentz spaces? For instance, provide some results with training times and memory utilization.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730774565896

thread_type:
official_review

<thread_object>
creator: reviewer_wtgh

<content>
summary:
This paper proposes a new personalized federated learning approach called FlatLand. Unlike current methods that commonly operate in Euclidean space, FlatLand embeds client data in a Tailored Lorentz space to more effectively represent the clients’ data distributions. Building on this foundation, the authors introduce a parameter decoupling strategy to aggregate shared parameters efficiently. Extensive experiments demonstrate the effectiveness of FlatLand.

strengths:
1. The paper is well-written and easy to follow. 
2. The idea of increasing dimensionality to embed client data in Lorentz space is novel.
3. Experiments are extensive.

weaknesses:
1. Techniques that personalize client models by creating localized models or layers are highly relevant to this paper. However, the related work section includes somewhat outdated references. It would benefit from incorporating more recent approaches, such as FedCAC [1] and GPFL [2]. Additionally, the experiments lack comparisons with these types of methods.

2. This paper primarily focuses on graph datasets. Can FlatLand effectively perform on commonly used benchmarks in PFL, such as CV datasets (e.g., CIFAR-100) and text datasets (e.g., AGNEWS)?

[1] Bold but cautious: Unlocking the potential of personalized federated learning through cautiously aggressive collaboration, ICCV 2023
[2] Gpfl: Simultaneously learning global and personalized feature information for personalized federated learning, ICCV 2023

questions:
Please see the weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730973397706

thread_type:
official_review

<thread_object>
creator: reviewer_4tqh

<content>
summary:
This paper studies the federated graph learning problem, and proposes to model data heterogeneity across different clients as different Lorentz spaces with different curvatures. The authors adapt the Lorentz network into the FL framework, and propose to maintain dedicated trainable client-specific parameters to learn different client curvatures, which makes the new method FlatLand different from previous baselines. Simulation results and ablation studies are also presented in the paper to support the better performance of the new method.

strengths:
Although there are already many works on federated hyperbolic learning, I found the idea in this paper still interesting and inspiring. The proposed method is concise yet intuitive, and the simulation results seem to suggest that it also works well in practice for graph-related problems.

weaknesses:
The presentation of this paper can be improved. 

- Firstly, some of the notations are not consistent in the paper and can make readers confusing. For example, line 286-290, if $x^{(l)}$ is in $\mathcal{L}^n_K$, then $x_s^{(l)}\in\mathbb{R}^n$ instead of $\mathbb{R}^d$; $m^{(l)}$ should be in $\mathbb{R}^{m}$ and $M^{(l)}$ should be in $\mathbb{R}^{m\times n}$. This type of inconsistency exists throughout the paper. 

- Secondly, some of the statements in the paper are left unexplained, which makes me very confused. For example, line 167-168 "For instance, for dropout, the operation function is f(Wx, v) = W dropout (x).", where is the usage of the bias v? Another example is Equation (4), it seems that the first row of $\hat{M}^{(l)}$ is not used at all in each LT layer. Is something missing here? I assume $K$ is a separate learnable scalar to represent the curvature?

- Thirdly, how the authors use the proposed method to solve federated **graph** learning problem is not clearly explained in the main text, as the proposed method is suitable for transforming the node feature, but is not easy to handle the adjacency matrix directly. So it would be confusing if the author just apply the method over graph datasets as presented in the simulation section. In appendix, I find that "For the node classification task, we employ 2-layer GCN for Euclidean models, 2-layer LGCN Chen et al. (2021) for FlatLand... For graph classification, we use 3-layer GIN Xu et al. (2018) as the Euclidean encoder, and the same 3-layer hyperbolic encoders as node classification for hyperbolic models" How LGCN is combined with the proposed method? What does the "3-layer hyperbolic encoders" mean in this case? I think the authors should first provide a clear explanation of these details in the main text before diving into the detailed results. Also, the paper can be reorganized to cover more important algorithmic details (i.e., pseudocode) in the main text.

- Lastly, I find that the authors provide the implementation of the method with a text link in the doc (line 404, "anonymous repository") instead of uploading it as a supplementary file. I nearly missed it. I would recommend the authors to upload the code directly whenever possible.

questions:
1. Is the parameter $K$ in equation (4) just a learnable scalar? Do we need to have any kind of regularization/restrictions on $K$?

2. Line 501, what does "no parameter decoupling strategy" mean? Does it mean the server will aggregate all parameters including the client-learned curvatures?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731188267570

thread_type:
official_review

<thread_object>
creator: reviewer_ntmn

<content>
summary:
The paper proposes a new personalized federated learning algorithm called Flatland employing hyperbolic geometry. One of the key challenges in federated learning is the data heterogeneity among clients and understanding the similarity between clients data distributions can be helpful. The paper tries to achieve this goal by projecting each client data to a higher dimension to better capture similarities between clients data. The paper more focused on graph data. Experiments are carried out to showcase the effectiveness of the approach.

strengths:
- The motivation and contribution of the paper is clear and it seems that leveraging hyperbolic geometric approaches can be helpful for personalized federated learning.
- The paper provides analysis in Section 6 to support their approach.
- Experimental study of the paper looks convincing.

weaknesses:
Reading of the paper, I felt that some arguments in the paper are vague and needs more explanations to become more clear. For example "*PFL approaches include segmenting models into generic and personalized components (McMahan et al., 2017; Tan et al., 2023), leveraging model weights and gradients to map client relationships (Xie et al., 2021), or integrating additional modules to facilitate customization (Baek et al., 2023)*." Or for example in "Moreover, embedding data from various clients into a fixed space complicates the interpretability of model parameters, making it difficult to segment the model into meaningful components (Arivazhagan et al., 2019) and often expensive to assess similarities between client models", the reasons are not clear.

Furthermore, the paper states that embedding data from various clients into a fixed space often expensive to assess similarities between client models. However, using the Flatland itself in step 4 of algorithm 2, each clients needs to project its data to another space. I believe this can add more computations compared to other methods which does not employ Lorentz space and hyperbolic geometry-based approaches. Therefore, to me the above argument seems a bit contradictory.

questions:
Looking into algorithm 2 of the paper, it seems that using Flatland clients and the server collaborate using conventional federated averaging methods such as Fedavg. Is Flatland compatible with other personalized federated learning algorithms? Is it possible to use other methods like PerFedAvg to orchestrate the collaboration between the server and clients? Maybe this can further improve personalization properties of Flatland.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732506757904

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Follow-up on Review Comments

comment:
Dear Reviewers,

&nbsp;

Thank you very much for your valuable feedback and comments. As the rebuttal period is ending, please don't hesitate to reach out if you have any further questions or concerns. We would also appreciate it if you could update your ratings if we have resolved your concerns appropriately.

&nbsp;

Best regards,

Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733220775129

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Summary

comment:
&nbsp;

We thank all the reviewers for their time and for providing constructive comments to enhance the paper. **We appreciate the reviewers' positive comments:**

* The motivation and contribution of the paper are clear (*Reviews nTmN, 4tQH, Wtgh, Dc6W*). The ideas that integrating geometry information to solve the heterogeneity in federated learning are novel, interesting, and inspiring. (*All reviewers*)
* The proposed method is elegant (*Reviewer Dc6W*) , concise yet intuitive (*Reviewer 4tQH*), and easy to integrate. (*Reviewer KA7L*). 
* Convincing evaluation and ablation study of proposed method in graph datasets. (*Reviewers nTmN, 4tQH, Wtgh, KA7L*)
* Solid theoretical support and sufficient analyses (*Reviewers nTmN, Dc6W*).
* This paper can inspire more ideas in federated hyperbolic learning problems (*Reviewer 4tQH, Dc6W*).

&nbsp;

---

We include the following revisions in the Rebuttal PDF according to reviewers' constructive suggestions to improve the quality of our manuscript. Due to the time constraints, we will continue to refine our paper later. Specifically, the added contents are as follows:

1. **Clearer arguments and typos.** We refined our explanations  (*Reviewers nTmN, 4tQHm*) and added experimental details about the curvature calculation  (*Reviewers 4tQH, Dc6W, 4tQH*) and graph classification settings in Appendix C2 (*Reviewer 4tQH*).

2. **Time complexity analysis.** We added a complexity analysis section in Appendix B6 to show that exponential maps do not lead to much extra computation overhead (*Reviewers nTmN, Dc6W*).

3. **Convergence analysis.** We added a convergence analysis in Appendix B5 to show that our method does not negatively impact the convergence rate (*Reviewer Dc6W*), which also enhances our theoretical analysis part (*Reviewers KA7L*).

4. **Experiments on image datasets.** We added a section in Appendix C3 to show some experiment results on the image dataset compared with some strong baselines that are not utilized in graph datasets (*Reviewers Dc6W, Wtgh, KA7L*).

5. **Experiments in a partial participation setting.** We added a section in Appendix C4 to analyze the performance with lower participation rate (*Reviewer KA7L*). 

We update the whole results in Cora (50 clients) as follows; we will integrate and update it in our final version:


| Participate Rate |  0.1  |  0.3  |  0.5  |  0.7  | **1.0** |
| :--------------: | :---: | :---: | :---: | :---: | :-----: |
|      FedAvg      | 18.14 | 36.64 | 34.30 | 33.61 |  30.20  |
| FedPer       | 56.03 | 60.22 | 74.79 | 73.48 | 68.11 |
| FedHGCN      | 54.38 | 77.08 | 69.87 | 73.35 | 62.52 |
|   **FlatLand**   | **81.82** | **81.11** | **79.27** | **79.42** |  **77.98**  |

&nbsp;

---

We would like to give clarification in the question:

&nbsp;

**The selection of evaluation in graph datasets.**

The core idea of our approach lies in the correlation between data heterogeneity and hyperbolic curvature. We propose a novel method to leverage geometry information to decouple data heterogeneity. This relationship is particularly intuitive and supported by solid theoretical guarantees in graph data, where hyperbolic geometry has been more extensively developed. Therefore, we primarily use graph data to evaluate the rational and effectiveness of our proposed method.


&nbsp;

**Future work and improvement.**

We acknowledge that the performance of our method across all data types remains an open question for further investigation. Due to time constraints, our exploration of other data types is somehow limited, though we supplemented our study with experiments on image datasets. However, **this does not conflict with the motivation, contribution, methodology, and experimental results presented in this paper.** This work can inspire more ideas. Below, we outline possible areas for further exploration following this work:

1. It is the truth that hyperbolic space may NOT be universally optimal for all data distributions as it is still a rapidly evolving field. One intuition is that some data may exhibit positive curvature. Therefore, we could follow the idea of this paper to model more complex data structures in **mixed-curvature spaces** that also contain manifolds that have positive curvature.

2. Inspired by Reviewer KA7L, we could investigate the **unique potential** of geometric information when *dealing with the rapid updates of parameters for newly arrived users*. A novel strategy could be designed, for example, by *leveraging curvature values to directly identify clients with similar data properties*. This would allow us to quickly retrieve personalized parameters from these clients to help initialize the personalized parameters for the new user. This is a task that traditional Euclidean-space methods struggle with.

3. As hyperbolic space continues to show promising results across data in other domains, we could *expand our analysis and validation to include more complex and challenging datasets and tasks.*
</content>
</thread_object>

</thread>


Metareview:
metareview:
**(a) Summary of Scientific Claims and Findings:**  
The innovations proposed by the authors are: i) use fully lorentz neural networks instead of just the hyperbolic graph neural networks, and ii) use client-specific curvatures. Their justifications are
- Fully Lorentz Neural Networks are even better than HGNNs at modeling hierarchical data due to the exponential volume expansion in hyperbolic space, enabling efficient representation of scale-free graphs with fewer dimensions than Euclidean models.

- Different clients may need different curvatures since the levels of complexity/hierarchy can be heterogenous. So the authors propose to learn personalized client-specific curvatures.

**(b) Strengths:**  
- The idea of using client specific curvatures is well motivated and unique.
- The authors do a pretty convincing ablation study to demonstrate the effectiveness of their two innovations.

**(c) Weaknesses and Missing Elements:**  
- Fully Lorentz Neural Networks are not novel - this is a well known technique in the centralized setting (Chen et al. 2022).
- While the results are great, the setup is currently quite limited. The number of clients and datasets are quite small. The effect of model architectures, client sampling, etc. is not explored.
- Theoretical contributions are limited and lack depth. Most the results are known or direct extensions.

**(d) Decision and Key Reasons for Rejection:**  
The paper is rejected due to limited theoretical depth, lack of robust empirical validation across diverse datasets and realistic scenarios.