Paper: Federated Generalization via Information-Theoretic Distribution Diversification (ICLR-2024)

Reviewer comment (Experiments):
The considered setup does not well capture the FL setup.

Author response:
*About the distributed learning.* The Joint Self-Information Weighted Expected Risk, as defined in Definition 2, is equivalent to the scenario where all the data sources are stored in a centralized parameter server. In this setup, the server conducts centralized training using these data sources, as the expectation is computed based on the joint distribution of these data sources. Consequently, the Information-Theoretic Generalization Gap, defined in Definition 3, essentially quantifies the disparity between the information-theoretic expected risks calculated on the joint distribution of data sources and the marginal distribution of data sources. This gap signifies, to some extent, the difference between centralized learning and distributed learning (when $\mathcal{I}=\mathcal{I}_p$. i.e., all the clients participate in FL).

Reviewer comment (Theory):
The proof techniques are simple and the resulted bounds “seem” to be loose.

Author response:
*About the techniques used in this paper.* The techniques employed in this paper are indeed based on general generalization theory and information theory. On one hand, for instance, in the proof of Theorem 2, we utilize techniques commonly employed in the field of information theory. Specifically, on page 19 of the supplementary material, we apply the concept of “conditioning reduces entropy” and leverage the properties of conditional entropy and mutual information to establish our proof. Additionally, in the analysis of the In-domain scenario, we utilize the entropy rate of a stochastic process to assess the complexity of the data sources and its influence on the generalization gap. Therefore, the techniques utilized in this paper are not limited to “purely algebraic manipulations.” Furthermore, the scope of this paper focuses on developing algorithms to enhance the out-of-distribution generalization of FL based on the derived generalization bound. This implies that the proposed generalization framework serves as a means to provide us with intuition and insights into understanding the OOD generalization of FL. Refining the presented generalization bound further is an avenue for our future research.

Reviewer comment (Theory):
Another major limitation is that, similar to some previous works, this work assumes that the FL algorithms used by the participating clients manage to find the "global minimizer" of the self-information weighted semi-empirical risk. This consequently means that the FL algorithm (considered for the participating clients) is "equivalent" to a centralized algorithm that has all training data of all participating clients in one place. In a sense, the paper assumes that there is no such thing as "distributed learning". While much of the work in the literature shows the different behavior of these architectures. In other words, it seems that the paper is not really about the federated learning setup, but rather about the "mismatch" between test and training data.

Author response:
**Distributed setup:** In our setup, the term “generalization” encompasses both “out-of-sample generalization” (which focuses on concentration) and “out-of-distribution generalization”. However, our primary focus in this paper is on the latter aspect. Equation 3 effectively quantifies the disparity between the information obtained from the entire data source  $Z^{\mathcal{I}}$  and the partial participating data sources  $Z^{\mathcal{I}_p}$. Therefore, in the context of Equation 3, the term “generalization” refers to the process of generalizing the model trained on  $Z^{\mathcal{I}_p}$ to the unknown $Z^{\mathcal{I}}$. The sample complexity term measures the impact of “out-of-sample generalization”, which tends to diminish as the size of the total samples stored in the participating training set approaches infinity.

Reviewer comment (Theory):
All results appear to be in terms of the "information-theoretic-generalization gap". However, except for some "intuitions", it has not been shown concretely what the relation is between such a definition and the true "generalization error". In other words, what are the concrete implications of these results for the true generalization error? If this relation cannot be established, there is uncertainty about the usefulness of the results.

Author response:
*About the connection between the true generalization error and the information theoretic generalization gap.* Thanks for your valuable comments! If we consider fixing the hypothesis $h \in \mathcal{H}$ and analyze the true generalization error and the proposed information-theoretic generalization gap using the Hoeffding inequality, we can observe that with high probability $(1-\delta)$, the true generalization error is bounded as $|\frac{1}{n}\sum_{i=1}^n\ell(h,z_i)-\mathbb{E}[\mathcal{L}(h)]|\leq\sqrt{\frac{b^2\log(2/\delta)}{2n}}$. Additionally, the information-theoretic generalization gap is bounded as $|\frac{1}{n}\sum_{i=1}^n\ell(h,z_i)\log\frac{1}{p(z_i)}-\mathbb{E}\ell(h,z)\log\frac{1}{p(z)}|\leq\sqrt{\frac{{(b\log(1/p_m))}^2\log(2/\delta)}{2n}}$, where $\ell$ is bounded by $b$, and $p_m$ represents the minimum probability of an outcome in $\mathcal{Z}$, with the condition $p_m\neq0$. The above bounds indicate that the information-theoretic generalization gap is dominated when there are outcomes with extremely high or low probabilities.

