Initialization Matters: Unraveling the Impact of Pre-Training on Federated Learning

Thank you for appreciating the rigor of our analysis, extension of Proposition 1 to the federated setting and experimental results. We address the weaknesses and questions below.

Re:Weaknesses

1. Focus on two-layer CNN can limit the broader impact of the findings

The key finding in our analysis is that learning the signal (for misaligned filters) becomes more challenging as data heterogeneity increases. However, this challenge can be mitigated to a large extent by starting from a pre-trained model where most filters are already aligned with the signal. We believe this finding can be extended to deeper CNN architectures as well. Intuitively, while deeper architectures provide the capacity to learn more complex hierarchical signals, aligning a filter with the signal will also progressively become harder across layers with the alignment of the L-th layer potentially depending on the alignment of the (L−1)-th layer. In non-IID FL, this suggests that feature learning deteriorates as we move from the base layer toward the final layer. This intuition is supported by empirical evidence (see experiment in Table 1 in [a]) and explains why pre-trained initialization offers greater benefits for deeper networks compared to shallower ones. We leave a more rigorous analysis with formal guarantees as future work.

2. The definition of data heterogeneity implies a binary classification problem. This means that for a multi-class problem, we would need to cast it into multiple binary classification problems to apply the data heterogeneity measure, which is less practical.

While modeling multi-class classification as multiple binary classification is one approach, we believe there can be better alternatives by changing the data generation model and subsequent theory to explicitly incorporate multi-class classification. In essence, the definition of data heterogeneity in Eq. (1) aims to capture how far the label distribution at each client is from the the global label distribution, which assigns equal probability to each label.

3. The assumption of the noise vector being orthogonal to the signal can hardly hold in more complicated datasets.

We assume orthogonality of signal and noise vectors for simplicity of analysis. This assumption can be easily relaxed as done in [b]. Our theoretical insights will remain the same with the only difference being that we need a slightly stronger condition on the dimension of the filters (Condition C2).

4. Are there any insights into determining what types of pre-trained weights are best suited for what particular tasks?

Our analysis in Section 3.4 shows that the smaller the $\ell_2$ distance between the signal in the pre-training data $\mu^{\text{pre}}$ and the signal in the fine-tuning data $\mu$, the higher will be the alignment of the pre-trained weights. Intuitively, this suggests that pre-trained weights derived from data similar to the fine-tuning task will yield better performance. An interesting direction for future work is to see if we can use empirical measures of alignment to rank the suitability of different pre-trained models for a given task.

5. In what situations might the theory not hold?

Our theory assumes a setting where data can be decomposed into signal and noise components, which is usually applicable to vision tasks. However, this assumption is less straightforward for language tasks, where decomposing data into signal and noise is more challenging, and pre-training is typically done autoregressively on text sequences. Therefore, we believe that explaining the benefits of pre-training for language models, such as Transformers, requires a different theoretical model. We are happy to discuss this in more detail if the reviewer has other situations in mind.

Re:Questions

1. The definitions of alignment and misalignment in Definition 1 are very interesting. Can this be quantified, for example, in an experimental setup where the data is MNIST, the model is VGG, and the task is odd and even number classification?

We note that even in this simple setup we cannot explicitly characterize the signal information being learned by every filter in the VGG model and hence cannot directly use Definition 1 to measure alignment. Nonetheless we can use our empirical definition of alignment introduced in Eq. (15) to measure alignment. The results are summarized in the table below, with additional experimental details and plots provided in Appendix F 2.4.

[Table 1: Results on training VGG11 on MNIST to classify odd and even digits with different initializations]

We observe that the percentage of misaligned filters for random initialization in this task is lower compared to our experiment on CIFAR-10 in Fig. 5 where it was around $25\\%$. Intuitively, this suggests that even random features generated by deep CNNs are sufficient to achieve reasonably good test accuracy on MNIST. Nonetheless, pre-trained initialization still achieves higher accuracy, as it results in a lower percentage of misaligned filters.

2. Could the authors elaborate further on Figure 4, as I do not see a clear trend over r? Additionally, the green color appears to be missing in Figures 4a and 4d, or is it on top of other colors?

