SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yuyang Deng
Pennsylvania State University  
yzd82@psu.edu
Mohammad Mahdi Kamani
Wyze Labs
mmkamani@alumni.psu.edu
Pouria Mahdavinia
Pennsylvania State University  
pxm5426@psu.edu
Mehrdad Mahdavi
mzm616@psu.edu
This paper advocates a new paradigm Personalized Empirical Risk Minimization (PERM) to facilitate learning from heterogeneous data sources without imposing stringent constraints on computational resources shared by participating devices. In PERM, we aim to learn a distinct model for each client by learning who to learn with and personalizing the aggregation of local empirical losses by effectively estimating the statistical discrepancy among data distributions, which entails optimal statistical accuracy for all local distributions and overcomes the data heterogeneity issue. To learn personalized models at scale, we propose a distributed algorithm that replaces the standard model averaging with model shuffling to simultaneously optimize PERM objectives for all devices. This also allows us to learn distinct model architectures (e.g., neural networks with different numbers of parameters) for different clients, thus confining underlying memory and compute resources of individual clients. We rigorously analyze the convergence of the proposed algorithm and conduct experiments that corroborate the effectiveness of the proposed paradigm.
Recently federated learning (FL) has emerged as an alternative paradigm to centralized learning to encourage federated model sharing and create a framework to support edge intelligence by shifting model training and inference from data centers to potentially scattered—and perhaps self-interested—systems where data is generated [1]. While undoubtedly being a better paradigm than centralized learning, enabling the widespread adoption of FL necessitates foundational advances in the efficient use of statistical and computational resources to encourage a large pool of individuals or corporations to share their private data and resources. Specifically, due to heterogeneity of data and compute resources among participants, it is necessary, if not imperative, to develop distributed algorithms that are i) cognizant of statistical heterogeneity (data-awareness) by designing algorithms that effectively deal with highly heterogeneous data distributions across devices; and ii) confined to learning models that meet available computational resources of participant devices (system-awareness).
To mitigate the negative effect of data heterogeneity (non-IIDness), two common approaches are clustering and personalization. The key idea behind the clustering-based methods [2, 3, 4, 5] is to partition the devices into clusters (coalitions) of similar data distributions and then learn a single shared model for all clients within each cluster. While appealing, the partitioning methods are limited to heuristic ideas such as clustering based on the geographical distribution of devices without taking the actual data distributions into account and lack theoretical guarantees or postulate strong assumptions on initial models or data distributions [4, 5]. In personalization-based methods [6, 7, 8, 2, 9, 10, 11], the idea is to learn a distinct personalized model for each device alongside the global model, which
can be unified as minimizing a bi-level optimization problem [12]. Personalization aims to learn a model that has the generalization capabilities of the global model but can also perform well on the specific data distribution of each participant suffers from a few key limitations. First, as the number of clients grows, while the number of training data increases, the number of parameters to be learned increases which limits to increase in the number of clients beyond a certain point to balance data and overall model complexity tradeoff-- a phenomenon known as incidental parameters problem [13]. Moreover, since the knowledge transfer among data sources happens through a single global model, it might lead to suboptimal results. To see this, consider an extreme example, where half of the users have identical data distributions, say  $\mathcal{D}$ , while the other half share a data distribution that is substantially different, say  $-\mathcal{D}$  (e.g. two distributions with same marginal distribution on features but opposite labeling functions). In this case, the global model obtained by naively aggregating local models (e.g., fixed mixture weights) converges to a solution that suffers from low test accuracy on all local distributions which makes it preferable to learn a model for each client solely based on its local data or carefully chosen subset of data sources.
Focusing on system heterogeneity, most existing works require learning models of identical architecture to be deployed across the clients and server (model-homogeneity) [14, 15], and mostly focus on reducing number [15] or size [16, 17, 18] of communications or sampling handling chaotic availability of clients [19, 20]. The requirement of the same model makes it infeasible to train large models due to system heterogeneity where client devices have drastically different computational resources. A few recent studies aim to overcome this issue either by leveraging knowledge distillation methods [21, 22, 23, 24] or partial training (PT) strategies via model subsampling (either static [25, 26], random [27], or rolling [28]). However, KD-based methods require having access to a public representative dataset of all local datasets at server and ignore data heterogeneity in the distillation stage to a large extent. The focus of PT training methods is mostly on learning a single server model using heterogeneous resources of devices and does not aim at deploying a model onto each client after the global server model is trained (which is left as a future direction in [28]). The aforementioned issues lead to a fundamental question: "What is the best strategy to learn from heterogeneous data sources to achieve optimal accuracy w.r.t. each data source, without imposing stringent constraints on computational resources shared by participating devices?"
We answer this question affirmatively, by proposing a new data&system-aware paradigm dubbed Personalized Empirical Risk Minimization (PERM), to facilitate learning from massively fragmented private data under resource constraints. Motivated by generalization bounds in multiple source domain adaptation [29, 30, 31, 32], in PERM we aim to learn a distinct model for each client by personalizing the aggregation of empirical losses of different data sources which enables each client to learn who to learn with using an effective method to empirically estimate the statistical discrepancy between their associated data distributions. We argue that PERM entails optimal statistical accuracy for all local distributions, thus overcoming the data heterogeneity issue. PERM can also be employed in other learning settings with multiple heterogeneous sources of data such as domain adaptation and multi-task learning to entail optimal statistical accuracy. While PERM overcomes the data heterogeneity issue, the number of optimization problems (i.e., distinct personalized ERMs) to be solved scales linearly with the number of data sources. To simultaneously optimize all objectives in a scalable and computationally efficient manner, we propose a novel idea that replaces the standard model averaging in distributed learning with model shuffling and establish its convergence rate. This also allows us to learn distinct model architectures (e.g., neural networks with different number of parameters) for different clients, thus confining to underlying memory and compute resources of individual clients, and overcoming the system heterogeneity issue. This addresses an open question in [28] where only a single global model can be trained in a model-heterogeneous setting, while PERM allows deploying distinct models for different clients. We empirically evaluate the performance of PERM, which corroborates the statistical benefits of PERM in comparison to existing methods.
2 Personalized Empirical Risk Minimization
In this section, we formally state the problem and introduce PERM as an ideal paradigm for learning from heterogeneous data sources. We assume there are  $N$  distributed devices where each holds a distinct data shard  $S_{i} = \{(x_{i,j},y_{i,j})\}_{j = 1}^{n_{i}}$  with  $n_i$  training samples that are realized by a local source distribution  $\mathcal{D}_i$  over instance space  $\Xi = \mathcal{X}\times \mathcal{Y}$ . The data distributions across the devices are not independently and identically distributed (non-IID or heterogeneous), i.e.,  $\mathcal{D}_1\neq \mathcal{D}_2\neq$
$\ldots \neq \mathcal{D}_N$ , and each distribution corresponds to a local generalization error or true risk  $\mathcal{L}_i(h) = \mathbb{E}_{(\boldsymbol{x},y)\sim \mathcal{D}_i}[\ell (h(\boldsymbol{x}),y)], i = 1,2,\ldots ,N$  on unseen samples for any model  $h\in \mathcal{H}$ , where  $\mathcal{H}$  is the hypothesis set (e.g., a linear model or a deep neural network) and  $\ell :\mathcal{V}\times \mathcal{V}\to \mathbb{R}^{+}$  is a given convex or non-convex loss function. We use  $\widehat{\mathcal{L}}_i(h) = (1 / n_i)\sum_{(\boldsymbol{x},y)\in S_i}\ell (h(\boldsymbol{x}),y))$  to denote the local empirical risk or training loss at ith data shard  $S_{i}$  with  $n_i$  samples.
We seek to collaboratively learn a model or personalized models that entail a good generalization on all local distributions, i.e. minimizing true risk  $\mathcal{L}_i(\cdot), i = 1,\dots ,N$  for all data sources (all-for-all [33]). A simple non-personalized solution, particularly in FL, aims to minimize a (weighted) empirical risk over all data shards in a communication-efficient manner [34]:
$$
\arg \min  _ {h \in \mathcal {H}} \sum_ {i = 1} ^ {N} p (i) \widehat {\mathcal {L}} _ {i} (h) \text {w i t h} \boldsymbol {p} \in \Delta_ {N}, \tag {WERM}
$$
where  $\Delta_N = \{\pmb{p} \in \mathbb{R}_+^N \mid \sum_{i=1}^N p(i) = 1\}$  denotes the simplex set.
It has been shown that a single model learned by WERM, for example by using fixed mixing weights  $p(i) = n_i / n$ , where  $n$  is total number of training samples, or even agnostic to mixture of distributions [35, 36], while yielding a good performance on the combined datasets of all devices, can suffer from a poor generalization error on individual datasets by increasing the diversity among distributions [37, 38, 39, 40]. To overcome this issue, there has been a surge of interest in developing methods that personalize the global model to individual local distributions. These methods can be unified as the following bi-level problem (a similar unification has been made in [12]):
$$
\arg \min  _ {h _ {1}, h _ {2}, \dots , h _ {m} \in \mathcal {H}} \widehat {\mathcal {F}} _ {i} \left(h _ {i} \oplus h _ {*}\right) \quad \text {s u b j e c t t o} \quad h _ {*} = \arg \min  _ {h \in \mathcal {H}} \sum_ {j = 1} ^ {N} \alpha (j) \widehat {\mathcal {L}} _ {j} (h) (\text {B E R M})
$$
where  $\oplus$  denotes the mixing operation to combine local and global models, and  $\widehat{\mathcal{F}}_i$  is a modified local loss which is not necessarily same as local risk  $\widehat{\mathcal{L}}_i$ . By carefully designing the local loss  $\widehat{\mathcal{F}}_i$  and mixing operation  $\oplus$ , we can develop different penalization schemes for FL including existing methods such as linearly interpolating global and local models [11, 2], multi-task learning [10] and meta-learning [9] as special cases. For example, BERM reduces to zero-personalization objective WERM when  $h_i \oplus h_* = h_*$ , and  $\widehat{\mathcal{F}}_i = \widehat{\mathcal{L}}_i$ . At the other end of the spectrum lies the zero-collaboration where the  $i$ th client trains its own model without any influence from other clients by setting  $h_i \oplus h_* = h_i$ ,  $\widehat{\mathcal{F}}_i = \widehat{\mathcal{L}}_i$ . The personalized model with interpolation of global and local models can be recovered by setting  $h_i \oplus h_* = \alpha h_i + (1 - \alpha)h_*$ , and  $\widehat{\mathcal{F}}_i = \widehat{\mathcal{L}}_i$ . While more effective than a single global model learned via WERM, personalization methods suffer from three key issues: i) the global model is still obtained by minimizing the average empirical loss which might limit the statistical benefits of collaboration, ii) overall model complexity increases linearly with number of clients, and iii) a same model space is shared across servers and clients.
To motivate our proposal, let us consider the empirical loss  $\sum_{i=1}^{N} \alpha(i) \widehat{\mathcal{L}}_i(h)$  in WERM (or the inner level objective in BERM) with fixed mixing weights  $\alpha \in \Delta_N$ , and denote the optimal solution by  $\widehat{h}_{\alpha}$ . The excess risk of the learned model  $\widehat{h}_{\alpha}$  on  $i$ th local distribution  $\mathcal{D}_i$  w.r.t. the optimal local model  $h_i^* = \arg \min_{h \in \mathcal{H}} \mathcal{L}_{\mathcal{D}_i}(h)$  (i.e., all-for-one) can be bounded by (informal) [31]
$$
\mathcal {L} _ {i} \left(\widehat {h} _ {\boldsymbol {\alpha}}\right) \leq \mathcal {L} _ {i} \left(h _ {i} ^ {*}\right) + \sum_ {j = 1} ^ {N} \alpha (j) \mathrm {R} _ {j} (\mathcal {H}) + 2 \sum_ {j = 1} ^ {N} \alpha (j) \operatorname {d i s c} _ {\mathcal {H}} \left(\mathcal {D} _ {j}, \mathcal {D} _ {i}\right) + C \sqrt {\sum_ {j = 1} ^ {N} \frac {\alpha (j) ^ {2}}{n _ {j}}} \quad (\text {G E N})
$$
where  $\mathrm{R}_j(\mathcal{H})$  is the empirical Rademacher complexity  $\mathcal{H}$  w.r.t.  $S_{j}$ , and  $\mathrm{disc}_{\mathcal{H}}(\mathcal{D}_i,\mathcal{D}_j)$  is a pseudodistance on the set of probability measures on  $\Xi$  to assess the discrepancy between the distributions  $\mathcal{D}_i$  and  $\mathcal{D}_j$  with respect to the hypothesis class  $\mathcal{H}$  as defined below [29]:
Definition 1. For a model space  $\mathcal{H}$  and  $\mathcal{D},\mathcal{D}'$  two probability distributions on  $\Xi = \mathcal{X}\times \mathcal{Y}$
$$
\operatorname {d i s c} _ {\mathcal {H}} \left(\mathcal {D}, \mathcal {D} ^ {\prime}\right) = \sup  _ {h \in \mathcal {H}} \left| \mathbb {E} _ {\xi \sim \mathcal {D}} \left(\ell (h, \xi)\right) - \mathbb {E} _ {\xi^ {\prime} \sim \mathcal {D} ^ {\prime}} \left(\ell (h, \xi^ {\prime})\right) \right|
$$
Intuitively, the discrepancy between the two distributions is large, if there exists a predictor that performs well on one of them and badly on the other. On the other hand, if all functions in the hypothesis class perform similarly on both, then  $\mathcal{D}$  and  $\mathcal{D}'$  have low discrepancy. The above metric which is a special case of a popular family of distance measures in probability theory and mathematical
statistics known as integral probability metrics (IPMs) [41], can be estimated from finite data by replacing the expected losses with their empirical counterparts (i.e.  $\mathcal{L}_i$  with  $\widehat{\mathcal{L}}_i$ ).
From GEN, it can be observed that a mismatch between pairs of distributions limits the benefits of ERM on all distributions. Indeed, the generalization risk w.r.t.  $\mathcal{D}_i$  will significantly increase when the distribution divergence terms  $\mathrm{disc}_{\mathcal{H}}(\mathcal{D}_j,\mathcal{D}_i)$  are large. It leads to an ideal sample complexity  $1 / \sqrt{n}$  where  $n = n_1 + n_2 + \ldots +n_N$  is the total number of samples, which could have been obtained in the IID setting with  $\alpha (j) = 1 / N$  when the divergence is small as the pairwise discrepancies disappear. Also, we note that even if the global model achieves a small training error over the union of all data (e.g., over parametrized setting) and can entail a good generalization error with respect to average distribution, the divergence term still remains which illustrates the poor performance of the global model on all local distributions  $\mathcal{D}_i,i = 1,2,\dots ,N$ . This implies that even personalization of the global model as in BERM can not entail a good generalization on all local distributions as there is no effective transfer of positive knowledge among data sources in the presence of high data heterogeneity among local distributions (similar impossibility results even under seemingly generous assumptions on how distributions relate have been made in multisource domain adaptation as well [42]).
Interestingly the bound suggests that seeking optimal accuracy on all local distributions requires choosing a distinct mixing of local losses for each client  $i$  that minimizes the right-hand side of GEN. This indicates that in an ideal setting (i.e. all-for-all), we can achieve the best accuracy for each local distribution  $\mathcal{D}_i$  by personalizing the WERM, i.e., (i) first estimating  $\alpha_{i}, i = 1,2,\dots,N$  for each client individually, then (ii) solving a variant of WERM for each client with obtained mixing parameters:
$$
\arg \min  _ {h \in \mathcal {H} _ {i}} \sum_ {j = 1} ^ {N} \alpha_ {i} (j) \widehat {\mathcal {L}} _ {j} (h) \quad \text {f o r} \quad i = 1, 2, \dots , N. \tag {PERM}
$$
By doing this each device achieves the optimal local generalization error by learning who to learn with based on the number of samples at each source and the mismatch between its data distribution with other clients. We also note that compared to WERM and BERM, in PERM since we solve a different aggregated empirical loss for each client, we can pick a different model space/model architecture  $\mathcal{H}_i$  for each client to meet its available computational resources.
While this two-stage method is guaranteed to entail optimal test accuracy for all local distributions  $\mathcal{D}_i$ , however, making it scalable requires overcoming two issues. First, estimating the statistical discrepancies between each pair of data sources (i.e.,  $\alpha_i, i = 1, \dots, N$ ) is a computing burden as it requires solving  $O(N^2)$  difference of (non)-convex functions in a distributed manner and requires enough samples form each source to entail good accuracy on estimating pairwise discrepancies [41]. Second, we need to solve  $N$  variants of the optimization problem in PERM, possibly each with a different model space, which is infeasible when the number of devices is huge (e.g., cross-device federated learning). In the next section, we propose a simple yet effective idea to overcome these issues in a computationally efficient manner.
3 PERM at Scale via Model Shuffling
In this section, we propose a method to efficiently estimate the empirical discrepancies among data sources followed by a model shuffling idea to simultaneously solve  $N$  versions of PERM to learn a personalized model for each client. We first start by proposing a two-stage algorithm: estimating mixing parameters followed by model shuffling. Then, we propose a single loop unified algorithm that enjoys the same computation and communication overhead as BERM (twice communication of FedAvg). For ease of exposition, we discuss the proposed algorithms by assuming all the clients share the same model architecture and later on discuss the generalization to heterogeneous model spaces. Specifically, we assume that the model space  $\mathcal{H}$  is a parameterized by a convex set  $\mathcal{W} \subseteq \mathbb{R}^d$  and use  $f_i(\boldsymbol{w}) \coloneqq \widehat{\mathcal{L}}_i(\boldsymbol{w}) = \sum_{(\boldsymbol{x}, y) \in S_i} \ell(\boldsymbol{w}; (\boldsymbol{x}, y))$  to denote the empirical loss at  $i$ th data shard.
3.1 Warmup: a two-stage algorithm
We start by proposing a two-stage method for solving  $N$  variants of PERM in parallel. In the first stage, we propose an efficient method to learn the mixing parameters for all clients. Then, in stage two, we propose a model shuffling method to solve all personalized empirical losses in parallel.
Stage 1: Mixing parameters estimation. In the first stage we aim to efficiently estimate the pairwise discrepancy among local distributions to construct mixing parameters  $\alpha_{i}, i = 1,2,\ldots N$ . From
generalization bound GEN and Definition 1, a direct solution to estimate  $\alpha_{i}$  is to solve the following convex-nonconcave minimax problem for each client:
$$
\boldsymbol {\alpha} _ {i} ^ {*} = \arg \min  _ {\boldsymbol {\alpha} \in \Delta_ {N}} \sum_ {j = 1} ^ {N} \alpha (j) \max  _ {\boldsymbol {w} \in \mathcal {W}} | f _ {i} (\boldsymbol {w}) - f _ {j} (\boldsymbol {w}) | + \sum_ {j = 1} ^ {N} \alpha (j) ^ {2} / n _ {j} \tag {1}
$$
where we estimate the true risks in pairwise discrepancy terms with their empirical counterparts and drop the complexity term as it becomes identical for all sources by fixing the hypothesis space  $\mathcal{H}$  and bounding it with a computable distribution-independent quantity such as VC dimension [43], or it can be controlled by choice of  $\mathcal{H}$  or through data-dependent regularization. However, solving the above minimax problem itself is already challenging: the inner maximization loop is a nonconcave (or difference of convex) problem, so most of the existing minimax algorithms will fail on this problem. To our best knowledge, the only provable deterministic algorithm is [44], and it is hard to generalize it to stochastic and distributed fashion. Moreover, since we have  $N$  clients, we need to solve  $N$  variants of (1), which makes designing a scalable algorithm even harder.
To overcome aforementioned issues, we make two relaxations to estimate the per-client mixing parameters. First, we optimize an upper bound of pairwise empirical discrepancies  $\sup_{\boldsymbol{w}} |f_i(\boldsymbol{w}) - f_j(\boldsymbol{w})|$  in terms of gradient dissimilarity between local objectives  $\| \nabla f_i(\boldsymbol{w}) - \nabla f_j(\boldsymbol{w})\|$  [45], which quantifies how different the local empirical losses are and widely used in analysis of learning from heterogeneous losses as in FL [46]. Second, given that the discrepancy measure based on the supremum could be excessively pessimistic in real-world scenarios, and drawing inspiration from the concept of average drift at the optimal point as a right metric to measure the effect of data heterogeneity in federated learning [47], we propose to measure discrepancy at the optimal solution obtained by solving WERM, i.e.,  $\boldsymbol{w}^* := \arg \min_{\boldsymbol{w} \in \mathcal{W}} (1/N) \sum_{i=1}^{N} f_i(\boldsymbol{w})$ . By doing this, the problem reduces to a simple minimization for each client, given the optimal global solution. These two relaxations lead to solving the following tractable optimization problem to decide the per-client mixing parameters:
$$
\arg \min  _ {\boldsymbol {\alpha} \in \Delta_ {N}} g _ {i} (\boldsymbol {w} ^ {*}, \boldsymbol {\alpha}) := \sum_ {j = 1} ^ {N} \alpha (j) \| \nabla f _ {i} (\boldsymbol {w} ^ {*}) - \nabla f _ {j} (\boldsymbol {w} ^ {*}) \| ^ {2} + \lambda \sum_ {j = 1} ^ {N} \alpha (j) ^ {2} / n _ {j} \tag {2}
$$
where we added a regularization parameter  $\lambda$  and used the squared of gradient dissimilarity for computational convenience. Thus, obtaining all  $N$  mixing parameters requires solving a single ERM to obtain optimal global solution and  $N$  variants of (2). To get the optimal solution in a communication-reduced manner, we adapt the Local SGD algorithm [48] (or FedAvg [14]) and find the optimal solution in intermittent communication setting [49] where the clients work in parallel and are allowed to make  $K$  stochastic updates between two communication rounds for  $R$  consecutive rounds. The detailed steps are given in Algorithm A1 in Appendix B for completeness. After obtaining the global model  $\boldsymbol{w}^R$  we optimize over  $\alpha$  in  $g_i(\boldsymbol{w}^R, \boldsymbol{\alpha})$  using  $T_{\alpha}$  iterations of GD to get  $\hat{\alpha}_i$ . Actually, we will show that as long as  $\boldsymbol{w}^R$  converge to  $\boldsymbol{w}^*$ ,  $\hat{\alpha}_i, i = 1, \dots, N$  converges to solution of (2) very fast. Our proof idea is based on the following Lipschitzness observation: