Generalization in Federated Learning: A Conditional Mutual Information Framework

We thank you sincerely for your valuable feedback on our paper. Our responses follow.

• While the mathematical methods are solid, including more straightforward explanations or diagrams could help clarify the concepts for readers who may not be experts in the field.

• Another suggestion is to enhance the clarity of the presentation, perhaps through more diagrams or flowcharts, which would help readers unfamiliar with CMI techniques.

Response. Thank you for your suggestion. We have included a visualization of our superclient and supersample construction, and we will also provide additional background on CMI techniques along with more accessible explanations of these concepts in the revised version.

• It would be beneficial if the paper discussed any limitations of these theoretical results, such as their dependence on the bounded loss function assumption.
• Consider including a more detailed discussion of potential limitations or assumptions (e.g., reliance on bounded loss functions or the i.i.d. assumption in some cases).
• How sensitive are the bounds to deviations from the bounded loss assumption in practice?

Response. Regarding the boundedness assumption, as noted in the paragraph immediately following Theorem 6.2 (Lines 361–362) in the paper, this assumption can be relaxed to a sub-Gaussian condition, where the loss function may be unbounded but exhibits Gaussian-like tail behavior. Under this relaxed assumption, all the theoretical results in our paper remain valid. In the case of heavy-tailed losses, additional techniques such as truncation would be required to ensure the results hold. We will elaborate further on these points in the revised version.

• The experimental work supports the theory; however, the paper could be improved by testing on more diverse datasets or different federated learning scenarios to provide a broader view of the approach’s effectiveness.

Response. Thank you very much for the suggestion. We will include an additional dataset in the revised version, and if the reviewer has any specific dataset in mind, we would be happy to incorporate it.

• While the paper references key works in CMI and FL, a discussion comparing the new bounds with very recent developments in PAC-Bayesian approaches or alternative stability-based analyses could further contextualize the contributions.

Response. Compared to PAC-Bayesian and stability-based bounds, our e-CMI bounds are significantly easier to estimate in practice. Moreover, existing PAC-Bayesian and stability-based bounds do not account for the participation gap and do not exhibit fast-rate behavior for overall generalization. We will include a more detailed discussion comparing our results with these existing bounds in the revised version.

• Could you provide further insight into how the e-CMI bounds might be estimated in high-dimensional settings beyond the one-dimensional case?

Response. We note that the e-CMI bounds in our paper always involve mutual information between one-dimensional random variables by construction. Specifically, e-CMI refers to the mutual information between a single loss value (treated as a one-dimensional random variable, since the loss function maps inputs to real values) and a Bernoulli mask, which is a binary random variable. This low-dimensional structure makes e-CMI particularly easy to estimate, which is one of its key advantages.

• How do the derived bounds perform in non-i.i.d. client scenarios that are common in cross-device FL?

Response. Our bounds hold under a non-i.i.d. setting in the sense that each client may have a different data distribution, which is typical in cross-device FL. However, we do assume that clients are sampled independently, and that within each client, data points are drawn i.i.d. If either of these assumptions does not hold, the generalization bounds would need to be refined using additional techniques to account for the dependencies, for example, by invoking graph-based dependencies among clients or modeling temporal (mixing) dependencies in the data. We will include this discussion in the revised version.

We thank you sincerely for your valuable feedback on our paper. Our responses follow.

-... the bounds established do not capture any characteristic of FL ...

Response. We respectfully disagree with the reviewer's claim that our bounds fail to capture key characteristics of FL, and the argument that the overall algorithm is ultimately treated as a black box is a misinterpretation of our results in Sections 4 & 5. To analyze FL in the two-level generalization setting using CMI, the introduction of the superclient is essential—a construction not needed in centralized learning and one that highlights the fundamental differences between the two settings. Moreover, treating the FL algorithm as a global mapping and construct a supersample of size $2nK$, would break the symmetry required for CMI analysis due to the non-i.i.d. nature.

Instead, we believe the more accurate interpretation is that our results in Sections 4 & 5 are designed to apply to any FL algorithm, without being restricted to a particular instantiation. The goal of these sections is to establish a general framework for analyzing FL as a learning problem. This generality should not be viewed as a weakness.

In Section 6, we go beyond this generality by considering specific aggregation strategies and loss functions, allowing us to derive sharper insights such as faster learning rates and the benefits of interpolation in local training. Although we focus on the one-round setting, as noted in Remark 6.1, extending the results to the multi-round case is straightforward. We plan to include the multi-round extension in the Appendix, as introducing it in the main text would further complicate the already dense notation.

• Main take-home message ... | main question ...|Q1. My main concern ...