Thank you for pointing this out. We have now modified the experiment setting and plots to make the trend over $r$ more clear. In the IID case, we see that both signal learning (Fig. 4a) and noise memorization (Fig. 4b) is similar across all $r$ and grows approximately linearly with the number of local steps. In the non-IID case however while noise memorization (Fig. 4e) grows linearly for all $r$, signal learning (Fig. 4b) saturates for misaligned filters. This in turn affects the ratio of signal learning to noise memorization which is almost constant across all $r$ in the IID case (Fig. 4e) but decays with the number of local steps in the non-IID case for misaligned filters as shown in Eq. (10). Thus Figure 4 highlights the importance of filter alignment for non-IID FL and verifies our theoretical results.

Thank you again for your review. We are happy to answer any further questions that you may have. If your concerns have been resolved, we would be really grateful if you could consider increasing your score.

References

[a] Yu, Yaodong, et al. "TCT: Convexifying federated learning using bootstrapped neural tangent kernels." Advances in Neural Information Processing Systems 35 (2022): 30882-30897.

[b] Kou, Yiwen, et al. "Benign overfitting in two-layer ReLU convolutional neural networks." International Conference on Machine Learning. PMLR, 2023.

Experiments on FL setup with varying levels of data heterogeneity:

We extend the experiment from Figure 5 of our paper, originally conducted on CIFAR-10 with $\alpha = 0.1$ Dirichlet heterogeneity to three other levels of heterogeneity: $\alpha = 0.05$ (high heterogeneity), $\alpha = 0.3$ (medium heterogeneity) and $\alpha = 10$ (low heterogeneity). The results are summarized in the table below, with additional experimental details and plots provided in Appendix F 2.2.

[Table 2: Results on training ResNet-18 with FedAvg with different heterogeneities on CIFAR-10]

First, we observe that the percentage of misaligned filters remains approximately $25\\%$ with random initialization and $10\\%$ with pre-trained initialization, regardless of the level of heterogeneity. However, as heterogeneity increases, the improvement in test accuracy provided by pre-trained initialization becomes more pronounced. This trend is consistent with our theoretical analysis in Theorem 2, which suggests that the percentage of misaligned filters will have a greater impact on test performance as data heterogeneity increases.

Thank you again for your review. We are happy to answer any further questions that you may have. If your concerns have been resolved, we would be really grateful if you could consider increasing your score.

References

[a] Yu, Yaodong, et al. "TCT: Convexifying federated learning using bootstrapped neural tangent kernels." Advances in Neural Information Processing Systems 35 (2022): 30882-30897.

[b] Le, Ya, and Xuan Yang. "Tiny imagenet visual recognition challenge." CS 231N 7.7 (2015): 3.

[c] Weyand, Tobias, et al. "Google landmarks dataset v2-a large-scale benchmark for instance-level recognition and retrieval." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020.

[d] Karimireddy, Sai Praneeth, et al. "Scaffold: Stochastic controlled averaging for federated learning." International Conference on Machine Learning. PMLR, 2020.

[e] Wang, Jianyu, et al. "Tackling the objective inconsistency problem in heterogeneous federated optimization." Advances in Neural Information Processing Systems 33 (2020): 7611-7623.

• C6) Learning rate is sufficiently small: This is a standard condition to ensure that gradient descent does not diverge. The conditions are derived from ensuring that the signal and noise coefficients remain bounded in the first stage of training and that the loss decreases monotonically in every round in the second stage of training.

We have added this discussion to Appendix C where we discuss our proof sketch. We note that all these assumptions have appeared in previous work (references Cao et al. (2022); Kou et al. (2023) in the paper) and are standard in the analysis of two-layer CNN. Finally, we also clarify that these conditions do not contradict each other and can be simultaneously satisfied. The relationship between C1 (restricted number of updates) and C6 (small learning rate) is not contradictory since C1 has a factor of $1/\eta$ in its definition. So if we make the learning rate $\eta$ smaller, we are also allowed a larger number of updates.

3. I would like to see more discussion on how results might apply to real-world datasets and FL setups with varying levels of data heterogeneity

Thank you for the suggestion. We have added more experiments on real-world datasets and FL setups with varying levels of data heterogeneity as we discuss below.

Experiments on real-world datasets:

We extend the experiment from Figure 5 of our paper, originally conducted on CIFAR-10, to evaluate the number of misaligned filters at initialization, on more challenging and real-world datasets which include

a) TinyImageNet [b]: $100k$ datapoints, $200$ classes, data partitioned across $20$ clients with $\alpha = 0.3$ heterogeneity

b) Google Landmarks v2 23k [c]: $23k$ datapoints, $203$ classes, $233$ clients, data naturally grouped by photographer to achieve a federated partitioning

