SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we compare our collaborative learning procedure with other popular least squares techniques based on imputation and comment on the statistical efficacy and communication cost differences. We summarize our analysis in Table 1. The proofs of the theorems are in Appendix C. For the sake of brevity, our results here will center around estimating  $\theta$ . Analogous results centered around estimating  $T_{i}\theta$  follow directly. These results can be connected to the full-feature and local test error by the same argument we made earlier to derive (5).
Local imputation w/ collaboration. Suppose a coordinating server collected covariance information  $\Sigma_{i}$  from each agent and then distribute  $\Sigma$  back to each of them. Then one intuitive strategy is to use this information to impute each agent's local data by replacing  $X_{i+}$  with  $\mathbb{E}[X_{i} \mid X_{i+}] = X_{i+}T_{i}$ , before performing local linear regression. In other words, instead of computing  $\hat{\theta}_{i}$ , compute
$$
\hat {\theta} _ {i} ^ {\mathrm {i m p}} = (T _ {i} ^ {\top} X _ {i +} ^ {\top} X _ {i +} T _ {i}) ^ {\dagger} T _ {i} ^ {\top} X _ {i +} ^ {\top} y _ {i}
$$
to send back to the coordinating server. Note that we use Moore-Penrose inverse here as  $T_{i}^{\top}X_{i + }^{\top}X_{i + }T_{i}$  is in general of rank  $d_{i}$ , and  $\hat{\theta}_{i}^{\mathrm{imp}}$  is then the min-norm interpolant of agent  $i$ 's data. Similar to COLLAB, we can use weighted empirical risk minimization parameterized by  $W_{i}\in \mathbb{R}^{d\times d}$  and to aggregate  $\hat{\theta}^{\mathrm{imp}}$  via
$$
\hat {\theta} ^ {\mathrm {i m p}} = \hat {\theta} (W _ {1}, \dots , W _ {m}) := \underset {\theta} {\operatorname {a r g m i n}} \sum_ {i = 1} ^ {m} \left\| T _ {i} ^ {\top} (T _ {i} T _ {i} ^ {\top}) ^ {- 1} T _ {i} \theta - \hat {\theta} _ {i} ^ {\mathrm {i m p}} \right\| _ {W _ {i}} ^ {2}.
$$
The next theorem, in conjunction with Theorem 3.1, implies that under the WERM optimization scheme, aggregation of least squares estimators on imputed local data does not bring additional statistical benefit. In fact, the local imputation estimator is a linearly transformed on local OLS  $\hat{\theta}_i$ .
Theorem 4.1. For  $\hat{\theta}_i^{\mathrm{imp}}$  from agent  $i$ , we have  $\hat{\theta}_i^{\mathrm{imp}} = T_i^\top (T_i T_i^\top)^{-1} \hat{\theta}_i$ . Given any weighting matrices  $W_i \in \mathbb{R}^{d \times d}$ , the aggregated imputation estimator  $\hat{\theta}^{\mathrm{imp}}$  is consistent and asymptotically normal
$$
\sqrt {n} \left(\hat {\theta} ^ {\mathrm {i m p}} - \theta\right) = \mathsf {N} (0, C ^ {\mathrm {i m p}} (W _ {1}, \dots , W _ {m})).
$$
Using the same weights  $W_{i}^{\star} \in \mathbb{R}^{d_{i} \times d_{i}}$  as in Theorem 3.1 for aggregated  $\hat{\theta}^{\mathrm{imp}}$ , we have under p.s.d. cone order, for weights  $W_{i}$ ,  $C^{\mathrm{imp}}(W_1, \dots, W_m) \succeq C^{\star}$ , where  $C^{\star} = (\sum_{i=1}^{m} T_i^\top W_i^{\star} T_i)^{-1}$ . In addition, the equality holds when  $W_{i} = T_{i}^\top W_{i}^{\star} T_{i}$ .
As we will see in Sec. 5 where we provide minimax lower bound for weak observation models, the fact that the weighted imputation does not outperform our COLLAB approach is because the WERM on local OLS without imputation is already optimal. In fact, having access to the features will not achieve better estimation rate for both the global parameter  $\theta$  and local parameters  $T_{i}\theta$ .
In terms of communication cost, this local imputation method requires more communication than COLLAB, as a coordinating server needs to communicate  $\Sigma$  to all the agents. Each agent must first send  $d_{i}^{2}$  real numbers corresponding to  $\Sigma_{i}$  to the coordinating server, and then the coordinating server will send back  $d^{2} - d_{i}^{2}$  real numbers corresponding to the entries of  $\Sigma$  that agent  $i$  has not observed. This amounts to  $d^{2}$  total communication cost per agent. The per-agent communication cost for this method scales with  $d$ , which is not desirable for the reason outlined in the introduction.
Global imputation. Finally, we analyze the setting where we allow each agent to send the coordinating server all of their data  $(X_{i+}, y_i)$  for  $i = 1, \dots, m$  instead of their local estimators,  $\hat{\theta}_i$  or  $\hat{\theta}_i^{\mathrm{imp}}$ . Having all the data with structured missingness available, a natural idea is to have the coordinating server first impute the data, replacing  $X_{i+}$  with  $\mathbb{E}[X_i \mid X_{i+}] = X_{i+}T_i$ , and then perform weighted OLS on all of the  $nm$  data points, before sending a  $d_i$  dimensional local model back to each agent. Namely for scalars  $\alpha_1, \dots, \alpha_m > 0$ , we take
$$
\hat {\theta} ^ {\mathrm {i m p - g l b}} = \hat {\theta} ^ {\mathrm {i m p - g l b}} (\alpha_ {1}, \dots , \alpha_ {m}) := \left(\sum_ {i = 1} ^ {m} \alpha_ {i} T _ {i} ^ {\top} X _ {i +} ^ {\top} X _ {i +} T _ {i}\right) ^ {- 1} \left(\sum_ {i = 1} ^ {m} \alpha_ {i} T _ {i} ^ {\top} X _ {i +} ^ {\top} y _ {i}\right).
$$
Surprisingly, in spite of the additional power,  $\hat{\theta}^{\mathrm{imp - glb}}$  still does not beat  $\hat{\theta}$  in Theorem 3.1.
Theorem 4.2. For any scalars  $\alpha_{1},\dots ,\alpha_{m} > 0$ $\hat{\theta}^{\mathrm{imp - glb}}$  is consistent and asymptotically normal
$$
\sqrt {n} \left(\hat {\theta} ^ {\operatorname {i m p - g l b}} - \theta\right) = \mathsf {N} (0, C ^ {\operatorname {i m p - g l b}} (\alpha_ {1}, \dots , \alpha_ {m})).
$$
Recall the lower bound matrix  $C^\star \coloneqq (\sum_{i=1}^m T_i^\top W_i^\star T_i)^{-1}$  in Theorem 3.1 If  $X_i \sim \mathsf{N}(0, \Sigma)$ , we have under p.s.d. cone order and any  $\alpha_i > 0$ ,  $C^{\mathrm{imp - glb}}(\alpha_1, \dots, \alpha_m) \succeq C^\star$ . In addition, the equality holds when  $\alpha_i = 1 / (\| \theta_{i-} \|_{\Gamma_{i-}}^2 + \sigma^2)$ .
The communication cost for this method is significantly larger than the other methods we discussed. Each agent must send all of its data to a coordinating server. Factoring in the additional communication cost of receiving a new local model from the server, this amounts to a total  $nd_{i} + d_{i}$  communication cost per agent. The fact that communication cost for this method scales with  $n$  is a significant disadvantage for the reason we outlined in the introduction.
5 Asymptotic Local Minimax Lower Bounds
In this section, we prove asymptotic local minimax lower bounds that show COLLAB is (nearly) optimal. We work in the partially-fixed-design regime. For every sample  $x \in \mathbb{R}^d$ ,  $x_{i+} \in \mathbb{R}^{d_i}$  is a fixed vector. We draw  $x_{i-}$  from  $\mathsf{N}(\mu_{i-}, \Gamma_{i-})$  where  $\mu_{i-}$  and  $\Gamma_{i-}$  is the conditional mean and variance of  $x_{i-}$  given  $x_{i+}$ . Here  $\Gamma_{i-}$  is also the Schur complement. We draw  $x_{i-}$  from  $\mathsf{N}(\mu_{i-}, \Gamma_{i-})$ . The samples  $x_{i+} \in \mathbb{R}^{d_i}$  comprise the matrices  $X_{i+} \in \mathbb{R}^{n \times d_i}$ . For all  $i \in [m]$ , we will assume we have an infinite sequence (w.r.t.  $n$ ) of matrices  $X_{i+}$ . This partially-fixed-design scheme gives the estimators knowledge of the observed features and the distribution of the unobserved features, which is consistent with knowledge that COLLAB has access to. In this section we fix  $\theta \in \mathbb{R}^d$ . The corresponding label  $y = x_{i+}\theta_{i+} + x_{i-}\theta_{i-} + \xi$ , where  $\xi \in \mathbb{R}$  is drawn from i.i.d.  $\mathsf{N}(0, \sigma^2)$ . We use  $y_j \in \mathbb{R}^n$  to denote its vector form for the agent  $j$ . To model the estimator's knowledge about the labels, we will have two observation models—one weaker and one stronger—which we will specify later when we present our results.
For each observation model, we will have two types of results. The first type of result is a minimax lower bound for full-featured data; i.e., how well can estimator perform on a fresh sample without
missing features. This type of result will concern the full-feature asymptotic local minimax risk
$$
\liminf_{n\to \infty}\mathfrak{M}_{m,\varepsilon}(\{X_{i + }\}_{i\in [m]};\mathcal{P}_{n},u):= \liminf_{n\to \infty}\inf_{\bar{\theta}}\sup_{P\in \mathcal{P}_{n}}n\mathbb{E}_{Z\sim P}\langle u,\bar{\theta} (Z,\{X_{i + }\}_{i\in [m]}) - \theta \rangle^{2}.
$$
We will show that there exists a  $B \in \mathbb{R}^{d \times d}$  such that the local minimax risk in the previous display is lower bounded by  $u^T B u$  for all  $u \in \mathbb{R}^d$ . In other words, we have lower bounded the asymptotic covariance of our estimator with  $B$  (with respect to the p.s.d. cone order). The second type of result is an agent specific minimax lower bound; i.e., what is the best prediction error an estimator (for the given observation model) can possibly have on a fresh sample for a given agent. This type of result will deal with the missing-feature asymptotic local minimax risk, defined as
$$
\liminf_{n\to \infty}\mathfrak{M}^{i +}_{m,\varepsilon}(\{X_{i +}\}_{i\in [m]};\mathcal{P}_{n},u):= \liminf_{n\to \infty}\inf_{\bar{\theta}}\sup_{P\in \mathcal{P}_{n}}n\mathbb{E}_{Z\sim P}\langle u,\bar{\theta} (Z,\{X_{i +}\}_{i\in [m]}) - T_{i}\theta \rangle^{2}.
$$
Similar to the first minimax error definition, we will again show that there exists a  $B_i \in \mathbb{R}^{d_i \times d_i}$  such that the local minimax risk we just defined is lower bounded by  $u^T B_i u$  for all  $u \in \mathbb{R}^{d_i}$ . Recall (2) for discussion surrounding why  $T_i \theta$  is the right object to compare against.
5.1 Weak Observation Model: Access only to local models and features
Recall the local least squares estimator  $\hat{\theta}_i = (X_{i+}^\top X_{i+})^{-1} X_{i+}^\top y_i$ . Let  $P_\theta^\hat{\theta}$  be a distribution over  $\hat{\theta}_1, \ldots, \hat{\theta}_m$  induced by  $\theta$  and  $(\xi_1, \ldots, \xi_m) \stackrel{\mathrm{iid}}{\sim} \mathbb{N}(0, \sigma^2 I_n)$ . We define the following family of distributions  $\mathcal{P}_{n,c}^\hat{\theta} := \{P_{\theta'}^\hat{\theta} : \| \theta' - \theta \|_2 \leq cn^{-1/2}\}$  which defines our observation model. Intuitively, in this observation model, we are constructing a lower bound for estimators which have access to the features  $X_{1+}, \ldots, X_{m+}$ , the population covariance  $\Sigma$ , and access to  $\hat{\theta}_1, \ldots, \hat{\theta}_m$ . In comparison, our estimator COLLAB only uses  $\Sigma$  and  $\hat{\theta}_1, \ldots, \hat{\theta}_m$ . We present our first asymptotic local minimax lower bound result here. The proof of this result can be found in Appendix D.1.
Theorem 5.1. Recall that  $C^{\mathrm{g}} \coloneqq (\sum_{i=1}^{m} T_i^\top W_i^{\mathrm{g}} T_i)^{-1}$ . For all  $\in [m]$  and  $n$  let the rows of  $X_{i+}$  be drawn i.i.d. from  $\mathsf{N}(0, \Sigma_{i+})$ . Then for all  $u \in \mathbb{R}^d$ , with probability 1, the full-feature asymptotic local minimax risk for  $\mathcal{P}_{n,c}^{\hat{\theta}}$  is bounded below as,
$$
\liminf_{c\to \infty}\liminf_{n\to \infty}\mathfrak{M}_{m,\varepsilon}(\{X_{i + }\}_{i\in [m]};\mathcal{P}_{n,c}^{\hat{\theta}},u)\geq u^{\top}C^{\mathrm{g}}u.
$$
For all  $u \in \mathbb{R}^{d_i}$ , with probability 1, the missing-feature asymptotic local minimax risk for  $\mathcal{P}_{n,c}^{\hat{\theta}}$  is bounded below as
$$
\liminf_{c\to \infty}\liminf_{n\to \infty}\mathfrak{M}^{i +}_{m,\varepsilon}(\{X_{i + }\}_{i\in [m]};\mathcal{P}^{\hat{\theta}}_{n,c},u)\geq u^{\top}T_{i}C^{\mathrm{g}}T_{i}^{\top}u.
$$
This exactly matches the upper bound for COLLAB we presented in Corollary 3.2
5.2 Strong Observation Model: Access to features and labels
Define the family of distributions  $\mathcal{P}_{n,c}^{y} \coloneqq \{P_{\theta'}^{y} : \| \theta' - \theta \|_2 \leq cn^{-1/2}\}$  as the observation model. Intuitively, in this model, we are constructing a lower bound for estimators having access to all of the features  $X_{1+}, \ldots, X_{m+}$  and access to  $y_1, \ldots, y_m$ . This observation model is stronger than the previous observation model because estimators now have access to the labels  $y$ . We note again that our estimator COLLAB only uses  $\Sigma$  and  $\hat{\theta}_1, \ldots, \hat{\theta}_m$ . The quantities our estimator rely on do not scale with  $n$ , making our estimator much weaker than other potential estimators in this observation model, as estimators are allowed to depend on  $y_i$ , which grows in size with  $n$ . We present our second asymptotic local minimax lower bound result here, starting with defining the strong local lower bound matrix  $C^s \coloneqq (\sum_{i=1}^{m} 2\Sigma / (\| \theta_{i-} \|_{\Gamma_{i-}}^2 + \sigma^2))^{-1}$ . The proof of this result is in Appendix D.2.
Theorem 5.2. For all  $i \in [m]$  and  $n$  let the rows of  $X_{i+}$  be drawn i.i.d. from  $\mathsf{N}(0, \Sigma_{i+})$ . Then for all  $u \in \mathbb{R}^d$ , with probability 1, the full-feature asymptotic local minimax risk for  $\mathcal{P}_{n,c}^y$  is bounded below as
$$
\liminf_{c\to \infty}\liminf_{n\to \infty}\mathfrak{M}_{m,\varepsilon}(\{X_{i + }\}_{i\in [m]};\mathcal{P}_{n,c}^{y},u)\geq u^{\top}C^{s}u.
$$
For all  $u \in \mathbb{R}^{d_i}$ , with probability 1, the missing-feature asymptotic local minimax risk for  $\mathcal{P}_{n,c}^y$  is bounded below as
$$
\liminf_{c\to \infty}\liminf_{n\to \infty}\mathfrak{N}_{m,\varepsilon}^{i + }(\{X_{i + }\}_{i\in [m]};\mathcal{P}_{n,c}^{y},u)\geq u^{\top}T_{i}C^{\mathrm{s}}T_{i}^{\top}u.
$$
(a) Prediction Error for large  $n$
(b) Prediction Error for small  $n$
(c) Small  $n$  omitting Impute algs.
Figure 1: Experimental results for US Census Experiment
In view of the lower bound in the strong observation model and that of the weak observation model in Theorem 5.1, it is clear that the lower bound in the strong observation setting is in general smaller as
$$
\Sigma - T _ {i} ^ {\top} \Sigma_ {i +} T _ {i} = \Pi_ {i} ^ {\top} \left[ \begin{array}{c c} 0 & 0 \\ 0 & \Gamma_ {i -} \end{array} \right] \Pi_ {i} \succeq 0,
$$
which further implies  $C^{\mathrm{g}} \succeq \left(\sum_{i=1}^{m} \Sigma / \left(\| \theta_{i-}\|_{\Gamma_{i-}}^{2} + \sigma^{2}\right)\right)^{-1} \succeq C^{\mathrm{s}}$ .
We argue that the two lower bounds are comparable in the missing completely at random [16]. Consider for every agent  $i$ , each coordinate is missing independently with probability  $p$ . In this case,  $(d_i, \Sigma_{i+}, T_i)$  are i.i.d. random triplets parameterized by  $p$ .
Corollary 5.3. Under the random missingness setup with missing probability  $p$ , let the eigenvalue of  $\Sigma$  be  $\lambda_1(\Sigma) \geq \dots \geq \lambda_d(\Sigma) > 0$  and define its condition number  $\kappa = \lambda_1(\Sigma) / \lambda_d(\Sigma)$ . Suppose  $p \leq \frac{1}{2} \kappa^{-1} (1 + \| \theta \|_{\Sigma}^2 / \sigma^2)^{-1}$ , we have the limits  $\lim_{m \to \infty} mC^{\mathrm{g}}$  and  $\lim_{m \to \infty} mC^{\mathrm{s}}$  exist and
$$
4 \lim  _ {m \rightarrow \infty} m C ^ {\mathrm {s}} \succeq \lim  _ {m \rightarrow \infty} m C ^ {\mathrm {g}} \succeq \lim  _ {m \rightarrow \infty} m C ^ {\mathrm {s}}.
$$
6 US Census Experiment
We experiment on real US census data modified from the ACSTTravelTime dataset from the folktables package [7]. The code can be found at https://github.com/garyxcheng/collab. We (artificially) construct  $m = 5$  datacenters (agents), each containing data from only one of California, New York, Texas, Florida, and Illinois. The goal is to collaboratively learn a model for each datacenter in a communication efficient way. This setup models potentially real settings where state governments are interested in similar prediction tasks but may not be allowed to directly transfer data about their constituents directly to one another due to privacy or communication constraints. In this setup, there is covariate shift across agents, and the feature and error distributions are not Gaussians—presenting significant deviations from our theoretical assumptions which serve as a proving ground to test how well COLLAB works in practice. For an experimental setup that more closely matches the assumptions we made in our theory, see Appendix A.2
After dataset preprocessing, described in the Appendix A.1, we have  $d = 37$  features. We compute the feature covariance matrix from training data across all the datacenters, which we plot in Appendix A.1. We assume we are able to do this because this computation can be done in a distributed manner, without communicating training data points or labels. The California datacenter will have access to 37 features, New York to 36, Texas to 35, Florida to 30, and Illinois to 27. This models the feature heterogeneity which varies across geography. Each datacenter will have  $n$  datapoints, which we vary in this experiment. The objective is to predict people from Illinois's travel time to work given all 37 features. This task models the setting where the datacenter of interest does not have access to labeled full-featured, data to use to predict on full-featured test data.
We compare our method COLLAB against methods we call Naive-Local, Naive-Collab, Optimized-Naive-Collab, Imputation, and RW-Imputation. We briefly describe each method here; Appendix A.1 contains a more detailed description of each method. Naive-Local refers to each agent locally perform OLS to construct  $\hat{\theta}_i$ . Naive-Collab does an equal-weighted average of the agent OLS models— $\sum_{i \in [m]} \Pi_{i+}^{\top} \hat{\theta}_i / m$ . Optimized-Naive-Collab uses gradient descent to optimize the choices of weights used in Naive-Collab. Optimized-Naive-Collab uses fresh labeled samples without any missing
features during gradient descent, so in this sense, Optimized-Naive-Collab is more powerful than our method. Imputation refers to the global imputation estimator  $\hat{\theta}^{\mathrm{imp - glb}}$  with  $\alpha_{i} = 1 / m$ . RW-Imputation is Imputation but with the optimal choice of weights  $\alpha_{i}$ . We also compare against Naive-Local trained with  $5n$  datapoints. We choose  $5n$  to model the hypothetical scenario setting where all of the other datacenters available contain data (albeit with missing features) from Illinois. For each method that we test, we run 80 trials to form  $95\%$  confidence intervals.
We see that for  $n \leq 800$  in Figures I(b) and I(c), COLLAB performs the best; the imputation methods do the worst, and have much higher variance. In this small  $n$  regime, even the Naive-Local method with 5 times the data does worse than COLLAB. For  $n \geq 2000$  in Figure I(a), the aggregation methods do worse than the imputation methods, and Naive-Local method with 5 times the data is the best performing method. However, COLLAB remains better than Optimized-Naive-Collab and Naive-Collab. The performance of Naive-Collab and Optimized-Naive-Collab being much closer to the performance of COLLAB here than in the synthetic data experiment in Appendix A.2 is not surprising, as the covariance of the features here is much more isotropic, meaning that the naive aggregation methods will not incur nearly as much bias.
7 Discussion and Future Work
Optimal weights beyond Gaussianity.  $\mathbb{E}\left[x_{i + }\theta_{i - }^{\top}z_{i + }z_{i + }^{\top}\theta_{i - }x_{i + }^{\top}\right]$  has a nice closed form in Gaussian setting because  $z_{i + }$  and  $x_{i + }$  are independent—which is in general not true without Gaussianity. If we can directly sample from the feature distribution  $\mathcal{P}$  (e.g., unlabeled data), then we can empirically estimate  $\mathbb{E}\left[x_{i + }\theta_{i - }^{\top}z_{i + }z_{i + }^{\top}\theta_{i - }x_{i + }^{\top}\right]$  by sampling from  $\mathcal{P}$  and using any consistent plug-in estimate  $\hat{\theta}$  (e.g., run COLLAB with weights  $W_{i} = I_{d_{i}}$ ). This will return a good estimate of the optimal weights. An interesting future direction is to prove lower bounds without the Gaussianity assumption.
Generalization to non-linear models. Recall in the Gaussian setting, the optimal weights in COLLAB are  $W_{i}^{\mathrm{g}} = \Sigma_{i + } / (\mathbb{E}_{x,y}[(\langle x_{i + },\hat{\theta}_{i}\rangle -y)^{2}])$ . Then, the optimal loss function in Eq. (3) becomes
$$
\sum_ {i = 1} ^ {m} \left\| \theta_ {i +} + \Sigma_ {i +} ^ {- 1} \Sigma_ {i \pm} \theta_ {i -} - \hat {\theta} _ {i} \right\| _ {W _ {i} ^ {\mathrm {g}}} ^ {2} = \sum_ {i = 1} ^ {m} \frac {\mathbb {E} _ {x _ {i +}} [ (\langle x _ {i +} , \hat {\theta} _ {i} \rangle - \langle x _ {i +} , T _ {i} \theta \rangle) ^ {2} ]}{\mathbb {E} _ {x , y} [ (\langle x _ {i +} , \hat {\theta} _ {i} \rangle - y) ^ {2} ]}.
$$
This hints at a generalization to non-linear models. Suppose the local agents train on models  $f^{i}(x_{i +};\theta_{i}), \mathbb{R}^{d_{i}} \times \mathbb{R}^{d_{i}} \mapsto \mathcal{V}$  and the global model  $f(x;\theta), \mathbb{R}^d \times \mathbb{R}^d \mapsto \mathcal{V}$  satisfies for some mapping  $T_{i}: \mathbb{R}^{d} \to \mathbb{R}^{d_{i}}$ ,  $f(x;T_i\theta) = f^i (x_{i +};\theta_i)$ . Consider a loss function  $\ell (\cdot ,\cdot):\mathcal{V}\times \mathcal{V}\rightarrow [0,\infty)$ . Then we can consider the following way of aggregation inspired by COLLAB for linear models
$$
\hat {\theta} := \underset {\theta} {\operatorname {a r g m i n}} \sum_ {i = 1} ^ {m} \frac {\mathbb {E} _ {x _ {i +}} \ell \left(f ^ {i} \left(x _ {i +}; \hat {\theta} _ {i}\right) , f \left(x _ {i +}; T _ {i} \theta\right)\right)}{\mathbb {E} _ {x _ {i +} , y} \ell \left(f ^ {i} \left(x _ {i +}; \hat {\theta} _ {i}\right) , y\right)}. \tag {6}
$$
We can consistently estimate the denominators (weights) using training time loss. An interesting future direction is to investigate the performance of this general approach for non-linear problems.
Application to Federated Learning. Federated Learning algorithms generally consist of repeating the following steps until termination: Step 1. server sends  $\hat{\theta}$  to all agent, Step 2. each agent  $i$  does some local training initialized at  $\hat{\theta}$  to form  $\hat{\theta}_i$ , and Step 3. server collects  $\hat{\theta}_i$  and aggregates them to form a new  $\hat{\theta}$ . Most federated learning algorithms assume the agent model sizes are the same across agents because aggregating models (Step 3) of different sizes is not obvious [19]. However, now [6] suggests a solution to this issue: in step 3, we can solve [6], albeit with population expectations replaced with empirical averages, to aggregate models of different sizes. Note that because data is assumed to stay with each agent in Federated Learning, this optimization problem may have to be solved via a federated learning algorithm.