Federated Generalization via Information-Theoretic Distribution Diversification

We thank the reviewer for the time to read our paper and the reviewer's hard work on our paper.

Questions:

• About the discrete alphabet. Thanks for your constructive comments! In fact, our paper adheres to the classic settings of Shannon information theory, specifically focusing on the discrete alphabet setup. Firstly, we emphasize that the information entropy of data sources serves as an upper bound for the proposed information-theoretic generalization gap. If we were to consider the continuous alphabets setting, the information entropy would be replaced by the concept of differential entropy. However, it is important to note that the differential entropy may assume negative values in certain cases, which can pose challenges when developing algorithms in such situations.

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

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

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

Thanks for your detailed reply. Here is my further explanation.

Discrete Alphabet: The key point of rate distortion theory is quantization, and the physical meaning of differential entropy is to measure the “difficulty” or the “degree of adjustment” involved in quantizing continuous sources. The vanilla definition of information entropy is unable to accurately measure the uncertainty of continuous sources. Therefore, we firmly believe that utilizing a discrete alphabet directly is sufficient to support our theoretical findings and uphold the spirit of Shannon’s information theory.

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.

In connection with the generalization gap: By examining their respective upper bounds, we observe that in order to achieve the same approximate error $\epsilon$ , the necessary sample size $n$ for the self-information-weighted generalization gap is greater than that for the true generalization error, when there are outcomes with extremely high or low probabilities.

Regarding suitability of information-theoretic analysis: Thank you for providing further clarification. In fact, we formulated Definition 3 to ensure consistency in the expected risk form between the joint distribution-based joint self-information weighted expected risk defined in Equation 2 and the marginal distribution-based self-information weighted expected risk defined in Equation 1. This alignment of definitions aims to ensure that all of them are formulated based on a unified information-theoretic framework.

I would like to thank the authors for their time in writing the rebuttal. Unfortunately, I am not convinced that the paper at this stage is ready for publication. I think the paper has valuable ideas that can be further investigated for future submission.

• Discrete Alphabet: I’m not sure to agree about this. In information theory, indeed continuous alphabets (for both source and channel coding) have been vastly explored. Besides, concerning the source coding part, while the lossless compressibility for the discrete alphabet is measured by Shannon discrete entropy, the counterpart for the continuous sources is measuring the lossy compressibility via rate-distortion theory (and not differential entropy).

• Distributed setup: After clarification, unless I missed something, there seems to be another issue: there does not exist any dependence on the training dataset in Definition 3, and calling this a generalization gap is not appropriate. Moreover, the dependence of the bounds (e.g. Theorem 1) on the dataset size is purely synthetic. Then, why one cannot let $n_i \to \infty$ and get rid of the last two terms?

• In connection with the generalization gap: We cannot compare two terms by comparing their corresponding upper bounds.

• Regarding suitability of information-theoretic analysis: Just to make it precise: exactly what I meant by “worst case” is considering the maximum entropy distributions, which e.g. for the finite alphabet case result in uniform distributions. Adhering to this principle, why in Definition 3, we do not consider such distributions for the set $Z^I$. In addition, why do we have still the term $\log(1/(P(Z^I)$? The justification is that this term for the “participating clients” is supposed to guarantee a good performance for all clients (for arbitrary distributions). But why we have $\log(1/(P(Z^I)$?

This does not answer my question. As I understand it, the object defined in Definition 3 does not "by definition" depend on the training data and thus on the sample size. However, the authors proposed a bound on this object that depends on the training size. This just doesn't make sense.

Thanks very much for your detailed comments! The following are our responses to the comments accordingly:

Weaknesses:

• Unaccounted-for notations in this paper. Thank you for reminding us! In this section, $F_i$ denotes the usual empirical risk and $g_i^t$ indicates the local gradient with respect to the usual empirical risk $F_i$. It is important to note that the self-information weighted expected risk and the corresponding information-theoretic generalization gap considered in this paper provide insights to guide the algorithm design for improving the out-of-distribution generalization performance of federated learning. This is because the oracle probability $P(z)$ of outcomes is unknown in practice.

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

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

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

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

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.

Yinan Zou, Zixin Wang, Xu Chen, Haibo Zhou, and Yong Zhou. Knowledge-guided learning for transceiver design in over-the-air federated learning. IEEE Transactions on Wireless Communications, 22(1):270–285, 2023. doi: 10.1109/TWC.2022.3192550.12

Dun Zeng, Xiangjing Hu, Shiyu Liu, Yue Yu, Qifan Wang, and Zenglin Xu. Stochastic clustered federated learning, 2023.

We sincerely appreciate the supportive and constructive comments received! In the following sections, we will address the weaknesses and questions raised by the reviewer regarding our paper.

Weakness:

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

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

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

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

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.

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

Thanks for the detailed comments and reminders presented by the reviewer. We would like to respond to the weaknesses and questions to support our paper.

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

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

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