FedJETs: Efficient Just-In-Time Personalization with Federated Mixture of Experts

We want to thank the reviewer for the detailed comments. Below, we respond to the points raised one by one. We hope our responses resolve further concerns and are available for other questions.

-Regarding lack of novelty: We understand your point of view; however, we believe our paper has a different focus than the references pointed out in the following ways:

-MMoe [1]: [1] proposes to learn task relationships by sharing ALL expert models across ALL tasks while using different gating functions that learn task similarities across all the experts. FedJET's main advantage is that it dynamically selects the most relevant experts (according to the client distribution) and only sends these models across the network, avoiding the communication costs of sending all the experts as FedMix. On the other hand, FedJETs uses a single gating function to learn the differences across different domains. At the same time, MMoe utilizes multiple gates (one per task) that learn how to model the different relationships across the experts and weigh them so they can be used differently.

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

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

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

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

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

Thanks for your response. I will keep my score.

We want to thank the reviewer for the detailed comments. We address each point raised and hope our responses will alleviate further concerns. We remain available to answer any additional questions you may have.

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

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

-Regarding expert/clients ratio: The experts in FedJETs are agnostic to the total number of clients. Experiments presented in Table 1 demonstrate the behavior of FedJETs for various $M$ values (total number of experts) with a fixed number of clients $K$. Each expert should have one "anchor" client that matches its specialization range to ensure that the gating function learns to assign clients to the right experts in the early stages of training. Appendix A, Figure 6, shows an example of how we randomly assigned specialization ranges to the experts. We hope this information clarifies the reviewer's concern.

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

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

We want to thank the reviewer for the detailed comments. Below, we respond to the points raised by the reviewer, one by one. We hope our responses will resolve any further concerns, and we are available for any other questions.

-Regarding more datasets: First, we agree with the reviewer that including more datasets can help to illustrate the differences among different domains better. For this reason, we provide new results on the EMNIST dataset (table below) using the same setup as Table 1 in the paper, which further demonstrates the advantage of our method; note that completing the existing experiments in the paper took a significant amount of GPU time, given the excessive ablation study results we present. Moreover, we point out that the CIFAR data suite is still highly relevant for current benchmarks on Federated Learning. Table 1 shows a noticeable difference in behavior between the two datasets; CIFAR10 has more difficulty reaching the desired performance due to its fewer classes, exacerbating the model drift problem in all baselines. FedJETs is the only method able to surpass the initial accuracy of the Common expert, whereas the rest only degraded during training.

-Regarding anchor clients ratio: The ablation study examines the importance of the ``anchor clients" during sampling, i.e., the clients that share the same expert for the target task. It shows that we need at least 30% of anchor clients to guarantee that our method can surpass the initial accuracy of the pretrained common expert. Still, Figure 5 also explores two other scenarios: $i)$ How does our approach perform if we randomly sample clients across the training iterations without any control over the expert distribution? And $ii)$ How does our method benefit from increasing the proportion of anchor clients to 50-50? While the latter ensures faster convergence of our method, it is not very realistic, as enforcing higher ratios of anchor clients in federated learning may affect the efficiency/deployment costs and privacy. Appendix B (Figure 9) shows the detailed performance of our method under $i$ and $ii)$ setups for both datasets.