Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: In this paper, the authors propose a domain-aware representation learning method (FedDAR) for the non-iid FL problem. The FedDAR assumes data on clients are from multiple domains and learns a classifier head for each domain. A representation module is shared for all classifier heads and updated by the vanilla FedAvg. The authors also proposed a second-order aggregation method to update domain-aware classifier heads, whose effectiveness is validated by ablation studies. Experiments on both synthetic data and real-world datasets validate the effectiveness of the proposed method.
 
 
 
 


strength_and_weaknesses: Strength
1. solid theoretical analysis on convergence.

2. extensive experiments on both synthetic data and real-world datasets validate the effectiveness of the proposed method.

Weaknesses
1. The domain shift problem in FL has been extensively discussed in various works in recent years, such as [1-3]. Specially, [4] proposed a similar architecture where a representation module is shared and multiple classifiers heads are applied for target domains. As these methods are not mentioned or compared, it is uncertain if the contributions of the proposed FedDAR hold and the experiments are not convincing.

2. This work assumes that there is labeled information about the domain a sample belongs to, which limits the proposed FedDAR’s application as it is usually not the case in FL. Besides, as there are no comparisons with those unsupervised methods[1,3] handling the domain shift problem in FL, it is unclear how many improvements the labeled information brings.

3. As discussed in Sec.1 and Sec 2., personalized models have been recognized to have superior performance to vanilla FL models. Then comparisons in Sec.3.1, Sec.5.1, and Sec.5.2 would be less persuasive as they compare the FedDAR with vanilla FedAvg rather than other personalized methods. For example, will FedDAR outperform a vanilla model fine-tuned on a client’s local data?

4. The second-order aggregation shares similar ideas as those personalization methods with a meta-learner [5]. Besides, as sharing and updating a Hessian matrix would be communication expensive and unstable, it requires a client to have sufficient training data and limits FedDAR to have simple classifier heads. Experiments in Appendix B.2. also indicates this problem. 

