SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zihan Chen<sup>1</sup>, Howard H. Yang<sup>2</sup>, Tony Q.S. Quek<sup>1</sup>, and Kai Fong Ernest Chong<sup>1*</sup>
$^{1}$ Singapore University of Technology and Design (SUTD)
<sup>2</sup>Zhejiang University/University of Illinois Urbana-Champaign Institute, Zhejiang University zihan_chen@sutd.edu.sg
Personalized federated learning (PFL) has been widely investigated to address the challenge of data heterogeneity, especially when a single generic model is inadequate in satisfying the diverse performance requirements of local clients simultaneously. Existing PFL methods are inherently based on the idea that the relations between the generic global and personalized local models are captured by the similarity of model weights. Such a similarity is primarily based on either partitioning the model architecture into generic versus personalized components, or modeling client relationships via model weights. To better capture similar (yet distinct) generic versus personalized model representations, we propose spectral distillation, a novel distillation method based on model spectrum information. Building upon spectral distillation, we also introduce a co-distillation framework that establishes a two-way bridge between generic and personalized model training. Moreover, to utilize the local idle time in conventional PFL, we propose a wait-free local training protocol. Through extensive experiments on multiple datasets over diverse heterogeneous data settings, we demonstrate the outperformance and efficacy of our proposed spectral co-distillation method, as well as our wait-free training protocol.
With the rapid rise in mainstream popularity of artificial intelligence (AI) models such as ChatGPT [1] and LoRA [2], there has been an increasing shift towards the development of personalized AI assistants [3]. Hence, in a future where personalized AI services become mainstream, training AI models on personal data while preserving data privacy would become increasingly important [4], and maintaining the quality for such models would require collaborative training across multiple models. Personalized federated learning (PFL) emerges as a promising privacy-preserving distributed learning paradigm that is well-equipped to meet such requirements [3]. As an extension of federated learning (FL), PFL aims to train a customized machine learning model for each client or each group of clients with similar preferences [5]. When faced with inconsistencies in the objective functions of different clients, conventional FL fails to generalize well with just a single model, while in contrast PFL promises to generalize well across all clients, even in the presence of data heterogeneity (e.g., label distribution skew and label quantity skew) [6-14].
To tackle the challenges of personalization, numerous works have focused on designing new PFL systems, or enhancing the performance of personalized models from different aspects, such robustness, fairness, and model convergence[15-18]. Under federated settings, personalization is achieved through capturing the (dis-)similarity of the local versus globally shared model representations. In
practical FL/PFL applications of collaboratively training deep neural networks (DNNs), only the model parameters (e.g., model weights or gradients) are exchanged between the clients and the server [8, 19]. Existing DNN-based PFL methods capture this (dis-)similarity either by decoupling the model architecture into groups of layers/channels[20-25], or by designing local optimization methods with regularization based directly on model weights [17, 15]. Unfortunately, the motivations for such approaches are based on empirical observations, without an overarching theory to explain model (dis-)similarity in relation to training dynamics.
In deep learning theory, the training dynamics of DNNs have been studied from the lens of Fourier analysis [26]. A crucial insight from this analysis is that there is an implicit self-regularization effect arising from the training process itself. Given a target function  $f$  to learn, the model tends to learn the lower frequencies of the Fourier spectrum of  $f$  first before learning the respective higher frequencies. Such a bias in this training process is called spectral bias [27, 28]. Informally, spectral bias describes the commonly encountered phenomenon that DNNs first learn low-level features before learning high-level features.
Motivated by this insight, we can distinguish different levels of features in a model representation by looking at its Fourier spectrum. Intuitively, diverse personalized models would still share the same low-level features, and a global generic model would contain the same low-level features. Hence, despite any inconsistencies in the objective functions of different clients, there would be no conflict in learning low-level features for both the generic and personalized models. Consequently, with the expected similarity in the lower frequency components of the Fourier spectra of both the generic and personalized models, we can distill the knowledge of the lower Fourier coefficients to boost the performance of the generic model. Dually, the entire Fourier spectrum of the generic model, which includes the "averaged" high-level features across all clients, would benefit the training of the personalized models. By combining both perspectives, we shall propose a co-distillation framework for PFL that captures (dis-)similarity in models via spectral information.
Typically, when designing PFL systems, a compute-and-wait protocol is implicitly assumed for local training [15, 23]. This means that the locally updated generic models would be sent by the clients to the server after all local computation tasks have been completed. Such a protocol would yield a period of idle waiting where clients have to wait for the next aggregated model to be broadcasted. By circumventing this compute-and-wait protocol, we shall utilize the local idle time for training to reduce the total PFL runtime.
- We propose a spectral co-distillation framework for PFL. In particular, this is the first ever use of spectral distillation in PFL to capture the (dis-)similarity of the generic and personalized models. Also, this is the first ever bi-directional knowledge distillation directly between the generic and personalized models.
- We propose a wait-free local training protocol for our spectral co-distillation framework, where we utilize the idle time during global communication so as to reduce the total PFL runtime.
- Through extensive experiments on multiple datasets with heterogeneous data settings, we demonstrate the outperformance and efficacy of our proposed spectral co-distillation framework with the wait-free communication protocol design for PFL, with respect to model generalizability and the total PFL runtime.
PFL. In PFL, prior efforts have focused on training multiple personalized models via leveraging the similarity and relationships between the global generic model and the local personalized models, such as via model interpolation/mixture [29], model decoupling [22], and personalized optimization with customized regularizers [15]. In DNN-based FL applications, decoupling-based approaches divide the model into a private part (kept at the local side) and a shared part (exchanged between the server and clients) [3, 25, 23]. In particular, FedPer [22] and FedRep [24] share the shallow layers and train personalized deep layers, while in contrast, LG-Fed [20] and  $\mathrm{CD^2}$ -pFed maintain personalized shallow layers and channels [21], respectively. Moreover, Fed-RoD proposes a framework to achieve state-of-the-art (SOTA) performance for generic and personalized models simultaneously, based on the "two-loss, two-predictor" design[23]. APFL [5] and L2GD [29] consider using a mixture of local
Figure 1: Spectral co-distillation framework with wait-free local training for PFL, in which the generic model (GM) training and the personalized model (PM) training are carried out via spectral distillation in two different stages.
and global models to achieve personalization, in which the mixture weight controls the personalization level. Personalized local training methods have been recently explored, which include local fine-tuning in FedBABU [30], bi-level optimization in Ditto [15], feature alignment in FedPAC [31], and personalized model sparsification in FedMask [16, 32] and PerFedMask [33]. More broadly, meta-learning [34, 35], gaussian processes [36], and hyper-network-based approaches [37] have been investigated in PFL. Specifically, there is another type of PFL that aims to train personalized models at the level of clusters of clients with similar preferences [38-40].
Knowledge Distillation (KD) in FL. KD has been widely explored in knowledge transfer scenarios, which usually is used to transfer knowledge from the pre-trained teacher model to the student model via minimizing the distance from the latent or logit outputs of the two models [41, 42]. KD-based FL frameworks have been developed with diverse setups, such as FedMD [43] and FedDF [44]. On the other hand, knowledge-transfer-based PFL frameworks are investigated in [45, 46] with different model structures at the local clients, which could address the system heterogeneity and improve communication efficiency. However, such methods rely on the assumption of having access to a public labeled/unlabelled dataset, which may not be a realistic assumption in FL applications [3]. Moreover, co-distillation methods have been investigated in communication-efficient decentralized scenarios to improve generalizability [45].
The main goal of this work is to train a generic global model and multiple personalized models simultaneously. As summarized in Sec. 1, our proposed framework consists of three major components: spectral distillation-based personalized model training, spectral co-distillation-based generic model training, and the wait-free sequential computation-communication protocol. In this section, we first provide the preliminary and problem formulation for PFL and model spectrum in Sec. 3.1. Next, we present our proposed spectral distillation approach for PFL in Sec. 3.2, co-distillation-based generic model training in Sec. 3.3, and the wait-free local training protocol in Sec. 3.4, accordingly. Moreover, the summarized algorithm is given in Sec. 3.5.
Problem formulation for FL and PFL. Consider an FL system consisting of a server and  $N$  clients, in which client  $i$  has a loss function  $f_{i}:\mathbb{R}^{d}\to \mathbb{R}$  used for training on its local private dataset  $\mathcal{D}_i = \{(x_i^j,y_i^j)\}_{j = 1}^{n_i}$ , where  $n_i = |\mathcal{D}_i|$  denotes the size of the local dataset of client  $i$ . In conventional FL, the objective of all the participating clients in this system is to find a global model  $w\in \mathbb{R}^d$  that
solves the following minimization problem [19]:
$$
\underset {w \in \mathbb {R} ^ {d}} {\text {m i n i m i z e}} \left\{F (w) := \sum_ {i = 1} ^ {N} \frac {n _ {i}}{n} f _ {i} (w) \right\}, \tag {1}
$$
where  $n = \sum_{i=1}^{N} n_i$  is the total number of training samples across the  $N$  clients. In a typical communication round  $t$ , a subset  $S_t$  of clients is selected to conduct local training, starting from the latest global model weights  $w_G^t$ . Let  $w_i^t$  denote the weights of client  $i$ 's model after local training. At the end of communication round  $t$ , the server would collect local models from the selected clients to update the global model via Federated Averaging (FedAvg), i.e.  $w_G^{t+1} \gets \sum_{i \in S_t} p_i^t w_i^t$ , in which  $p_i^t = n_i / \sum_{k \in S_t} n_k$  represents the ratio of the local data samples in client  $i$  over the total number of data samples in the selected subset  $S_t$  of clients for communication round  $t$ .
There are two general types of PFL: a) training  $N$  personalized models for all  $N$  clients; and b) training 1 generic model and  $N$  personalized models simultaneously. In this work, we investigate the latter one, which we term as "PFL+". This means each client  $i$  has a local personalized model  $w_{\mathrm{p},i}$  for its private dataset  $\mathcal{D}_i$ , and all clients jointly participate in the training of the generic model  $w_{\mathrm{G}}$ . After local training at client  $i$ , the updated generic model is denoted by  $w_{\mathrm{G},i}$ . Thus, PFL can be formulated using a regularized loss function with regularization term  $R_{\mathrm{p}}(w_{\mathrm{p},i}, w_{\mathrm{G},i})$ . For example,  $R_{\mathrm{p}}(w_{\mathrm{p},i}, w_{\mathrm{G},i})$  could represent the similarity/divergence between the global and local models' features, such as model weights, feature centroids, and prototypes. In our method,  $R_{\mathrm{p}}(w_{\mathrm{p},i}, w_{\mathrm{G},i})$  represents cross-model distillation during the training of client  $i$ 's personalized model. Therefore, the objective of personalized model training in PFL+ can be formally formulated as a bi-level optimization problem [17]:
$$
\text {P 1):} \quad \underset {w _ {\mathrm {p}, i} \in \mathbb {R} ^ {d}} {\text {m i n i m i z e}} \quad \left\{f _ {\mathrm {p}, i} \left(w _ {\mathrm {p}, i}\right) := f _ {i} \left(w _ {\mathrm {p}, i}\right) + \lambda_ {\mathrm {p}} R _ {\mathrm {p}} \left(w _ {\mathrm {p}, i}, w _ {\mathrm {G}, i}\right) \right\} \quad \text {f o r e a c h c l i e n t} i \tag {2}
$$
$$
\text {s u b j e c t} \quad w _ {\mathrm {G}, i} \leftarrow \text {u p d a t e d g e n e r i c m o d e l f r o m} w _ {\mathrm {G}}, \tag {3}
$$
where the regularization coefficient  $\lambda_{\mathrm{p}}$  is used to control the level of personalization. For client  $i$ , when referring to a specific communication round  $t$ , we shall denote the untrained personalized model and updated generic model by  $w_{\mathrm{p},i}^{t-1}$  and  $w_{\mathrm{G},i}^{t}$ , respectively.
3.2 Personalized local model training
Motivated by both theoretical and empirical insights of the spectral bias inherent in the training dynamics of DNNs, we explore the use of the Fourier spectrum of the generic model for knowledge distillation to enhance the training of personalized local models. In particular, we propose a distillation regularization term representing the divergence between the full model spectra of the generic and personalized models.
First, we introduce some notation. Given vectors  $p = (p_1,\dots ,p_d)$ ,  $q = (q_{1},\ldots ,q_{d})$  in  $\mathbb{R}^d$ , define the divergence function  $\mathfrak{D}(p\| q)\coloneqq \sum_{i = 1}^{d}p_{i}\log p_{i} - p_{i}\log q_{i}$ . (By convention,  $0\log 0\coloneqq 0$ .) Note that when  $p$  and  $q$  are stochastic vectors representing parameter vectors of multinomial distributions  $P$  and  $Q$ , then  $\mathfrak{D}(p\| q)$  is identically the Kullback-Leibler (KL) divergence from  $P$  to  $Q$ . Next, let DFT:  $\mathbb{C}^d\to \mathbb{C}^d$  denote discrete Fourier transform, let  $\varrho :\mathbb{C}^d\to \mathbb{R}^d$  be the map given by  $(z_{1},\ldots ,z_{d})\mapsto (\| z_{1}\| ,\ldots ,\| z_{d}\|)$ , and define the function  $s:\mathbb{R}^d\rightarrow \mathbb{R}^d$  by  $s\coloneqq \varrho \circ \mathrm{DFT}$ . For an input vector of the weights of a DNN model, the output vector after applying  $s$  shall be called the spectrum vector of that model [28]. Thus, in communication round  $t$ , the spectrum vectors of the personalized model  $w_{\mathrm{p},i}^{t - 1}$  of client  $i$  and updated generic model  $w_{\mathrm{G},i}^{t}$  are written as  $s(w_{\mathrm{G}}^{t})$  and  $s(w_{\mathrm{p},i}^{t - 1})$ , respectively. We shall represent the divergence of the personalized and generic models by  $\mathfrak{D}(s(w_{\mathrm{p},i}^{t - 1})||s(w_{\mathrm{G},i}^{t}))$ , the divergence of their spectrum vectors.
Concretely, we define  $R_{\mathrm{p}}(w_{\mathrm{p},i},w_{\mathrm{G},i})\coloneqq \mathfrak{D}(s(w_{\mathrm{p},i}^{t - 1})\| s(w_{\mathrm{G},i}^{t}))$ , and let  $f_{i}$  be the cross-entropy loss  $\mathcal{L}_{\mathrm{CE}}$  for all  $i$ . Then the personalized objective function  $f_{\mathrm{p},i}$  of client  $i$  in communication round  $t$  (cf. (2)) is given by:
$$
\mathcal {L} ^ {\mathrm {p}} \left(w _ {\mathrm {p}, i} ^ {t - 1} \mid w _ {\mathrm {G}, i} ^ {t}\right) := \mathcal {L} _ {\mathrm {C E}} \left(w _ {\mathrm {p}, i} ^ {t - 1} \mid \mathcal {D} _ {i}\right) + \lambda_ {\mathrm {p}} \mathfrak {D} \left(s \left(w _ {\mathrm {p}, i} ^ {t - 1}\right) \| s \left(w _ {\mathrm {G}, i} ^ {t}\right)\right). \tag {4}
$$
Figure 2: A comparison of the (a) conventional compute-and-wait protocol with the (b) proposed wait-free training protocol.
For simplicity, we use a common time-invariant  $\lambda_{\mathrm{p}}$  for all clients throughout training. Since we are distilling the knowledge of the spectrum vector  $s(w_{\mathrm{G},i}^{t})$  in (4), we term our approach as spectral distillation.
3.3 Generic model training
Given a  $\mathrm{PFL + }$  training framework, it is natural to connect the roles of generic and personalized models to the roles of the teacher and student models in distillation, where the training of one model is guided by the knowledge distilled by the other. Co-distillation extends this idea. Intuitively, the role of each model alternates between teacher and student for knowledge distillation during training. In  $\mathrm{PFL + }$ , since we are concurrently training both the generic and personalized models, either of them could be used for knowledge distillation. The key challenge for applying co-distillation to  $\mathrm{PFL + }$  is that it is not obvious what knowledge should be distilled from the personalized models to enhance the training performance of the generic model.
In the theory of deep learning, it is well-known that when training a DNN, there is a learning bias towards the lower frequencies of its Fourier spectrum [27, 28]. In fact, the lower-frequency components of this spectrum are robust to random weight perturbations. Hence, with diverse personalized models, we would still expect the lower-frequency components of the spectra of all models (both generic and personalized) to be similar. Consequently, we could use such lower-frequency components for knowledge distillation to enhance generic model training.
Motivated by this, we propose a truncated spectrum-based distillation loss as the regularizer for generic model training. Given  $0 < \tau \leq 1$ , let  $\iota_{\tau}:\mathbb{R}^{d}\to \mathbb{R}^{\lceil \tau d\rceil}$  be the projection map onto the first  $\lceil \tau d\rceil$  entries, and define  $\widehat{s} \coloneqq \iota_{\tau}\circ s$ . Then the loss function for generic model training, which depends on the truncated spectrum vectors  $\widehat{s}(w_{\mathrm{G},i}^{t})$  and  $\widehat{s}(w_{\mathrm{p},i}^{t-1})$ , is given by:
$$
\mathcal {L} ^ {\mathrm {G}} \left(w _ {\mathrm {G}, i} ^ {t} \mid w _ {\mathrm {p}, i} ^ {t - 1}\right) := \mathcal {L} _ {\mathrm {C E}} \left(w _ {\mathrm {G}, i} ^ {t} \mid \mathcal {D} _ {i}\right) + \lambda_ {\mathrm {g}} \mathfrak {D} \left(\widehat {s} \left(w _ {\mathrm {G}, i} ^ {t}\right) \| \widehat {s} \left(w _ {\mathrm {p}, i} ^ {t - 1}\right)\right), \tag {5}
$$
where the regularization term  $R_{\mathrm{G}}(w_{\mathrm{G},i}^{t},w_{\mathrm{p},i}^{t - 1})\coloneqq \mathfrak{D}(\widehat{s} (w_{\mathrm{G},i}^{t})\| \widehat{s} (w_{\mathrm{p},i}^{t - 1}))$  depends on the hyperparameter  $\tau$ , and  $\lambda_{\mathrm{g}}$  is the coefficient of this regularization term. Analogous to (P1), the objective of generic model training in PFL+ could be formulated as the following bi-level optimization problem:
$$
\text {P 2):} \quad \underset {w _ {\mathrm {G}} \in \mathbb {R} ^ {d}} {\operatorname {m i n i m i z e}} \quad \left\{f \left(w _ {\mathrm {G}}\right) := \sum_ {i = 1} ^ {N} \frac {n _ {i}}{n} \left(f _ {i} \left(w _ {\mathrm {G}}\right) + \lambda_ {\mathrm {G}} R _ {\mathrm {G}} \left(w _ {\mathrm {G}}, w _ {\mathrm {p}, i}\right)\right) \right\} \tag {6}
$$
$$
\text {s u b j e c t} \quad w _ {\mathrm {p}, i} \leftarrow \text {o u t p u t o f (P 1) f o r c l i e n t} i, \text {f o r} i = 1, \dots , N. \tag {7}
$$
Overall, by combining the two spectral distillation approaches together, we get a training framework for PFL+, which we shall call spectral co-distillation.
Algorithm 1 Spectral Co-Distillation with Wait-free Training for PFL+
Inputs:  $N,T,\eta_{\mathrm{p}},\eta_{\mathrm{G}},w_{\mathrm{G}}^{0},\{w_{\mathrm{p},i}^{0}\}_{i = 0}^{N},E_{\mathrm{G}},E_{\mathrm{p}}$    
Outputs: Generic model  $\boldsymbol {w}_{\mathrm{G}}^T$  , personalized models  $\{w_{\mathrm{p},i}^{T}\}_{i = 1}^{N}$    
1: for  $t = 1$  to  $T$  do   
2: for each client  $k = 1$  to  $N$  in parallel do // Generic model training and update   
3:  $w_{\mathrm{G},k}^{t}\gets \mathrm{GMUPDATE}(w_{\mathrm{G}}^{t - 1},w_{\mathrm{P},k}^{t - 1})$    
4: Upload weights  $w_{\mathrm{G},k}^{(t)}$  to server // Personalized model training (Task 1: Line 5)   
5:  $w_{\mathrm{p},k}^{t}\gets \mathrm{PMUPDATE}(w_{\mathrm{p},k}^{t - 1},w_{\mathrm{G},k}^{t})$    
6: do sequentially (Task 2: Lines 7-8)   
7: Uplink communication of generic model to server   
8: Generic model aggregation to obtain  $w_{\mathrm{G}}^{t}$  at server   
9: return  $w_{\mathrm{G}}^{T}$ $\{w_{\mathrm{p},i}^{T}\}_{i = 1}^{N}$    
function GMUPDATE  $(w_{\mathrm{G}}^{t - 1},w_{\mathrm{p},k}^{t - 1})$  Require:  $w_{\mathrm{G}}^{t - 1},w_{\mathrm{p},k}^{t - 1}$  are the latest generic model and personalized model.   
1:  $w_0\gets w_G^{t - 1}$    
2: for  $j = 1$  to  $E_{\mathrm{G}}$  do   
3:  $w_{j}\leftarrow w_{j - 1} - \eta_{\mathrm{G}}\nabla \mathcal{L}^{\mathrm{G}}(w_{j - 1}|w_{\mathrm{p},k}^{t - 1})$  // Using truncated low frequency spectrum information   
4: return  $w_{j}$    
function PMUPDATE  $(w_{\mathrm{p},k}^{t - 1},w_{\mathrm{G},k}^{t})$  Require:  $w_{\mathrm{p},k}^{t - 1},w_{\mathrm{G},k}^{t}$  are the latest personalized model and updated generic model of client  $k$  .   
1:  $w_0\gets w_{p,k}^{t - 1}$    
2: for  $j = 1$  to  $E_{\mathrm{p}}$  do   
3:  $w_{j}\leftarrow w_{j - 1} - \eta_{\mathrm{p}}\nabla \mathcal{L}^{\mathrm{p}}(w_{j - 1}|w_{\mathrm{G},k}^{t})$  // Using full model spectrum information   
4: return  $w_{j}$
3.4 Wait-free Local Training Protocol
In the context of federated computing, the total runtime, which includes both local computation and communication time throughout the entire training process, is a direct indicator of communication efficiency. However, current PFL frameworks adopt a compute-and-wait protocol for local training. This means that in each round, the client performs both generic and personalized model updates only after all local computation tasks have been completed, and resumes local training upon receiving the latest global model broadcasted from the server. In consequence, there is idle waiting time between model update and model broadcast; see Fig. 2(a).
To improve the communication efficiency of PFL training with respect to the total runtime, we propose a wait-free local training protocol, as depicted in Fig. 2(b). In our protocol, the client updates the generic model according to the conventional generic FL training and trains the personalized model during the global communication time period. Unlike existing PFL frameworks, local clients would send the updated generic model to the server before the start of the personalized model training. Thus, our protocol eliminates idle waiting time, thereby dramatically reducing total runtime. Furthermore, it could be easily incorporated into existing PFL frameworks, such as Ditto [15], to further improve the efficiency; see Tab. 4.