Reviewer comment (Theory):
Considering Theorem 1, this result is obtained by almost purely algebraic manipulations (e.g., proofs of Lemmas 1 and 2) and excessive use of the triangle inequality. Besides the simplicity of the techniques used (except perhaps the proof of Lemma 3, which is inspired by and similar to previous work), the bounds using such techniques are susceptible to being very loose. Similar comments apply to Theorem 2. Therefore, considering the terms appearing in the bound as optimization proxies does not seem very justified.

Author response:
*About the techniques used in this paper.* The techniques employed in this paper are indeed based on general generalization theory and information theory. On one hand, for instance, in the proof of Theorem 2, we utilize techniques commonly employed in the field of information theory. Specifically, on page 19 of the supplementary material, we apply the concept of “conditioning reduces entropy” and leverage the properties of conditional entropy and mutual information to establish our proof. Additionally, in the analysis of the In-domain scenario, we utilize the entropy rate of a stochastic process to assess the complexity of the data sources and its influence on the generalization gap. Therefore, the techniques utilized in this paper are not limited to “purely algebraic manipulations.” Furthermore, the scope of this paper focuses on developing algorithms to enhance the out-of-distribution generalization of FL based on the derived generalization bound. This implies that the proposed generalization framework serves as a means to provide us with intuition and insights into understanding the OOD generalization of FL. Refining the presented generalization bound further is an avenue for our future research.

Reviewer comment (Theory):
The authors mentioned that “Our hypothesis is grounded in the notion that a model exhibiting proficiency with low-probability examples from training distributions might demonstrate adaptability to unfamiliar testing distributions.” In a sense, intuitively, the goal here is to consider the "worst case" scenarios. However, it seems that the information-theoretic approach is not appropriate for this goal. In essence, information theory relies heavily on the probabilistic behavior of the learned model and data distribution, and the associated concentration of measures. Could the authors intuitively justify their information-theoretic approach and consider "self-information-weighted expected risk"?

Author response:
*About "the information-theoretic approach is not appropriate for this goal".* In fact, the scenario considered in this paper is commonly encountered in practice and does not represent the "worst case" scenario. For example, in the AIoT scenario, certain IoT devices may only collect data in specific areas, while the global model trained by these devices is expected to provide spatial-related services for all devices across the entire area. Consequently, certain outcomes within the sample space may have a relatively low probability within the training domain, but exhibit a higher probability within the test domain. During this period, the test domain is assumed to possess maximum entropy distributions, indicating a greater degree of uncertainty, as mentioned on Page 3 of the paper. It is important to note that considering maximum entropy distributions with significant uncertainty is a natural approach within the field of information theory.

Reviewer comment (Theory):
Describing the first term in the RHS of equation (4) as overfitting could be misleading. In fact, if, for example, $\hat{h}^*$ is the minimizer of the self-information weighted expected risk for the set $\mathcal{I}$, this term can still be large. What exactly is the intuition behind this term?

Author response:
*About  the intuition behind the first term in the RHS of equation (4).* Thank you for your valuable comments! The first term $\sup_{h \in \mathcal{H}}L\Vert\hat{h}^*-h\Vert H(Z^{\mathcal{I}})$ in the upper bound in Theorem 1 indicates the distance between the optimal hypothesis $\hat{h}^*=\arg\inf_{h}\mathcal{L}_{\mathcal{I}_p}(h)$ on self-information-weighted semi-empirical risk and other $h$ in the hypothesis space. Therefore, this first term essentially reflects the effect of model complexity on the considered generalization gap and the difficulty in obtaining this optimal model.

Reviewer comment (Theory):
Remark 1 simply rewrites the RHS of equation (5). Could you please give some intuition about the behavior described and why it makes sense?