The results are summarized in the table below, with additional experimental details and plots provided in Appendix F 2.1.

[Table 1: Results on training ResNet-18 with FedAvg across various datasets]

Observe that as the complexity of the signal in the data increases (CIFAR-10 < TinyImageNet < Google Landmarks), we see a corresponding sharp increase in the percentage of misaligned filters ($25\\%$ to $40\\%$). In contrast, with pre-trained initialization, the percentage of misaligned filters remains less than $15\\%$ across datasets leading to a larger improvement in test accuracy for harder datasets. Thus results align well with our theoretical findings: as the percentage of misaligned filters increases, the benefits of pre-training become more pronounced.

Thank you for appreciating our detailed theoretical analysis, intuitive understanding of results and experimental support. We address the weaknesses below.

1. The analysis relies on a simplified two-layer CNN, raising concerns about the transferability of the derived bounds to more complex architectures

The key finding in our analysis is that learning the signal (for misaligned filters) becomes more challenging as data heterogeneity increases. However, this challenge can be mitigated to a large extent by starting from a pre-trained model where the filters are already aligned with the signal. We believe this finding can be extended to deeper CNN architectures as well. Intuitively, while deeper architectures provide the capacity to learn more complex hierarchical signals, aligning a filter with the signal will also progressively become harder across layers with the alignment of the L-th layer potentially depending on the alignment of the (L−1)-th layer. In non-IID FL, this suggests that feature learning deteriorates as we move from the base layer toward the final layer. This intuition is supported by empirical evidence (see experiment in Table 1 in [a]) and explains why pre-trained initialization offers greater benefits for deeper networks compared to shallower ones. We leave a more rigorous analysis with formal guarantees as future work.

2. My main concern is on the unclear assumptions and rationale

The assumptions are primarily used to ensure that the model is sufficiently overparameterized, i.e., training loss can be made arbitrarily small, and that we do not begin optimization from a point where the gradient is already zero or unbounded. We provide a more intuitive reasoning behind each of the assumptions below:

• C1) Bounded number of communication rounds: This is needed to ensure that the magnitude of filter weights remains bounded throughout training since they grow logarithmically with the number of updates (Theorem 3). We note that this is quite a mild condition since the max rounds can have polynomial dependence on $1/\epsilon$ where $\epsilon$ is our desired training error.

• C2) Dimension $d$ is sufficiently large: This is needed to ensure that the model is sufficiently overparameterized and the training loss can be made arbitrarily small. Recall from our data generation model in Section 2 that our input $\mathbf{x}$ consists of a signal component $\mu \in \mathbb{R}^d$ that is common across all datapoints and noise component $\xi \in \mathbb{R}^d$ that is independently drawn from $\mathcal{N}(0,\sigma_p^2 \cdot \mathbf{I})$. Having a sufficiently large $d$ ensures that the correlation between any two noise vectors, i.e. $\langle \mathbf{\xi}, \mathbf{\xi}' \rangle/ \|\| \mathbf{\xi}\|\|^2$ is not too large (Lemma 4). Otherwise if the correlation between two noise vectors is large and negative, then minimizing the loss on one data point could end up increasing the loss on another training point which complicates convergence and prevents loss from becoming arbitrarily small.

• C3) Training set size and network width is sufficiently large: The condition ensures that a sufficient number of filters get activated at initialization with high probability (Lemma 6 and Lemma 7) and prevents cases where the initial gradient is zero. The condition on training set size also ensures that there are a sufficient number of datapoints with negative and positive labels (Lemma 8).

• C4) Standard deviation of Gaussian random initialization is sufficiently small: This condition is needed to ensure that the magnitude of the initial correlation between the filter weights and the signal and noise components, i.e $|\langle \mathbf{w}\_{j,r}^{(0)}, \mathbf{\mu} \rangle|$, $|\langle \mathbf{w}\_{j,r}^{(0)}, \xi \rangle|$ is not too large. This simplifies the analysis and prevents cases where none of the filters get activated at initialization (Lemma 21). It also ensures that after some number of rounds all filters get aligned with the signal (Lemma 30).

• C5) Norm of signal is larger than noise variance: This condition is needed to ensure that all misaligned filters at initialization eventually become aligned with the signal after some rounds (Lemma 30). This allows us to derive a meaningful bound on test performance that is not dominated by noise memorization.

I would like to thank the authors for the detailed response. I have raised my positive score from 6 to 8.

