SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chen Cheng*
chencheng@stanford.edu
Gary Cheng*
chenggar@stanford.edu
John Duchi
jduchi@stanford.edu
We study the problem of collaboratively learning least squares estimates for  $m$  agents. Each agent observes a different subset of the features—e.g., containing data collected from sensors of varying resolution. Our goal is to determine how to coordinate the agents in order to produce the best estimator for each agent. We propose a distributed, semi-supervised algorithm COLLAB, consisting of three steps: local training, aggregation, and distribution. Our procedure does not require communicating the labeled data, making it communication efficient and useful in settings where the labeled data is inaccessible. Despite this handicap, our procedure is nearly asymptotically, local-minimax optimal—even among estimators allowed to communicate the labeled data such as imputation methods. We test our method on US Census data. We also discuss generalizations of our method to non-Gaussian feature settings, non-linear settings, and Federated Learning.
Consider a set of agents that collect data to make predictions, where different agents may collect different features—because of different sensor availability or specialization—but wish to leverage shared structure to achieve better accuracy. Concretely, suppose we have  $m$  agents, where each agent  $i \in [m]$  observes  $n$  samples of  $(x_{i+}, y)$  where  $x_{i+} \in \mathbb{R}^{d_i}$  is some subset of  $x \in \mathbb{R}^d$ . We set this as a regression problem where the data  $(x, y)$  has the linear relationship  $y = \langle x, \theta \rangle + \xi$  for some noise variable  $\xi$ . For example, these agents could be a network of satellites, each collecting data with a distinct set of sensors of varying resolution and specialization, with the purpose of estimating quantities like crop-yields [22], biomass [18], and solar-flare intensity [12]. Or these agents could be a group of seismic sensors, using acoustic modalities or accelerometers to predict whether an earthquake will occur [2]. Other examples may include networks of hospitals or phones [13]. In these settings, the agents can share information to collaboratively train a model; however, they are limited by communication bandwidth constraints, a situation satellites and seismic sensors often face due to radio frequency spectrum scarcity and interference [6, 20]. Without being too rigorous, we will define a communication efficient algorithm as one where the per-agent communication cost does not scale with  $n, d$ , or  $m$ . No dependence on the dataset size  $n$  is suited for applications with significant data volume but limited communication resources. No dependence on the overall dimension  $d$  of the data ensures that agents are not incentivized to turn away agents with access to a richer set of features that would improve statistical performance but increase  $d$ . Finally, no dependence on  $m$  ensures that the algorithm is scalable to large collectives of agents. Can we construct a statistically optimal and communication efficient procedure to estimate  $\theta$ ?
We answer in the affirmative and introduce our estimator COLLAB. COLLAB consists of three steps: local training on all agents, aggregation on a coordinating server, and distribution back to all agents. Our algorithm is communication-efficient: each agent  $i \in [m]$  syncs twice with a coordinating server
*Equal contribution, authors ordered alphabetically by (last name, first name)
and incurs communication cost scaling like  $\Theta (d_i^2)$ . We prove local minimax lower bounds which prove that COLLAB is (nearly) instance-optimal. We choose to study this problem in a stylized linear setting so that we can provide stronger guarantees for the algorithms we make. Indeed, our results which pair the exact asymptotic covariance of our estimator COLLAB with matching asymptotic local minimax lower bounds heavily rely on the linearity of our problem and would not be possible without strong structural assumptions. Having said this, the theory we develop for linear models does hint at potential methods for non-linear settings, which we discuss in Section 7. We also acknowledge privacy considerations are important for real world systems such as hospitals. We choose to focus instead on sensor settings where privacy is less of a concern. We leave adapting our results to privacy-sensitive settings to future work.
We compare our methods to single-imputation methods theoretically and empirically. We choose to baseline against imputation methods for three reasons. First, if we ignore communication constraints, our problem is a missing data problem, where formally the data is "missing at random" (MAR) [16]. MAR problems are well studied, so we know that imputation methods work well theoretically and in practice [25, 15]. Second, because we have instance-optimal lower bounds, we know that imputation methods are also optimal for our problem. Finally, because imputation methods use more information than the method we propose, imputation will serve as a "oracle" baseline of sorts.
Contributions. We briefly summarize our contributions.
1. We design a communication-efficient, distributed learning algorithm COLLAB which performs a weighted de-biasing procedure on the ordinary least squares estimator of each agent's data.
2. We show COLLAB is asymptotically locally minimax optimal among estimators which have access to the ordinary least squares estimator of each agent's data. We also show that with some additional assumptions, COLLAB is also asymptotically locally minimax optimal among estimators that have access to all of the training data of all agents.
3. We propose and develop theory for various baseline methods based on imputation. We compare the statistical error and communication cost of COLLAB against these baseline methods both theoretically and empirically on real and synthetic data.
4. We discuss generalizations of COLLAB for non-Gaussian feature settings and non-linear settings. We highlight open problems and identify possible directions for future work.
Missing data. If we ignore the communication and computational aspects of our problem, the problem we study reduces to one of estimation with missing data. There has been a lot of work on this topic; please see [16] for an overview. The data in our problem is missing at random (MAR)—the missing pattern does not depend on the value of the data and is known given agent  $i$ . There are many approaches to handling missing data such as weighting and model-based methods [24]. Most related to our work are methods on single imputation. Schafer and Schenker [25] shows imputation with conditional mean is nearly optimal with special corrections applied. More recently, Chandrasekher et al. [3] show that single imputation is minimax optimal in the high dimensional setting. Another closely related popular approach is multiple imputation [23, 1]. Previous work [27, 29] has shown that multiple imputation in low dimensional settings produces correct confidence intervals under a more general set of assumptions compared to single imputation settings. However, we choose to focus on single imputation methods for two reasons. First, we are interested in estimation error and not confidence intervals, and our lower bounds show that single imputation has optimal estimation error for our setting. Second, in our problem context, multiple imputation would require more rounds of communication and consequently higher communication cost. Other methods for missing data include weighting and model-based methods.
Distributed learning. Learning with communication constraints is a well studied practical problem. We provide a couple of examples. Suresh et al. [26] study how to perform mean estimation with communication constraints. Duchi et al. [8] develop communication-constrained minimax lower bounds. Distributed convex optimization methods like Hogwild [21] have also been well studied. However, the works mentioned all concern the no-missing-data regime. A more relevant subfield of distributed learning is federated learning. In federated learning, a central server coordinates a
collection of client devices to train a machine learning model. Training data is stored on client devices, and due to communication and privacy constraints, clients are not allowed to share their training data with each other or the central server [13]. In the no-missing-features regime, optimization algorithms for federated optimization are well studied. There is also more theoretical work, which focus on characterizing communication, statistical, and privacy tradeoffs, albeit for a more narrow set of problems such as mean and frequency estimation [4]. More related to the missing data regime we consider is cross-silo federated learning [13] or vertical federated learning [30]. In this paradigm, the datasets on client machines are not only partitioned by samples but also by features. Researchers have studied this problem in the context of trees [5], calculating covariance matrices [14], k-means clustering [28], support vector machines [31], and neural nets [17]. Most related to our work is Gascon et al. [9], Hardy et al. [11]; they study how to privately perform linear regression in a distributed manner. However, unlike our work, these works focus more on developing algorithms with privacy guarantees rather than statistical ones.
2 Mathematical model
We assume we have  $m$  agents that observe a subset of the dimensions of the input data  $x \in \mathbb{R}^d$ . Each agent  $i$  has a "view" permutation matrix  $\Pi_i^\top := \left[\Pi_{i+}^\top \quad \Pi_{i-}^\top\right] \in \mathbb{R}^{d \times d}$ .  $\Pi_{i+} \in \mathbb{R}^{d_i \times d}$  describes which feature dimensions the agent sees, and  $\Pi_{i-} \in \mathbb{R}^{(d-d_i) \times d}$  describes the dimensions the agent does not see. For a feature vector  $x \in \mathbb{R}^d$  and corresponding label  $y \in \mathbb{R}$ , the  $i$ -th agent observes  $(x_{i+}, y)$  where  $x_{i+} := \Pi_{i+} x \in \mathbb{R}^{d_i}$ . Each agent has  $n$  fresh, independent observations (independent across agents) denoted as a matrix  $X_{i+} \in R^{n \times d_i}$  and vector  $y_i \in \mathbb{R}^n$ . We let  $X_{i-} \in \mathbb{R}^{n \times (d-d_i)}$  denote the unobserved dimensions of the input data  $x$  drawn for the  $i$ -th agent, and we let  $X_i \in \mathbb{R}^{n \times d}$  denote the matrix of input data  $x$  drawn for the  $i$ -th agent, including the dimensions of  $x$  unobserved by the  $i$ -th agent. To simplify discussions in the following sections, for any vector  $v \in \mathbb{R}^d$  we use the shorthand  $v_{i+} = \Pi_{i+} v$  and  $v_{i-} = \Pi_{i-} v$ . Similarly for any matrix  $A \in \mathbb{R}^{d \times d}$  we denote by
$$
A _ {i +} = \Pi_ {i +} A \Pi_ {i +} ^ {\top},
$$
$$
A _ {i -} = \Pi_ {i -} A \Pi_ {i -} ^ {\top},
$$
$$
A _ {i \pm} = \Pi_ {i +} A \Pi_ {i -} ^ {\top},
$$
$$
A _ {i \mp} = \Pi_ {i -} A \Pi_ {i +} ^ {\top}.
$$
For a p.s.d. matrix  $A$ , we let  $\|x\|_A = \langle x, Ax \rangle$ .
We assume the data from the  $m$  agents follow the same linear model. The feature vectors  $x$  comprising the data matrices  $X_{1},\ldots ,X_{m}$  are i.i.d. with zero mean and covariance  $\Sigma \succ 0$ . We will assume that each agent has knowledge of  $\Sigma_{i + }$  -e.g., they have a lot of unlabeled data to use to estimate this quantity. We also assume for each index tuple pair  $(j,k)$  for  $1\leq j,k\leq d$ , there exists some agent that observes dimensions  $j$  and  $k$ . This ensures the coordinating server can construct the complete covariance matrix  $\Sigma$  from individual  $\Sigma_{i + }$ . The labels generated follow the linear model
$$
y _ {i} = X _ {i} \theta + \xi_ {i}, \quad \xi_ {i} \stackrel {{\mathrm {i d}}} {{\sim}} \mathsf {N} (0, \sigma^ {2} I _ {n}).
$$
Throughout this work we consider a fixed ground truth parameter  $\theta$ .
Objectives. We are interested in proposing a method of using the data of the agents to form an estimate  $\hat{\theta}$  which minimizes the full-feature prediction error on a fresh sample  $x\in \mathbb{R}^d$
$$
\mathbb {E} _ {x} \left[ \left(\langle x, \hat {\theta} \rangle - \langle x, \theta \rangle\right) ^ {2} \right] = \| \hat {\theta} - \theta \| _ {\Sigma} ^ {2}. \tag {1}
$$
We are also interested in forming an estimate  $\hat{\theta}_i$  which minimizes the missing-feature prediction error of a fresh sample  $x_{i+} \in \mathbb{R}^{d_i}$  for agent  $i$ -i.e.,  $x_{i+} = \Pi_{i+} x$  where  $x \in \mathbb{R}^d$  is fresh. Define  $T_i := [I_{d_i} \Sigma_{i+}^{-1} \Sigma_{i\pm}] \Pi_i$  and the Schur complement  $\Gamma_{i-} := \Sigma \backslash \Sigma_{i+} := \Sigma_{i-} - \Sigma_{i\mp} \Sigma_{i+}^{-1} \Sigma_{i\pm}$ . The local test error is then
$$
\mathbb {E} _ {x} \left[ \left(\left\langle x _ {i +}, \hat {\theta} _ {i} \right\rangle - \left\langle x, \theta \right\rangle\right) ^ {2} \right] = \| \hat {\theta} _ {i} - T _ {i} \theta \| _ {\Sigma_ {i +}} ^ {2} + \| \theta_ {i -} \| _ {\Gamma_ {i -}} ^ {2} \tag {2}
$$
We see that  $\| \theta_{i - }\|_{\Gamma_{i - }}^{2}$  is irreducible error. The role of the operator  $T_{i}$  is significant, as (2) shows that  $T_{i}\theta$  is the best possible linear predictor for the  $i$ th agent. Through this paper, we will also highlight the communication costs of the methods we consider. Recall that we would like our methods to have communication cost that do not scale with  $n$  or  $d$ .
2. Maybe surprisingly,  $T_{i}\theta$  is better than naively selecting the subset of  $\theta$  corresponding to the features observed by agent  $i$  —i.e.,  $\Pi_{i}\theta$ . This is because  $T_{i}\theta$  leverages the correlations between features.
We begin with some intuition of why collaboration should help. Then we outline an approach of solving this problem for general feature distributions. The general approach is not immediately usable because it requires some knowledge of  $\theta$ , so we need to do some massaging. In Section 3.3, we show how to circumvent this issue in the Gaussian feature setting and introduce our method COLLAB. Adapting the general approach to other non-Gaussian settings is an open problem, but we discuss some potential approaches in Section 7.
3.1 Intuition
In order to build some intuition for how our approach leverages collaboration, we begin with some basic analysis of the no-collaboration solution: where each agent  $i$  only performs ordinary least squares on their local data
$$
\hat {\theta} _ {i} = X _ {i +} ^ {\dagger} y _ {i} \stackrel {\mathrm {(i)}} {=} (X _ {i +} ^ {\top} X _ {i +}) ^ {- 1} X _ {i +} ^ {\top} y _ {i} \stackrel {\mathrm {(i i)}} {=} \theta_ {i +} + \widehat {\Sigma} _ {i +} ^ {- 1} \widehat {\Sigma} _ {i \pm} \theta_ {i -} + (X _ {i +} ^ {\top} X _ {i +}) ^ {- 1} X _ {i +} ^ {\top} \xi_ {i},
$$
here  $A^\dagger$  denotes the Moore-Penrose inverse for a general matrix  $A$ , and (i) holds whenever  $\mathrm{rank}(X_{i+}) \geq d_i$ . Because we focus on the large sample asymptotics regime ( $n \gg d_i$ ), (i) will hold with probability 1. Expanding out the  $y_i$  gives equality (ii). Ignoring the zero-mean noise term, we see that  $\theta$  must satisfy  $\theta_{i+} + \widehat{\Sigma}_{i+}^{-1}\widehat{\Sigma}_{i\pm}\theta_{i-} - \hat{\theta}_i \approx 0$  for all  $i \in [m]$ . Our approach essentially tries to find such a  $\theta$ .
3.2 Our General approach
In particular, after each agent  $i$  performing ordinary least squares on their local data to produce  $\{\hat{\theta}_i\}_{i\in [m]}$ , we recover an estimate of  $\theta$  by performing a form of weighted empirical risk minimization parameterized by the positive definite matrices  $W_{i}\in \mathbb{R}^{d_{i}\times d_{i}}$
$$
\hat {\theta} = \hat {\theta} \left(W _ {1}, \dots , W _ {m}\right) := \underset {\theta} {\operatorname {a r g m i n}} \sum_ {i = 1} ^ {m} \left\| \theta_ {i +} + \Sigma_ {i +} ^ {- 1} \Sigma_ {i \pm} \theta_ {i -} - \hat {\theta} _ {i} \right\| _ {W _ {i}} ^ {2}. \tag {3}
$$
We know by first order stationarity that  $\hat{\theta} = \left(\sum_{i=1}^{m} T_i^\top W_i T_i\right)^{-1} \left(\sum_{i=1}^{m} T_i^\top W_i \hat{\theta}_i\right)$ . We can show that  $\hat{\theta}$  is a consistent estimate of  $\theta$  regardless the choice of weighting matrices  $W_i$ . In fact, if the features  $X_i, \xi_i$  are Gaussian,  $\hat{\theta}$  is also unbiased. See Lemma B.1 in the Appendix for this result. While Lemma B.1 shows that the choice of weighting matrices  $W_i$  does not affect consistency, the choice of weighting matrices  $W_i$  does affect the asymptotic convergence rate of the estimator. In the next theorem, we show what the best performing choice of weighting matrices are. The proof is in Appendix B.2
Theorem 3.1. For any weighting matrices  $W_{i}$ , the aggregated estimator  $\hat{\theta} = \hat{\theta}(W_1, \dots, W_m)$  is asymptotically normal
$$
\sqrt {n} \left(\hat {\theta} - \theta\right) = \mathsf {N} (0, C (W _ {1}, \dots , W _ {m})),
$$
with some covariance matrix  $C(W_{1},\dots ,W_{m})$  . The optimal choice of weighting matrices is
$$
W _ {i} ^ {\star} := \Sigma_ {i +} (\mathbb {E} \left[ x _ {i +} \theta_ {i -} ^ {\top} z _ {i +} z _ {i +} ^ {\top} \theta_ {i -} x _ {i +} ^ {\top} \right] + \sigma^ {2} \Sigma_ {i +}) ^ {- 1} \Sigma_ {i +},
$$
where  $z_{i+} = x_{i-} - \sum_{i\mp}\Sigma_{i+}^{-1}x_{i+}$ . In particular, for all  $W_i$ ,  $C(W_1,\dots ,W_m)\succeq C(W_1^\star ,\dots ,W_m^\star) = \left(\sum_{i=1}^{m}T_i^\top W_i^\star T_i\right)^{-1}$ .
The main challenge of using Theorem 3.1 is in constructing the optimal weights  $W_{i}^{\star}$ , as at face value, they depend on knowledge of  $\theta$ . While we will discuss high level strategies of bypassing this issue in non-Gaussian data settings in Section 7, we will currently focus our attention on how we can make use of Gaussianity to construct our estimator COLLAB.
Algorithm 1: COLLAB algorithm
end
Data:  $m$  agents with training data  $(X_{1 + },y_{1}),\ldots ,(X_{m + },y_{m})$  each with  $n$  datapoints for Each agent  $i = 1,\dots ,m$  in parallel do Compute  $\hat{\theta}_i = (X_{i + }^\top X_{i + })^{-1}X_{i + }^\top y_i;$  Compute  $\Sigma_{i}$  with (labeled and unlabeled) data; Compute  $R_{i} = \frac{1}{n}\| X_{i + }\hat{\theta}_{i} - y\|_{2}^{2};$  Send  $\hat{\theta}_i,\Sigma_i,R_i$  to coordinating server;
Coordinating server constructs  $\Sigma$  from  $\{\Sigma_j\}_{j = 1}^m$    
Coordinating server constructs  $\hat{W}_i^{\mathrm{g}}\coloneqq \Sigma_{i + } / R_i$    
Coordinating server computes  $\hat{\theta}_i^{\mathrm{clb}} = T_i\hat{\theta} (\hat{W}_1^{\mathrm{g}},\dots$  agents;
3.3 COLLAB Estimator - Gaussian feature setting
If  $X_{i}$  are distributed as  $\mathsf{N}(0,\Sigma)$ ,  $W_{i}^{\star}$  has an explicit closed form as
$$
W _ {i} ^ {\star} = W _ {i} ^ {\mathrm {g}} := \frac {\Sigma_ {i +}}{\| \theta_ {i -} \| _ {\Gamma_ {i -}} ^ {2} + \sigma^ {2}} = \frac {\Sigma_ {i +}}{\mathbb {E} _ {x , y} [ (\langle x _ {i +} , \hat {\theta} _ {i} \rangle - y) ^ {2} ]},
$$
where  $\Gamma_{i-} = \Sigma_{i-} - \Sigma_{i\mp}\Sigma_{i+}^{-1}\Sigma_{i\pm}$  is the Schur complement. Recall we assume that each agent has enough unlabeled data to estimate  $\Sigma_{i+}$ . The denominator  $\mathbb{E}_{x,y}[(\langle x_{i+},\hat{\theta}_i\rangle -y)^2]$  cannot be computed but we can use  $\frac{1}{n}\|X_{i+}\hat{\theta}_i - y\|_2^2$ , which is a consistent estimator of  $\mathbb{E}_{x,y}[(\langle x_{i+},\hat{\theta}_i\rangle -y)^2]$ , in its place. Thus, each agent is able to construct estimates of  $W_i^g$  by computing
$$
\hat {W} _ {i} ^ {\mathrm {g}} := \frac {\Sigma_ {i +}}{\frac {1}{n} \| X _ {i +} \hat {\theta} _ {i} - y \| _ {2} ^ {2}}
$$
Now we construct our global and local COLLAB estimators defined respectively as
$$
\hat {\theta} ^ {\mathrm {c l b}} := \hat {\theta} \left(\hat {W} _ {1} ^ {\mathrm {g}}, \dots , \hat {W} _ {m} ^ {\mathrm {g}}\right), \quad \quad \quad \hat {\theta} _ {i} ^ {\mathrm {c l b}} := T _ {i} \hat {\theta} \left(\hat {W} _ {1} ^ {\mathrm {g}}, \dots , \hat {W} _ {m} ^ {\mathrm {g}}\right). \tag {4}
$$
We summarize the COLLAB algorithm in Algorithm 1. At a high level,  $\hat{\theta}^{\mathrm{clb}}$  is an estimate of  $\theta$  which also minimizes the full-feature prediction error (1) and  $\hat{\theta}_i^{\mathrm{clb}}$  minimizes the missing-feature prediction error for agent  $i$  (2). Now we show that using the collective "biased wisdom" of local estimates  $\hat{\theta}_i$ , our collaborative learning approach returns an improved local estimator. The proof is in Appendix B.3.
Corollary 3.2. Let  $X_{i} \sim \mathsf{N}(0,\Sigma)$  and define  $C^{\mathrm{g}} \coloneqq (\sum_{i=1}^{m} T_{i}^{\top} W_{i}^{\mathrm{g}} T_{i})^{-1}$ . The global COLLAB estimator  $\hat{\theta}^{\mathrm{clb}}$  and the local  $\hat{\theta}_{i}^{\mathrm{clb}}$  on agent  $i$  are asymptotically normal
$$
\sqrt {n} \left(\hat {\theta} ^ {\mathrm {c l b}} - \theta\right) \xrightarrow {d} \mathsf {N} \left(0, C ^ {\mathfrak {g}}\right) \quad a n d \quad \sqrt {n} \left(\hat {\theta} _ {i} ^ {\mathrm {c l b}} - T _ {i} \theta\right) \xrightarrow {d} \mathsf {N} \left(0, T _ {i} C ^ {\mathfrak {g}} T _ {i} ^ {\top}\right).
$$
The following are true
(i)  $W_{i}^{\mathrm{g}}$  are the optimal choice of weighting matrices i.e.,particular,  $C(W_{1},\dots ,W_{m})\succeq C(W_{1}^{\mathrm{g}},\dots ,W_{m}^{\mathrm{g}}) = C^{\mathrm{g}}.$
(ii) On agent  $i$ , we have  $\sqrt{n} (\hat{\theta}_i - T_i\theta) \stackrel{d}{\to} \mathsf{N}(0, (W_i^{\mathrm{g}})^{-1})$ . The asymptotic variance of  $\hat{\theta}_i$  is larger than that of the COLLAB estimator  $\hat{\theta}_i^{\mathrm{clb}}$  —i.e.,  $(W_i^{\mathrm{g}})^{-1} \succeq T_i C^{\mathrm{g}} T_i^\top$ .
Remark 3.3. In Corollary 3.2, we can replace  $\Sigma_{i+}$  in  $\hat{W}_i^g$  with local sample covariances  $\hat{\Sigma}_{i+}$  for the plug-in estimator in Eq. (4). We prove this stronger statement in the proof of Corollary 3.2. However, as the coordinating server still need to use exact operators  $T_i$  to compute the COLLAB estimator, relaxing this condition does not bring practical benefit.
Method	Full-feature asymptotic covariance	Missing-feature asymptotic covariance	Communication cost for agent i
Local OLS - θi	-	(Wi^g)^{-1}	0
Local imputation w/ collaboration - θi^imp	(Σm i=1 T_i^T W_i^g T_i)^{-1}	Ti(Σm i=1 T_i^T W_i^g T_i)^{-1} Ti^T	d^2
Global imputation - θi^imp-glb	(Σm i=1 T_i^T W_i^g T_i)^{-1}	Ti(Σm i=1 T_i^T W_i^g T_i)^{-1} Ti^T	nd_i + di
COLLAB - θi^clb	(Σm i=1 T_i^T W_i^g T_i)^{-1}	Ti(Σm i=1 T_i^T W_i^g T_i)^{-1} Ti^T	di^2 + 2di + 1
Table 1. Full and Missing feature asymptotic covariance and communication cost for agent  $i$  . Communication cost is measured by how many real numbers are received and sent from agent  $i$  .
This result characterizes the error distribution of the COLLAB estimator; guarantees about prediction error, defined in (1) and (2), directly follow. In particular, letting  $z \in \mathbb{R}^d$  be distributed as  $\mathsf{N}(0,C^{\mathrm{g}})$ , by the continuous mapping theorem we have
$$
n \mathbb {E} _ {x} \left[ \left(\langle x, \hat {\theta} ^ {\mathrm {c l b}} \rangle - \langle x, \theta \rangle\right) ^ {2} \right] = n \| \hat {\theta} ^ {\mathrm {c l b}} - \theta \| _ {\Sigma} ^ {2} \xrightarrow {d} \| z \| _ {\Sigma} ^ {2}. \tag {5}
$$
For large dimension  $d$ ,  $\|z\|_{\Sigma}^{2}$  concentrates around its expectation  $\mathbb{E}[||z||_{\Sigma}^{2}] = \mathrm{Tr}(C^{\mathrm{g}})$ . The same arguments can be applied to make a statement about the local test error (2).
Communication Cost In the first round of communication, agent  $i$  sends  $d_i^2 + d_i + 1$  real numbers to the coordinating server. In the second round of communication, the server sends back the updated local model which is  $d_i$  real numbers. In total, this amounts to  $d_i^2 + 2d_i + 1$  communication cost per agent. The communication cost of COLLAB does not scale with  $n$ ,  $d$ , or  $m$ , satisfying our desiderata.
4 Comparison with other methods