Author response:
*About  the Remark 1.* The key point in Equation 5 is the participation gap term $3bH(Z^{\mathcal{I}})-b\sum_{i}\alpha_iH(Z_i), i \in \mathcal{I}_p$. This term signifies that a smaller joint entropy of all possible data sources and a larger information entropy of participating data sources can lead to a reduction in the derived generalization bound. This observation is in line with our intuition. In our setting, we plan to train a global model by partial participating clients. This trained global model is intended to provide services for all possible clients. Consequently, the corresponding test domain becomes "complex" and generalizing the trained model to this test domain becomes challenging when the joint entropy $H(Z^{\mathcal{I}})$ is large. On the other hand, if the information entropy of the training domain is large, it implies that the trained model is powerful and capable of performing well on unknown test domains. This is because a larger information entropy indicates that the model has encountered a sufficient number of outcomes during the training stage.

Reviewer comment (Writing):
This notational ambiguity makes me unable to understand Algorithm 2 fully. At line 5, especially, what local loss function does client i try to minimize?

Author response:
**Notational ambiguity in Alg.2.** Thank you once again for bringing this to our attention! The local loss function in Line 5 of Algorithm 2 refers to the standard loss function commonly used in FL, i.e., the usual mini-batch loss $\frac{1}{|B_i|}\sum_j \ell(z_j), z_j \in B_i, B_i \subset S_i$, where $B_i$ is the mini-batch of i-th local dataset $S_i$.

Reviewer comment (Writing):
In Section 4.2.2, again, I do not fully understand what the convex hull is with respect to.

Author response:
**About the convex hull.** The detailed procedure of the proposed convex hull-based client selection method is described in the supplementary material and we will discuss it in this response below. 
The formal work process of the proposed convex hull construction-based client selection  is outlined as follows. a) The server initiates the execution of the quickhull algorithm to construct the convex hull of the local gradients stored in the gradient table. b) Once  the vertices of the considered point set, which comprises 
 all the stored gradients are identified, the server proceeds to select the corresponding clients whose gradients are located on these discovered vertices. These selected clients are then required to participate in the subsequent round of FL.

Reviewer comment (Experiments):
Empirically speaking, it is quite rare to update clients’ parameters through only one round of gradient descent as the authors propose in Algorithm 2, due to the cost of communication relative to local computation.

Author response:
**One round of gradient descent in local update.** The derived generalization bound is based on the uniform convergence-type generalization bound used in the classic study of federated generalization in (Hu et al., 2023). This theoretical technique focuses on the algorithm-independent generalization bound, which means that the derived generalization bound is not based on one specific algorithm, such as FedSGD or FedAvg. In other words, the information-theoretic generalization framework can be applied in one-round or multi-round local optimization settings since it is based on the algorithm-independent generalization bound.

Reviewer comment (Theory):
The authors aim to optimize the participation-dependent term in the generalization bound. However, the relationship between this problem and the formulations in (11) and (12) is insufficiently clear. For example, what do the authors mean by the claim that cosine similarity is “suitable for FL”? Similarly, I do not see any proof for the equivalence of this optimization problem and that in (12).

Author response:
**Questions: The relationship between optimizing the participation dependent term and the formulations in (11) and (12):** The cross entropy term $\sum_{j}H(P_i,P_j), i\in\mathcal{I}_t, j \in \mathcal{I}_p \setminus \mathcal{I}_t $ plays a crucial role in the generalization bound presented in Theorem 2. It indicates that the distribution discrepancy between the selected data sources and the unselected data sources can significantly impact the out-of-distribution generalization performance of federated learning.  Furthermore, it is worth noting that in the field of federated learning, the distance between different local gradients $\Vert \nabla F_i-\nabla F_j\Vert^2, \forall i,j \in \mathcal{I}_p$ can be considered as an approximate metric of distribution discrepancy, as discussed in the study (Zou et al, 2023). This distance can be further expressed as $ \Vert\nabla F_i-\nabla F_j\Vert^2=\Vert \nabla F_i\Vert^2-2<\nabla F_i,\nabla F_j>+\Vert \nabla F_j\Vert^2$. In the context of federated learning, optimizing the cosine similarity $\frac{<\nabla F_i,\nabla F_j>}{\Vert \nabla F_i\Vert\Vert \nabla F_j\Vert}$ provides a more stable approach (Zeng et al, 2023).