3. Explicitly clarify whether your theorems apply to centralized settings, federated learning (FL) settings, or both. Provide clear reasoning to support each case.

Theorems 1 and 2 are applicable to both centralized and federated learning (FL) settings, as the centralized setting can be viewed as a special case of FL where clients perform only a single local step (i.e., $\tau = 1$). However, the effects of data heterogeneity and misaligned filters on the test error bound in Theorem 2 become evident only when $\tau > 1$, which is the typical scenario in FL. Thus Theorem 2 highlights that initialization plays a more critical role in federated settings compared to centralized ones, as also seen in practice.

4. In line 357, the authors mention: "We focus on centralized pre-training, but our discussion here can be extended to federated pre-training as well." It is unclear why this extension holds.

The key insight in our discussion is that misaligned filters eventually align with the signal in both centralized and FL settings. In FL, however, the number of rounds taken for all the filters to get aligned will increase as the data heterogeneity increases. This is demonstrated empirically by our new experiments in Appendix F 2.2 where we extend the experiment from Figure 5 of our paper, originally conducted on CIFAR-10 with $\alpha = 0.1$ Dirichlet heterogeneity to three other levels of heterogeneity: $\alpha = 0.05$ (high heterogenity), $\alpha = 0.3$ (medium heterogeneity) and $\alpha = 10$ (low heterogeneity).

Since since pre-training is usually done on a public dataset, we focus on centralized pre-training in the discussion in Section 3.4. As such, Lemma 2 gives an upper bound on the number of iterations needed for all filters to get aligned in the centralized setting. As part of our proof for Theorem 2 we also derive the federated counterpart of this lemma- Lemma 30 in Appendix B - which gives an upper bound on the the number of rounds required for all filters to get aligned in the FL setting. We have now added an explanation after line 357 to clarify this extension.

5. Minor Issues

Thank you for the careful proofreading! We have fixed the points mentioned in points (a), (c), (d). We also clarify that the statement in Nguyen (2022) refers to an FL setup.

Thank you again for your review. We are happy to answer any further questions that you may have.

References

[a] Yu, Yaodong, et al. "TCT: Convexifying federated learning using bootstrapped neural tangent kernels." Advances in Neural Information Processing Systems 35 (2022): 30882-30897.

[b] Le, Ya, and Xuan Yang. "Tiny imagenet visual recognition challenge." CS 231N 7.7 (2015): 3.

[c] Weyand, Tobias, et al. "Google landmarks dataset v2-a large-scale benchmark for instance-level recognition and retrieval." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020.

Thank you for appreciating the clarity of our writing and novel contributions. We address the weaknesses and questions below.

1. I would appreciate if the authors could discuss how the theorems introduced in the paper might change under generalization to deeper CNN architectures

The key finding in our analysis is that learning the signal (for misaligned filters) becomes more challenging as data heterogeneity increases. However, this challenge can be mitigated to a large extent by starting from a pre-trained model where the filters are already aligned with the signal. We believe this finding can be extended to deeper CNN architectures as well. Intuitively, while deeper architectures provide the capacity to learn more complex hierarchical signals, aligning a filter with the signal will also progressively become harder across layers with the alignment of the L-th layer potentially depending on the alignment of the (L−1)-th layer. In non-IID FL, this suggests that feature learning deteriorates as we move from the base layer toward the final layer. This intuition is supported by empirical evidence (see experiment in Table 1 in [a]) and explains why pre-trained initialization offers greater benefits for deeper networks compared to shallower ones. We leave a more rigorous analysis with formal guarantees as future work.

2. Additional experiments on larger datasets would significantly strengthen the paper’s contributions.

We extend the experiment from Figure 5 of our paper, originally conducted on CIFAR-10, to evaluate the number of misaligned filters at initialization, on more challenging and real-world datasets which include

a) TinyImageNet [b]: $100k$ datapoints, $200$ classes, data partitioned across $20$ clients with $\alpha = 0.3$ heterogeneity

b) Google Landmarks v2 23k [b]: $23k$ datapoints, $203$ classes, $233$ clients, data naturally grouped by photographer to achieve a federated partitioning

The results are summarized in the table below, with additional experimental details and plots provided in Appendix F 2.1.

[Table 1: Results on training ResNet-18 with FedAvg across various datasets]

