Task-Agnostic Machine-Learning-Assisted Inference

Thanks for your reply! W2: We would like to clarify that our framework does not require the consistent and asymptotically normal estimators for ML algorithms that are commonly expected in statistics. In fact, our approach can accommodate any "black box" ML algorithm. In our setting, the ML algorithm is used to impute (predict) labels in unlabeled data. These predictions are then used as input to a statistical method (algorithm), such as linear regression, to solve a statistical problem, such as estimating the effect of DNA on height. The requirement for consistent and asymptotically normal estimators applies only to the statistical algorithm, not to the ML algorithm used for label prediction. We have updated our paper to explicitly state that we did not impose any constraints on the ML algorithms used for label prediction.

**Q (continued):** * Even for M-estimation tasks that have already been implemented, PSPS offers the additional advantage of relying solely on summary statistics. The “task-specific derivations” mentioned throughout our paper were not only referring to statistical tasks, but also scientific tasks. Real-world data analysis in any scientific discipline often involves conventions and nuisances that require careful consideration. For example, our work is partly motivated from genome-wide association studies (GWAS). Statistically, GWAS is a linear regression that regresses an outcome (e.g., height) on many genetic variants. While the regression-based statistical foundation is simple, conducting a valid GWAS requires accounting for numerous technical issues, such as sample relatedness (i.e., study participants may be genetically related) and population structure (i.e., unrelated individuals of the same ancestry are both genetically and phenotypically similar, creating confounded associations in GWAS). Sophisticated algorithms and software have been developed to address these complex issues (`Mbatchou et al. 2021`). It will be very challenging if all these important features need to be reimplemented in an ML-assisted GWAS framework. With our PSPS protocol, researchers can utilize existing algorithms and software that are highly optimized for genetic applications to perform ML-assisted GWAS. This adaptability is not just limited to GWAS but is a major feature of our approach across scientific domains. A main result of this paper is that PSPS enables researchers to conduct ML-assisted inference using well-established data analysis pipelines. **L:** Following this suggestion, we have conducted experiments to compare the computational efficiency of our method (PSPS) against existing methods. Utilizing a dataset with 500 labeled and 10,000 unlabeled data points, PSPS required 1.62 seconds for linear regression and 8.27 seconds for logistic regression using 200 bootstrap resampling. The computation of one-step de-biasing using summary statistics alone took 0.032 seconds for linear regression and 0.033 seconds for logistic regression. Current methods, which estimate asymptotic variance via the closed form derived by the Central Limit Theorem instead of resampling, ranged from 0.024 to 0.049 seconds for linear regression and 0.032 to 0.077 seconds for logistic regression. Although PSPS is slower due to its resampling nature, the overall runtime remains relatively short. | Method | Linear regression | Logistic regression | | ------------- | ------------- |------------- | | PSPS | 1.62s | 8.27s | | PPI | 0.024s | 0.032s | | PPI++ | 0.031s | 0.077s | | POP-Inf | 0.049s |0.034s | We also note that we designed PSPS to utilize summary statistics, aiming to integrate seamlessly with existing computationally efficient software routinely used in data analysis. For example, Regenie is a software employed in GWAS that allows for fast computation of tens of millions of linear regressions using tens of thousands of samples (`Mbatchou et al. 2021`). Our protocol involves initially generating summary statistics using such a software tool, followed by their integration, allows high computational efficiency. **Final note**: We are excited that you find our work tackling an important problem and appreciate our theoretical study in statistical efficiency. If you have any further questions, please do not hesitate to let us know. If our responses have resolved your concerns, we kindly request you to consider increasing your score and championing our paper.

