SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\left\| \widehat {\mathbf {B}} ^ {(l r)} - \mathbf {B} ^ {*} \right\| _ {F} \leq c _ {2} \left(\sqrt {\frac {r \gamma (Q + p)}{n _ {\operatorname* {m i n}}}} + \sqrt {r} \| \Xi \| _ {\mathrm {o p}}\right)
$$
This theorem recovers precisely the same behavior with respect to  $\gamma$  and  $\Xi$  as Theorem 3.1. As  $\gamma \rightarrow 1$ , we achieve the minimax rate of estimation for low-rank regression as derived in [27] as long as  $\Xi = 0$ .
4 Tuning parameter selection with Lepski's method
A key challenge of applying penalized regression models to summary statistics is that model tuning based on data splitting (e.g., training and validation) is no longer an option. Model selection methods based on information criteria require knowing the log squared loss  $\log \| \mathbf{Y}^{(q)} - \mathbf{X}^{(q)}\beta \| _2^2$  which cannot be recovered from  $\widehat{\mathbf{S}}^{(q)}$  and  $\widetilde{\Sigma}^{(q)}$  [1]. To address this, we propose to use a tuning scheme based on Lepski's method [14], a classical tool of nonparametric statistics for adaptive estimation with unknown tuning parameters. The authors of [5] apply the ideas of Lepski to the LASSO, providing a fast algorithm for model tuning with non-asymptotic guarantees. In this section, we extend the methods in [5] to tune the multi-task estimators described in the present work.
The results and ideas in this section apply to both  $\widehat{\mathbf{B}}^{(sp)}$  and  $\widehat{\mathbf{B}}^{(lr)}$ , so without loss of generality, let  $(\widehat{\mathbf{B}}_{\lambda},\mathcal{P})$  denote a generic estimator-regularizer pair with tuning parameter  $\lambda$ , which may refer to either  $(\widehat{\mathbf{B}}_{\lambda}^{(sp)},\|.\|_{2,1})$  or  $(\widehat{\mathbf{B}}_{\lambda}^{(lr)},\|.\|_{*})$ . Additionally, let  $\mathcal{P}^*$  denote the dual of  $\mathcal{P}$ , meaning that
$$
\mathcal{P}^{*}(X) = \sup_{Y:\mathcal{P}(Y)\leq 1}\langle X,Y\rangle
$$
Finally, we let  $\mathcal{L}$  denote the loss function for both estimators and let  $\nabla \mathcal{L}$  denote its gradient.
The intuition behind the adaptive tuning procedure is that the tuning parameter should be chosen large enough to control fluctuations in the gradient of the loss function, but not too large such that too much bias is incurred. [26] articulates that the performance of regression estimator with a convex penalty is contingent on the following event occurring with high probability:
$$
\mathcal {A} (\lambda) = \left\{\mathcal {P} ^ {*} (\nabla \mathcal {L} (\mathbf {B} ^ {*})) \leq \frac {\lambda}{2} \right\}
$$
where we use our problem's notation for continuity. The proofs of Theorem 3.1 and 3.2 involve showing that  $\mathcal{A}(\lambda)$  holds with high probability under our stated conditions. It is straightforward to prove the following proposition, which states that conditional on  $\mathcal{A}$ , the gradient of the loss function at our generic estimator  $\widehat{\mathbf{B}}$  is close to the gradient at the true parameter  $\mathbf{B}^*$ .
Proposition 4.1. Let  $(\widehat{\mathbf{B}}_{\lambda},\mathcal{P})$  denote a generic estimator-regularizer pair. Conditional on the event  $\mathcal{A}(\lambda)$ , there exists a constant  $C > 0$  such that the following inequality is satisfied almost surely:
$$
\mathcal {P} ^ {*} \left(\nabla \mathcal {L} \left(\widehat {\mathbf {B}} _ {\lambda}\right) - \nabla \mathcal {L} \left(\mathbf {B} ^ {*}\right)\right) \leq C \lambda
$$
This proposition motivates the following definition, which we adopt from [5].
Definition 4.1. Let  $\Lambda = \{\lambda_1, \lambda_2, \dots, \lambda_M\}$  denote a grid of potential tuning parameters ordered such that  $0 < \lambda_1 < \lambda_2 < \dots < \lambda_M < \infty$ . Fix  $\delta \in (0,1)$ . The oracle tuning parameter  $\lambda_{\delta}^{*}$  is defined as
$$
\lambda_{\delta}^{*} = \operatorname *{arg  min}_{\lambda \in \Lambda}\left\{\mathbb{P}\left\{\mathcal{A}(\lambda)\right\} \geq 1 - \delta \right\}
$$
The oracle tuning parameter provides the tightest upper bound in Proposition 4.1, but is unknowable in practice, since we do not observe  $\mathbf{B}^*$  and hence cannot verify  $\mathcal{A}$ . The aim of our Lepski-type method is to mimic the performance of  $\lambda_{\delta}^{*}$  in an entirely data-driven fashion. Letting  $\Lambda$  denote our ordered grid of potential tuning parameters, we follow [5] and choose the tuning parameter  $\hat{\lambda} \in \Lambda$  that satisfies
$$
\hat {\lambda} = \arg \min  _ {\lambda \in \Lambda} \left\{\max  _ {\lambda^ {\prime}, \lambda^ {\prime \prime} \in \Lambda , \lambda^ {\prime}, \lambda^ {\prime \prime} \geq \lambda} \mathcal {P} ^ {*} (\nabla \mathcal {L} (\widehat {\mathbf {B}} _ {\lambda^ {\prime}}) - \nabla \mathcal {L} (\widehat {\mathbf {B}} _ {\lambda^ {\prime \prime}})) \leq \bar {C} (\lambda^ {\prime} + \lambda^ {\prime \prime}) \right\} \tag {5}
$$
where  $\bar{C}$  is a constant chosen by the statistician. The following theorem states that  $\hat{\lambda}$  recovers the behavior of  $\lambda_{\delta}^{*}$  with high probability, as long as  $\bar{C}$  is sufficiently large.
Theorem 4.1. Let  $C$  denote the constant in Proposition 4.1. If  $\hat{\lambda}$  is chosen as in Equation 5 with  $\bar{C} \geq C$ , then the following inequalities hold simultaneously with probability at least  $1 - \delta$ :
1.  $\hat{\lambda} \leq \lambda_{\delta}^{*}$
2.  $\mathcal{P}^* (\nabla \mathcal{L}(\widehat{\mathbf{B}}_{\hat{\lambda}}) - \nabla \mathcal{L}(\mathbf{B}^*))\leq C^*\lambda_\delta^*$
where  $C^* \geq \bar{C}$ .
This theorem is a generalization of Theorem 3 in [5], adapted to our setting. The primary advantage of this Lepski-style tuning scheme is that it can be performed using only the gradient of the loss function, which in our setting consists only of summary-level statistics. This is a marked improvement over other summary statistic-based estimators, which typically require an additional set of individual-level data for tuning.
The adaptive tuning scheme does hold some disadvantages. First of all, it requires a choice of constant  $\bar{C}$ , which should be taken to be as close to the constant in Proposition 4.1 as possible. Remark 10 in [5] offers some guidance as to how to choose  $\bar{C}$  but unfortunately their analysis corresponds only to the LASSO. Deriving the exact constant in Proposition 4.1 may be possible under stronger assumptions on the data-generating process (i.e. Gaussianity), and we view this as an area of future work. Furthermore, Theorem 4.1 offers only a bound on  $\nabla \mathcal{L}(\widehat{\mathbf{B}}) - \nabla \mathcal{L}(\mathbf{B}^{*})$ . Translating this to a bound on  $\widehat{\mathbf{B}} - \mathbf{B}^{*}$  will require strong element-wise conditions on each of the matrices  $\widetilde{\Sigma}^{(q)}$ , which we do not explore in the present work. Nevertheless, our simulations in Section 5 indicate that the adaptive tuning method performs well in terms of the MSE of  $\widehat{\mathbf{B}}$ , suggesting that adaptive tuning is a good option for model selection when only summary statistics are available.
5 Numerical experiments and real data application
We validate our theory and demonstrate the effectiveness of multi-task learning in proxy data settings via extensive experiments. In each experiment, we take the proxy dataset to be well-specified;
in other words, we assume that  $\Xi = 0$ . When  $\Xi$  is nonzero, this predictably leads to worse performance, which we demonstrate in the supplement. The code, further implementation details, and additional simulations which explore the use of our adaptive tuning procedure are also available in the supplement.
First, we consider the effect of varying proxy data size on empirical MSE per task. We generate synthetic Gaussian data with  $n_{\mathrm{min}} = 100$ ,  $p = 100$ ,  $\tilde{n}_{\mathrm{min}} = \tau n_{\mathrm{min}}$  for  $\tau \in \{0.5, 1, 2, 5, 10\}$ , and  $\tilde{\rho}_q = 0$  for each  $q$ . The number of tasks was fixed at 8. Furthermore, we generate a row-sparse  $\mathbf{B}^*$  matrix with 10 nonzero rows and a  $\mathbf{B}^*$  with rank 2 for the sparse and low-rank multi-task estimators, respectively. We then fit the proxy data multi-task learning estimator and compare the prediction MSE per task to the estimator that has access to all of the individual level data and to the estimator that uses the true covariance matrix  $\Sigma$ . The results of this simulation are given in Figure 1.
Figure 1: Average prediction MSE per task after 100 repetitions plotted against  $\tau = \tilde{n} / n$ . The left hand side corresponds to the sparse estimator, and the right hand side is the low-rank estimator. The red boxes indicate the estimator that uses all of the individual-level data (IL_Cov), the blue boxes indicate the estimator that uses the true covariance matrix of the features (true_Cov), and the green boxes correspond to the estimator that uses just the proxy data (Proxy_Cov).
We observe a performance gap between the estimators that use the true covariance matrix and the individual level estimators, as predicted by our theory in Section 3. The performance of the proxy data estimators increase with increasing proxy sample size, but are unable to match the performance of the individual level estimator, as expected.
Next we study the effect of varying the proportion of overlapping samples between the discovery and proxy datasets. Similarly, we generate synthetic data with  $n = \tilde{n} = 100$ , and vary  $\tilde{\rho}$ , which indicates the proportion of proxy data points that are also in the discovery dataset. With  $Q = 8$ , we generate  $\mathbf{B}^*$  in the same way as in the previous simulation. These results are given in Figure 2.
Once again, we observe the expected performance gap between the estimators with the true covariance and the individual-level estimators. As the proportion of overlap between the proxy dataset and the discovery dataset grows, we see that the performance of the proxy data estimator converges to that of the individual level estimator. This is anticipated by Theorems 3.1 and 3.2.
Figure 2: Average MSE per task after 100 repetitions plotted against  $\tilde{\rho}$ . The orientation and colors are the same as in Figure 1.
Finally, we have applied our method to analyze real genetic data to demonstrate the real-world applicability of our method. We use a multi-site data obtained from the electronic Medical Records and Genomics (eMERGE) network [28], which includes individual-level genotype data from multiple research sites in the United States. Our goal is to predict levels of low-density lipoprotein (LDL) across five adult sites, treating the data from each site as a separate task. We split the data (with sample sizes  $n_1 = 3813$ ,  $n_2 = 546$ ,  $n_3 = 2666$ ,  $n_4 = 1435$ ,  $n_5 = 525$ ) at each task into a training and test set (with a test set data size of 100 for each task) and evaluate the performance of our method using the prediction MSE on the test set. The training data from each site is used to construct the discovery summary statistics  $\widehat{\mathbf{S}}^{(q)}$  for each task. For approximating  $\boldsymbol{\Sigma}^{(q)}$ , we choose two different approaches: one is to use the half of the genotype data from each site (this approach is labeled as Proxy_MTL1); the other approach is to use  $\mathbf{X}_1$  (genotype data from site 1) to approximate  $\boldsymbol{\Sigma}^{(q)}$  for all the sites. This approach is labeled as Proxy_MTL2. We use these two approaches to demonstrate a potential trade-off in the construction of the reference panel: Proxy_MTL1 uses a well-specified reference dataset with a smaller sample size; Proxy_MTL2 uses a larger reference dataset that may suffer from a distribution shift. For comparison, we also fit a multi-task learning estimator that uses all of the individual level training data for each task, which is labeled 'Individual_MTL', and we fit a ridge regression estimator that models each task separately, for comparison to our multi-task learning approach. The ridge estimator uses the proxy sample covariance instead of the individual level covariance matrix for a fair comparison with our method. We repeat the train-test split process 10 times, which admits the distribution of prediction MSE values that we report in the figure. We use the nuclear norm penalized multi-task learning estimators in this application, because we believe that the genetic effects in this dataset are dense.
Figure 3: Prediction MSE per task after 10 splits of the eMERGE data.
Our results in Figure demonstrate that our method is highly practical when only summary-level information is available, as the prediction MSE of our method is nearly the same as the estimator which uses the individual-level data, despite a slight cost in performance. Furthermore, all multi-task learning estimators outperform the ridge-estimator, confirming that multi-task learning is a strong approach when there is shared structure between tasks.
6 Discussion, Limitations, and Broader Impacts
We have described a flexible multi-task framework incorporating summary statistics from distinct sources with a general data-driven tuning scheme for selecting tuning parameters. Our theoretical analysis sheds light on the intrinsic price of using summary-level information from distinct sources for statistical analysis, and suggests that more overlap between the sources, less distributional shift, and larger proxy data sample sizes can alleviate this cost. Our data-driven tuning scheme allows
models to be trained without sample splitting, making it more applicable to real-world settings with only summary statistics available.
The limitations of our work are summarized as follows. First of all, our methods depend on a linear relationship between the covariates and the outcomes. This assumption is often satisfied in our target application of genetic risk prediction [3], but it does limit the applicability of our method to other domains. To extend our framework to non-linear models, we may use the second-order Taylor approximation of the loss function as in [16]. However, the summary statistics used by such an algorithm are not found in existing literature or publicly available databases. Additionally, our theoretical results provide only upper bounds on the estimation error of the two estimators that we consider in this work. To fully characterize the cost of using summary statistics for multi-task learning, lower bounds resembling Theorem 2.2 of [15] are needed. We conjecture that our estimators converge at a minimax optimal rate, and we view the proof of this conjecture as an important future direction. Finally, we may also extend the framework of [10] to our summary statistic based setting to adjust for potential differences between tasks.
Nevertheless, our results have important implications beyond high-dimensional statistical theory. The trade-off between proxy data sample size and discovery-proxy overlap may inform how polygenic risk models are built in real-world applications: Practitioners should prioritise alignment between the sources of summary statistics that they use to build these models, rather than optimizing for large sample sizes. This guidance may lead to more accurate polygenic scores, which have emerged as an important predictive tool in the field of precision medicine.
We recognize that the development of polygenic risk scores, if done without care, may worsen existing health disparities [23]. This is a potential negative societal impact of our work. We hope that our multi-task learning framework may be used to incorporate data from diverse populations to improve generalizability and transportability of genetic risk predictions to overcome these negative impacts.
Rui Duan is supported by National Institute of General Medical Sciences (NIGM) R01GM148494. Parker Knight is supported by an NSF Graduate Research Fellowship. We would like to thank Rajarshi Mukherjee for an insightful discussion of Lepski's method, and thank the reviewers for their comments and feedback which greatly improved the paper.