Adaptive Optimization for Prediction with Missing Data

TL;DR

This work reframes prediction with missing data as a two-stage adaptive optimization problem and introduces adaptive linear regression models whose coefficients depend on the observed feature pattern. It establishes a hierarchy of adaptive models (static, affine, polynomial, finite, and fully adaptive) and proves finite-sample generalization bounds, while providing practical implementation details. The paper further connects adaptive linear models to joint impute-then-regress strategies and extends to non-linear predictors via an alternating-minimization heuristic, enabling joint learning of imputation rules and downstream models. Through extensive synthetic and real-data experiments, the authors demonstrate that their approaches perform comparably to or better than traditional impute-then-regress pipelines, particularly when data are NMAR or adversarially missing, highlighting robustness to the underlying missingness mechanism. Collectively, the methods offer a principled, adaptable framework for prediction with missing data that can generalize to non-linear models and improve out-of-sample accuracy in challenging missingness regimes.

Abstract

When training predictive models on data with missing entries, the most widely used and versatile approach is a pipeline technique where we first impute missing entries and then compute predictions. In this paper, we view prediction with missing data as a two-stage adaptive optimization problem and propose a new class of models, adaptive linear regression models, where the regression coefficients adapt to the set of observed features. We show that some adaptive linear regression models are equivalent to learning an imputation rule and a downstream linear regression model simultaneously instead of sequentially. We leverage this joint-impute-then-regress interpretation to generalize our framework to non-linear models. In settings where data is strongly not missing at random, our methods achieve a 2-10% improvement in out-of-sample accuracy.

Figures (15)

• Figure 1: Graphical representation of the 4 experimental designs implemented in our benchmark simulations with real-world design matrix $\bm{X}$. Solid (resp. dashed) lines correspond to correlations explicitly (resp. not explicitly) controlled in our experiments.
• Figure 2: Out-of-sample $R^2$ of mean-impute-then-regress, joint impute-then-regress, and optimal $\bm{\mu}$-impute-then-regress (synthetic data, linear signal).
• Figure 3: Out-of-sample $R^2$ of mean-impute-then-regress, joint impute-then-regress, and optimal $\bm{\mu}$-impute-then-regress (synthetic data, linear signal), as the number of observations $n$ increases.
• Figure 4: Performance of adaptive linear regression (a) and joint impute-then-regress models (b) on real-world design matrix $\bm{X}$.
• Figure 5: Comparison of adaptive linear regression and joint impute-then-regress methods vs. mean-impute-then-regress on real-world design matrix $\bm{X}$. We also compare to the "Oracle" which has access to the fully observed data and the "Complete Feature" regression which only uses features which are never missing.

Theorems & Definitions (6)

• Theorem 2.1
• Lemma 1.1
• proof : Proof of Lemma \ref{['lemma:taylor']}
• Lemma 1.2
• proof : Proof of Lemma \ref{['lemma:taylor.exp']}
• proof : Proof of Theorem \ref{['prop:finite']}