Thank you for your valuable comments! **W1:** We have included in the manuscript the formal regularity conditions required for consistent bootstrap variance estimation. These conditions are detailed in Theorem 3.10 (i) from `Shao, Jun, and Dongsheng Tu. The jackknife and bootstrap. Springer Science & Business Media, 2012.`, which proves the consistency of the bootstrap variance estimator. Below is the detailed theorem: We assume that $X _1, \cdots, X _n$ are i.i.d random p-dimensional vectors from distribution $F$. Let $T _n = T _n(X _1, \cdots, X _n)$ be a n estimator of an unknown parameter $\theta$, and $\Re _n = \sqrt{n}(T _n - \theta) \sim N(0, \sigma_n^2)$. Let $\{X _1^*, \ldots, X _n^*\}$ be a bootstrap sample from the empirical distribution $F _n$ based on $X _1, \cdots, X _n$, $T_n^*=T_n\left(X_1^*, \ldots, X_n^*\right)$ and $\Re_n^*=\sqrt{n}\left(T_n^*-T_n\right)$. Denote the bootstrap variance estimator $v _{\text {boot}} = \text{Var}(T_n^*)$. Theorem 3.10 (i) Let $T _n=T\left(F _n\right)$, assume that $T$ is $\rho _{\infty}$-Fréchet differentiable at $F$, and $ \max _{i_1, \ldots, i _n}\left|T _n\left(X _{i _1}, \ldots, X _{i _n}\right)-T _n\right| / \tau _n \rightarrow _{a . s .} 0 $ where the maximum is taken over all integers $i _1, \ldots, i _n$ satisfying $1 \leq$ $i _1 \leq \cdots \leq i _n \leq n$, and $\{ \tau _n \}$ is a sequence of positive numbers satisfying $\liminf _n \tau _n>0$ and $\tau _n=O\left(e^{n^q}\right)$ with a $q \in (0, \frac{1}{2})$., then $v _{\mathrm{boot}} / \sigma _n^2 \rightarrow _p 1$, where $\sigma _n^2=n^{-1} E[\phi _F(X _1)]^2>0$ We have also added a remark in our paper to refer the reader to the theoretical results in a more recent paper `Hahn, Jinyong, and Zhipeng Liao. "Bootstrap standard error estimates and inference." Econometrica 89.4 (2021): 1963-1977.` for bootstrap variance estimation. This paper (THEOREM 1) shows that common bootstrap based standard error in fact leads to a valid but (potentially conservative) inference. **W3:** Thank you for bringing this to our attention. Since we have decided to remove Section 4 from the paper and keep only the remarks in Section 3, we will no longer include the theoretical results related to this section in the paper. Instead, we will empirically verify the performance of ML-assisted FDR control (as we have done in Section 5 of the paper). As the reviewer pointed out, this does not affect the main contribution of our paper, which is the compelling feature of "task-agnostic" ML-assisted inference. **Task-agnostic feature of PSPS:** In Section 3.1, we have added a paragraph specifically highlighting why PSPS is considered "task-agnostic." This addition includes the three bullet points from our previous rebuttal that clearly delineate the relevant scenarios. **Runtime comparison:** We have added the runtime comparison into the Section 5 (Numerical experiments and real data application) and the relevant discussions in the Section 6 (Conclusion) in the paper. We hope our responses address your concerns, and please do not hesitate to let us know if you have any further questions.

**W6:** Examples of PRDS p-values include independent p-values and p-values from test statistics that are jointly normally distributed, if all correlations between test statistics are positive (`Wang, Ruodu, and Aaditya Ramdas. "False discovery rate control with e-values." Journal of the Royal Statistical Society Series B: Statistical Methodology 84.3 (2022): 822-852.`). **W7:** We appreciate these suggestions. We have thoroughly addressed these issues in the revised manuscript. **Q1:** The “privacy-preserving feature” of PSPS refers to the fact that we only require summary statistics as input for inference, rather than individual-level raw data (features $X$ and label $Y$). This terminology is commonly used in human genetics, healthcare, and multicenter electronic health record analysis to describe methods that enable statistical inference without directly accessing personal data. This approach is analogous to federated learning (`Kairouz, Peter, et al. 2021`) in the machine learning literature. For example, consider a scenario where labeled data is in one center and unlabeled data in another, yet researchers cannot access individual-level data from both centers simultaneously. Under such conditions, current ML-assisted inference, which relies on accessing both labeled and unlabeled data to minimize a joint loss function, is not feasible. However, PSPS circumvents this issue by aggregating summary statistics from multiple centers, thus performing statistical inference while upholding the privacy of individual-level data. We also acknowledge that in the machine learning literature, "privacy-preserving" often specifically refers to techniques like differential privacy. To avoid confusion, we have revised the terminology from “privacy-preserving” to "federated inference" and provided an example to illustrate this use case more clearly. **Q2:** First, we apologize for the error that we mistakenly used Figure 2e for Figure 2f due to a typo in the code for making the figure. The correct Figure 2f is attached in the rebuttal pdf. However, this does not change our main results related to statistical efficiency: PSPS is more efficient than PPI++ and other existing methods. As we have previously explained in our comparison of PSPS and PPI++, PSPS leverages information across multiple coordinates in a multi-dimensional parameter estimation task, resulting in asymptotic variances that are less than or equal to those produced by PPI++. Regarding the concern about the choice of tuning, the POP-Inf method, which tunes the element-wise variance, also shows wider confidence intervals than PSPS in linear regression. **Q3:** Our prediction pipeline comprises two components: prediction for unlabeled data and prediction for labeled data. To predict bone mineral density in unlabeled data, we first selected predictive features by 1) calculating the correlation of bone mineral density with 466 other variables (sample size > 200,000 from UK Biobank) using labeled data and 2) selecting the top 50 variables with the highest correlations as inputs for the SoftImpute algorithm to predict bone mineral density in the unlabeled data. For the labeled data, we employ a similar approach but incorporate 10-fold cross-validation to prevent overfitting. We select the predictive variables and train the SoftImpute model using 90% of the labeled data. We then perform predictions on the remaining 10% in each fold and repeat this process 10 times across all folds. We have included these details in the Appendix of the manuscript. **Final note**: We are excited that you find our method flexible, simple and elegant, and the real data application compelling. If you have any further questions, please do not hesitate to let us know. If our responses have resolved your concerns, we kindly request you to consider increasing your score and championing our paper.