Surrogate-assisted optimal sampling for risk prediction under measurement constraints

Technical Summary: Surrogate-assisted optimal sampling for risk prediction under measurement constraints

Problem Context
In many risk prediction scenarios, such as electronic health record (EHR) analysis, covariates and a response surrogate (e.g., diagnostic codes) are available for a large target population, but the true response is costly to ascertain and observed only for a limited subset. This creates a design problem: how to allocate a fixed measurement budget to select which observations should have their true responses measured to build an optimal prediction model.

Existing subsampling methods often face limitations in this setting. Simple Random Sampling (SRS) fails to exploit heterogeneity in informativeness. Methods optimizing parameter estimation (e.g., Wang et al., 2018) require observed responses to design sampling probabilities, which are unavailable for the unlabeled group. Other approaches (e.g., Shen & Ning, 2025; Lee et al., 2023) focus on minimizing Mean Squared Error (MSE) of parameters or predictions, which may not align directly with binary risk prediction performance. Furthermore, many existing methods assume surrogates perfectly identify positive cases, a condition that may not hold in practice.

Methodology
The authors propose a surrogate-assisted optimal sampling framework designed specifically for binary risk prediction under measurement constraints.

1. Setting:

   • Let $y_i \in \{0, 1\}$ be the true response and $s_i$ be a surrogate indicator where $s_i=1 \implies y_i=1$.
   • Observations with $s_i=1$ are confirmed positive cases with known responses.
   • Observations with $s_i=0$ are unlabeled (containing both true positives and negatives). The goal is to select a subset of these $s_i=0$ observations for response measurement.
   • A measurement budget $C$ limits the total number of sampled observations from the $s_i=0$ group.

2. Estimation:

   • A prediction model $p(x, \beta)$ estimates $P(Y=1|x)$.
   • The model parameter $\beta$ is estimated by minimizing an Inverse Probability Weighted (IPW) Cross-Entropy Loss:
     $$L_N(\beta) = -\sum_{i=1}^N \frac{\delta_i}{\pi_T(x_i)} [y_i \log p(x_i, \beta) + (1-y_i) \log(1-p(x_i, \beta))]$$
     where $\delta_i$ is the sampling indicator and $\pi_T(x_i)$ is the total sampling probability.

3. Optimal Sampling Design:

   • The authors derive the leading term of the expected out-of-sample cross-entropy loss, which depends on the sampling probability $\pi(z)$ for the $s_i=0$ group.
   • The optimal sampling probability is derived by minimizing the empirical counterpart of this leading term subject to the budget constraint.
   • The solution takes the form:
     $$\hat{\pi}_i \propto \sqrt{a_i(\hat{\beta})}$$
     where $a_i(\hat{\beta})$ represents the contribution of the $i$-th observation to the asymptotic variance of the cross-entropy loss, and $\hat{\beta}$ is a preliminary estimator.
   • Two-Step Procedure: Since $\hat{\beta}$ is unknown initially, the authors propose a two-step algorithm:
     1. Obtain a preliminary estimator $\tilde{\beta}$ using a pilot sample or simple random sampling.
     2. Compute optimal sampling probabilities $\hat{\pi}$ based on $\tilde{\beta}$ and re-sample/measure responses to fit the final model.

4. Theoretical Properties:

   • Under regularity conditions (compact parameter space, smoothness, invertibility of the Hessian), the proposed estimator $\hat{\beta}_{\pi}$ is shown to be consistent and asymptotically normal.
   • Theorem 3 establishes that the proposed design minimizes the leading-order term of the expected out-of-sample cross-entropy loss among all admissible sampling strategies with the same budget.

Key Contributions

1. General Applicability: The method applies to a general class of smooth parametric risk prediction models, not restricted to logistic regression or specific Generalized Linear Models (GLMs).
2. Feasibility under Constraints: The design does not require true responses from the target unlabeled observations at the design stage, relying only on covariates, the surrogate, and a preliminary estimator.
3. Prediction-Oriented Criterion: Unlike existing methods that minimize parameter estimation MSE or prediction MSE, this framework directly optimizes the expected cross-entropy loss, aligning the sampling strategy with the goal of binary risk prediction.

Results
The paper validates the method through extensive simulations and two real-data applications:

• Simulations:

  • Correct Surrogates: When surrogates correctly identify positive cases ($P(Y=1|S=1)=1$), the proposed Cross-Entropy (CE) method achieves the lowest cross-entropy loss and competitive MSE, sensitivity, and specificity compared to MSE-based optimal sampling (Lee et al., 2023), OSCA, and SRS.
  • Misspecified Surrogates: When the assumption $P(Y=1|S=1)=1$ is violated (i.e., false positives exist in the surrogate group), the CE method demonstrates superior robustness. It maintains lower cross-entropy loss and better classification stability compared to competitors, whose performance degrades significantly under false-positive contamination.
  • Rare Outcomes: In a scenario with low prevalence ($\approx 5\%$) and no surrogate information, the CE method outperforms SRS and OSCA (which reduces to SRS in this case) and avoids the numerical degeneracy observed in the MSE-based method.

• Real Data Applications:

  • Depression Prediction (MIMIC-III): Using ICU admission records, the CE method consistently achieved lower cross-entropy loss and higher AUC than competing methods across various sampling budgets ($C=100$ to $400$).
  • Stroke Prediction (Kaggle): In a low-prevalence setting without surrogate indicators, the CE method provided the most stable predictions and highest AUC among subsampling methods, closely approaching the performance of the full-data benchmark.

Significance and Claims
The paper claims that the proposed framework offers a theoretically grounded and practically effective solution for risk prediction under resource constraints. By directly targeting the cross-entropy loss, the method ensures that the sampling design prioritizes observations most informative for predictive accuracy rather than parameter estimation. The authors emphasize the method's robustness to surrogate misspecification and its ability to function effectively even in rare outcome settings where traditional optimal sampling strategies may fail. The work bridges the gap between measurement-constrained experimental design and binary classification performance, providing a viable strategy for scenarios where labeling is expensive and surrogates are imperfect.