Observe that as the complexity of the signal in the data increases (CIFAR-10 < TinyImagenet < Google Landmarks), we see corresponding a sharp increase in the percentage of misaligned filters ($25\\%$ to $40\\%$). In contrast, with pre-trained initialization, the percentage of misaligned filters remains less than $15\\%$ across datasets leading to a larger improvement in test accuracy for harder datasets. These results align well with our theoretical findings: as the percentage of misaligned filters increases, the benefits of pre-training become more pronounced.

2. It seems to me that the presence of aligned or misaligned filters alone does not fully explain why a model benefits from pre-training in heterogeneous setting. Could the authors comment on this point?

Great question! Our analysis for the test performance of the CNN model at round $T$ ultimately depends on bounding the ratio of the signal learning to noise memorization across all the filters. Noise memorization is independent of initialization as we discuss in Section 4 (see Eq. (14)). For signal learning, (measured by $\\langle \\mathbf{w}_{j,r}^{(T)} , j\mathbf{\mu} \\rangle$), we have the following decomposition

$\langle \mathbf{w}\_{j,r}^{(T)}, j\mathbf{\mu} \rangle = \langle \mathbf{w}\_{j,r}^{(0)}, j\mathbf{\mu} \rangle \hspace{10pt} + \underbrace{\Gamma\_{j,r}^{(0, T)}}_{\text{Signal learning coefficient starting from round 0}}.$

Based on our conditions for model initialization in Section 3.3, we have that the initial correlation between the signal and filter weights at round $0$ will not be very large, i.e, $\langle \mathbf{w}\_{j,r}^{(0)} , j\mathbf{\mu} \rangle \approx 0$. Intuitively this implies that the model’s test performance begins at a near-random baseline, which is true even for pre-trained initialization. In this case we only care about the sign of $\langle \mathbf{w}\_{j,r}^{(0)} , j\mathbf{\mu} \rangle$, i.e. alignment, since $\Gamma\_{j,r}^{(0, T)}$ grows much faster for aligned filters as shown in Eqs. (12) and (13).

On the other hand, if we consider the same analysis starting from an intermediate round $T' < T$, then

