Semi-supervised linear regression with missing covariates

1. Problem Statement

The paper addresses the challenge of estimating regression parameters ($\beta^*$) in a linear model where:

1. Missing Covariates: The labelled dataset contains missing values in the predictor variables ($X$). The missingness can be unstructured (sporadic, independent across variables) or structured (blockwise, where variables are missing in groups, e.g., due to different data modalities like MRI vs. PET scans).
2. Semi-Supervised Setting: The available data consists of a labelled dataset ($L$) with missing covariates and an unlabelled dataset ($U$) containing covariates but no response ($Y$).
3. Goal: To develop estimators that leverage the unlabelled data to improve estimation accuracy and to characterize the fundamental statistical limits (minimax rates) of this problem, comparing them to fully supervised settings.

The authors aim to answer:

• How can unlabelled data be utilized when labelled samples have missing covariates?
• How do missingness patterns (structured vs. unstructured) affect optimal performance?
• Can procedures be optimal for both settings?

2. Methodology

The paper proposes distinct estimators for low-dimensional and high-dimensional settings, handling both Ordinary Semi-Supervised (OSS) and Ideal Semi-Supervised (ISS, where the population covariance $\Sigma$ is known) scenarios.

A. Low-Dimensional Setting

• Estimator Construction: The authors propose a weighted least squares estimator based on imputation.
  • For a missingness pattern $k$ observing covariates $O_k$, they impute the missing variables $M_k$ using the conditional expectation: $\hat{X}_{M_k} = \hat{\Sigma}_{M_k O_k} \hat{\Sigma}_{O_k O_k}^{-1} X_{O_k}$.
  • The estimator $\hat{\beta}$ minimizes a weighted sum of squared errors across all patterns:
    $$\hat{\beta} = \arg \min_{\beta} \sum_{k=1}^K \sum_{i \in I_k} \hat{D}_k (Y_i - X_{O_k, i}^T \hat{P}_k \beta)^2$$
    where $\hat{P}_k$ is the imputation matrix derived from the estimated covariance $\hat{\Sigma}$, and $\hat{D}_k$ are reweighting factors.
• Reweighting: A key innovation is the reweighting step. Naive imputation (setting weights to 1) is suboptimal. The optimal weight $D_k^*$ is the ratio of the variance of a complete case to the variance of a case with missingness pattern $k$. The authors provide a procedure to estimate these weights from the data.
• Covariance Estimation:
  • In the OSS setting, $\hat{\Sigma}$ is estimated using the unlabelled data (and potentially labelled data via cross-fitting).
  • In the Supervised setting ($N=0$), they employ two-fold cross-fitting to estimate $\Sigma$ and the weights to avoid overfitting and ensure independence between the estimation of nuisance parameters and the final estimator.

B. High-Dimensional Sparse Setting

• Estimator Construction: The authors adapt the Dantzig Selector to handle missing data and unlabelled samples.
  • They construct unbiased estimators for $\gamma = \Sigma \beta^*$ and $\Sigma$ using the available partial observations and unlabelled data.
  • The estimator solves the linear program:
    $$\hat{\beta} \in \arg \min_{\beta} \{ \|\beta\|_1 : \|\hat{\Sigma}\beta - \hat{\gamma}\|_\infty \leq \lambda \}$$
• Regularization: The tuning parameter $\lambda$ is carefully chosen to account for the effective sample size of the missingness patterns and the presence of unlabelled data.

3. Key Contributions

1. Novel Estimators:

   • Introduced a weighted imputation estimator for low-dimensional settings that outperforms naive imputation and standard supervised methods.
   • Extended the Dantzig selector to the semi-supervised setting with blockwise and unstructured missingness.