Reviewer comment (Theory):
First of all, it is unclear what it the meaning of bounding the self-information weighted expected risk. In practice, we care about the accuracy over the generalization dataset, which is different from the information-weighted risk. Therefore, the fundamental question is: why bounding the information theoretic-generalization gap means good generalization accuracy? I didn't found in this paper an answer to this question or any solid explanations on this connection.

Author response:
* **The meaning of bounding the self-information weighted expected risk.** The motivation behind defining the self-information weighted expected risk is described in Equation 1. To further support this concept, we will provide a concrete example. During the training stage, the target distributions are unknown. It is natural to assume that these unknown target distributions have greater uncertainty. For instance, we can consider a uniform distribution of discrete labels as an example. In such cases, it is reasonable to incorporate the self-information of outcomes to formulate the expected risk.  This is because the considered model $h\in\mathcal{H}$  may perform poorly on certain outcomes $z \in \mathcal{Z}$ with low probability in the training distributions, while these outcomes may have higher probabilities in the unknown target distributions. Here we provide a concrete example: Internet of Things (IoT) devices utilize spatially correlated data for federated learning, where the data collected by devices exhibits spatial heterogeneity. Devices participating in federated learning have limited data collected from certain regions, resulting in the possibility of the global model obtained through federated learning providing poorer services to devices that have collected a larger amount of data from that region.

Reviewer comment (Theory):
The upper bound in Theorem 1 lacks examples to explain each term in concrete cases. Given its current form, it is nearly impossible to compare with existing bounds and reasoning about tightness.

Author response:
* **Concrete cases of each term in the upper bound in Theorem 1.** The first term $\sup_{h \in \mathcal{H}}L\Vert\hat{h}^*-h\Vert H(Z^{\mathcal{I}})$ in the upper bound in Theorem 1 reflects the discrepancy between the optimal global model $\hat{h}^*=\arg\inf_{h}\mathcal{L}_{\mathcal{I}_p}(h)$, considering self-information of outcomes, and other models $h$. Therefore, this first term essentially captures the impact of model complexity on the considered generalization gap.  It is also impacted by the joint entropy $H(Z^{\mathcal{I}})$ of all the data sources,  indicating that the complexity of hypothesis space $\mathcal{H}$ and data sources $Z^{\mathcal{I}}$ affect the generalization. The second term $3bH(Z^{\mathcal{I}}) -b \sum_i H(Z_i), i \in \mathcal{I}_p$ measures the difference between the joint entropy of all the data sources and the sum of information entropy of participating data sources. The insight behind this crucial term is that reducing the uncertainty of all the data sources, including unknown target data sources, and increasing the information of participating data sources will help decrease the considered generalization bound. The third term, which incorporates model complexity (VC dimension) and sample complexity, is commonly found in generalization theory (Hu et al., 2023).

Reviewer comment (Theory):
VC dimension is usually too large for modern neural networks to make it a useful complexity measure. Unfortunately, theorem 1 involves VC dimension.

Author response:
* **VC dimension used in the proposed generalization bound.** First of all, we follow the setting of the classic generalization study of federated learning (Hu et al., 2023), which considers unparticipating clients. Therefore, we directly employ the techniques they used to bound the generalization gap, namely the VC dimension. Besides, alternative metrics such as Rademacher complexity and covering number can also be used in the uniform convergence-type generalization bound discussed in this paper, refining the generalization bound further. However, this paper’s focus is on designing algorithms to improve the out-of-distribution generalization of federated learning, inspired by the proposed information-theoretic generalization framework. Enhancing traditional generalization techniques to refine the generalization bound is beyond the scope of this paper.

Reviewer comment (Experiments):
The experiment section seems also quite limited in that only some simple and small-scale datasets are tested.

Author response:
* **Experimental results.**  Experiment setup of this paper follows previous works on federated generalization studies (Yuan et al, 2021; Hu et al, 2023). We believe the key points of this paper are fully considered by experiments. It will be appreciated if the reviewer provides direct requirements of the further experiment study. We are also willing to conduct further experiments if suggested.

Reviewer comment (Theory):
What is the connection of information theoretic-generalization gap and generalization accuracy?