$\langle \mathbf{w}\_{j,r}^{(T)}, j\mathbf{\mu} \rangle = \langle \mathbf{w}\_{j,r}^{(T')}, j\mathbf{\mu} \rangle \hspace{10pt} + \underbrace{\Gamma\_{j,r}^{(T', T)}}_{\text{Signal Learning Coefficient starting from round $T'$}}.$

In this case, in addition to the signal learning coefficient we also have to account for the magnitude of filter-signal correlation at initialization, i.e. $\langle \mathbf{w}\_{j,r}^{(T')}, j\mathbf{\mu} \rangle \gg 0$, since the model has now trained on the signal for some rounds. For large enough $T'$, while the number of aligned filters will be same for pre-trained and random initialization (growth of $\Gamma\_{j,r}^{(T', T)}$ will be similar), we can show that the initial magnitude of signal learning across all the filters i.e, $\sum\_{j,r} \langle \mathbf{w}\_{j,r}^{(T')}, j\mathbf{\mu} \rangle$ will be much higher for pre-trained initialization since the filters get aligned faster in this case, as also seen in Figure 5 left plot. Therefore even though all the filters may be aligned at round $T'$, pre-trained initialization will continue to have a higher value of signal learning for all subsequent rounds leading to better accuracy compared to random initialization.

Thank you again for your review. We are happy to answer any further questions that you may have. If your concerns have been resolved, we would be really grateful if you could consider increasing your score.

References:

[a] Song, C., Granqvist, F., & Talwar, K. (2022). Flair: Federated learning annotated image repository. Advances in Neural Information Processing Systems, 35, 37792-37805.

[b] Van Horn, Grant, et al. "The inaturalist species classification and detection dataset." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.

[c] Hsu, Tzu-Ming Harry, Hang Qi, and Matthew Brown. "Measuring the effects of non-identical data distribution for federated visual classification." arXiv preprint arXiv:1909.06335 (2019).

[d] Wang, Jianyu, et al. "On the unreasonable effectiveness of federated averaging with heterogeneous data." Transactions of Machine Learning Research (2024).

[e] Venkateswara, Hemanth, et al. "Deep hashing network for unsupervised domain adaptation." Proceedings of the IEEE conference on computer vision and pattern recognition. 2017.

[f] Nguyen, John, et al. "Where to begin? On the impact of pre-training and initialization in federated learning." International Conference of Learning Representations (2023).

Thank you for appreciating the clarity of our paper, notion of aligned/misaligned filters and results regarding the harmful effect of local steps and heterogeneity on test error. We address the weaknesses below.

1. On measuring data heterogeneity....could the authors explain their choice of measurement of data heterogeneity?

We believe the reviewer here is referring to heterogeneity in domain space, i.e., while two clients can have the same label $y$, the corresponding images or $x$ values can be quite different. While domain heterogeneity can exist in FL systems, we find that label heterogeneity is the more prevalent source of heterogeneity in FL systems. This can be seen both in real world FL datasets such as Flair[a], iNaturalist[b] as well techniques to simulate heterogeneity from centralized datasets such as restricting client data to certain labels (as done in [B]) or the Dirichlet sampling technique [c].

At the same time, recent works [d] have shown that the definition of data heterogeneity appearing in [A,B] is quite pessimistic in practice since they require the difference between local and global gradients to be bounded over the entire parameter space and do not accurately reflect heterogeneity in real world FL training. In this aspect, we consider our definition of heterogeneity to be a more transparent and realistic measure of heterogeneity in FL systems.

Lastly, we also conduct the following experiment to show that while heterogeneity in domain space can cause the gradients to become diverse, it does not seem to affect test performance significantly. To simulate domain heterogeneity, we consider the Office-Home dataset [e] which consists of images of $65$ objects in $4$ different domains - Art, Clipart, Product and Real World. Each domain has around $20-100$ images of every object. We split the data across $4$ clients in the following ways:

1. IID: Data across all domains is split IID across clients, i.e., each client will have images corresponding to every domain and every label.
2. Domain Heterogeneity: Each client only has images corresponding to a single domain
3. Label Heterogeneity: Data is split across clients with $\alpha = 0.1$ Dirichlet label heterogeneity, i.e, each client will have images corresponding to all domains but only certain labels.

[Table 1: Results on training a randomly initialized ResNet18 using FedAvg for the 3 types of data heterogeneity. Additional details can be found in Appendix F.]

The table above shows while gradient diversity in the domain heterogeneity setting is higher than in the IID case, it does not significantly affect test performance of FedAvg unlike the label heterogeneity setting. We conjecture that the impact of domain heterogeneity is mitigated due to standard pre-processing data augmentations such as rotation and cropping which have a regularizing effect of enabling clients to learn similar features across domains. Thus, this experiment shows that label heterogeneity is the more challenging form of heterogeneity in FL systems and explains why most FL works, including ours, focuses on the impact of label heterogeneity on FL training.

We also take this opportunity to respond to several claims made by the AC in the Additional Discussion:

Claim 1:

"Most reviews of this paper are relatively light."

We are not sure how the AC comes to this conclusion, as each review provides detailed strengths, weaknesses, and constructive suggestions for improvement. If the AC genuinely believed the reviews were insufficient, we would have hoped the AC could solicit additional experts to review and foster proper discussion among reviewers.

Claim 2:

...analysis here is specific to a superficial CNN model

We disagree with the claim that the two-layer CNN model is "superficial." While it is a simplification of a standard CNN model, it has been widely used in theoretical ML research to gain insights into deep networks. Several recent works have demonstrated that the two-layer CNN model captures key phenomena observed in practical settings, including:

• The generalization gap between Adam and SGD [11]
• The emergence of benign overfitting [12,13]
• The benefits of FedAvg over local training [14]
• The effectiveness of data augmentation techniques like Cutout and CutMix [15]
• The role of local steps in FL for learning heterogeneous features [16]

By adopting the two-layer CNN and corresponding data model, we eliminate the need for common assumptions in FL convergence, such as smoothness and bounded heterogeneity [17, 18]. Instead, this setup allows us to introduce a more interpretable notion of heterogeneity based on label proportion mismatch across clients (Equation 1), establish convergence to a global minimum for a non-convex problem and provide generalization guarantees.

Claim 3:

"No reviewer has mentioned this important point in the review, so I suspect that the reviewers have not read the theoretical setup and analysis in detail."

This assertion is unfounded and unfair. A lack of concern from the reviewers does not imply that they did not carefully read the paper. We would have hoped that the AC could consult the reviewers to verify this before overriding their decision.

In conclusion, we sincerely appreciate the reviewers' thoughtful evaluations and their unanimous recognition of the novelty and significance of our work. However, as outlined above, we are disappointed by the inconsistencies and claims made in the AC’s assessment without sufficient justification. We hope our response fosters a broader discussion on the fairness and transparency in the ML review process.

Best,

Divyansh, Pranay, Zheng, Gauri