2. Minimax Theory (Upper and Lower Bounds):

   • Non-Asymptotic Bounds: Provided rigorous upper bounds on the risk (MSE) for both estimators.
   • Matching Lower Bounds: Derived new minimax lower bounds that match the upper bounds up to constant factors, proving rate optimality.
   • Resolution of Conjectures: Resolved a conjecture by Wang et al. (2019) regarding the minimax rate in the high-dimensional supervised case with unstructured missingness.

3. Characterization of Unlabelled Data Utility:

   • Demonstrated that unlabelled data improves the convergence rate by effectively increasing the observation probability.
   • Low-Dimensional: The rate improves from $O(\frac{1}{\rho n_L})$ to $O(\frac{1}{\rho n_L} + \frac{1}{\rho^2 n_L + N})$, effectively boosting the observation rate from $\rho$ to $\rho^{1/2}$ when $N$ is large.
   • Dimensionality Reduction: In structured missingness (e.g., monotonic patterns), unlabelled data allows for a reduction in effective dimensionality. If $p_0$ variables are missing in a large portion of labelled data but observed in unlabelled data, the rate depends on $p_0$ rather than the full dimension $p$.

4. Safety: The proposed estimators are "safe," meaning they perform at least as well as the best supervised estimator in the worst case.

4. Key Results

Low-Dimensional Rates

The minimax risk scales as:
$$\text{Risk} \asymp \sigma^2 \sum_{i=1}^p \frac{1}{\alpha_i} + \|\beta^*\|_2^2 \zeta$$
Where $\alpha_i$ represents the effective sample size for variable $i$, and $\zeta$ depends on the ability to estimate the covariance matrix (scaling with $1/N$ in OSS).

• Unstructured: Rate is $O(\frac{p}{\rho n_L} + \frac{p}{\rho^2 n_L + N})$.
• Structured: Rate depends on modality sizes $h_g$ and cross-modality sizes $n_{g,h}$. Unlabelled data reduces the penalty for missing blocks.

High-Dimensional Rates

The minimax risk scales as:
$$\text{Risk} \asymp \frac{s \log p}{\rho n_L} + \frac{s \log p}{\rho^2 n_L + N}$$

• This confirms that the presence of unlabelled data ($N$) improves the rate from depending on $\rho n_L$ to $\rho^2 n_L + N$, effectively doubling the "information" regarding the covariance structure.
• For structured missingness, the rate depends on the specific modality sizes, showing that unlabelled data can mitigate the impact of missing blocks.

Empirical Performance

• Simulations: Extensive simulations on synthetic data show the proposed estimators significantly outperform:
  • Complete Case (CC) analysis.
  • Naive imputation (Mean, kNN, Random Forest).
  • Existing semi-supervised methods (e.g., Azriel et al., 2022) that ignore missingness.
  • Semiparametric efficient estimators (Robins, Rotnitzky, Zhao) which often fail to leverage unlabelled data effectively in this specific context.
• Real Data: Applied to the California Housing dataset with synthetic missingness. The OSS estimator reduced the Mean Squared Error (MSE) significantly compared to complete case analysis, particularly when the number of complete cases was small.

5. Significance

• Bridging a Gap: This work fills a critical theoretical gap by jointly addressing missing covariates and semi-supervised learning. Previous literature treated these as separate problems.
• Practical Relevance: The findings are highly relevant for modern data science, particularly in healthcare (EHR data with blockwise missingness due to different tests) and multi-modal learning, where obtaining labels is expensive but covariates are abundant.
• Theoretical Insight: The paper provides the first matching upper and lower bounds for these settings, clarifying exactly how much unlabelled data helps and under what conditions (e.g., the necessity of observing pairs of variables together).
• Methodological Advancement: The introduction of the reweighting step in low-dimensional settings and the adaptation of the Dantzig selector for semi-supervised missing data provide robust, rate-optimal tools for practitioners.

In summary, the paper establishes that unlabelled data is not just a bonus but a fundamental requirement for achieving optimal rates in linear regression with missing covariates, effectively transforming the problem's difficulty by improving the effective sample size and reducing the impact of missingness patterns.