FedGTST: Boosting Global Transferability of Federated Models via Statistics Tuning

We thank the reviewer for their comments. Our answers are as follows.

1. We additionally conducted experiments on DomainNet, and the results are presented below.

   Following FedSR, we applied a leave-one-out strategy, where one domain is treated as the target domain and the other five are source domains.

• The target domains listed in the table are: C (Clipart), I (Infograph), P (Painting), Q (Quickdraw), R (Real), and S (Sketch). |Model | C | I | P | Q | R | S | AVG | |--------|----------|-----------|----------|----------|----------|----------|------| | FedAVG | 59.3±0.7 | 16.5±0.9 | 44.2±0.7 | 10.8±1.8 | 57.2±0.8 | 49.8±0.4 | 39.6 | | FedSR | 61.0±0.6 | 18.6±0.4 | 45.2±0.5 | 13.4±0.6 | 57.6±0.2 | 51.8±0.3 | 41.3 | |FedGTST| 63.9±0.5 | 20.7±0.3 | 47.8±0.4 | 15.2±0.5 | 59.5±0.6 | 54.3±0.2 | 43.6 |

2. The proposed method is effective when the feature extractors are also optimized during FL.

   In our setting, the classifier is not the only component optimized during federated communication. In addition to the classifier, the feature extractor is also updated. As stated in Line 152, "Let $h^*=g^*\circ f^*$ be an optimal solution for the source objective 2", which indicates that both the feature extractor f and the classifier g are optimized during federated pretraining.

3. Our current setting indeed assumes multiple source domains.

   As mentioned in Lines 148–149, "...federated pretraining where the source domain is a union of the agents’ local domains," we use multiple domains for pretraining. Furthermore, in Lines 302–308 (Experiments), we explained the construction of the agents’ local domains: "We … to construct non-iid local domains". In Appendix G, Line 578, we reiterate, "For the source domains MNIST and CIFAR-10, we only allow each local client to have access to…".

   Additionally, in our response to the reviewer’s question about adding results on DomainNet, our setting also involves letting clients have access to different domains, enabling multiple source domains.

We appreciate the reviewer’s feedback, and we are happy to discuss further with the reviewer to address any remaining questions or concerns.

We thank the reviewer for their comments. Our answers are below.

1. Additional Baseline. We discuss Scaffold, a baseline suggested by the reviewer (and will cite this and other suggested works). They are as follows:

• 10 clients: | | MNIST → MNIST-M (LeNet)| MNIST → MNIST-M (ResNet) | CIFAR10 → SVHN (LeNet)| CIFAR10 → SVHN (ResNet) | Average| |--------|-------|---------|------|------|------| | Scaffold | 75.6 ± 0.8 | 80.8 ± 0.3| 66.0 ± 0.5 | 71.1 ± 0.4| 73.3 | | FedGTST | 76.2 ± 0.9 | 82.3 ± 0.5| 70.1 ± 0.8 | 74.5 ± 0.3| 75.8 |

• 100 clients: | | MNIST → MNIST-M (LeNet)| MNIST → MNIST-M (ResNet) | CIFAR10 → SVHN (LeNet)| CIFAR10 → SVHN (ResNet) | Average| |--------|-------|---------|------|------|------| | Scaffold | 52.3 ± 0.5 | 63.1 ± 0.3| 45.5 ± 0.1 | 55.5 ± 0.3| 54.1 | | FedGTST | 57.5 ± 0.3 | 67.6 ± 0.2| 52.4 ± 0.1 | 63.1 ± 0.2| 60.2 |

2. Intuitive explanations for Theorem 2

• Tuning Cross-Client Statistics Improves Transferability.

  • Larger $|J_p|$ in initial stages enhances transferability. By preventing small averaged gradients, we prevent clients from becoming trapped in individual's local minima or overfitting locally. As training progresses, the averaged gradient norm naturally decreases due to convergence.
  • Smaller $\sigma_p$ induces similarity across different domains. This helps avoid local overfitting and improves global transferability.

