Collaboratively Learning Linear Models with Structured Missing Data
Authors: Chen Cheng
Gary Cheng
John Duchi
Venue: NeurIPS 2023 poster (2023)
Keywords: Collaborative Learning
Missing Data
Sensors
Linear Regression

Abstract: 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.

Collaboratively Learning Linear Models with
Structured Missing Data
Chen Cheng"
Gary Cheng"
John Duchi Stanford University
Stanford University
Stanford University chenggardstannord.cco
chencheng@sttaffrdd.crm
j (duchi@stants ord. edu
Abstract
We study the problem of collaboratively learning least squares estimates for m agents. Each agent observes different subset of the features-c.g, containing data collected from sensors of varying resolution. Our goal is to determine how coordinate the agents in order to produce the best estimator for each agent. We propose a distributed, semi- supe-sspe 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 settings where the labeled data is inaccessible. Despite this handicap, our procedure is nearly asymptotically, local-minimax optimal-even among estimators allowed 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.
Introduction
Consider set of agents that collect data to make predictions, where different agents may collect dif- ferent because of different sensor availability or specialization but wish to leverage shared structure achieve better accuracy. Concretely, suppose we have agents, where each agent E [m] observes samples of (2++, y) where x++ Rdd: is some subset of Rd We set this regression problem where the data (xr, y) has the linear relationship y = (xx 0) + for some noise variable ç. For network of satellites, each collecting data with distinct of sensors
example, these agents could of varying resolution and specialization, with the purpose of estimating quantities like crop-yields (22), biomass [8]. and solar-flare intensity 112]. Or these agents could be 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 [33. 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 inter- ference 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 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 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 statistically optimal and communication efficient procedure to estimate 0?
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 communication--fffccee each agent i [m] synes twice with coordinating server
Equua contribution, authors ordered alphabetically by (last name, first name)
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

and incurs communication cost scaling like (d))). We prove local minimax lower bounds which prove that COLLAB is (nearly) instance-optimal. We choose 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 sensor settings where privacy is less of a concern. We leave adapting our results to privacy-sensidee settings to future work.
We compare our methods to single-inputation methods theoretically and empirically. We choose to baseline against imputation methods for three reasons. First, if we ignore communication constraints, our problem is 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 P5 |5 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 "oracle" baseline of sorts.
Contributions. We briefly summarize our contributions.
1. We design a communication-effcieenn, distributed learning algorithm COLLAB which per- forms 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, COLLABB 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.
1.1 Related 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 [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 weighting and model-based methods
Zi! Most related
to our work are methods on single imputation. Schafer and Schenker (5 shows imputation with conditional mean is nearly optimal with special corrections applied. More recently, Chandrasekher et al. [ show that single imputation is minimax optimal in the high dimensional setting. Another
closely related popular approach is multiple imputation [2 T. Previous work [27] 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 well studied practical problem. We provide a couple of examples. Suresh et al. [26] study how to perform mean estimation with minimax lower bounds. Distributed convex optimization methods like Hogwild [2] have also been well studied. However, the works mentioned all concern the no-missing-dan regime. A more relevant subfield of distributed learning is federated learning. In federated learning, central server coordinates
communication constraints. Duchi al. [ develop conmunication-conntiiiinn

---

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-feature 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 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 | or vertical federated learning [0) 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 [[4], k-means clustering [28]. support vector machines [1I, and neural nets [7]. Most related to our work is Gascón et al. 19) Hardy et al. [[] they study how to privately perform linear regression in 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 observes a subset of the dimensions of the input data & E IRd. Each agent i has a "view" permutation matrix II = [ _ Rdxd lh+ E &Rdxdd describes which feature dimensions the agent sees, and II,- E (ddd).) describes the dimensions the agent does not see. For a feature vector x E IR and corresponding label E R. the i-th agent observes (2++.8) y) where I2+ = 11+ E Rdd, Each agent has fresh, independent observations (independent across agents) denoted as a matrix Xx+ E Rmxd; and vector y4 KR We let X;- (d denote the unobserved dimensions of the input data r drawn for the i-th agent, and we let X; E Rnxd denote the matrix of input data æ drawn for the i-th agent, including the dimensions of æ unobserved by the i-th agent. To simplify discussions in the following sections, for any vector E Rd we use the shorthand 2i+
H1 v and 0 = -I Similarly for any matrix Rdxd we denote by
A++ = 11
A AII
A+ 11,
A I1,-AII. AII
For a p.s.d. matrix A, we let ||ا||| Ar).
We assume the data from the m agents follow the same linear model. The feature vectors comprising the data matrices Xj..... Xm are i.i.d with zero mean and covariance each agent has knowledge of E,,-e.r e.g., they have lot of unlabeled data to use to estimate this quantity. We also assume for each index tuple pair (j. k:) for S j.k S d, there exists some agent that observes dimensions j and k. This ensures the coordinating server can construct the complete covariance matrix from individual E,+. The labels generated follow the linear model
0. We will assume that
yi X,0 + &s; & sid N(0,o
Throughout this work we consider fixed ground truth parameter 6.
Objectives. We are interested in proposing a method of using the data of the agents to form an estimate 0 which minimizes the full-feature prediction error on a fresh sample æ E Rd
Eg [(x,, 0) (z, 0) = 0
(1)
We are also interested in forming an estimate 0 which minimizes the missing-feettd feature prediction error of fresh sample 244 E Rdd: for agent i-i.., Ii+ 11,4 where E Rd is fresh. Define T; = [La E II, and the Schur complement r,_ = E = Σ, Exp E. E The local test error is then
Ex
E. 1. 0:) (x. 0) 110, -T,0 + 110;-112
2
We see that |10,-112, is irreducible error. The role of the operator T; is significant, as shows that T:0 is the best possible linear predictor for the ith agen Through this paper, we will also highlight the communication costs of the methods consider. Recall that we would like our methods have communication cost that do not scale with n or d.
Maybe surprisingly. T:0 is better than naively selecting the subset of 0 corresponding to the features observed by agent -ii.., 11:0. This is because 7,0 leverages the correlations between features.
3

---

Method
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 0, so we need to do some massaging. In Section 3.3 we show how to circumvent this issue the Gaussian feature setting and introduce our method COLLAB. Adapting the general approach to other non- -on-Gassian 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
0. =x ①1)
Xx+
(i)
0ct + E E-1 + - -
here AT denotes the Moore-Penrose inverse for general matrix A, and (i) holds whenever (ank(X;+) 2 di. Because we focus on the large sample asymptotics regime (n > d:), (i) will hold with probability 1. Expanding out the yi gives equality (ii). Ignoring the zero-mean noise term, we see that & must satisfy 0:4 + E 0 0, z 0 for all i E [m]. Our approach essentially tries find such a 0.
3.2 Our General approach
In particular, after each agent performing ordinary least squares on their local data to produce [m]). we recover an estimate of f by performing form of weighted empirical risk minimization parameterized by the positive definite matrices W; E xdd
= @W, Wm) = argmin Σ
0,t + E,1
0 1
We know by first order stationarity that 0 IE W;Z:) W,T.) (Em, T: W,O.). We can show that i is consistent estimate of 0 regardless the choice of weighting matrices W,. In fact, if the features X,, & are Gaussian, is also unbiased. See Lemma B.J in the Appendix for this result, While Lemma B. shows that the choice of weighting matrices W, does not affect consistency, the choice of weighting matrices W; 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 Wa, the aggregated estimator Ô = Wm) is asymptotically normal
Vn (0 -0- C(W,........... ,,,,,,, Ww.)),
with some covariance matrix C(W,," ... ). The optimal choice of weighting matrices is W = Ei+ (E (24 O_ 2,4 2 x xJ -
where Zi4 xi- Eip D xe+ . In particular, for all W., C(W,, ....... Wm) 2 C(Wi, ... Wh)
EM T; W.T:)
The main challenge using Theorem 3.1 is in constructing the optimal weights W as at face value, they depend on knowledge of 0. 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 of Gaussianity to construct our estimator COLLAB.

---

Algorithm 1: COLLAB algorithm
Data: agents with training data (X1+, (X)+.y1)... .n.,..... (Xu+, Um) each with datapoints for Each agent i = 1,..., m in parallel do
Compute 02 (
Compute E; with (labeled and unlabeled) data;
Compute R 0, yll::
Send 6, Σ, R; to coordinating server; end
Coordinating server constructs Σ from {- Coordinating server constructs W : Ei+/Ik. Coordinating server computes 60th T.A(WF.... W and distributes them to respective
agents;
3.3
COLLAB Estimator Gaussian feature setting
If X; are distributed as N(O, 2), W has an explicit closed form as
Σx
E+ Eu.y(((z (11)) - 8)
W W 10 10, + +
where I;-- D; Emp 2, 2 is the Schur complement. Recall assume that each agent has sairs+,00. 0.)
enough unlabeled data to estimate Ej+. The denominator [e.,(((T:4, but we use 1 || A, v12.. which is consistent estimator of Ez,y [(#+++ 0.) - -y)? in its place. Thus, each agent is able construct estimates of W by computing
- y} cannot be computed
E++
W
1 11 8 v112
Now we construct our global and local COLLAB estimators defined respectively as
gelb @( ..l......... W
gell = T, ... W..)
(4)
We summarize the COL.IAB algorithm Algorithm At high level, jelh is an estimate of 0 which also minimizes the full-feature prediction error I and طelh minimizes the missing-feature prediction error for agent Now we show that using the collective "biased wisdom" of local estimates 0,, our collaborative learning approach returns an improved local estimator. The proof is in Appendix B.
Corollary 3.2. Let X; N(O, Σ) and define Cs = E T.T W The global COLLAB estimator jclb and the local 05 on agent i are asymptotically normal
/n (g.th - 4 CO, and /n @ Ti0) $ (O,T.CTTT)
The following are true
(i) W are the optimal choice of weighting matrices i.e.particular. C(W, Wm) 2 C(WE..... ...., ...) C
(ii) On agent i, we have vn(0, T,0) 4 N(O. (W?)-1) -). The asymptotic variance of 0, is larger than that of the COLLAB estimator gelh-ee. (W3)-1 = T.
Remark 3.3. In Corollary 3.2 we can replace Σ,+ in W with local sample covariances E for plug-in estimator in Eq. We prove this stronger statement in the proof of Corollary 3.2 However, as the coordinating server still need to use exact operators T; compute the COLLAB estimator, relaxing this condition does not bring practical benefit.
5

---

Full-feature asymplotic covariance
Missing-feature asymptotic covariance
Communication cost for agent
Method
Local OLS A
(W
Local imputation w/
EEE=? (TT:)))
T; (E T. W TT
d²
collaboration 6ppp
Global imputation Glmp-alb
Ei=1 T.T:) EE T. (.
T, Ei m T, W T.¹
nd; + d;
COLLAB
thth
T, (E T, W.T) T.¹
df + 2d, + 1
Table Full and Missing feature asymptotic covariance and communication cost for agent i. Commu- nication cost measured by how many real numbers are received and sent from agent
This result characterizes the error distribution of the COLLAB estimator; guarantees about prediction error, defined in ① and 2 directly follow. In particular, letting = E IR be distributed as N(0, C*), by the continuous mapping theorem have
nE,1((a, 0 - (z, 0))} - n n 0
d
|
(5)
For large dimension d, || concentrates around its expectation ||| [3] Tr(Cs). The same arguments can be applied to make a statement about the local test error
Communication Cost
In the first round of communication, agent i sends df + d, + 1 real numbers to the coordinating server. In the second round of communication, the server sends back the updated local model which is d; real numbers. In total, this amounts to df + 2d, + communication cost per agent. The communication cost of COLLAB does not scale with n, d, or m, satisfying our desiderata.
Comparison with other methods
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 D The proofs of the theorems are in Appendix C For the sake brevity, our results here will center around estimating 0. Analogous results centered around estimating 7,0 follow directly. These results can be connected to the full-feature and local test error by the same argument we made earlier derive
Local imputation w/ collaboration. Suppose coordinating server collected covariance informa- tion E, from each agent and then distribute back each of them. Then one intuitive strategy is to use this information to impute each agent's s local data by replacing X,+ with E[X, X,, xX, before performing local linear regression. In other words, instead of computing 0, compute
jimp -T; X X.) X
to send back to the coordinating server. Note that we use T X,+ is in general of rank d,, and 0TmP is then the min-norm interpolant of agent i's data. Similar to COLLAB, we can use weighted empirical risk minimization parameterized by W; Rdxd and aggregate gimp via
inverse here as
@imp Ō(W1, W..): == argmin Σ T. (TTT))--TT 0.mp)
w:
The theorem, in conjunction with Theorem 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 Ú,.
Theorem 4.1. For 0i from agent i, have 0. T. (T;T,T)-'0, Given any weighting matrices W3 ddddd the aggregated imputation estimator gimp is consistent and asymptotically normal
V (jimp 0) N(O, C ....... ..... Wn)))

---

Using the same weights W E Rd; as in Theorem 3.1 for aggregated @imp, we have under p.s.d. - In
cone order, for weights W,, C Wm) = C*, where C* = (Z addition, the equality holds when W,- = T. W
1==1 T
As we will see in Sec. 3 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 0 and local parameters T,0.
In terms of communication cost, this local imputation method requires more communication than COLLAB, as a coordinating server needs to communicate E to all the agents. Each agent must first send df real numbers corresponding to 2, the coordinating server, and then the coordinating server will send back d² - df real numbers corresponding to the entries of Σ that agent has not observed. This amounts to de 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 coordi- nating server all of their data (X++, 1/)) for = 1,... , instead of their local estimators, 0; or 0TTP. Having all the data with structured missingness available, natural idea is to have the coordinating server first impute the data, replacing X++ with E[X; X++] = X1+T, and then perform weighted OLS on all of the nm data points, before sending d, dimensional local model back to each agent. Namely for scalars 01,
m
amm > 0, we take
m Σ
Gimp-gth gimp (ai, amm))
am))
a XT
Σ T X
Surprisingly, in spite of the additional power, jump-gib still does not beat Ô in Theorem 3.1 Theorem 4.2. For any scalars o1. a0m > 0, gimp-gth is consistent and asymptotically normal
vn (gimp-Etb 0) N(0, Chimp-eth (a1.... ,am)).
Recall the lower bound matrix -== ( m
=
i=1 T:T:) in Theorem 3.1 Xf N(O, M), we
have under p.s.d. cone order and any aq > 0, Cimp-glb (amp) (ar1,".. ... (m)) 2 C In addition, the equality holds when o; 1/(110;-11)) IIF: r²).
The communication cost for this method is significantly larger than the other methods we discussed. Each agent must send all its data to coordinating server. Factoring in the additional communication cost of receiving a new local model from the server, this amounts total nd,; + d, 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.
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 cartiallyfixed-desiee regime. For every sample x E Rd Fz+ E Rd is fixed vector. We draw h₁- from N(,,-, r,-) where 14- and r;_ is the conditional mean and variance of I;-- given Zz4- Here r is also the Schur complement. We draw Z;- from ((,,.T;-)) The samples 244 e dd comprise the matrices X,+ E Raxdd For all i E [m]. we will assume we have infinite sequence (w.r.t. n) of matrices X1+. This partially- sertialll-fixee-eee 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 0 E Rd The corresponding label y = 0446,+ + zz= 0: + E, where & E R is drawn from i.i.d. N(O, o2). We use yj R to denote its vector form for the agent j. To model the estimator's 's knowledge about the labels, we will have two observation models one weaker and one stronger- we will specify later when present our results.
For each observation model, we will have two types of results. The first type of result is minimax lower bound for full-featured data; i..., how well can estimator perform fresh sample without
7

---

missing features. This type of result will concern the full-feature asymptotic local minimax risk lim inf 1--000
Jeliml Pn, u) == 11-000 0 PEP.
inf sup nL 0(2, {Xi+Jemm)) 0)².
We will show that there exists B E Rdxd such that the local minimax risk in the previous display lower bounded by u Bu for all E TR. 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.c. what is the best prediction error an estimator (for the given observation model) can possibly have fresh sample for given agent. This type of result will deal with the missing-feature asymptotic local minimax risk, defined as
lim inf M ((Xe)) [m]]m] P n-+00
lim inf inf sup nEZzr(( (Z, (X+ hE(m)) - n-+00 PEP,
Similar to the first minimax enror definition, we will again show that there exists a B, E RR4 x such that the local minimax risk just defined is lower bounded by u B,u for all E R Recall for discussion surrounding why T,0 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 0: = (x+ Vr. Let Pd be a distribution over iid ~ N(O, T²T,)) We define the following family of distri-
0,...., 0m induced by and (LI,..., Em) butions pá (Pd ||0' A1|2 S en 1/3 which defines observation model. Intuitively, in this observation model, we are constructing a lower bound for estimators which have access to the features _ Xm+, the population covariance Σ, and access to
In comparison, our estimator COLLAB only uses E and é1.... ......... Am. We present our first asymptotic local minimax lower bound result here. The proof of this result can be found in Appendix D.I
Theorem 5.1. Recall that C = (Ei T. W For all E [m]] and n let the rows of Xit be drawn i..d from N(O, E++). Then for all u E IRd, with probability 1, the full-feature asymptotic local minimax risk for ph is bounded below as,
lim inf lim inf Wm.. ((Xiee((Xeeeee P Pa,ar u) 2 uTCuu.
c-000
n-+00
For all u E Rd, with probability 1, the missing-feature asymptotic local minimax risk for pb is bounded below as
lim inf lim inf wi ((X++ Jeemm: ne u) 2 T TT.TTT
20000
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 PK, == (P* - S cn-1/23 as the observation model. Intuitively, in this model, we are constructing lower bound for estimators having access to all of the features X1+..... Xm++ and access to y1,.... .. ym This observation model is stronger than the previous observation model because estimators now have access to the labels y. We note again estimator rely do not
that our estimator COLLAB only uses Σ and 0 The quantities scale with n, making our estimator much weaker than other potential estimators in this observation model, estimators are allowed to depend on Vi, which grows in size with We present second asymptotic local minimax lower bound result here, starting with defining the strong local lower bound matrix C : (Em) 22/(110,-11 + +o²))=
0m.
The proof of this result is in Appendix D.2
Theorem 5.2. For all E [m]] and n let the rows of Xi+ be drawn i..dd. from N(O, Ex+). Then for all E Rd, with probability 1, the full-feature asymptotic local minimax risk for nie is bounded below as
lim inf lim inf Mnm. ((X++ hiemm]; Ph.ss u) 2 u T. n-000
For all u E Rdd, with probability 1, the missing-feature asymptotic local minimax risk for Py is bounded below as
lim inf lim inf wnt.tt(fX0 (1X++ Fie(on] PH u) 2 u T T.C.T. u.
c0000 n no0
8

---

Ulioots Full Peature Prediction Ereer (Lange n)
Iliowis Full Festure Prodiction Erve Frnall m)
Illinii Full Fratio Prediction Frnu (imalli s)
Tred Itr.
Nove
Proc [rr.
Sig
Dumber samples (n)
Trmmmer camples (n)
number samples
(a) Prediction Error for large n
(b) Prediction Error for small
(c) Small n omitting Impute algs.
Figure I: 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.I it is clear that the lower bound in the strong observation setting is in general smaller as
0 0 T 11,
Σ- T, E4 Ti II!
2,
which further implies C = (E = 2/ 11,, + o + -²))
We argue that the two lower bounds comparable in the missing completely at random [L.. Consider for every agent i, each coordinate missing independently with probability p. In this case, (d,, Ei+, T.) are i.i.d. random triplets parameterized by p.
Corollary 5.3. Under the random missingness setup with missing probability let the eigenvalue (E) > 0 and define its condition number K (x())aa((E) Suppose we have the limits limm-soc mC and lim mC exist and
of Σ be (Σ) S 24
1 (1 +
lim mC* lim mC > lim mC. m 100
m0-000
m0000
6 US Census Experiment
We experiment on real US census data modified from the ACSTravelTime dataset from the folktables package 77). The code can be found bettps:/ggttte ctw commgaryccheng//// collab. We (artificially) construct m 5 datacenters (agents), each containing data from only one of California, New York, Texas, Florida, and Illinois. The goal 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 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 cressianss-reeenttii
significant deviations from our theoretical assumptions which serve 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.I have 37 features. We compute the feature covariance matrix from training data across all the datacenters, which we plot in Appendix A. We assume we are able do this because this computation can be done in 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 27. This models the feature heterogeneity which varies across geography. Each datacenter will have 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-Colllab, Imputation, and RW-Imputation. We briefly describe each method here; Appendix A. contains a more detailed description of each method. Naive-Local refers to each agent locally perform OLS to construct 0, Naive-Collab does an equal-weighted average of the agent OLS models- Eie(m) ITT_ 0./m Optimized-Naive--Cnn uses gradient descent to optimize the choices of weights used in Naive-Collab. Optimized-Nar -aadd--aaiiilllaa uses fresh labeled samples without any missing

---

features during gradient descent, so in this sense, method. Imputation refers to the global imputation estimator fimp-glb with a, = 1/m. RW-Imputation Imputation but with the optimal choice of weights a,, 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 test, run 80 trials to form 95% confidence intervals.
is more powerful than our
We see that for n < 800 in Figures [[b] and [[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 2 2000 in Figure T(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-Colllo Naive-Collab. The performance of Naive-Collab and Seiiiiize-Naaie-CCae to the performance COLLAB here than the synthetic data experiment Appendix AA.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.
and being much closer
Discussion and Future Work
Optimal weights beyond Gaussianity. E [xx++ 24 =[00 has nice closed form in Gaus- sian setting because Zi4- and de+ are independent-whice is in general not true without Gaussianity. If can directly sample from the feature distribution P (e.g., unlabeled data), then can empirically estimate E [ 2,42 T 0:-2 by sampling from P and using any consistent plug-in estimate (e.g., run COLLAB with weights W; = Ia,) ). This will return a good estimate of the optimal weights. An interesting future direction prove lower bounds without the Gaussianity assumption.
Generalization to non-linear models. Recall in the Gaussian setting, the optimal weights in COLLAB are W = E (Bz,, [((x)+ 0.) 1)) Then, the optimal loss function in Eq. 3 becomes
Emg [((x4 0.) - (x14, T,0))
Σ 0++ + E E - 8, 81llv
Er.y(((0et. (,0e4+0.)) 0.) y)
This hints at a generalization to non-lincar models. Suppose the local agents train on models f' 0,), Rd. x Rd. + )) and the global model fa; 0), R x H + y satisfies for some mapping
T, : + Rd, f(z T,0) = f (z++; 0:). Consider a loss function e(-,-) 3 x y > [0, co). Then
we can consider the following way of aggregation inspired by COLLAB for linear models
Exe (( J ((I:):710)))
argmin
(6)
Bz4.y (1 0.).v)
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 Ô to all agent, Step 2. each agent i does some local training initialized at 0 to form Ô., and Step server collects 0; and aggregates them to form new 9. 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 [1]] However, now K suggests solution 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.
10

---

References
[1] Melissa J. Azur, Elizabeth A. Stuart, Constantine Frangakis, and Philip J. Leaf. Multiple imputation by chained equations: what and how does it work? International Journal of Methods in Psychiatric Research, 20, 2011.
[2] David Chas Bolton, Parisa Shokouhi, Bertrand Philippe Gerard Rouct-Leduc, Claudia Hulbert, Jacques Rivière, Chris Marone, and Paul Allan Johnson. Characterizing acoustic signals and searching for precursors during the laboratory seismic cycle using unsupervised machine learning. Seismological Research Letters, 2019.
[3] Kabir Aladin Chandrasekher, Ahmed El Alaoui, and Andrea Montanari. Imputation for high- dimensional linear regression. arXiv preprint arXivv2001.091 09000, 2020.
[4] Wei-Ning Chen, Peter Kairouz, and Ayfer Özgür. Breaking the communication wwwpraattoinnvvvvceeeeeeee privacy-accuracy trilemma. IEEE Transactions on Information Theory, 69:1261-1222 12811, 2020.
[5] Kewei Cheng, Tao Fan, Yilun Jin, Yang Liu, Tianjian Chen, and Qiang Yang. Secureboost: lossless federated learning framework. IEEE Intelligent Systems, 36:87-98, 2019.
[6] Giacomo Curzi, Dario Modenini, and Paolo Tortora. Large constellations of small satellites: survey of near future challenges and missions. Aerospuce, 2020.
[7] Frances Ding, Moritz Hardt, John Miller, and Ludwig Schmidt. Retiring adult: New datasets for fair machine learning. In Neural Information Processing Systems, 2021.
[8] John C. Duchi, Michael I. Jordan, Martin J. Wainwright, and Yuchen Zhang. Optimality guarantees for distributed statistical estimation. arXiv: Information Theory, 2014.
[9] Adrià Gascón, Phillipp Schoppmann, Borja Balle, Mariana Raykova, Jack Doerner, Samee Zahur, and David Evans. Privacy-preserving distributed linear regression high-dimensional data. Proceedings on Privacy Enhancing Technologies, 2017:345 364, 2017.
[10] Elisabeth Gassiat. Revisiting the van Trees inequality in the spirit of Hájek and Le Cam. 2014.
[11] Stephen Hardy, Wilko Henecka, Hamish Ivey-Law, Richard Nock, Giorgio Patrini, Guillaume Smith, and Brian Thome. Private federated learning vertically partitioned data via entity resolution and additively homomorphic encryption. ArXiv, abs/ abss/7711.1067) 2017.
[12] Zhenbang Jiao, Hu Sun, Xiantong Wang, Ward B. Manchester, Tamas I. I. Gombosi, Alfred O. Hero, and Yang Chen. Solar flare intensity prediction with machine learning models. Space Weather, 18, 2019.
[13] Peter Kairouz, H. B. McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Keith Bonawitz, Zachary B. Charles, Graham Cormode, Rachel Cummings, Rafael G. I.. D' Oliveir,,, Salim Y. El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, 7aïd Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson, Justin Hsu, Martin Jaggi. Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konecný, Aleksandra Korolova, Farinaz Koushanfar, Oluwasanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür. R. Pagh, Mariana Raykova, Hang Qi, Daniel Ramage, Ramesh Raskar, Dawn Xiaodong Song, Weikang Song, Sebastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramer, Praneeth Vepakomma, Jianyu Wang. Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning. Found. Trends Mach. Learn., 14:1-210, 2019.
[14] Alan F. Karr, Xiaodong Lin, Ashish P. Sanil, and Jerome P. Reiter. of vertically partitioned data using secure matrix products. Journal of Official Statistics, 25: 125-138, 2009.
analysis
[15] Gayaneh Kyureghian, Oral Capps, and Rodolfo M. Nayga. A missing variable imputation methodology with an empirical application. 2011.

---

[16] Roderick J. A. Little and Donald B. Rubin, Statistical analysis with missing data, third edition. Wiley Series in Probability and Statistics, 2019.
[17] Yang Liu, Yan Kang, Chaoping Xing, Tianjian Chen, and Qiang Yang. secure federated transfer learning framework. IEEE Intelligent Systems, 15,70-82, 2020.
[18] Lamin R. Mansaray, Adam Sheka Kanu, Lingbo Yang. Jingfeng Huang, and Fumin Wang. Evaluation of machine learning models for rice dry biomass estimation and mapping using quad-source optical imagery. GlScience & Remote Sensing, 57:785 -96, 2020.
[19] H. B. McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Agüera y Ar- cas. Comminicationn-fficcent learning of deep networks from decentralized data. In In- ternational Conference on Artificial Intelligence and Statistics, 2016. URL https://api
semanticscholar arg/CorpusIDi :4925533::4995545 14955348
[20] Board on Physics. A strategy for active remote sensing amid increased demand for radio spectrum. 2015.
[21] Benjamin Recht, Christopher Ré, Stephen J. Wright, and Feng Niu. Hogwild: A lock-free approach to parallelizing stochastic gradient descent. In NIP,, 2011.
[22] Ritvik Sahajpal, Lucas Fontana, Pedro Lafluf, Guillermo Leale, Estefania Puricelli, Dan O'Neill, Mehdi Hosseini, Mauricio Varela, and Inbal Becker-Reshef. Using machine-learning models for field-scale crop yield and condition modeling in argentina. 2020.
[23] Joseph L Schafer. Multiple imputation: primer. Statistical methods in medical research, 8(1):
3 15, 1999.
[24] Joseph L. Schafer and John W. Graham. Psychological methods, 2:147-77, 2002.
Missing data: our view of the state of the art.
[25] Joseph L Schafer and Nathaniel Schenker. Inference with imputed conditional means. Journal of the American Statistical Association, 95(449):144-154, 2000.
[26] Ananda Theertha Suresh, Ziteng Sun, Jae Hun Ro, and Felix X. Yu. Correlated quantization for distributed mean estimation and optimization. ArXiv, abs/2203.04925, 2022.
[27] Anastasios A. Tsiatis. Semiparametric theory and missing data. 2006.
[28] Jaideep Vaidya and Chris Clifton. tioned data. In Knowledge Discovery and Data Mining, 2003.
k-means clustering over vertically parti-
[29] Naisyin Wang and James M. Robins. Large-sample theory for parametric multiple imputation procedures. Biometrika,
1998.
[30] Qiang Yang. Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. arXiv: Artificial Intelligence, 2019.
[31] Hwanjo Yu, Jaideep Vaidya, and Xiaoqian Jiang. Privacy-preserving svm classification on vertically partitioned data. In Pacific-Asia Conference on Knowledge Discovery and Data Mining, 2006.
12


Reviews:
[{"review_id": "QGxtp0VAsF", "reviewer": "Reviewer_jiy8", "summary": "This paper studied the problem of collaboratively learning least squares estimates with multiple agents, each of which only observes a different subset of the features. The authors proposed a distributed, semi-supervised algorithm called Collab consisting of three steps: 1) local training, 2) aggregation, and 3) distribution. The authors showed that the proposed Collab algorithm is nearly asymptotically local minimax optimal. The authors conducted experiments to verify their algorithm on real and synthetic data.", "questions": "1. In Step 2 of the main loop in the Collab algorithm, what happens if the estimate $\\hat{\\Sigma}_i$ is inaccurate? Could the authors analyze the impact of such errors?\n\n2. Although the paper is well written in general, there are some minor typos. For example, in Line 162, it appears that the global estimator should be $\\hat{\\theta}^{clb}$.", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. This paper provides deep theoretical insights into their proposed Collab algorithms. The results in Theorems 4.1 and 4.2 that the performance of Collab is no worse than local imputation with collaboration and global imputation are novel and surprising.\n\n2. The asymptotic local minimax lower bound is interesting and could be of independent interest.", "weaknesses": "Although the results in this paper are interesting, as the authors themselves admitted, they are only limited to the linear model and Gaussian features. But the authors did provide some interesting discussions in Section 7 on future directions. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. \n\n**Response to weaknesses**: Our theory is indeed limited to Gaussian features. We did experiments on non-Gaussian data in our Folktables experiment. Though it's only a preliminary experiment, we hope that it shows our method does not overfit to the Gaussian data setting.\n\n**Response to Q1**: noisy estimate of covariance: This is an interesting point.  We first want to point out a potential confusion that could be caused by a typo in the definition of $\\hat{W}_i^g$ (line 155). The numerator is supposed to be the sample sub-covariance $\\hat{\\Sigma}\\_{i+} = X\\_{i+}^\\top X\\_{i+}/n$ instead of the exact sub-covariance matrix $\\Sigma\\_{i+}$. In fact, we use $\\hat{\\Sigma}\\_{i+}$ in Algorithm 1 and in the proof of Corollary 3.2 in the submitted supplementary materials. It is not clear what is the optimal procedure if we do not have a consistent estimate for the population covariance, which essentially boils down to the harder problem of distributional shift. This could be a future direction and we will include it in our discussion.\n\n**Response to Q2: typos**:  Thanks for pointing this out. We will fix the typo you mentioned, the typo in the definition of $\\hat{W}\\_i^g$ (line 155), and other typos in the camera ready version of the paper."}, {"review_id": "6LuN3HD5dA", "reviewer": "Reviewer_Mws6", "summary": "Summary\n-------\n\nThe paper studies collaborative linear regression when m agents attempt to collaboratively\nestimate a linear model, under communication constraints. Each agent i only observes di\nof the d features. A central server designs a protocol to elicit sufficient information\nfrom each agent and compute the parameter theta of the linear model so as to minimize\ncommunication costs and the estimation error.\nThe authors present, what appears to be a near-complete solution, for the case where the\ncovariates are distributed as a Gaussian and when the covariance matrix is known. This\nincludes asymptotic normality results for the proposed estimator and lower bounds which\nmatch as n goes to infinity. The authors also compare their method against other baselines\nbased on imputing data and show that their estimation errors are no worse but at\nsignificantly lower communication cost.\n\nDecision: While the paper is well-presented and pleasant to read, I am not an expert in\nthis topic and did not have the time to go through the proofs in detail. As such,\nI am unable to evalute the technical merit of the paper (challenges of the problem,\nnovelty of proof techniques). I have given a positive score with low confidence to reflect\nthis but will heed to more expert reviewers during the discussion.\n\n\nDetailed comments\n-----------------\n\nThe authors have used local and global imputation\nbaselines in Table 1 to show the communication benefits, and have gone on to show that\ntheir method does no worse theoretically than these methods.\nHowever, I am not sure if these are particularly strong baselines to compete with; for\ninstance, I would not have expected communicating all data points to be necessary.\nIn fact, at the outset, my intuition suggested a solution which computes local\ncoefficients by each agent which are then aggregated by the server in an appropriate\nfashion.  It would have been helpful if the authors had better illustrated the challenges\nin doing so. For instance, is the method in 3.1 the most natural way to solve this\nproblem, or are there other naive ways to aggregate the coefficients that yield\nsub-optimal solutions?\n\nThe same applies to the non-Gaussian case and the setting when the co-variance matrices\nare unknown. The authors could have done a better job of illustrating the challenges.\nIn general, I did not get a sense for how challenging this problem setting was to\nappreciate the contributions by the authors.\n\nDo the results in Theorems 3.1-3.2 implicitly capture the difficulty of the problem in\nterms of the number of covariates each agent has access to?\n- For instance, in the worst case there could be only n samples (when there is a perfect\n    partition of the covariates among the agents, and they have the same points),\n    but in the best case there could be mn\n    samples (when all agents have all the covariates and have distinct points)\nIf so, is it possible to make this more explicit in the results.\n\nThe paper is largely well-written. The problem is well-motivated, the setting is described\nclearly, and the method/results are organized well. I would have however liked to see a\nsketch of the main proof ideas and how they differ from similar results in the linear\nregression literature. A breakdown of the key proof challenges that the authors had to\novercome would have also been useful.\n\nWhat was the reason for the discussion around o(n) communication complexity? It appears\nthat local imputation methods and the method of the authors are able to achieve consistent\nestimation with communication cost that does not depend on n.\n\nIn the experiments, the imputation-based methods outperform the method of the authors when\nthere are more samples. Intuitively, I would have expected this since they are more\ncommunication-heavy. However, this is not the case in the low-sample regime. Can you\nexplain why this is the case?\n", "questions": "See above.", "limitations": "See above.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "See above.", "weaknesses": "See above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. \n\n**Response to baseline strength**: Our method is asymptotically instance-optimal (shown by our lower bounds), meaning that no algorithm could perform statistically better on any specific problem instance. In some sense, our lower-bound is the ultimate theoretical “baseline”. The motivation for comparing baseline algorithms which are “stronger” than our method in terms of communication cost is to ultimately show in the experiments section that Collab is not overfit to the assumptions of our theory and performs well against communication-ignorant, conventionally-adopted methods, like imputation.\n\n**Response to the suggested intuitive solution**: If we understand correctly, our method falls in the scope of your intuition—Collab takes the locally computed coefficients and debiases them using covariance information. We note that we need covariance information from the agents, otherwise the debiasing procedure would return a biased estimate of the parameter. Having said this, we do want to point out that our aggregation approach is novel to the best of our knowledge and not standard in the missing data literature. As we discuss in the related work, imputation-type algorithms are more standard, which is also why we chose to baseline against them in our theory and experiments. \n\n**Response to non-Gaussian/unknown covariance case**: We first want to point out a potential confusion that could be caused by a typo in the definition of $\\hat{W}\\_i^g$ (line 155). The numerator is supposed to be the sample sub-covariance $\\hat{\\Sigma}\\_{i+} = X\\_{i+}^\\top X\\_{i+}/n$ instead of the exact sub-covariance matrix $\\Sigma\\_{i+}$. In fact, we use $\\hat{\\Sigma}\\_{i+}$ in Algorithm 1 and in the proof of Corollary 3.2 in the submitted supplementary materials. We discuss the non-gaussian setting in lines 147-150 and section 7; to summarize, in the gaussian setting, we can estimate $W_i^\\star$ from the data we have access to. In the non-gaussian setting we cannot. We are happy to answer any specific questions you may have about this.\n\n**Response to the implicit difficulty in Theorem 3.1 and Corollary 3.2**: Yes, you are right. Theorem 3.1 and Corollary 3.2 (and our lower bounds) capture the difficulty of the problem in terms of the number of covariates each agent has access to. In fact, if we just look at the lower bound $C^g$ in the Gaussian setting and consider each of the summand. Note that\n$T\\_i^\\top W\\_i^g T\\_i = \\frac{\\Sigma - \\Pi\\_i^\\top \\begin{bmatrix} 0 & 0 \\\\\\ 0 & \\Gamma\\_{i-} \\end{bmatrix} \\Pi\\_i}{\\\\|\\theta\\_{i-}\\\\|\\_{\\Gamma\\_{i-}}^2 + \\sigma^2}$. If we have strictly more coordinates observed for the $i$-th agent, then the Schur complement $\\Pi\\_i^\\top \\begin{bmatrix} 0 & 0 \\\\\\ 0 & \\Gamma\\_{i-} \\end{bmatrix} \\Pi\\_i$ will be a smaller matrix and also $\\\\|\\theta_{i-}\\\\|\\_{\\Gamma_{i-}}^2$ will be a smaller quantity, resulting in an overall larger $T\\_i^\\top W\\_i^g T\\_i $ and thus smaller uncertainty $C^g$ and therefore also smaller test error. We will add this discussion in the camera ready version of the paper.\n\n**Response to o(n) bandwidth**: Thank you for bringing this up. You raise a good point about our storytelling with your observation that the local imputation baseline also has o(n) communication cost. While hopefully we were convincing in the introduction of why o(n) bandwidth constraints are important in real settings like satellites and seismic sensors, we agree that we did not adequately justify why a communication reduction from $d^2$ to $d\\_i^2$ is important. We present this justification here now, and we will add it to the camera ready version. One property of a good collaborative algorithm is that agents which are a part of the collective should be incentivized to welcome new agents. In Collab, adding new agents to the collective never increases the communication cost to any of the existing agents a part of the collective. On the other hand, in the local imputation baseline algorithm, new agents with data collected from new/different sensors increase the communication cost of other agents in the collective (i.e., because $d$ becomes larger). This implicitly incentivizes homogeneity of sensors within the collective, which is antithetical to the idea of leveraging diverse data to make better predictions.\n\n**Response to the low-sample imputation underperformance**: This is a good observation. We believe the phenomenon boils down to numerical instability also known as the double descent for linear regression when the number of samples $n$ and underlying dimension $d$ is comparable. We point out that our theory is asymptotic and thus the statistical optimality holds for $n \\\\gg d$, and therefore not predictive in the low-sample regime. Though not within the scope of our current theoretical setup, it is nontheless an interesting future direction to investigate what is the optimal procedure when $n$ is comparable to $d$ (for instance, whether adding regularization would help). We will add this to the discussion section of our work.\n"}, {"review_id": "U05wbcf2iA", "reviewer": "Reviewer_vMLQ", "summary": "The paper studies statistical inference in learning a linear regression model in the cross-silo or vertical federated learning setting. The paper formulates the problem as a missing data problem and chooses single imputation methods to deal with the associated inference of the common parameter. The paper shows the theoretical properties of the proposed estimator under some assumptions. Finally the paper compares their results with other existing methods.", "questions": "What is the motivation to consider minimizing a weighted empirical loss in Equation (3)?", "limitations": "The paper considers single imputation, claiming that this is okay given that the goal is estimation error instead of confidence intervals. However, many theoretical results given in the paper are about the asymptotical distributions of the proposed estimators. Do these asymptotical distributions properly consider the uncertainty associated with the missing data imputation?", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper studies a very important problem and relatively easy to follow.", "weaknesses": "1. The paper is not carefully written with confusing notations and problem formulation:\n\n1(a). In Section 2, it says that “the $i^{th}$ agent has data $(x_{i+}, y)$” and then later uses $y_i$ to denote the $i^{th}$ agent labeled data. It is unclear whether these labels are the same or not. If they are the same, this is an unrealistic assumption in practice and obviously contrary to the common assumption in the literature that only one active party has access to the label data. \n\n1(b). Equation (1) defines a weighted MSE, with expectation assessed with respect to the feature $x$, which is problematic. Why not consider the typical unweighted MSE with expectation evaluated with respect to $y$ in regression setting? \n\n2. It seems that the contribution of the paper is to apply existing single imputation based missing data method to vertical federated learning setting, with the exception that here they assume each agent has access to its own label data (which again is problematic).\n\n3. The paper does compare their methods with several other methods in the experiment. However, there are not enough discussions of these methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. \n\n**Based on your comments, we believe there is a misunderstanding**. We are not doing vertical federated learning [Liu et al.]. Unlike vertical federated learning, agents in our framework are not measuring the same underlying set of users. Vertical federated learning is not the right way to model settings like the satellite and seismic sensor estimation problem examples we discussed in the introduction, as this would mean, each agent (i.e., the sensor) would take (different) measurements of the same locations at the same time. In the setting we study, each agent has their own input data and labels. Thank you for pointing this out; we will clarify this in the camera ready version.\n\n**Response to 1a**: We believe this confusion stems from some notational issues. In line 102, $(x, y)$ refers to a single sample drawn from the generating distribution. In line 104, $y\\_i\\in \\mathbb{R}^n$ refers to the vector of labels the agent observes. To be clear, as said in line 103, each agent gets $n$ draws of $(x, y)$ from the generating distribution and observes $X\\_{i+}$ and $y\\_i$: the labels are not the same across agents. Furthermore, just to make sure there is no confusion, the features observed by agent $i$ and agent $j$ ($j\\neq i$) are *not* different subsets of the same $n$ feature vectors: each agent draws $n$ fresh covariates and observes a subset of the covariates. Finally, as we are not doing vertical federated learning, we do not have an “active party” constraint on the labeled data. We will make this very clear in the camera ready version.\n\n**Response to 1b**: We are confident that what we wrote is the standard notion of prediction/generalization error for a fresh sample $(x \\in \\mathbb{R}^d, y \\in \\mathbb{R})$. In the regression setting, weighted MSE with respect to the data distribution is the same up to an additive constant as the test error on a fresh sample: $\\mathbb{E}\\_{x, y}[(\\langle x, \\hat{\\theta}\\rangle) - y)^2] = \\mathbb{E}\\_x[(\\langle x, \\hat{\\theta} \\rangle - \\langle x, \\theta \\rangle)^2] + \\sigma^2 = \\\\|\\hat{\\theta} - \\theta \\\\|\\_{\\Sigma}^2+ \\sigma^2$. This formulation is standard (e.g., see section 2.1 in Hastie et al.).\n\n**Response to 2**: First, Collab is not doing single imputation at all. In fact, in section 4, we compare our method Collab against traditional single imputation methods. Second, we are not doing vertical federated learning as we discussed above. \n\n**Response to 3**: Can you provide more details about what discussion you feel is missing? Is there anything you want us to clarify? We are open to suggestions. We want to emphasize that while we chose to compare against imputation methods in Section 4, our method is asymptotically instance-optimal (as shown by our lower bounds), meaning that no algorithm could theoretically perform statistically better. In this sense, our lower-bound is the ultimate theoretical “baseline”. \n\n**Additional References**\n\nHastie, Trevor J. et al. “Surprises in High-Dimensional Ridgeless Least Squares Interpolation.” Annals of statistics 50 2 (2019): 949-986 .\n\nLiu, Yang et al. “Vertical Federated Learning.” ArXiv abs/2211.12814 (2022): n. pag.\n"}, {"review_id": "X1JXvxQ77T", "reviewer": "Reviewer_N12y", "summary": "The authors investigate collaborative learning of least squares estimates for multiple agents with varying feature subsets. The goal is to coordinate the agents efficiently to achieve optimal estimators without exchanging labeled data. To address this, the authors propose the distributed algorithm Collab, consisting of local training, aggregation, and distribution steps. Despite not sharing labeled data, Collab approaches near-asymptotic local minimax optimality, outperforming methods that do utilize labeled data. They validate our approach through experiments on real and synthetic datasets.", "questions": "In Section 7, is that the numerator in *generalizing to non-linear models* requires sharing global information for $x_{i+}$? What are the possible ways to alleviate this constraint? ", "limitations": "I do not foresee any potential societal impact of this work. ", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "Distributed learning with heterogeneous data sources is a problem of broad interest. In this study, the authors tackle this problem in a simplified setting and provide robust theoretical guarantees. Their theoretical results are strong, demonstrating the solid foundation of their approach. In addition, the point of minimizing communication resources is an interesting angle. Furthermore, the experimental results are compelling, further supporting the effectiveness of their method. ", "weaknesses": "- Settings need more justification: the authors discuss a setting where a linear regression problem is running on satellites with completely different features. This seems restrictive and can the authors elaborate more on the motivation of their study?\n- Random X with zero mean. If X is not zero-mean, then the regression with partial information is not consistent anymore (different features may have correlations). I can imagine this is a very common scenario and can the authors provide some discussions on this? \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n**Response to justifications for our setting**: Regarding the satellite application, we agree it is a stylized example and have expanded the introduction to discuss other potential applications like sensor networks, weather stations, and hospitals. The key aspects we wish to capture are: 1) agents have different features due to heterogeneous data sources, 2) agents wish to leverage correlations and train collaborative models, 3) communication constraints exist. We believe these aspects apply broadly. The high level motivation for our work is to study how practitioners should handle hardware heterogeneity in communication-constrained collaborative learning environments. Works in federated learning have thought about how to do training with heterogeneous computational hardware; e.g., learning with client phones that have different processing power [Yang et al.]. But little attention has been focused on settings on heterogeneous measurement devices; e.g., what would we do if the features collected by each agent were heterogeneous? As we show, this problem is challenging, even in the simplest linear regression setting.\n\n**Response to zero-mean data**: Thank you for bringing up the point about zero-mean assumption. It is standard practice in machine learning to first center the data by subtracting the mean. This pre-processing step results in zero-mean features. This is why many other theory papers study the zero mean setting (e.g., Hastie et al., [3]). In fact, we followed this approach in the real census data experiment by centering each feature before model training and evaluation. As seen in Figure 1, COLLAB either outperforms or is competitive against baselines, indicating it is robust even when the zero-mean assumption is violated in practice.\n\n**Response to the question**: This is a good question. We want to clarify that the formula is under the expectation of $x\\_{i+}$ which is a population quantity that each agent can estimate consistently and locally. So, no, the numerator does not require sharing individual data samples $x\\_{i+}$, as we assume that each agent has $n$ labeled samples with $n$ growing to infinity (over the subset of the features they can observe). Thus, each agent can estimate the numerator individually with the samples they have access to and with their own local model.\n\n**Additional References**\n\nHastie, Trevor J. et al. “Surprises in High-Dimensional Ridgeless Least Squares Interpolation.” Annals of statistics 50 2 (2019): 949-986 .\n\nYang, Chengxu et al. “Characterizing Impacts of Heterogeneity in Federated Learning upon Large-Scale Smartphone Data.” Proceedings of the Web Conference 2021 (2021): n. pag.\n"}, {"review_id": "g2WlySCRif", "reviewer": "Reviewer_93Dr", "summary": "The paper discusses the idea of estimating least squares collaboratively when each agent has access to a different set of features of the same data. The aim is to design an effective and efficient algorithm in terms of communication cost (various agents transferring/communicating information/data). The paper introduces a new algorithm *COLLAB* that is efficient and also applicable in security fields where features can not be transferred between agents/sensors/machines/systems. The authors theoretically prove that the proposed algorithm *COLLAB* is minimax optimal and perform a set of experiments to showcase the capability of the proposed algorithm *COLLAB*.", "questions": "* The paper assumes that all the agents have the same linear model. How restrictive is it? One of the agents has a $R^{16}$ dimension feature vector, whereas the other might have a $R^2$ dimension feature vector. So, the same complex model does not seem logical here.\n* Second assumption is that the agent has enough data to estimate $\\Sigma$. How will it impact the algorithm if there is little data and $\\Sigma$ is unreliable/biased?\n* A toy example to plot various parameter $\\theta_i$ of the agents, how they combine to give $\\theta_{global}$, and how they compare with the oracle value of $\\theta$.\n* A setup where the computational benefits can be shown empirically as well?\n", "limitations": "The authors do not discuss them, and I do not see a limitation directly of this work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper discusses the setup typical in many fields (especially security applications) where data can not be transferred between agents due to security reasons or input-output constraints (network bottleneck). In the setup discussed, when the agent has a linear model, ordinary least squares (OLS) is a logical way to solve and get the parameter value. The idea behind *COLLAB* is logical and intuitive however needs more clarity in the presentation. The comparison against the imputation methods is also logical, as these are go-to models for such setups. The authors derive local minimax lower bounds to showcase that the proposed algorithm *COLLAB* is very close to the optimal, which is interesting. For the correctness of this section, I would rely on other reviewers as I was not able to follow the derivations clearly. The experiments are also a good mixture of real-world and synthetic data sets where the capability of the proposed algorithm is shown. ", "weaknesses": "* A nice idea, but some assumptions are very hard/limited, and the evaluations are limited.\n* Some details (especially Sections 3.1 and 3.2) are presented convolutedly and are hard to grab. Some minor clarifications:\n  - L110, is the `x` written same as L99? It probably should be `X`?\n  - L132, L133: I am still not sure what `(i)` refers to here.\n  - In general, I believe it is convenient to stick to the notation where bold **`X`** is a matrix, bold **`x`** is a vector and `x` is a scalar. It makes it easy to follow the equations.\n* Evaluations are limited. More empirical evaluations are to be performed to showcase the results of the proposed algorithm against other methods. A limitation here is the linear model in the agents, which limits the modelling capability of the model.\n* Considering non-linear cases as well would make the paper more solid. Then the agents are more flexible, and the authors can experiment with more complex data sets.\n* Computation benefits should be shown empirically as well. Currently, it is only shown theoretically.   \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n**Response to W1: limited assumptions and evaluations**: We admit that our current theory relies heavily on Gaussianity assumption.These generalizations have various technical challenges, which make giving strong theoretical guarantees (i.e., what we were able to do in our linear Gaussian model) very difficult. In our work, we opt to provide these strong guarantees in lieu of more generality. We point out that the experiments are on real data, suggesting potential generalization of our theory beyond the Gaussian case. Removing linearity/Gaussianity would be an interesting future direction that we have discussed in Section 7. Having said that, we are happy to  explain in more detail if you want further clarification on  particular assumptions/evaluations.\n\n**Response to W2: minor clarifications**: Thank you for the suggestions, we will make clarifications in the camera-ready version for the following points.\n2.1: Yes, the $x\\in \\mathbb{R}^d$ in line 110 is the same as the $x\\in\\mathbb{R}^d$ in line 99.\n2.2: The $(i)$ refers to the second equality in the equation between lines 131 and 132.\n2.3: This is a good suggestion, and we will try to make the notation more clear in the camera-ready version.\n\n**Response to W3: limited evaluations**: We partially respond in Response to weakness 1. We want to emphasize that we believe our main contribution is theory. The focus of our preliminary experiments is to show that our method does not overfit to the Gaussian data setting.\n\n**Response to W4: non-linear cases**:  We respond in Response to weakness 1.\n\n**Response to W5: empirical justification of computational benefits**: Thank you for bringing this up.e think there is some confusion, as our focus is on communication cost instead of computational cost. Please correct us if we misunderstood. We note that the communication costs in Table 1 are actually not asymptotic, local imputation requires communicating $d^2$ real numbers, global imputation $nd_i$ real numbers, and our method $d_i^2$ real numbers; we will clarify this in the camera-ready version. For this reason, we believe benchmarking real communication costs is not necessary.\n\n**Response to Q1: dimensional differences among agents**: Indeed, our model setup assumes the same underlying model with agent-dependent, unobserved dimensions. We want to clarify that it is not restrictive to dimensional differences among agents. Our setup allows one agent to observe data in $\\mathbb{R}^{16}$ dimensions while another observes data in $\\mathbb{R}^{2}$ dimensions; the second agent would just observe fewer dimensions than the first agent. It would be very challenging to develop theory when the underlying models are different for each agent, as there is no shared global model we can analyze. We feel this is beyond the scope of the discussion in our paper. \n\n**Response to Q2: unreliable/biased estimate of the covariance**: This is an interesting point. We first want to point out a potential confusion that could be caused by a typo in the definition of $\\hat{W}\\_i^g$ (line 155). The numerator is supposed to be the sample sub-covariance $\\hat{\\Sigma}\\_{i+} = X\\_{i+}^\\top X\\_{i+}/n$ instead of the exact sub-covariance matrix $\\Sigma\\_{i+}$. In fact, we use $\\hat{\\Sigma}\\_{i+}$ in Algorithm 1 and in the proof of Corollary 3.2 in the submitted supplementary materials. It is not clear what is the optimal procedure if we do not have a consistent estimate for the population covariance, which essentially boils down to the harder problem of distributional shift. This could be a future direction and we will include it in our discussion.\n\n**Response to Q3: suggestion of a toy example**: Thank you for this experimental suggestion. As $\\theta_i$ are in general high dimensional (>3), the plots we can visually make would probably be $\\ell\\_2$ error against the ground-truth, which might be less clear, so we opted for doing real data experiments due to page-limit. Please let us know if you have any suggestions of how to visualize a toy example.\n\n**Response to Q4: computational benefits**: See our response to W5."}]

Author rebuttal:
We want to clear up some possible confusion due to a typo we made. Our method Collab only needs the sample covariance $\hat{\Sigma}\_{i+} = X\_{i+}^\top X\_{i+}/n$ --- **not** the population covariance $\Sigma\_{i+}$ --- for our results to hold (see Algorithm 1 for the correct pseudocode). In other words, Collab does not need to know additional population information relative to the baselines. This potential confusion is likely caused by a typo in the definition of $\hat{W}\_i^g$ (line 155). The numerator is supposed to be the sample sub-covariance $\hat{\Sigma}\_{i+}$ instead of the exact sub-covariance matrix $\Sigma\_{i+}$. In fact, we use $\hat{\Sigma}\_{i+}$ in Algorithm 1 and in the proof of Corollary 3.2 in the submitted supplementary materials.

Meta-review: Reviewers had mixed scores on this paper.  The paper has some merits: viz. the minimax optimality, but also suffers from the limitations reviewers have pointed out (viz. applying to a limited Gaussian covariates setting). However, overall it was felt that the merits outweigh the negatives, so the paper is accepted to NeurIPS. The authors are advised to carefully revise their paper with the feedback provided by the reviewers for the final version.
Decision: Accept (poster)