[1] Disentangled Federated Learning for Tackling Attributes Skew via Invariant Aggregation and Diversity Transferring (https://arxiv.org/abs/2206.06818)

[2] FedBN: Federated Learning on Non-IID Features via Local Batch Normalization (https://arxiv.org/abs/2102.07623)

[3] Federated Adversarial Domain Adaptation (https://arxiv.org/abs/1911.02054)

[4] Diurnal or Nocturnal? Federated Learning of Multi-branch Networks from Periodically Shifting Distributions (https://openreview.net/forum?id=E4EE_ohFGz)

[5] Personalized Federated Learning: A Meta-Learning Approach (https://arxiv.org/abs/2002.07948)


clarity,_quality,_novelty_and_reproducibility: quality
The proposed method is well analyzed and discussed, but there are insufficient comparisons and experiments to support the claims.

clarity 
The paper is well organized.
 
originality
The work is less original as there are numerous works that share similar ideas with the proposed method.

summary_of_the_review: This paper proposes a domain-aware representation learning method (FedDAR) for the non-iid FL problem. The authors provide solid convergence analysis and validate FedDAR’s performance on both synthetic and real-world datasets. However, assumptions on domain labels and simple classifier heads limit the application of FedDAR. Besides, as there are limited comparisons with existing methods sharing similar ideas or settings, the paper’s contributions could be limited, and the results are less persuasive.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The authors address the problem of personalized (to a client) federated learning in which the distribution from each client is assumed to be drawn from a mixture of predefined domains. The proposed model learns a shared data representation with independent prediction modules (heads) for each client. Moreover, the authors show that in the linear case, the proposed approach enjoyed linear convergence and empirically, that on both artificial and real-world data, the proposed approach outperforms existing federated learning methods.
strength_and_weaknesses: The theoretical results in Theorem 5.1 are interesting and highlight the sample requirements scaling with k^2, which is convenient, specially, if k << d as the authors argue. However, it is also important to recognize that modern applications involving (medical) images and text leverage encoders with representations whose dimensionality is in the high hundreds or lower thousands, which may be prohibitive for data-poor clients and in applications where sample size is limited, e.g., rare diseases and hard to obtain data such as pathology. Moreover, the large number of clients scenario described by the authors in which the cost decreases as n^-1 is not that realistic in healthcare domains.

The toy example in Section 3.1 is welcome, however, it is not necessarily clear or discussed why in the FedRep scenario, where heads are specified, one for each client, the representation is not flexible enough to account for the differences in domain. One can certainly devise a toy example where it is possible.

The need and motivation for u_m (inverse probability weighting) is not provided. Moreover, no intuition is provided about the difference or relationship between \hat{\cal R}_m and \hat{\cal R}_i, which only differ by u_m.

The justification for only using second order aggregation for the heads (with low enough dimensionality) is reasonable (accounting for communication costs), however, its relative benefit compared to only considering simpler aggregation for the encoder (which is understandable), is not provided, though explored in the experiments.

Though the use of the artificial data in Section 5.1 is justified, the considered generative process seems too simplistic to resemble any realistic scenario. Further, the results in Table 1 may be more impactful if the experiment is repeated multiple times to better understand the variation around the performance statistics (max, min and average accuracy). It will be also important to note the value of k used for the results in Table 1.

It is not necessarily clear whether AUC is a proper metric to assess fairness across domains. Perhaps the authors can elaborate on the reasoning behind their fairness claim.

The results in Table 3 are impressive, however, it is not clear why the authors used a cross-validation scheme considering that one of the claims is that the proposed approach is not data hungry. Further, the dimensionality of the EXAM data is not specified.
clarity,_quality,_novelty_and_reproducibility: The paper is in general clear and well written, though it can use some proofreading to clean out some glaring typos. The novelty consisting on the domain-aware formulation and the theoretical results are satisfactory. Even though some of the aspects of the formulation are not properly justified the general concept is reasonable. The authors provide enough details in the main paper and supplementary material to make the experiments at least conceptually reproducible.
summary_of_the_review: The authors presented a domain-aware approach for federated learning together with a theoretical analysis of the (per iteration) sample complexity of the methodology and empirical results using both synthetic (sections 5.1 and 5.2) and convincing real-world data (Section 5.3). The contributions of each of the components of the model are justified with an ablation study in Section 5.3. Though the theoretical results are interesting, they may not be applicable to many real-world scenarios; and though the weights u_m seem reasonable, they are not justified or motivated.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper considers a federated learning problem with multiple heterogeneous clients, where each client’s data distribution is a distinct mixture of predefined domain distributions. It is assumed that the domains of the data samples are known. The goal is to federatively learn a model that can perform well in every domain. For this purpose, the paper proposes FedDAR algorithm, which learns a shared representation across all domains and a separate predictor head for each domain. Training of the encoder and the head predictors are decoupled. Learning a shared representation with a small number of outputs allows domain heads to be accurately learned even with a small set of per-domain samples. Experiments show that FedDAR achieves significant performance gains over methods that do not take domain labels into account.
strength_and_weaknesses: Strengths: 

-Shared representation structure enables the algorithm to learn low-dimensional classifiers for each domain, which benefits domains with a small number of data instances. 

-Experiments demonstrate that the proposed algorithm comes with significant performance gains compared with client-centric or domain-agnostic federated learning methods. 

Weaknesses: 

-How FedDAR optimizes the learning objective in Equation 1 is unclear. Minimizing the expected risk averaged over all domains may not lead to a model which can perform well in every domain. This paper does not address fairness issues, nor does it achieve minimization of risks on a per-client basis, as each client has a data distribution that is a different mixture of predefined domains. While it is mentioned that the algorithm uses sample re-weighting, it is not clear how this is achieved with the current objective function. 

-The paper does not study the privacy implications of sharing model parameters.

clarity,_quality,_novelty_and_reproducibility: A lot of details regarding the convergence analysis are deferred to the appendix. It will be good to provide a sketch of the proof in the main paper. Please elaborate on techniques used in convergence analysis, highlighting the technical novelty of the results.
summary_of_the_review: This work tackles the important problem of domain-aware federated learning. The proposed algorithm is intuitive, and its applicability is fortified by convergence analysis in a simplified setting and experimental results in a complex setting. While there are some issues related to privacy and computational costs, my first impression of the paper is positive.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper investigates the problem of data heterogeneity in cross-silo federated learning. It assumes that datasets in each client are composed of a mixture of known domains, and domain metadata is available for each sample.

Building on a fair learning formulation across domains (where each domain plays an equal role), the authors propose to train a neural network composed of two parts: a representation shared across domains and linear heads specific for each domain, called, FedDAR. Throughout each round, each part is trained alternatively, with the other part freezed. For the local heads, the authors consider two variants, FedDAR-WA (equivalent to FedAvg) and FedDAR-SA (second-order, requires to transmit hessians).

The paper proves on a toy linear model that the shared trunk converges to the true one linearly (Theorem 1).

Numerical experiments on synthetic data, real data with controlled distribution (fairface), and real data without controlled distribution (EXAM healthcare dataset) demonstrate that the method improves over previous personalization methods. An ablation study helps to understand the role of the different components.
strength_and_weaknesses: # Strengths

- The method proposed by the authors reaches the best performance on all datasets
- Well written overall

# Weaknesses

- The main weakness of the paper stems from the way baselines are handled
  - Harmonized baselines across examples (for synthetic and fairface, separate FedAvg is not proposed although it's a natural baseline)
  - Unless mistaken, FedAvg with per-domain reweighting is not considered as a baseline, except in the top row of Table 2 (ablation study). I think it would be fairer to compare all methods with the same optimization problem, to really demonstrate the benefit of sharing the trunk in the representation (given the results of Table 2, FedDAR would still probably be the best method).
- Writing could be improved in some sections (cf detailed comments below)
- The approach is limited to known domain mixtures (acknowledged by the authors)
clarity,_quality,_novelty_and_reproducibility: # Clarity

- The paper is well written overall. There are minor typos (see below)
- I did not understand well Section 4.3. I think that some rewriting could help
  - After initially reading this section, I understood that FedDAR-SA did 1 (full-batch) local Newton step which was then re-aggregated via Eq (8), which made sense.
  - However, App B.3.3 indicates that 5 epochs of update for heads are done for FedDAR-SA. This would tend to indicate that, in fact, what is done in FedDAR-SA is a minibatch optimization of the 2nd order Taylor expansion of the loss (7). But in this case, is the hessian computed at each minibatch? And how is the aggregation done after 5 such epochs?
  - Similarly, the projection layers hinted at the last of this section are not very clear to me, compared to the rest of the paper which is very precise.

## Minor typos/comments
- Notations for $\hat{R}_{i, m}$ are not consistent in Algorithm 1 and in Eq (6) (comma vs composition operator $\circ$ )
- In algorithm 1, for the update of the trunk $\phi$, shouldn't the operator GRD apply to $\hat{R_{i}} (\phi_i)$ rather than $\hat{R_{i, m}}$?
- $\alpha$ is used both to denote the learning rate in Algorithm 1 as well as to control the Dirichlet mixtures in Sec 5, which can be confusing

# Quality

- The experiments are overall sound. My main comment is related to the choice of baselines, cf weaknesses above. I also have more minor comments:
- Regarding measurement of AUC, it is written in the last sentence of Appendix B.3.3 that AUCs are averaged over the last five communication rounds. What justifies this choice?
- Why are only cross-validation results reported for the EXAM dataset, and not a fixed train/test split?
- Would it be possible to provide error bars for the synthetic results as well as the ones with FairFace?

# Novelty

To the best of my knowledge, this is the first time that specific heads have been proposed with this formulation.

# Reproducibility

- How were parameters tuned?
summary_of_the_review: Although I have some concerns regarding the choice of baselines, I think the paper makes an overall compelling case for the proposed approach. I am willing to upgrade my score if the authors address my questions.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: The paper considers a specific non-iid federated learning environment where each client's distribution is a mixture of predefined domains. To prevent the performance drop of the FL method in this environment, the authors propose a method called FedDAR consisting of a shared encoder across all domains and a personalized head in each domain.

All reviewers, including this AC, somewhat agreed that the setting is important, idea is novel and the paper is well written in most places, and unanimously suggested accept. However, even through rebuttal, reviewers were not completely convinced on some issues such as fairness metric/connection to optimization objective. It would be nice if these concerns could be completely resolved in the final version. And also in the final version, I think a concept figure about your method is absolutely necessary to help readers understand.
justification_for_why_not_higher_score:  it's well-deserved acceptance, but it didn't show any more surprising results than that.
justification_for_why_not_lower_score:  n/a
Decision: Accept: poster