• Clarification on Gradient Norm Behavior.

  The RHS of the equation cannot be $-\infty$. The averaged gradient norm $\|J_p\|$ in term 2 of the RHS is bounded by the smoothness parameter $\alpha$. Specifically, the increment of local gradients is upper-bounded by $\|J_1 - J_2\| \leq \alpha \|x_1 - x_2\|$.

• Trade-Off Between Enlarging $\|J_p\|$ and Decreasing $\sigma$

  A larger $|J_p|$ leads to more substantial local model updates, increasing variance $\sigma_{p+1}$ in subsequent rounds. Thus, it's essential to balance this trade-off.

• Higher initial average Jacobian norms do not necessarily decrease generalization error.

  Let's begin by comparing the definitions of generalization error and transferability as defined in our manuscript.

  • Generalization error is defined as $L_{D_T}(h_{\text{pretrained}}) - L_{D_S}(h_{\text{pretrained}})$.

  • Transferability measure in our paper is $L_{D_T}(h_{\text{finetuned}})$.

  The key differences between these two concepts are: a) Whether $h$ is finetuned; b) Whether source error is subtracted.

  From Theorem 2, we can see that cross-client statistics help the transferability defined by us by controlling its upper bound, and we provide an intuitive explanation above.

  However, tuning of statistics does not necessarily decrease generalization error. Specifically:

  1. $L_{D_T}(h_{\text{pretrained}})$ is not necessarily smaller, since even if the pretrained feature extractor $f_{\text{pretrained}}$ transfers well to the target domain, the pretrained classifier $g_{\text{pretrained}}$, without any finetuning on the target domain, can result in a high $L_{D_T}(h_{\text{pretrained}})$.

  2. $L_{D_S}(h_{\text{pretrained}})$ is not necessarily larger, as tuning the statistics intuitively prevents clients from getting stuck in their individual local minima or overfitting, thus reducing cross-client source loss.

  To summarize, tuning the statistics may lead to a higher $L_{D_T}(h_{\text{pretrained}})$ and a lower $L_{D_S}(h_{\text{pretrained}})$, and therefore, does not decrease generalization error.

3. Assumptions Clarification

   • The single-step assumption is commonly used in FL theoretical analyses - see FedSplit [2] and FEDAC [3].

     [2] Pathak, Reese, and Martin J. Wainwright. "FedSplit: An algorithmic framework for fast federated optimization." Advances in neural information processing systems 33 (2020): 7057-7066.

     [3] Yuan, Honglin, and Tengyu Ma. "Federated accelerated stochastic gradient descent." Advances in Neural Information Processing Systems 33 (2020): 5332-5344.

   • GD assumption

     The reviewer's point is correct. However, GD can be extended to stochastic learning with batch sampling. The additional randomness in sampling would require incorporating the variance of batch sampling into the generalization bound. This variance term, being independent of the algorithm design, was omitted in our theoretical analysis.

4. Minor clarifications

   • We appreciate the reviewer pointing out typo issues, and will make sure to fix them.

   • Convexity assumptions are solely made for theoretical tractability, and such simplifications were also considered in SCAFFOLD [1], FedSplit [2], etc.

     [1] Karimireddy, Sai Praneeth, et al. "Scaffold: Stochastic controlled averaging for federated learning." International conference on machine learning. PMLR, 2020.

     [2] Pathak, Reese, and Martin J. Wainwright. "FedSplit: An algorithmic framework for fast federated optimization." Advances in neural information processing systems 33 (2020): 7057-7066.

5. By epoch in experiments, we mean step. In the theoretical part of the paper, we dealt with steps, while in practical evaluations, we referred to epoch. Sorry for the confusion, we will clarify this in the revision.

6. We use SGD for experimental evaluation.

7. Ablation Study: removing the surrogate Jacobian norm.

   What the reviewer is asking is part of our additional baseline cross-examination involving Scaffold, which uses Jacobian regularization but does not consider surrogate Jacobian norm. Please see the results presented in the first bullet point.