Author response:
* **Connection of information theoretic-generalization gap and generalization accuracy.**  As mentioned earlier, the proposed information-theoretic generalization gap is formulated by considering the self-information of outcomes. Therefore, it is well-suited for situations where the target distribution is unknown during the training stage. Consequently, if the information-theoretic generalization gap is decreased, the out-of-distribution generalization accuracy tends to decrease as well.

Reviewer comment (Theory):
What are some examples to explain each term of the upper bound in Theorem 1, in concrete cases?

Author response:
* **Examples to explain each term of the upper bound in Theorem 1, in concrete cases.**  We have provided detailed examples in the weaknesses section of this comment.

Reviewer comment (Theory):
How to compare with existing bounds and reasoning about tightness on Theorem 1?

Author response:
* **About how to compare with existing bounds and reasoning about tightness on Theorem 1.** The generalization bound proposed in this paper is derived based on the considered self-information weighted expected risk. The theoretical framework of this paper differs from other studies on federated generalization. Therefore, comparing the proposed generalization bounds with existing bounds is beyond the scope of this paper. Our focus is on developing algorithms inspired by the proposed information-theoretic generalization framework.

Reviewer comment (Theory):
Speaking about OOD generalization, is there anywhere in the framework that captures the distribution shift? I thought something like mutual information should appear?

Author response:
* **The framework captures the distribution shift.** In fact, both the difference between the joint distribution of the sum of information entropy of participating data sources $3bH(Z^{\mathcal{I}})-\sum_i H(Z_i), i \in \mathcal{I}_p$ in Theorem 1 and the cross entropy $ \sum_j H(P_i ,P_j), i\in \mathcal{I}_t, j \in \mathcal{I}_p \setminus \mathcal{I}_t$ in Theorem 2 capture the distribution shift. The mutual information $I( Z^{\mathcal{I}_t};Z^{\mathcal{I}_p \setminus \mathcal{I}_t})$ also affects the generalization bound. Please refer to the proof of Theorem 2 on page 19 of the supplementary material for more details.

Reviewer comment (Evaluation):
In Table 1, why is your methods' ID accuracy on CIFAR-10 worse than MaxSim, Power-of-Choice, and Full Sampling?

Author response:
* **The reason of proposed methods' ID accuracy on CIFAR-10 show worse.** The algorithm designed in this paper is based on the proposed information-theoretic generalization framework. The core idea is to consider the information entropy of individual data sources and the distribution discrepancy between different sources to achieve distribution diversification. However, baseline methods such as MaxSim, Power-of-Choice, and Full Sampling aggregate local models from data sources with lower information content and smaller distribution differences. Consequently, when these aggregated global models are tested on a target domain that shares the same distribution as the source domain, they tend to perform better, resulting in lower ID accuracy. However, it is important to note that the focus of this paper is to design methods for improving the out-of-distribution generalization performance of federated learning by considering the proposed information-theoretic generalization framework. Therefore, the focus is on the OOD accuracy rather than the ID accuracy.

Reviewer comment (Theory):
The tightness of the information-theoretic generalization error bound is unknown. Thus, the weighted aggregation and client selection strategies based on the generalization error bound are unclear. Also, several notions in the derived information-theoretic generalization error bound are unclear.

Author response:
* **About the tightness of the information-theoretic generalization error bound.** Our paper follows the same setting as the classic study of federated generalization by Hu et al. (2023). In this setting, there are unparticipating clients in FL, while the global model trained by participating clients still needs to provide service for these unparticipating clients. Therefore, we directly utilize the theoretical techniques used in their study, specifically the uniform convergence-type generalization bounds.  In these generalization bounds, various metrics are used to measure the complexity of the model. While VC dimension is one commonly used metric, it is important to note that alternative metrics such as Rademacher complexity or covering number can also be employed to refine the proposed generalization bound. However, exploring these alternative metrics falls beyond the scope of our paper. Our primary focus is on designing algorithms to improve the out-of-distribution generalization performance of FL, drawing insights from the proposed information-theoretic generalization framework. We firmly believe that the information entropy term in Theorem 1 and the cross-entropy term in Theorem 2 can provide valuable insights towards achieving this goal.

Reviewer comment (Experiments):
The numerical experimentations may be inadequate.

