SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhiqin Yang $^{1,2*}$  Yonggang Zhang $^{2*}$  Yu Zheng $^{3}$  Xinmei Tian $^{5}$
Hao Peng $^{1,6^{\dagger}}$  Tongliang Liu $^{4}$  Bo Han $^{2}$
$^{1}$ Beihang University  $^{2}$ Hong Kong Baptist University  $^{3}$ Chinese University of Hong Kong
$^{4}$ Sydney AI Centre, The University of Sydney  $^{5}$ University of Science and Technology of China
$^{6}$  Kunming University of Science and Technology
Federated learning (FL) typically faces data heterogeneity, i.e., distribution shifting among clients. Sharing clients' information has shown great potentiality in mitigating data heterogeneity, yet incurs a dilemma in preserving privacy and promoting model performance. To alleviate the dilemma, we raise a fundamental question: Is it possible to share partial features in the data to tackle data heterogeneity? In this work, we give an affirmative answer to this question by proposing a novel approach called Federated Feature distillation (FedFed). Specifically, FedFed partitions data into performance-sensitive features (i.e., greatly contributing to model performance) and performance-robust features (i.e., limitedly contributing to model performance). The performance-sensitive features are globally shared to mitigate data heterogeneity, while the performance-robust features are kept locally. FedFed enables clients to train models over local and shared data. Comprehensive experiments demonstrate the efficacy of FedFed in promoting model performance. The code is publicly available at: https://github.com/tmlr-group/FedFed
Federated learning (FL), beneficial for training over multiple distributed data sources, has recently received increasing attention [1, 2, 3]. In FL, many clients collaboratively train a global model by aggregating gradients (or model parameters) without the need to share local data. However, the major concern is heterogeneity issues [4, 5, 6] caused by Non-IID distribution of distributed data and diverse computing capability across clients. The heterogeneity issues can cause unstable model convergence and degraded prediction accuracy, hindering further FL deployments in practice [1].
To address the heterogeneity challenge, the seminal work, federated averaging (FedAvg), introduces the model aggregation of locally trained models [4]. It addresses the diversity of computing and communication but still faces the issue of client drift induced by data heterogeneity [7]. Therefore, a branch of works for defending data heterogeneity has been explored by devising new learning objectives [6], aggregation strategies [8], and constructing shareable information across clients [9]. Among explorations as aforementioned, sharing clients' information has been considered to be a straightforward and promising approach to mitigate data heterogeneity [9, 10].
However, the dilemma of preserving data privacy and promoting model performance hinders the practical effectiveness of the information-sharing strategy. Specifically, it shows that a limited amount of shared data could significantly improve model performance [9]. Unfortunately, no matter for sharing raw data, synthesized data, logits, or statistical information [10, 11, 12, 7] can
† Corresponding author (penghao@act.buaa.edu.cn)
incur privacy concerns [13, 14, 15]. Injecting random noise to data provides provable security for protecting privacy [16, 17]. Yet, the primary concern for applying noise to data lies in performance degradation [18] as the injected noise negatively contributes to model performance. Consequently, effectively fulfilling the role of shared information necessitates addressing the dilemma of data privacy and model performance.
We revisit the purpose of sharing information to alleviate the dilemma in information sharing and performance improvements and ask:
Q.1 Is it possible to share partial features in the data to mitigate heterogeneity in  $FL$ ? This question is fundamental to reaching a new phase for mitigating heterogeneity with the information-sharing strategy. Inspired by the partition strategy of spurious feature and robust feature [19], the privacy and performance dilemma could be solved if the data features were separated into performance-robust and performance-sensitive parts without overlapping, where the performance-robust features contain almost all the information in the data. Namely, the performance-robust features that limitedly contribute to model performance are kept locally. Meanwhile, the performance-sensitive features that could contribute to model generalization are selected to be shared across clients. Accordingly, the server can construct a global dataset using the shared performance-sensitive features, which enables clients to train their models over the local and shared data.
Q.2 How to divide data into performance-robust features and performance-sensitive features? The question is inherently related to the spirit of the information bottleneck (IB) method [20]. In IB, ideal features should discard all information in data features except for the minimal sufficient information for generalization [21]. Concerning the information-sharing and performance-improvement dilemma, the features discarded in IB may contain performance-robust features, i.e., private information, thus, they are unnecessary to be shared for heterogeneity mitigation. Meanwhile, the performance-sensitive features contain information for generalization, which are the minimum sufficient information and should be shared across clients for heterogeneity mitigation. Therefore, IB provides an information-theoretic perspective of dividing data features for heterogeneity mitigation in FL.
Q.3 What if performance-sensitive features contain private information? This question lies in the fact that a non-trivial information-sharing strategy should contain necessary data information to mitigate the issue of data heterogeneity. That is, the information-sharing strategy unavoidably causes privacy risks. Fortunately, we can follow the conventional style in applying random noise to protect performance-sensitive features because the noise injection approach can provide a de facto standard way for provable security [16]. Notably, applying random noise to performance-sensitive features differs from applying random noise for all data features. More specifically, sharing partial features in the data is more accessible to preserve privacy than sharing complete data features, which is fortunately consistent with our theoretical analysis, see Theorem 3.3 in Sec. 3.3.
Our Solution. Built upon the above analysis, we propose a novel framework, named Federated Feature distillation (FedFed), to tackle data heterogeneity by generating and sharing performance-sensitive features. According to the question  $Q.1$ , FedFed introduces a competitive mechanism by decomposing data features  $\mathbf{x} \in \mathbb{R}^d$  with dimension  $d$  into performance-robust features  $\mathbf{x}_r \in \mathbb{R}^d$  and performance-sensitive features  $\mathbf{x}_s \in \mathbb{R}^d$ , i.e.,  $\mathbf{x} = \mathbf{x}_r + \mathbf{x}_s$ . Then following the question  $Q.2$ , FedFed generates performance-robust features  $\mathbf{x}_r$  in an IB manner for data  $\mathbf{x}$ . In line with  $Q.3$ , FedFed enables clients to securely share their protected features  $\mathbf{x}_p$  by applying random noise  $\mathbf{n}$  to performance sensitive features  $\mathbf{x}_s$ , i.e.,  $\mathbf{x}_p = \mathbf{x}_s + \mathbf{n}$ , where  $\mathbf{n}$  is drawn from a Gaussian distribution  $\mathcal{N}(0, \sigma_s^2\mathbf{I})$  with variance  $\sigma_s^2$ . To this end, the server can construct a global dataset to tackle data heterogeneity using the protected features  $\mathbf{x}_p$ , enabling clients to train models over the local private and globally shared data.
We deploy FedFed on four popular FL algorithms, including FedAvg [4], FedProx [6], SCAFFOLD [7], and FedNova [22]. Atop them, we conduct comprehensive experiments on various scenarios regarding different amounts of clients, varying degrees of heterogeneity, and four datasets. Extensive results show that the FedFed achieves considerable performance gains in all settings.
1. We pose a foundation question to challenge the necessity of sharing all data features for mitigating heterogeneity in FL with information-sharing strategies. The question sheds light on solving the privacy and performance dilemma, in a way of sharing partial features that contribute to data heterogeneity mitigation (Sec 3.1)
2. To solve the dilemma in information sharing and performance improvements, we propose a new framework FedFed. In FedFed, each client performs feature distillation—partitioning local data into performance-robust and performance-sensitive features (Sec 3.2)—and shares the latter with random noise globally (Sec 3.3). Consequently, FedFed can mitigate data heterogeneity by enabling clients to train their models over the local and shared data.
3. We conduct comprehensive experiments to show that FedFed consistently and significantly enhances the convergence rate and generalization performance of FL models across different scenarios under various datasets (Sec 4.2).
Federated Learning. Federated learning allows multiple clients to collaboratively train a global model parameterized by  $\phi$  without exposing clients' data [4, 1]. In general, the global model aims to minimize a global objective function  $\mathcal{L}(\phi)$  over all clients' data distributions:
$$
\min  _ {\phi} \mathcal {L} (\phi) = \sum_ {k = 1} ^ {K} \lambda_ {k} \mathcal {L} _ {k} \left(\phi_ {k}\right), \tag {1}
$$
where  $K$  represents the total number of clients,  $\lambda_{k}$  is the weight of the  $k$ -th client. The local objective function  $\mathcal{L}_k(\phi_k)$  of client  $k$  is defined on the distribution  $P(X_{k},Y_{k})$  with random variables  $X_{k},Y_{k}$ :
$$
\mathcal {L} _ {k} \left(\phi_ {k}\right) \triangleq \mathbb {E} _ {\left(\mathbf {x}, \mathbf {y}\right) \sim P \left(X _ {k}, Y _ {k}\right)} \ell \left(\phi_ {k}; \mathbf {x}, \mathbf {y}\right), \tag {2}
$$
where  $\mathbf{x}$  is input data with its label  $\mathbf{y}$  and  $\ell(\cdot)$  stands for the loss function, e.g., cross-entropy loss. Due to the distributed property of clients' data, the global objective function  $\mathcal{L}(\phi)$  is optimized round-by-round. Specifically, within the  $r$ -th communication round, a set of clients  $C$  are selected to perform local training on their private data, resulting in  $|C|$  optimized models  $\{\phi_k^r\}_{k=1}^{|C|}$ . These optimized models are then sent to a central server to derive a global model for the  $(r+1)$ -th communication round by an aggregation mechanism  $AGG(\cdot)$ , which may vary from different FL algorithms:
$$
\phi^ {r + 1} = A G G \left(\left\{\phi_ {k} ^ {r} \right\} _ {k = 1} ^ {| C |}\right). \tag {3}
$$
Differential Privacy. Differential privacy [16] is a framework to quantify to what extent individual privacy in a dataset is preserved while releasing the data.
Definition 2.1. (Differential Privacy). A randomized mechanism  $\mathcal{M}$  provides  $(\epsilon, \delta)$ -differential privacy  $(DP)$  if for any two neighboring datasets  $D$  and  $D'$  that differ in a single entry,  $\forall S \subseteq$  Range  $(\mathcal{M})$ ,
$$
\operatorname * {P r} (\mathcal {M} (D) \in S) \leq e ^ {\epsilon} \cdot \operatorname * {P r} (\mathcal {M} (D ^ {\prime}) \in S) + \delta .
$$
where  $\epsilon$  is the privacy budget and  $\delta$  is the failure probability.
The definition of  $(\epsilon, \delta)$ -DP shows the difference of two neighboring datasets in the probability that the output of  $\mathcal{M}$  falls within an arbitrary set  $S$  is related to  $\epsilon$  and an error term  $\delta$ . Similar outputs of  $\mathcal{M}$  on  $D, D'$  (i.e., smaller  $\epsilon$ ) represent a stronger privacy guarantee. In sum, DP as a de facto standard of quantitative privacy, provides provable security for protecting privacy.
Information Bottleneck. Traditionally, in the context of the Information Bottleneck (IB) framework, the goal is to effectively capture the information relevant to the output label  $Y$ , denoted as  $Z$ , while simultaneously achieving maximum compression of the input  $X$ .  $Z$  represents the latent embedding, which serves as a compressed and informative representation of  $X$ , preserving the essential information while minimizing redundancy. This objective can be formulated as:
$$
\mathcal {L} _ {\mathrm {I B}} = I (X; Y | Z), \quad s. t. I (X; Z) \leq I _ {\mathrm {I B}} \tag {4}
$$
where  $I(\cdot)$  denote the mutual information and  $I_{\mathrm{IB}}$  is a constant. This function is defined as a rate-distortion problem, indicating that IB is to extract the most efficient and informative features.
We detail the Federated Feature distillation (FedFed) proposed to mitigate data heterogeneity in FL. FedFed adopts the information-sharing strategy with the spirits of information bottleneck (IB). Roughly, it shares the minimal sufficient features, while keeping other features at clients<sup>1</sup>.
Appendix A displays FedFed's overview.
Figure 1: Have a Guessing Game! Question: Which one is the first image from? A or B or C?
Data heterogeneity inherently comes from the difference in data distribution among clients. Aggregating models without accessing data would inevitably bring performance degradation. Sharing clients' data benefits model performance greatly, but intrinsically violates privacy. Applying DP to protect shared data looks feasible; however, it is not a free lunch with paying the accuracy loss on the model.
Draw inspiration from the content and style partition of data causes [19], we investigate the dilemma in information sharing and performance improvements through a feature partition perspective. By introducing an appropriate partition, we can share performance-sensitive features while keeping performance-robust features locally. Namely, performance-sensitive features in the data are all we need for mitigating data heterogeneity!
We will elucidate the definitions of performance-sensitive features and performance-robust features. Firstly, we provide a precise definition of a valid partition (Definition 3.1), which captures the desirable attributes when partitioning features. Subsequently, adhering to the rules of a valid partition, we formalize the two types of features (Definition 3.2).
Definition 3.1. (Valid Partition). A partition strategy is to partition a variable  $X$  into two parts in the same measure space such that  $X = X_{1} + X_{2}$ . We say the partition strategy is valid if it holds: (i)  $H(X_{1},X_{2}|X) = 0$ ; (ii)  $H(X|X_1,X_2) = 0$ ; (iii)  $I(X_{1};X_{2}) = 0$ ; where  $H(\cdot)$  denotes the information entropy and  $I(\cdot)$  is the mutual information.
Definition 3.2. (Performance-sensitive and Performance-robust Features). Let  $X = X_{s} + X_{r}$  be a valid partition strategy. We say  $X_{s}$  are performance-sensitive features such that  $I(X;Y|X_s) = 0$ , where  $Y$  is the label of  $X$ . Accordingly,  $X_{r}$  are the performance-robust features.
Intuitively, performance-sensitive features contain all label information, while performance-robust features contain all information about the data except for the label information. More discussions about Definition 3.1 and Definition 3.2 can be found in Appendix C.1. Now, the challenge turns out to be: Can we derive performance-sensitive features to mitigate data heterogeneity? FedFed provides an affirmative answer from an IB perspective, boosting a simple yet effective data-sharing strategy.
3.2 Feature Distillation
Given the motivation above, we propose to distil features such that data features can be partitioned into performance-robust features depicting mostly data and performance-sensitive features favourable to model performance. We design feature distillation and answer Q.2 below.
We draw inspiration from the information bottleneck method [21]. Specifically, only minimal sufficient information is preserved in learned representations while the other features are dismissed [21]. This can be formulated as:
$$
\min  _ {Z} I (X; Y | Z), \text {s . t .} I (X; Z) \leq I _ {\mathrm {I B}}, \tag {5}
$$
where  $Z$  represents the desired representation extracted from input  $X$  with its label  $Y$ ,  $I(\cdot)$  is the mutual information, and  $I_{\mathrm{IB}}$  stands for a constant. Namely, given the learned representation  $Z$ , the mutual information between  $X$  and  $Y$  is minimized. Meanwhile,  $Z$  dismisses most information about the input  $X$  so that the mutual information  $I(X;Z)$  is less than a constant  $I_{\mathrm{IB}}$ .
Similarly, in FedFed, only minimal sufficient information is necessary to be shared across clients to mitigate data heterogeneity, while the other features are dismissed before sharing data, i.e., kept locally on the client. The differences between IB and FedFed have two folds: 1) FedFed aims to make
2We follow the formulation used in Tishby and Zaslavsky [21, 23].
the dismissed features close to the original features (i.e., depict most private data information), while IB focuses on the dissimilarity between the preserved and original features (i.e., contain minimal sufficient information about data); 2) FedFed dismisses information in the data space while IB does that in the representation space. More specifically, the objective of feature distillation is:
$$
\min  _ {Z} I (X; Y | Z), \text {s . t .} I (X; X - Z | Z) \geq I _ {\mathbb {F F}}, \tag {6}
$$
where  $Z$  represents the performance-sensitive features in  $X$ ,  $Y$  is the label of  $X$ ,  $(X - Z)$  stands for performance-robust features that are unnecessary to be shared, and  $I_{\mathrm{FF}}$  is a constant. In Eq. (6), the  $(X - Z)$  should represent data mostly conditioned on performance-sensitive features  $Z$ , while the  $Z$  should contain necessary information about the label  $Y$ . Consequently, learning with the objective can divide data features into performance-sensitive features and performance-robust features, achieving the goal of feature distillation.
To make the feature distillation tractable, we derive an objective equal to the original objective in Eq. (6) (see Appendix C.2 for more details) for client  $k$  as follows:
$$
\min  _ {\theta} - \mathbb {E} _ {(\mathbf {x}, \mathbf {y}) \sim P \left(X _ {k}, Y _ {k}\right)} \log p (\mathbf {y} | \mathbf {z} (\mathbf {x}; \theta)), \text {s . t .} \| \mathbf {z} (\mathbf {x}; \theta) \| _ {2} ^ {2} \leq \rho , \tag {7}
$$
where  $\theta$  is the parameter to be optimized for generating performance-sensitive features  $\mathbf{z}(\cdot ;\theta)$ ,  $(\mathbf{x},\mathbf{y})$  represents the input pair drawn from the joint distribution  $P(X_{k},Y_{k})$  of client  $k$ ,  $p(\mathbf{y}|\cdot)$  is the probability of predicting  $Y = \mathbf{y}$ , and  $\rho >0$  stands for a constant. The underlying insight of the objective function in Eq. (7) is intuitive. Specifically, the learned performance-sensitive features  $\mathbf{z}(\cdot ;\theta)$  are capable of predicting the label  $\mathbf{y}$  while having the minimal  $\ell_2$ -norm.
Unfortunately, the original formulation in Eq. (7) can hardly be used for dividing data features. This is because the feature distillation in Eq. 7 cannot make the preserved features  $\mathbf{x} - \mathbf{z}(\mathbf{x};\theta)$  similar to the raw features  $\mathbf{x}$ . Namely, the original Eq. (7) cannot guarantee the desired property of the preserved features. To solve the problem, we propose an explicit competition mechanism in the data space. Specifically, we model preserved features, i.e., performance-robust features, with  $q(\mathbf{x};\theta)$  explicitly and model the performance-sensitive features  $\mathbf{z}(\mathbf{x};\theta) \triangleq \mathbf{x} - q(\mathbf{x};\theta)$  implicitly:
$$
\min  _ {\theta} - \mathbb {E} _ {(\mathbf {x}, \mathbf {y}) \sim P \left(X _ {k}, Y _ {k}\right)} \log p (\mathbf {y} | \mathbf {x} - q (\mathbf {x}; \theta)), \text {s . t .} \| \mathbf {x} - q (\mathbf {x}; \theta) \| _ {2} ^ {2} \leq \rho . \tag {8}
$$
Thus, by Eq. 8, performance-sensitive features can predict labels and performance-robust features are almost the same as raw features.
The realization of Eq. (8) is straightforward. To be specific, we can employ a generative model parameterized with  $\theta$  to generate performance-robust features, i.e.,  $q(\mathbf{x};\theta)$  in Eq. (8). Meanwhile, we train a local classifier  $f(\cdot ;\mathbf{w}_k)$  parameterized with  $\mathbf{w}_k$  for client-  $k$  to model the process of predicting labels, i.e.,  $-\log p(\mathbf{y}|\cdot)$  in Eq. (8). Accordingly, the realization of feature distillation is formulated as follows with  $\mathbf{z}(\mathbf{x};\theta)\triangleq \mathbf{x} - q(\mathbf{x};\theta)$ :
$$
\min  _ {\theta , \mathbf {w} _ {k}} - \mathbb {E} _ {(\mathbf {x}, \mathbf {y}) \sim P \left(X _ {k}, X _ {k}\right)} \ell \left(f \left(\mathbf {z} (\mathbf {x}; \theta); \mathbf {w} _ {k}\right), \mathbf {y}\right), s. t. | | \mathbf {z} (\mathbf {x}; \theta) | | _ {2} ^ {2} \leq \rho , \tag {9}
$$
where  $\ell(\cdot)$  is the cross-entropy loss,  $q(\mathbf{x};\theta)$  stands for a generative model, and  $\rho$  represents a tunable hyper-parameter. Built upon Eq. (9), we can perform feature distillation, namely, the outputs of a generator  $q(\mathbf{x};\theta)$  serve as performance-robust features and  $\mathbf{x} - q(\mathbf{x};\theta)$  are used as performance-sensitive features. We merely share  $\mathbf{x} - q(\mathbf{x};\theta)$  to tackle data heterogeneity by training models over both local and shared data. Algorithm 1 summarizes the procedure of feature distillation.
3.3 Protection for Performance-Sensitive Features
Until now, we have intentionally overlooked the overlap between performance-sensitive features and performance-robust features. Since we prioritize data heterogeneity, overlapping is almost unavoidable in practice, i.e., performance-sensitive features containing certain data privacy. Accordingly, merely sharing performance-sensitive features can risk privacy. Thus, we answer Q. 3 below.
Why employ DP? The constructed performance-sensitive features may contain individual privacy, thus, the goal of introducing a protection approach is for individual privacy. In addition, the employed protection approach is expected to be robust against privacy attacks [24, 25]. According to the above analysis, differential privacy (DP) is naturally suitable in our scenario of feature distillation. Thus, we employ DP to protect performance-sensitive features before sending them to the server.
Algorithm 1 Feature Distillation
Server input: communication round  $T_{d}$  DP noise level  $\sigma_s^2$    
Client  $k$  s input: local epochs of feature distillation  $E_{d}$ $k$  -th local dataset  $\mathcal{D}^k$  and rescale to [0, 1] Initialization: server distributes the initial model  $\mathbf{w}^0,\theta^0$  to all clients   
Server Executes:   
for each round  $t = 1,2,\dots ,T_d$  do server samples a subset of clients  $C_t\subseteq \{1,\ldots ,K\}$  . server communicates  $\mathbf{w}^t,\theta^t$  to selected clients for each client  $k\in C_t$  in parallel do  $\mathbf{w}_k^{t + 1},\theta_k^{t + 1}\gets \mathrm{Local\_FeatDis}(\mathbf{w}^t,\theta^t,\sigma_s^2)$  end for  $\mathbf{w}^{t + 1},\theta^{t + 1}\gets \mathrm{AGG}(\mathbf{w}_k^t,\theta_k^t,k\in C_t)$    
end for   
 $\mathcal{D}^s = \{\mathcal{D}_k^s\}_{k = 1}^K\gets \mathrm{Collecting~}\mathbf{x}_p\mathrm{~generated~by~}k$  th client use Eq (9), where  $\mathbf{x}_p = \mathbf{x}_s + \mathbf{n}$  Local_FeatDis(w,t,  $\theta^t,\sigma_s^2$  ..   
for each local epoch e with  $e = 1,\dots ,E_d$  do  $\mathbf{w}_k^{t + 1},\theta_k^{t + 1}\gets \mathrm{SGD~update~use~Eq~(9)}.$    
end for   
Return  $\mathbf{w}_k^{t + 1},\theta_k^{t + 1}$  to server