We appreciate the reviewer’s feedback, and we are happy to discuss further to clarify any remaining questions.

We appreciate the reviewer's insightful suggestion on showcasing examples where improving transferability might reasonably hurt generalization, or vice versa.

Below, we provide: 1) some current results that show that Fed-GTST improves transferability at a slight cost of source domain generalization; and 2) reference to existing literature related to cases where promoting transferability may not always "align" with generalization quality.

We report results in two columns: a) Generalization: averaged accuracy evaluated on the source testing domains, where the model is pre-trained on source training domains, and b) Transferability: the accuracy on the target domain, where the model is fine-tuned on the target domain after being pre-trained on the source training domains.

• For MNIST → MNIST-M (100 clients, local models based on ResNet-18):

• For CIFAR10 → SVHN (100 clients, local models based on ResNet-18):

We observe that the transferable FL method (FedGTST) demonstrates better transferability (second column), while its generalization performance (first column) is comparable to or slightly lower than that of the heterogeneous FL method (e.g., Scaffold). Conversely, Scaffold is optimized for strong generalization on source test domains but is outperformed by Fed-GTST in terms of target domain performance (which we use to measure transferability).

Additionally, adversarially robust models offer another example where improving transferability may come at the expense of generalization. As shown in [1], adversarially robust models tend to have better transferability. However, since these models are designed to perform well against adversarial manipulations or perturbations, they do not necessarily exhibit lower generalization error on "clean" test data. In this scenario, transferability may not have any influence on generalization or even work against it.

[1] Salman, H., Ilyas, A., Engstrom, L., Kapoor, A., and Madry, A. "Do Adversarially Robust ImageNet Models Transfer Better?" arXiv preprint arXiv:2007.08489, 2020.

We thank the reviewer for their comments. Our answers are as follows.

1. We agree with the reviewer that there are other non iid sampling methods that may be more frequently used than ours. We therefore implemented the Dirichlet sampling. Following the reviewer’s suggested references, we set the concentration parameter to 0.5. The results clearly show that our method still outperforms others even when client sampling strategies change.

• Results for 10 clients: | Method | MNIST → MNIST-M (LeNet)| MNIST → MNIST-M (ResNet) | CIFAR10 → SVHN (LeNet) | CIFAR10 → SVHN (ResNet) | Average | |--------------|--------------|-------------|--------------|-------------|--------| | FedAVG | 74.1±0.6 | 82.2±0.3 | 65.2±0.4 | 72.8±0.9 | 73.5 | | FedSR | 75.5±0.8 | 82.0±0.2 | 66.1±0.5 | 72.1±0.3 | 73.9 | | FedIIR | 75.8±0.2 | 82.8±0.6 | 66.5±0.9 | 74.6±0.1 | 74.9 | | FedGTST| 77.3±0.8 | 82.7±0.4 | 70.8.0±0.7 | 75.2±0.2 |76.5 |

• Results for 100 clients: | Method | MNIST → MNIST-M (LeNet)| MNIST → MNIST-M (ResNet) | CIFAR10 → SVHN (LeNet) | CIFAR10 → SVHN (ResNet) | Average | |--------------|--------------|-------------|--------------|-------------|--------| | FedAVG | 50.4±0.1 | 63.0±0.3 | 43.3±0.2 | 54.6±0.5 | 52.8 | | FedSR | 52.7±0.2 | 62.9±0.3 | 44.7±0.1 | 56.5±0.3 | 54.2 | | FedIIR | 54.2±0.4 | 64.1±0.1 | 47.4±0.4 | 58.7±0.2 | 56.1 | | FedGTST | 59.5±0.3 | 69.2±0.2 | 55.1±0.1 | 65.6±0.2 | 62.4 |