Response. The main goal of our paper is to frame FL as a learning problem and analyze its learning rate in a two-level setting using CMI. Without assuming any specific FL strategy, our general two-step CMI bounds obtain several insights: (1) strong client privacy (both globally and locally) implies generalization; (2) in homogeneous settings, the bounds reduce to standard CMI analysis, with $I(W; V | \widetilde{Z}, U)$ capturing data heterogeneity; and (3) we identify conditions for fast learning rates.

When applied to specific algorithms like FedAvg under certain loss assumptions, we show that FL can achieve even faster rates, thanks to model averaging and favorable loss properties. We believe this behavior extends to broader loss classes, though a formal proof remains open.

Compared to existing bounds for FL, our e-CMI bound is notably easier to estimate and is the first MI-based bound to give fast-rate guarantees for both the participation and out-of-sample gaps.

Lastly, we refer the reviewer to our first response to Reviewer 4mA5 for a discussion on the practical implications of our results, and the difference between using generalization bounds as indicators of learnability v.s. as sufficient conditions for generalization performance.

• their proof techniques are rather standard ...

Response. We acknowledge that, following our superclient construction, our CMI bounds build on existing techniques, with some extensions to accommodate the two-level setting. However, developing a general CMI framework for FL does not require a complete overhaul of prior methods. As noted in the review guidelines, "originality may arise from creative combinations of existing ideas", and we believe this applies to our contribution.

• Q2. ... order-wise behavior is not correct.

**Response.**We believe the reviewer raises a valid point regarding order-wise behavior, and we also share the concern that many prior works overlook MI, KL, or other complexity terms in such analyses. In the absence of clear decay rates, we think the stated rates should be seen as reflecting worst-case behavior. For example, $O(1/\sqrt{K} + 1/\sqrt{Kn}) = O(1/\sqrt{K})$ implies that in highly heterogeneous settings, even with infinite local data ($n \to \infty$) for participating clients, FL may still fail to generalize to unseen clients—an intuitive outcome given the impact of extreme heterogeneity.

Regardless of the exact decay of the CMI term, removing the square-root dependency remains desirable for faster convergence. We will clarify this in the revision.

• Q4. It is claimed that the bound ...

Response. This seems to be a misinterpretation of our result. The CMI term in Theorem 6.2, unlike those in Sections 4 & 5, is a local CMI for each client, involving the local model $W_i$ rather than the global model $W$, and thus is not influenced by client heterogeneity.

However, the reviewer’s intuition holds in a multi-round extension of Theorem 6.2, where $W_i$ may depend on other clients through repeated communication. We will clarify this in the revised version.

We thank you sincerely for your valuable feedback on our paper. Our responses follow.

• ... why the MI-based ... will be useful and how ... lead to regularization ...
• ... uncertain about the practical role ...

Response. We note that our bounds do have practical implications, e.g., as mentioned in the Related Work section (Line 84-90), if the CMI terms in Theorem 6.2 are small, then the norm $||w-w_i||$ is also small, which is the regularization term used in FedProx. Upon reflection, we realize that this discussion would be more appropriately placed directly after Theorem 6.2, to make the implication more explicit.

We understand that the reviewer may have been looking for new regularization schemes for reducing participation gap and out-of-sample gap, rather than connections to well-known ideas such as norm-based capacity control. In response, we recommend focusing on Theorem C.2 (which, while looser than the CMI bounds, is more interpretable). This result will ultimately lead to a gradient-norm-based regularizer when SGD or SGLD is used, and hints at a potential feature alignment mechanism in the KL sense for clients. We will elaborate on these implications in the revised version.

Finally, regarding the practical implication, we would like to share a perspective based on our own research experience in learning theory. If the goal is to derive sharper generalization bounds with fast decay rates, i.e. to study learnability and learning rates of a problem, then practical applications may not follow directly. In this case, generalization bounds mainly serve to evaluate whether learning is theoretically possible. In the extreme case, the tightest generalization bound is the generalization error itself, which provides no new actionable insights for algorithm design. If the goal is to obtain actionable insights for algorithm design, then tightness of the bound becomes less critical. Generalization upper bounds, even if loose, can serve as a basis for designing regularizers. For example, penalizing weight norms is a widely used practice to improve generalization, despite the well-known fact that norm-based bounds are often vacuous and cannot explain generalization behavior in deep learning [R1, R2].

[R1] Vaishnavh Nagarajan, and J. Zico Kolter. "Uniform convergence may be unable to explain generalization in deep learning." NeurIPS 2019.

[R2] Yiding Jiang, et al. "Fantastic Generalization Measures and Where to Find Them." ICLR 2020.

• ... how the theoretical analysis ...?
• ... explain what extra challenges ...?

