Paper: FedJETs: Efficient Just-In-Time Personalization with Federated Mixture of Experts (ICLR-2024)

Reviewer comment (Novelty):
The architecture is similar to STAR model in paper [2] without anchor users.

Author response:
-*Star Topology* [2]: [2] proposes a method that constantly updates two models during training: a) Global model, which shares parameters across all personalized models, and b) In-domain personalized models (1 per domain). In FedJETs, the "pretrained model" remains frozen at all times, it is never updated during training or testing, and its only purpose is to serve as a feature extractor and feed its embeddings to the gating function so that it can rank the expert models. During inference, [2] combines these weighted networks and unifies them into a new one via an element-wise multiplication process. FedJETs ranks the subset of matching experts per sample level, capitalizing on a single expert more relevant to the sample distribution. Lastly, in [2], the domain ID needs to be explicitly fed into the network to facilitate the model learning; FedJETs dynamically identify the relevant experts by training the gating function.

Reviewer comment (Experiments):
There are couple of places mention that our approach lower communication costs but there is no experimental results that show that how much improvement is there as in FedMix paper.

Author response:
**Regarding lower communication cost**:
In our experiments, we demonstrate how our method is more communication efficient than FedMix by showing how our approach can achieve significantly better final performance using the same communication budget shown in Tables 1 and 5 (in the appendix). In Table 5, if we compare FedJETs using two experts with FedMix using five experts, we can see FedJETs achieve better final performance with a smaller communication budget.

Reviewer comment (Experiments):
There are no experiments for the cold-start problem as claimed in the paper (this is not same as unseen new users for testing. Testing, of course, should be unseen).

Author response:
**Regarding cold-start problem**:
Our paper does not use the term "cold-start" for our targeted scenario. Our method will adapt to new clients with unseen local dataset distributions and without labels during testing by dynamically selecting experts. We are not using any label from the test client, so we define this as zero-shot penalization.

Reviewer comment (Experiments):
The claim "we partition the data samples by classes to turn full datasets into non-i.i.d. subsets", how do you make sure that samples with different class labels with same data is non-i.i.d.?

Author response:
**Regarding data non-iidness**: 
We reproduce two main partitioning strategies proposed in [19] to simulate the non-iidness in our clients. The first partitioning is quantity-based label imbalance, and the second uses a more standard approach, distribution-based, which simulates the Dirichlet function across the total number of clients. All the details about the clients' partition can be found in Appendix A, and a more detailed evaluation of FedJETs under these distributions is included in Appendix B.

Reviewer comment (Novelty):
This work also very similar to multi-task learning, one of the main problems is conflicting gradients. Since you claim the data is non-i.i.d., have you ever encountered this problem as in these papers [3] [4]?

Author response:
**Regarding conflicting gradients**: 
During experiments, we discovered that anchor clients act as regularizers and help maintain consistency in the experts' updates when the gating function makes the wrong selection. The pretrained model - permanently frozen - serves as a feature extractor; these embeddings feed the gating function, ensuring that only the relevant experts are sent to each client and the expert updates are directed towards the same objectives, thus avoiding model drift. Additionally, there is an initial degree of randomness in the gating function. During the first couple of iterations, the gating function sends random top $K$ experts to each client while the experts learn to specialize in the different regions of the label space. We found a way to keep consistency during these initial rounds: through the anchor clients. By introducing at least 30\% anchor clients during each training round, we can ensure a balance between the wrong selection of the gating function by allowing them to act as regularizers in light of conflicting updates on the experts. Appendix C provides a detailed explanation of this phenomenon encountered during training.

Reviewer comment (Novelty):
My biggest concern is the novelty of the proposed method. The general framework of having individualized models and selecting a subset of experts for performing ensemble learning is a traditional topic. The specific setting of having a pre-trained model along with a gating function to select a subset of experts to update is new. I am not entirely familiar with the current federated learning literature, so I will leave other reviewers to decide on the novelty of the paper to the federated learning community.

Author response:
-**Regarding differences with ensemble learning**: There is a critical difference between ensemble learning and FedJETs: while the former combines multiple models to create a more accurate prediction during testing, FedJETS uses a weighted average per sample on all the selected $K$ experts, and utilizes the highest expert ranking score to fully capitalize on the specialization of a single expert at the sample level. Algorithm 1 (Testing) shows this is not a trivial adaptation of ensembles. Also, Appendix D provides further details on the gating function behavior during inference; if required, we can move some of this material into the main text to make ideas emerge more naturally. To the best of our knowledge, though, we are unaware of any works that learn mechanisms to select on-the-fly subsets of experts to be combined during inference; if there is such literature we are missing, we are eager to know.

Reviewer comment (Experiments):
In addition to the concern about the novelty of the work, another concern I have is the applicability of the method when expert models need to be very large. It seems to be inefficient to use the common expert (a large pre-trained model) to just perform expert selection. Would it be more reasonable, computation-wise at least, to not have individual expert models but different expert heads so that the pretrained common expert can be used to extract a common representation to pass into different experts?

Author response:
-**Regarding large-scale models, efficiency, and pretrained models**: First, we note that one of the main advantages of our method is that the ``pretrained common expert'' is considered a black box. This means that the architecture of the pretrained model is not strongly coupled with that on the expert's side. This allows us to leverage the knowledge of any pretrained model with capabilities over the different domains on the dataset. Further, this allows us to not rely on the model's full capacity; therefore, we assume this model is not extremely large, and neither has to be fully trained. Yet, it is noteworthy that during the early stages of our design, we experimented with the case where information (e.g., the whole model or part of the model, or even embeddings) from the pretrained model was available to be used along with the experts selected. Yet, we noticed that this often led to performance degradation, and the experts required further training to achieve some specialization.

Reviewer comment (Experiments):
Could the authors comment on the computation and memory costs of having individual experts? How big should the expert model be compared to the pretrained common expert model?

Author response:
-**Regarding computational/memory costs**: 
The architecture of the experts and that of the pretrained model are not necessarily coupled; instead, the pretrained model operates as a black box that decides which experts should be selected and activated. Further, the pretrained model does not require full training, as it is used as a feature extractor. Conversely, since the experts specialize in specific labels of the input space, the model size can be considered "small" compared to a monolithic model that covers all labels.

Reviewer comment (Experiments):
Other than using the pre-trained model for obtaining representation for the gating function, is it used in some other ways, e.g., is there a way to combine its output with the expert model?

Author response:
-**Regarding combining the expert model with that of pretrained model**:
While we do not exclude the possibility that the pretrained model can be somehow combined with expert models, as we mentioned above, we have observed a performance degradation in our initial attempts. This is an exciting research direction that we will be closely looking into in the near future.