2. Thank you for your suggestion. We will add a discussion of [7] into the revision as follows. A recent work, FedImpro [7], considers the domain generalization problem over different clients in the FL setting. Specifically, FedImpro assumes that the server can maintain an estimate of the aggregated feature distribution, which is used as an extra term in the client gradient computation to promote lower gradient dissimilarity and thus better generalization ability. Unlike our approach, FedImpro does not use cross-client statistics explicitly during client updates.

3. We would like to point out that the hyperparameters are described in Lines 321-333. If the reviewer would like additional information, please do let us know.

4. The optimal learning rate in assumption 4.3 can be directly computed in practice. An assumption that we made in 4.3 is that we can know the exact value of $\alpha$, as all other terms can be computed explicitly from the client statistics of round p-1. The assumption is only used for theoretical analysis. In practice we can always estimate $\alpha$. This type of assumption is also adopted in other papers, such as [Haddadpour, 19].

   [Haddadpour, 19] Farzin Haddadpour and Mehrdad Mahdavi. On the convergence of local descent methods in federated learning. arXiv preprint arXiv:1910.14425, 2019.

5. We appreciate the reviewer’s suggestion. A detailed conceptual explanation of Theorem 2 is follows:

• Tuning Cross-Client Statistics Improves Transferability:

  • Large $\|J_p\|$ Norm: A large $\|J_p\|$ intuitively enhances transferability. In the early stages of training, it is crucial to prevent the averaged gradient from becoming too small. Small local gradients can cause clients to become trapped in local minima or overfit. As training progresses toward its final stages, the averaged gradient norm will naturally decrease due to convergence.
  • Small Cross-Client Gradient Variance: A small cross-client gradient variance improves transferability by inducing similarity across different domains. This helps avoid local overfitting and promotes better global transferability.

• Clarification on Gradient Norm Behavior: It is impossible for the right-hand side (RHS) of the equation to be $-\infty$. The averaged gradient norm $\|J_p\|$ in term 2 of the RHS cannot approach $\infty$ because it is bounded by the smoothness parameter $\alpha$. Specifically, the increment of local gradients is upper-bounded by $\|J_1 - J_2\| \leq \alpha \|x_1 - x_2\|$.

• Trade-Off Between Enlarging $\|J_p\|$ and Decreasing $\sigma$: A larger $\|J_p\|$ results in more substantial updates of the local models, which can intuitively lead to an increased variance $\sigma_{p+1}$ in subsequent rounds. Therefore, Theorem 2 highlights that when implementing the approach, it is essential to balance this trade-off—i.e., increase $\|J_p\|$ while keeping $\sigma$ small.

We appreciate the reviewer’s feedback, and we are happy to discuss further with the reviewer to clarify any remaining questions or concerns.

We appreciate the reviewer's valuable feedback once again and thank them for reviewing our rebuttal response.

We thank the reviewer for their comments. Our answers are as below.

• Experimental Setting:

  1. Federated Pre-training was explained in the original manuscript. In Lines 148–149, we said: “...federated pretraining where the source domain is a union of the agents’ local domains.” Hence, we indeed used federated pretraining with the source domain comprising data from all local domains. In Appendix G, Line 579, we reiterated the claim: “For the source domains MNIST and CIFAR-10, we only allow each local client to have access to…”
  2. Target domain data or labels are available. Target domain data and labels were indeed available, as stated in Lines 147 and 151. One can see from Equations (1) and (3) that the target loss is defined based both on the data and labels.
  3. Yes, we are indeed performing inductive transfer learning.