Response. The bounding steps for the out-of-sample gap are indeed similar to those in standard centralized analysis, which is expected as each individual out-of-sample gap is close to an in-distribution generalization gap. The main challenges arise in bounding the participation gap. Compared to existing MI-based bounds, our contributions include: the construction of the superclient, the proof of its symmetry properties (Lemma 4.1), the shifted Rademacher representation of the weighted participation gap (Lemma E.1), and the leave-one-out argument for the participation gap (Lemma F.1). To clarify, the use of weighted generalization error and leave-one-out arguments is not new in the broader learning theory literature. However, their application in our setting is enabled by the superclient construction, which serves as the foundation for these results.

• ... on estimating the mutual information ...

Response. Please notice that we use e-CMI bounds in experiments to avoid the difficulties of estimating MI between high-dimensional R.V.'s. The second sentence in Section 7 states: "we estimate the fast-rate e-CMI bound for FL, as… Additionally, due to the challenges associated with estimating MI when dealing with high-dimensional random variables, we compute an evaluated version of the CMI bound from Theorem 4.1". Importantly, the e-CMI bound is computed between two one-dimensional variables (one of which is binary), making the estimation easy and eliminating the need for advanced MI estimators. Further details are in Appendix H.

• ... do not separately report ...
• ... the non-IID setting ... and the separated ...

Response. The pathological non-IID data partitioning follows McMahan et al. (2017): data are sorted by label, split into 200 shards of size 300, and each client is randomly assigned 2 shards. We will include more details in the revision.

As for separate plots of the participation gap (PG) and out-of-sample gap (OG), we would like to clarify that our experiments are based on the e-CMI bound in Theorem 5.1, which consists of three components weighted by jointly optimizable coefficients $C_1, C_2, C_3, C_4$. These coefficients are optimized using the SLSQP algorithm implemented in the SciPy package. As a result, the e-CMI bound is not a simple addition of the individual bounds for PG and OG, which makes it difficult to present separate plots for these two quantities.

We thank you sincerely for your valuable feedback on our paper. Our responses follow.

• you said in l.122-123 left column that your CMI framework is inspired by the meta-learning one of Hellstrom &Durisi 2022. Is it possible to precise what are the specificities of your derived results (e.g. those of Section 4) compared to theirs? It seems you control the same global true risk, due to the assumptions that all tasks are drawn according to $\mathcal{D}$.

Response. As mentioned in the paper, our construction is indeed inspired by the framework for meta-learning proposed by Hellström & Durisi (2022). However, our results are not directly comparable to theirs due to a key difference in problem setup: their meta-learning framework requires the meta-learner (i.e., the global model $W$) to be further trained on the test tasks (i.e., previously non-participating clients), whereas in FL, the global model is evaluated directly on unseen clients without any additional local fine-tuning. To enable a direct comparison with Hellström & Durisi (2022), the CMI framework presented in this paper would need to be extended to the personalized FL setting, where the global model is allowed further local adaptation. In that case, Lemma 4.1 would also need to be revised accordingly, as the hypothesis may no longer be invariant to the ordering of the "test data". We will include these discussions in the next revision.

• An important point: it seems that there is no comparison with existing bounds. For instance, how does your bound behave wrt those of, e.g., Sefidgaran et al. 2024? Is it challenging to plot their results?

Response. The PAC-Bayesian and rate-distortion bounds in Sefidgaran et al. (2024) are indeed challenging to compute numerically for more complex neural networks, as they require estimating KL or MI between high-dimensional random variables, even when the model parameters are quantized. Notably, in their paper, the generalization bounds are not plotted for their ResNet experiments on CIFAR-10; instead, they only present the behavior of the generalization error to support the insights behind their bounds. In contrast, a key advantage of our results lies in the ease of estimating the e-CMI bound in the FL setting, making it more practical for empirical evaluation.

• In most of your results, you have a $O(1/\sqrt{K})$ term. Would it be possbile to recover the influence of $n$ in such terms (maybe through the mutual information term)?

Response. Yes, the reviewer raises a valid point. Indeed, the participation gap term seems to follow an $O(1/\sqrt{K})$ behavior. At the same time, in the i.i.d. setting, increasing $n$ is also expected to reduce the participation gap. This effect is implicitly captured by the CMI term $I(W;V|\widetilde{Z},U)$, as demonstrated in Corollary 4.2, where all clients share the same data distribution (i.e., there is only one client distribution). In the non-i.i.d. case, however, it is more challenging to explicitly characterize the impact of $n$ on the participation gap, as its quantitative effect depends on the degree of data heterogeneity, which can vary significantly across scenarios. We will include this discussion in the revised version.