Author response:
* **About the numerical experimentations.** The experimental setting of this paper aligns with the existing theoretical works on federated generalization (Yuan et al., 2021; Hu et al., 2023). In their studies, Yuan et al. utilize the EMNIST and CIFAR datasets to verify their proposed two-level generalization framework, while Hu et al. employ the EMNIST dataset to validate their proposed uniform convergence-type generalization bound. In addition to these datasets, we also consider the Shakespeare dataset, which is commonly used in NLP, to evaluate the effectiveness of our proposed algorithms and to provide empirical support for our theoretical framework. By including this dataset in our experiments, we ensure that the core aspects of our paper are thoroughly examined, and the experimental results serve as strong evidence to support our proposed theoretical framework.

Honglin Yuan, Warren Morningstar, Lin Ning, and Karan Singhal. What do we mean by generalization in federated learning? arXiv preprint arXiv:2110.14216, 2021.

Xiaolin Hu, Shaojie Li, and Yong Liu. Generalization bounds for federated learning: Fast rates, unparticipating clients and unbounded losses. In International Conference on Learning Representations, 2023.

Reviewer comment (Theory):
It's unclear how tight the proposed joint self-information generalization error bounds in Theorems 1 and 2 are. Also, it is known in the literature that the VC-dimension-based generalization error bound, which is also used in Theorems 1 and 2, could be loose. Could the authors provide corresponding lower bounds for the proposed joint self-information-based generalization error bounds to show the tightness of the upper bounds? This paper could benefit tremendously from such insights and provide theoretical guarantees for the subsequent aggregation weighting and client selection strategies.

Author response:
* **About the tightness of the information-theoretic generalization error bound.** Our paper follows the same setting as the classic study of federated generalization by Hu et al. (2023). In this setting, there are unparticipating clients in FL, while the global model trained by participating clients still needs to provide service for these unparticipating clients. Therefore, we directly utilize the theoretical techniques used in their study, specifically the uniform convergence-type generalization bounds.  In these generalization bounds, various metrics are used to measure the complexity of the model. While VC dimension is one commonly used metric, it is important to note that alternative metrics such as Rademacher complexity or covering number can also be employed to refine the proposed generalization bound. However, exploring these alternative metrics falls beyond the scope of our paper. Our primary focus is on designing algorithms to improve the out-of-distribution generalization performance of FL, drawing insights from the proposed information-theoretic generalization framework. We firmly believe that the information entropy term in Theorem 1 and the cross-entropy term in Theorem 2 can provide valuable insights towards achieving this goal.

Reviewer comment (Theory):
In Definition 3, the information theoretic-generalization gap is based on the $\hat{h}^*$, which is the optimal model parameter corresponding to the proposed self-information weighted expected risk. A similar notion of $\hat{h}_t^*$ is used for the selected participating clients in Theorem 2. However, in practice, such optimal model parameters are rarely found due to the non-convexity of the problem. Instead, the model parameters in use are highly dependent on the problem setting (e.g., non-convexity, smoothness, etc.) and the optimization methods (e.g., a large number of FedAvg-type variants) and the associated hyper-parameters (e.g., learning rates, batch sizes, etc.). Thus, the generalization error bounds derived in Theorems 1 and 2 based on the optimal parameters might not be very meaningful. Could the authors further characterize generalization error bounds for commonly used FL algorithms (e.g., FedAvg-type)?

Author response:
* **About algorithm-dependent generalization bound.** Thank you very much for your valuable comment. Indeed, the factors mentioned in your comment, such as the optimization method, the problem setting (e.g., Lipschitz constant of the loss function), and related hyperparameters, can be studied using important techniques in generalization analysis known as “algorithm-dependent generalization bounds”. These techniques include algorithmic stability-based generalization bounds, PAC-Bayes-based generalization bounds, and others. However, the fundamental framework of generalization analysis employed in this paper is the uniform convergence-based generalization bound, which is algorithm-independent. In this framework, we focus on incorporating the self-information weighted expected risk and leveraging it to derive the corresponding generalization bound. We then design methods for improving the out-of-distribution generalization performance of FL motivated by our theoretical findings, which is the primary scope of our paper. In other words, deriving the algorithm-dependent generalization bound is beyond the scope of this paper. Furthermore, since the proposed generalization bound is algorithm-independent, it is well-suited for various federated optimization algorithms, such as FedAvg.