• Additional Dataset. We additionally run experiments on DomainNet, as suggested by another reviewer, to accommodate both requests. DomainNet is a widely used large dataset for both transfer learning and transferable federated learning, such as FedSR. Following FedSR we applied a leave-one-out strategy, where one domain is treated as the target domain and the other five are source domains. The target domains listed in the table are: C (Clipart), I (Infograph), P (Painting), Q (Quickdraw), R (Real), and S (Sketch). The results confirm that FedGTST still outperforms other methods from the literature. | Model | C | I | P | Q | R | S | AVG | |-|-|-|-|-|-|-|-| | FedAVG | 59.3±0.7 | 16.5±0.9 | 44.2±0.7 | 10.8±1.8 | 57.2±0.8 | 49.8±0.4 | 39.6 | | FedSR | 61.0±0.6 | 18.6±0.4 | 45.2±0.5 | 13.4±0.6 | 57.6±0.2 | 51.8±0.3 | 41.3 | | FedGTST | 63.9±0.5 | 20.7±0.3 | 47.8±0.4 | 15.2±0.5 | 59.5±0.6 | 54.3±0.2 | 43.6 |

• Baseline clarification.

  1. We indeed considered adaptations of existing transfer learning approaches to federated learning: FedSR, one of the baselines in our comparative study, is a direct adaptation of transfer learning. We will explain how FedSR performs this adaptation in the revision.
  2. FedSR also represents an example of using L2 norm regularization on mapped features. Also, we would appreciate it if the reviewer could clarify what reference [1] is, as we do not believe it corresponds to the reference in our submission. Other examples of using L2 regularizers in transfer learning include Li et al., Explicit inductive bias for transfer learning with CNs, etc.

• Transferability measure, theoretical bound and cross-client statistics.

  • Theorem: As in Lines 157-158, we use the optimal target loss as the transferability performance measure. More precisely, in Theorem 2, Line 251, the LHS represents the actual transferability measure. Proper tuning of the cross-client statistics influences transferability by controlling the RHS, which is an upper bound on the transferability metric on the LHS.

  • Intuition: Intuitively, a large averaged Jacobian norm in the early stages helps prevent clients from getting stuck in local minima or local overfitting, while a small cross-client Jacobian variance promotes similarity across domains, both of which enhance global transferability. Further clarifications regarding the Jacobian norm are: 1) the increment of the Jacobian is controlled as $\|J_1 - J_2\| \leq \alpha \|x_1 - x_2\|$, ensuring that the second term on the RHS of Theorem 2 cannot be $\infty$; 2) An average Jacobian norm will naturally decrease in the final stages due to convergence.

• We appreciate the reviewer pointing out the typos, phrasing and reorganization issues. We will correct these errors in the revised manuscript.

• Regarding the README file, we sincerely apologize for this oversight. This was an unfortunate copy-paste error that we can easily correct since we always had a ready-to-use instruction manual for the code. Due to space limitation, we can only repeat some parts of our intended README file below:

  Sample run command (i.e., 100 clients with resnet18):

  python main.py --num_clients 100 –back_bone resnet18

  Our evaluation metric is implemented in the code contained in main.py, Line 207.

• Please see the discussion below regarding gradient matching in transfer learning, and challenges of applying it directly to federated learning.

  • Discussion: Gradient matching in transfer learning are approaches of aligning gradients between source and target domains. The work Gradient Matching for Domain Generalization explores how matching gradients between the source and target domains can lead to better domain generalization by ensuring that the learned representations are robust to domain shifts. Fishr: Invariant Gradient Variances for Out-of-Distribution Generalization extends this idea by focusing on invariant gradient variances, which helps in maintaining generalization performance when faced with out-of-distribution samples. Understanding Hessian Alignment for Domain Generalization investigates the role of Hessian alignment in gradient-based methods, highlighting how aligning the Hessian matrices of the source and target domains can further improve generalization performance. All these and other suggested works will be cited in the revised manuscript.
  • Challenges: 1) Privacy leakage. Such approaches violate data privacy because source clients are given access to the target domain. 2) Local overfitting. Clients perform training solely on their local models, representations, and labels, thus intuitively end up overfitting in the local domain.
  • Our contribution: We ensure that our communication schemes prevent uncontrolled access to data, thereby maintaining data privacy, and that our approach enhances global transferability rather than focusing solely on local transferability.

We appreciate the reviewer’s feedback, and we are happy to provide additional clarifications.