SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Royson Lee $^{1,2*}$ , Minyoung Kim $^{2}$ , Da Li $^{2}$   
Xinchi Qiu $^{1}$ , Timothy Hospedales $^{2,3}$ , Ferenc Huszár $^{1}$ , Nicholas D. Lane $^{1,4}$
<sup>1</sup> University of Cambridge, UK  <sup>2</sup> Samsung AI Center, Cambridge, UK  <sup>3</sup> University of Edinburgh, UK  <sup>4</sup> Flower Labs
Federated learning (FL) research has made progress in developing algorithms for distributed learning of global models, as well as algorithms for local personalization of those common models to the specifics of each client's local data distribution. However, different FL problems may require different personalization strategies, and it may not even be possible to define an effective one-size-fits-all personalization strategy for all clients: depending on how similar each client's optimal predictor is to that of the global model, different personalization strategies may be preferred. In this paper, we consider the federated meta-learning problem of learning personalization strategies. Specifically, we consider meta-nets that induce the batch-norm and learning rate parameters for each client given local data statistics. By learning these meta-nets through FL, we allow the whole FL network to collaborate in learning a customized personalization strategy for each client. Empirical results show that this framework improves on a range of standard hand-crafted personalization baselines in both label and feature shift situations. $^{0}$
Federated learning (FL) is an emerging approach to enable privacy-preserving collaborative learning among clients who hold their own data. A major challenge of FL is to learn from differing degrees of statistical data heterogeneity among clients. This makes it hard to reliably learn a global model and also that the global model may perform sub-optimally for each local client. These two issues are often dealt with respectively by developing robust algorithms for learning the global model [33, 32, 39] and then offering each client the opportunity to personalize the global model to its own unique local statistics via fine-tuning. In this paper, we focus on improving the fine-tuning process by learning a personalization strategy for each client.
A variety of approaches have been proposed for client personalization. Some algorithms directly learn the personalized models [62, 43], but the majority obtain the personalized model after global model learning by fine-tuning techniques such as basic fine-tuning [46, 48], regularised fine-tuning [34, 59], and selective parameter fine-tuning [17, 11, 2, 37]. Recent benchmarks [45, 9] showed that different personalized FL methods suffer from lack of comparable evaluation setups. In particular, dataset- and experiment-specific personalization strategies are often required to achieve state-of-the-art performance. Intuitively, different datasets and FL scenarios require different personalization strategies. For example, scenarios with greater or lesser heterogeneity among clients, would imply different strengths of personalization are optimal. Furthermore, exactly how that personalization should be conducted might depend on whether the heterogeneity is primarily in marginal label shift, marginal feature shift, or conditional shift. None of these facets can be well addressed by a one size fits all personalization algorithm. Furthermore, we identify a previously understudied issue: even
*Corresponding Author: dsrl2@cam.ac.uk
Code is available at https://github.com/royson/fedl2p
for a single federated learning scenario, heterogeneous clients may require different personalization strategies. For example, the optimal personalization strategy will have client-wise dependence on whether that client is more or less similar to the global model in either marginal or conditional data distributions. Existing works that learn personalized strategies through the use of personalized weights are not scalable to larger setups and models [55, 14]. On the other hand, the few studies that attempt to optimize hyperparameters for fine-tuning do not sufficiently address this issue as they either learn a single set of personalization hyperparameters [65, 22] and/or learn a hyperparameter distribution without taking account of the client's data distribution [30].
In this paper, we address the issues above by considering the challenge of federated meta-learning of personalization strategies. Rather than manually defining a personalization strategy as is mainstream in FL [27, 9], we use hyper-gradient meta-learning strategies to efficiently estimate personalized hyperparameters. However, apart from standard centralised meta-learning and hyperparameter optimization (HPO) studies which only need to learn a single set of hyperparameters, we learn meta-nets which inductively map from local client statistics to client-specific personalization hyperparameters. More specifically, our approach FedL2P introduces meta-nets to estimate the extent in which to utilize the client-wise BN statistics as opposed to the global model's BN statistics, as well as to infer layer-wise learning rates given each client's metadata.
Our FedL2P thus enables per-dataset/scenario, as well as per-client, personalization strategy learning. By conducting federated meta-learning of the personalization hyperparameter networks, we simultaneously allow each client to benefit from its own personalization strategy, (e.g., learning rapidly, depending on similarity to the global model), and also enable all the clients to collaborate by learning the overall hyperparameter networks that map local meta-data to local personalization strategy. Our FedL2P generalizes many existing frameworks as special cases, such as FedBN [35], which makes a manual choice to normalize features using the client BN's statistics, and various selective fine-tuning approaches [37, 48, 11, 2], which make manual choices on which layers to personalize.
Existing FL works aim to tackle the statistical heterogeneity of learning personalized models by either first learning a global model [46, 39, 32, 48, 29] and then fine-tuning it on local data or directly learning the local models, which can often be further personalized using fine-tuning. Many personalized FL approaches include the use of transfer learning between global [56] and local models [44], model regularization [34], Moreau envelopes [59], and meta-learning [15, 10]. Besides algorithmic changes, many works also proposed model decoupling, in which layers are either shared or personalized [2, 11, 64, 48, 37, 17, 54] or client clustering, which assumes a local model for each cluster [42, 44, 53, 7]. These methods often rely on or adopt a fixed personalization policy for local fine-tuning in order to adapt a global model or further improve personalized performance. Although there exists numerous FL approaches that propose adaptable personalization policies [55, 40, 14], these works are memory intensive and do not scale to larger setups and models. On the other hand, our approach has a low memory footprint (Appendix C) and is directly applicable and complementary to existing FL approaches as it aims to solely improve the fine-tuning process.
Another line of work involves HPO for FL (or FL for HPO). These methods either learn one set of hyperparameters for all clients [22, 30, 65] or random sample from learnt hyperparameter categorical distributions which does not take into account of the client's meta-data [30]. Moreover, some of these methods [22, 65] search for a set of hyperparameters based on the local validation loss given the initial set of weights prior to the federated learning of the model [22, 65], which might be an inaccurate proxy to the final performance. Unlike previous works which directly learn hyperparameters, we deploy FL to learn meta-nets that take in, as inputs, the client meta-data to generate personalized hyperparameters for a given pretrained model. A detailed positioning of our work in comparison with existing literature can be found in Appendix B.
3.1 Background & Preliminaries
Centralized FL. A typical centralized FL setup using FedAvg [46] involves training a global model  $\theta_{g}$  from  $C$  clients whose data are kept private. At round  $t$ ,  $\theta_{g}^{t-1}$  is broadcast to a subset of clients selected,  $\tilde{C}$ , using a fraction ratio  $r$ . Each selected client,  $i \in \tilde{C}$ , would then update the model using
Figure 1: FedL2P's workflow. For each client, the meta-net (parameterized by  $\lambda$ ) takes the client meta-data (e.g., local data profile) and outputs hyperparameters. We update  $\lambda$  by optimizing the meta objective which is the validation loss of the model finetuned with the hyperparameters returned by the meta-net. The updated personalization strategies  $\lambda^{*}$  from clients are collected/aggregated (via FedAvg) in the server for the next round.
a set of hyperparameters and its own data samples, which are drawn from its local distribution  $P_{i}$  defined over  $\mathcal{X} \times \mathcal{Y}$  for a compact space  $\mathcal{X}$  and a label space  $\mathcal{Y}$  for  $e$  epochs. After which, the learned local models,  $\theta_{i}^{t - 1}$ , are sent back to the server for aggregation  $\theta_{g}^{t} = \sum_{i}^{\tilde{C}}\frac{N_{i}}{\sum_{i'}N_{i'}}\theta_{i}^{t - 1}$  where  $N_{i}$  is the total number of local data samples in client  $i$  and the resulting  $\theta_{g}^{t}$  is used for the next round. The aim of FL is either to minimize the global objective  $\mathbb{E}_{(x,y)}\sim_{P}\mathcal{L}(\theta_{g};x,y)$  for the global data distribution  $P$  or local objective  $\mathbb{E}_{(x,y)}\sim_{P_i}\mathcal{L}_i(\theta_i;x,y)$  for all  $i\in C$  where  $\mathcal{L}_i(\theta ;x,y)$  is the loss given the model parameters  $\theta$  at data point  $(x,y)$ . As fine-tuning is the dominant approach to either personalize from a high-performing  $\theta_{g}$  or to optimize  $\theta_{i}$  further for each client, achieving competitive or state-of-the-art results in recent benchmarks [45, 9], we focus on the collaborative learning of meta-nets which generates a personalized set of hyperparameters used during fine-tuning to further improve personalized performance without compromising global performance.
Non-III Problem Setup. Unlike many previous works, our method aims to handle both common label and feature distribution shift across clients. Specifically, given features  $x$  and labels  $y$ , we can rewrite the joint probability  $P_{i}(x,y)$  as  $P_{i}(x|y)P_{i}(y)$  and  $P_{i}(y|x)P_{i}(x)$  following [28]. We focus on three data heterogeneity settings found in many realistic settings: both label & distribution skew in which the marginal distributions  $P_{i}(y) \& P_{i}(x)$  may vary across clients, respectively, and concept drift in which the conditional distribution  $P_{i}(x|y)$  may vary across clients.
3.2 FedL2P: FL of Personalization Strategies
We now present our proposed method, FedL2P, to tackle collaborative learning of personalization strategies under data heterogeneity. Our main motivation is that the choice of the hyperparameters for the personalized fine-tuning, such as the learning rates and feature mean and standard deviations (SD) statistics of the batch normalization [26] (BN) layers, is very crucial. Although existing FL-HPO approaches aim to learn these hyperparameters directly[22, 65, 30], we aim to do it in a meta learning or hypernetwork-like fashion: learning a neural network (dubbed meta-net) that takes certain profiles of client data (e.g., summary statistics of personalized data) as input and outputs near-optimal hyperparameters. The meta-net is learned collaboratively in an FL manner without sharing local data. The returned hyperparameters from the meta-net are then deployed in client's personalized fine-tuning. The main advantage of this meta-net approach over the direct HPO is that a brand new client needs not do time-consuming HPO, but just gets their optimal hyperpameters by a single feed-forward pass through the meta-net. Our idea is visualized in Fig. 1. For each FL round, the latest meta-net is distributed to the participating clients; each client then performs meta learning to update the meta-net; the local updated meta-nets are sent to the server for aggregation. Details of our algorithm are described below.
Hyperparameter Selection. There is a wide range of local training hyperparameters that can be optimized, some of which were explored in previous FL works [22, 65, 30]. As local training is often a costly process, we narrowed it down to two main sets of hyperparameters based on previous works that showed promising results in dealing with non-IID data: BN hyperparameters and selective update hyperparameters.
Batch Normalization Hyperparameters. The first set of hyperparameters involves BN and explicitly deals with feature shift. Notably, FedBN [35] proposed keeping the BN layers local and the other layers global to better handle feature shift across clients; BN has found success at mitigating domain
To our knowledge, no previous FL-HPO works learn BN statistics, crucial for dealing with feature shifts.
shifts in various domain adaptation tasks [36, 8] and can be formulated as follows:
$$
g (x) = \frac {x - \hat {\mu}}{\sqrt {\hat {\sigma} ^ {2} + \epsilon}} * \gamma + \delta \tag {1}
$$
where  $g$  is a BN layer,  $x$  is its input features,  $(\hat{\mu},\hat{\sigma}^2)$  are the estimated running mean, variance, and  $\epsilon$  is used for numerical stability. During training, the batch statistics  $E(x)$  and  $Var(x)$  are used instead, keeping running estimates of them,  $\hat{\mu}$  and  $\hat{\sigma}^2$ . During inference, these estimates are used<sup>3</sup>. Both  $\gamma$  and  $\delta$  are learnable parameters used to scale and shift the normalized features.
Although deploying local BN layers is useful to counteract the drawbacks of feature shift, using global BN layers is often beneficial for local fine-tuning in cases where the feature shift is minimal as it helps speed-up convergence and reduces overfitting. Hence, we propose learning a hyperparameter  $\beta$  for each BN layer as follows:
$$
\hat {\mu} = (1 - \beta) * \hat {\mu} _ {p t} + \beta * \hat {\mu} _ {i}, \quad \hat {\sigma} ^ {2} = (1 - \beta) * \hat {\sigma} _ {p t} ^ {2} + \beta * \hat {\sigma} _ {i} ^ {2} \tag {2}
$$
where  $\hat{\mu}_{pt}$  and  $\hat{\sigma}_{pt}^2$  are the estimated running mean and variance by the given pretrained model;  $\hat{\mu}_i$  and  $\hat{\sigma}_i^2$  are the running mean and variance estimated by the client. When  $\beta \to 0$ , the client solely uses pretrained model's BN statistics and when  $\beta \to 1$ , it uses its local statistics to normalize features. Thus  $\beta$  indicates the degree in which the client should utilize its own BN statistics against pretrained model's, to handle the feature shift in its local data distribution.
Selective Update Hyperparameters. A variety of recent personalized FL works achieved promising results by manually selecting and fine-tuning a sub-module of  $\theta_{g}$  during personalization [48, 2, 37], (e.g., the feature extractor or the classifier layers), while leaving the remaining modules in  $\theta_{g}$  frozen. It is beneficial because it allows the designer to manage over- vs under-fitting in personalization. e.g., if the per-client dataset is small, then fine-tuning many parameters can easily lead to overfitting, and thus better freezing some layers during personalization. Alternatively, if the clients differ significantly from the global model and/or if the per-client dataset is larger, then more layers could be beneficially personalised without underfitting. Clearly the optimal configuration for allowing model updates depends on the specific scenario and the specific client. Furthermore, it may be beneficial to consider a more flexible range of hyperparameters that control a continuous degree of fine-tuning strength, rather than a binary frozen/updated decision per module.
To automate the search for good personalization strategies covering a range of wider and more challenging non-IID setups, we consider layer-wise learning rates,  $\eta$  for all learnable weights and biases. This parameterization of personalization encompasses all the previous manual frozen/updated split approaches as special cases. Furthermore, while these approaches have primarily considered heterogeneity in the marginal label distribution, we also aim to cover feature distribution shift between clients. Thus we also include the learn
Algorithm 1 FedL2P: FL of meta-nets for Personalization Hyperparameters
Input: Pretrained global model  $\theta_{g}$  with  $M$  layers and  $B$  BN layers, each has  $J$  input channels. Fraction ratio  $r$ . Total no. of clients  $C$ . No. of update iterations  $K$ . Training loss  $\mathcal{L}_T$  and validation loss  $\mathcal{L}_V$ .  $\zeta$  is the learning rate for  $\lambda$ .
1: initialize  $\lambda = \{\pmb{w}_{bn},\pmb{w}_{lr},\tilde{\pmb{\eta}}\}$
2: for round  $t = 1,\dots ,T$  do
3:  $\tilde{C} \gets$  Random sample  $Cr$  clients
4: for client  $i \in \tilde{C}$  do
5: send  $\theta_{g},\lambda$  to client
6: Forward pass of local dataset to compute 1)  $E(x_{m}),SD(x_{m})$  for  $1\leq m\leq M - 1$  2)  $\mu_{b,j},\sigma_{b,j}^2$  for  $b = 1,\dots ,B$  and  $j = 1,\ldots ,J$
7: Compute  $\xi_{b}$  for  $b = 1,\dots ,B$  using Eq.3
8: for iteration  $k = 1, \dots, K$  do
9:  $\theta_{i}\gets$  Finetune  $\theta_{g}$  using  $\mathcal{L}_T$  for  $e$  epochs
10:  $\lambda \gets \lambda -\zeta$  Hypergradient  $(\mathcal{L}_V,\mathcal{L}_T,\lambda ,\theta_i)$
11: end for
12: send  $\lambda$  and num of data samples  $N$  to server
13: end fo
14:  $\lambda \leftarrow \sum_{i}^{C}\frac{N_{i}}{\sum N_{i}}\lambda_{i}$
15: end for
Output:  $\lambda$
Algorithm 2 Hypergradient
Input: Validation Loss  $\mathcal{L}_V$  and training Loss  $\mathcal{L}_T$ .
Learning rate  $\psi$  and no. of iterations  $Q$ . Fixed point  $(\lambda', \theta^*(\lambda'))$ .
1:  $p = v = \partial_{\theta}\mathcal{L}_V|_{(\lambda ',\theta^* (\lambda '))}$
2: for iteration 1, ...,  $Q$  do
3:  $v\gets v - \psi v\partial_{\theta^{*}}^{\prime}\mathcal{L}_{T}$
4:  $p \gets p + v$
5: end for
Output:  $\left.\partial_{\lambda}\mathcal{L}_{V}\right|_{\left(\lambda^{\prime},\theta^{*}\left(\lambda^{\prime}\right)\right)} - p\left.\partial_{\lambda \theta}\mathcal{L}_{T}\right|_{\left(\lambda^{\prime},\theta^{*}\left(\lambda^{\prime}\right)\right)}$
ing rates for the BN parameters,  $\gamma$  and  $\delta$ , allowing us to further tackle feature shift by adjusting the means and SD of the normalized features.
3When the running mean & variance are not tracked, the batch statistics is used in both training and inference.
Hyperparameter Inference for Personalization. We aim to estimate a set of local hyperparameters that can best personalize the pretrained model for each client given a group of clients whose data might not be used for pretraining. To accomplish this, we learn to estimate hyperparameters based on the degree of data heterogeneity of the client's local data with respect to the data that the model is pretrained on. There are many ways to quantify data heterogeneity, such as utilizing the earth mover's distance between the client's data distribution and the population distribution for label distribution skew [63] or taking the difference in local covariance among clients for feature shift [35]. In our case, we aim to distinguish both label and feature data heterogeneity across clients. To this end, we utilize the client's local input features to each layer with respect to the given pretrained model. Given a pretrained model with  $M$  layers and  $B$  BN layers, we learn  $\eta = \eta_{1},\ldots ,\eta_{2M}^{4}$  and  $\beta = \beta_{1},\dots,\beta_{B}$ , using two functions, each of which is parameterized as a multilayer perceptron (MLP), named meta-net, with one hidden layer due to its ability to theoretically approximate almost any continuous function [12, 57]. We named the meta-net that estimates  $\beta$  and the meta-net that estimates  $\eta$  as BNNet and LRNet respectively. Details about the architecture can be found in Appendix. E.
To estimate  $\beta$ , we first perform a forward pass of the local dataset on the given pretrained model, computing the mean and SD of each channel of each input feature for each BN layer. We then measure the distance between the local feature distributions and the pretrained model's running estimated feature distributions of the  $b$ -th BN layer as follows:
$$
\xi_ {i, b} = \frac {1}{J} \sum_ {j = 1} ^ {J} \frac {1}{2} \left(D _ {K L} \left(P _ {j} \mid \mid Q _ {j}\right) + D _ {K L} \left(Q _ {j} \mid \mid P _ {j}\right)\right), \tag {3}
$$
where  $P_{j} = \mathcal{N}(\mu_{i,b,j},\sigma_{i,b,j}^{2})$ ,  $Q_{j} = \mathcal{N}(\hat{\mu}_{pt,b,j},\hat{\sigma}_{pt,b,j}^{2})$ ,  $D_{KL}$  is the KL divergence and  $J$  is the number of channels of the input feature.  $\xi$  is then used as an input to BNNet, which learns to estimate  $\beta$  as shown in Eq. 4.
Similarly, we compute the mean and SD of each input feature per layer by performing a forward pass of the local dataset on the pretrained model and use it as an input to LRNet. Following best practices from previous non-FL hyperparameter optimization works [49, 3], we use a learnable post-multiplier  $\tilde{\eta} = \tilde{\eta}_{1},\dots,\tilde{\eta}_{2M}$  to avoid limiting the range of the resulting learning rates (Eq 4).
$$
\begin{array}{l} \boldsymbol {\beta} = \operatorname {B N N e t} \left(\boldsymbol {w} _ {b n}; \xi_ {1}, \xi_ {2}, \dots , \xi_ {B - 1}, \xi_ {B}\right) \\ \boldsymbol {\eta} = \operatorname {L R N e t} \left(\omega_ {b h}, \zeta_ {1}, \zeta_ {2}, \dots , \zeta_ {B - 1}, \zeta_ {B}\right) \tag {4} \\ \end{array}
$$
where  $\odot$  is the Hadamard product,  $x_{m-1}$  is the input feature to the  $m$ -th layer, and  $\boldsymbol{w}_{bn}$  and  $\boldsymbol{w}_{lr}$  are the parameters of BNNet and LRNet respectively.  $\beta$  is used to compute the running mean and variance in the forward pass for each BN layer as shown in Eq. 2 and  $\eta$  is used as the learning rate for each weight and bias in the backward pass. We do not restrict  $\tilde{\eta}$  to be positive as the optimal learning rate might be negative [5].
Federated Hyperparameter Learning. We deploy FedAvg [46] to federatedly learn a set of client-specific personalization strategies. Specifically, we learn the common meta-net  $\lambda = \{\pmb{w}_{bn},\pmb{w}_{lr},\tilde{\eta}\}$  that generates client-wise personalization hyperparameters  $\{\beta_i,\eta_i\}$ , such that a group of clients can better adapt a pre-trained model  $\theta_{g}$  by fine-tuning to their local data distribution. So we solve:
$$
\begin{array}{l} \min  _ {\lambda} \mathcal {F} (\lambda , \theta_ {g}) = \sum_ {i = 1} ^ {C} \frac {N _ {i}}{\sum_ {i ^ {\prime}} N _ {i ^ {\prime}}} \mathcal {L} _ {i, V} \left(\theta_ {i} ^ {*} (\lambda), \lambda\right) \\ \text {s . t .} \theta_ {i} ^ {*} (\lambda) = \underset {\theta_ {i}} {\arg \min } \mathcal {L} _ {i, T} \left(\theta_ {i}, \lambda\right) \tag {5} \\ \end{array}
$$
where  $\theta_{i}^{*}$  is the set of optimal personalized model parameters after fine-tuning  $\theta_{g}$  for  $e$  epochs on the local dataset,  $\mathcal{L}_{i,V}(\theta ,\lambda) = \mathbb{E}_{(x,y)\sim V_i}\mathcal{L}_i(\theta ,\lambda ;x,y)$  and  $V_{i}$  is the validation set (samples from  $P_{i}$ ) for the client  $i$  - similarity for  $L_{i,T}$ .
For each client  $i$ , the validation loss gradient with respect to  $\lambda$ , known as the hypergradient, can be computed as follows:
$$
d _ {\lambda} \mathcal {L} _ {V} \left(\theta^ {*} (\lambda), \lambda\right) = \partial_ {\lambda} \mathcal {L} _ {V} \left(\theta^ {*} (\lambda), \lambda\right) + \partial_ {\theta^ {*} (\lambda)} \mathcal {L} _ {V} \left(\theta^ {*} (\lambda), \lambda\right) \partial_ {\lambda} \theta^ {*} (\lambda) \tag {6}
$$
To compute  $\partial_{\lambda}\theta^{*}$  in Eq. 6, we use the implicit function theorem (IFT):
4We assume all layers have weights and biases here.