Knockout: A simple way to handle missing inputs

## Abstract

Abstract Deep learning models can extract predictive and actionable information from complex inputs.
The richer the inputs, the better these models usually perform.
However, models that leverage rich inputs (e.g., multi-modality) can be difficult to deploy widely, because some inputs may be missing at inference.
Current popular solutions to this problem include marginalization, imputation, and training multiple models.
Marginalization can obtain calibrated predictions but it is computationally costly and therefore only feasible for low dimensional inputs.
Imputation may result in inaccurate predictions because it employs point estimates for missing variables and does not work well for high dimensional inputs (e.g., images).
Training multiple models whereby each model takes different subsets of inputs can work well but requires knowing missing input patterns in advance.
Furthermore, training and retaining multiple models can be costly.
We propose an efficient way to learn both the conditional distribution using full inputs and the marginal distributions.
Our method, Knockout, randomly replaces input features with appropriate placeholder values during training.
We provide a theoretical justification of Knockout and show that it can be viewed as an implicit marginalization strategy.
We evaluate Knockout in a wide range of simulations and real-world datasets and show that it can offer strong empirical performance.

## 1 Introduction

In many real-world applications of machine learning and statistics, not all variables might be available for every data point.
This issue, called missingness, is well-studied in the literature (Little and Rubin, ) and commonly arises in applications including healthcare, social sciences, and environmental studies.
From a Bayesian perspective, missingness can be viewed as a marginalization problem, where we would like a model to marginalize out the missing variables from the conditioning set.
However, during training, we often do not know which features will be missing at inference time.
In lieu of training multiple models for every missingness pattern, a common strategy is imputation, which uses a point estimate (usually the mean or mode) to impute the missing feature.
This can be seen as approximating the marginalization with a delta function.

In this work, we propose an augmentation strategy, called Knockout, for handling missing inputs, in a simple but theoretically justified way.
During training, features are augmented by randomly “knocking out” and replacing them with constant “placeholder” values.
At inference time, using the placeholder value corresponds mathematically to estimation with
the appropriate marginal distribution.
In particular, we demonstrate how Knockout can be seen as implicitly maximizing the likelihood of a weighted sum of the conditional estimators and all desired marginals in a single model.

In a suite of experiments, we demonstrate the broad applicability of Knockout.
We experiment on both synthetic and real-world data using tabular and image-based inputs.
Real world experiments include Alzheimer’s forecasting, noisy label learning, and multi-modal MR image segmentation and detection.
We demonstrate the effectiveness of Knockout in handling low and high-dimensional missing inputs, and compare it with appropriate baselines, including imputation and ensembling-based methods.

## 2 Method

### 2.1 Background

The goal of supervised machine learning (ML) is to learn the conditional distribution $p(Y|\bm{X})$ where $Y$ is the output (predictive target) and $\bm{X}\in\mathbb{R}^{N}$ are the vector of inputs or features.
The prediction for a new sample $\bm{x}$ is $\hat{y}=\operatorname*{arg\,max}_{Y}p(Y|\bm{X}=\bm{x})$.
However, in many practical applications, not all features may be present for a given input.
Consider the case when $X_{i}$ is missing, and denote the vector of non-missing features as $\bm{X}_{-i}$.
In general, multiple features may be missing at a time, and we can represent this with a missingness indicator set $\mathcal{M}$ and corresponding non-missing features as $\bm{X}_{-\mathcal{M}}$.
In this case, what we really want is $p(Y|\bm{X}_{-\mathcal{M}})$.

How can we account for missingness?
A simple approach is to train a separate model for $p(Y|\bm{X}_{-\mathcal{M}})$, i.e. a model that takes only the non-missing features $\bm{X}_{-\mathcal{M}}$ as inputs.
However, this is expensive because a separate model is needed for each missingness pattern.
Furthermore, there is no sharing of information between these separate models, even though they are theoretically related.

Another approach relies on rewriting $p(Y|\bm{X}_{-\mathcal{M}})$ in terms of the already available $p(Y|\bm{X})$:

$$ $p(Y|\bm{X}_{-\mathcal{M}})=\int_{\bm{X}_{\mathcal{M}}}p(Y,\bm{X}_{\mathcal{M}} |\bm{X}_{-\mathcal{M}})d\bm{X}_{\mathcal{M}}=\int_{\bm{X}_{\mathcal{M}}}p(Y| \bm{X})p(\bm{X}_{\mathcal{M}}|\bm{X}_{-\mathcal{M}})d\bm{X}_{\mathcal{M}}.$ (1) $$

The goal now is to obtain $p(\bm{X}_{\mathcal{M}}|\bm{X}_{-\mathcal{M}})$ and perform the integration over all possible $\bm{X}_{\mathcal{M}}$.

Imputation methods approximate Eq. (1) by replacing $p(\bm{X}_{\mathcal{M}}|\bm{X}_{-\mathcal{M}})$ with a delta function.
For example, “mean imputation” uses the mean of the missing features $\bm{X}_{\mathcal{M}}$, $\mathbb{E}[\bm{X}_{\mathcal{M}}]$, for $\bm{X}_{\mathcal{M}}$ itself.
In Eq. (1), this corresponds to approximating $p(\bm{X}_{\mathcal{M}}|\bm{X}_{-\mathcal{M}})\approx\delta(\mathbb{E}[\bm{X}_{
\mathcal{M}}])$, a delta function.
While convenient and commonly used, mean imputation ignores the dependency between $\bm{X}_{\mathcal{M}}$ and $\bm{X}_{-\mathcal{M}}$, and does not account for any uncertainty.

More sophisticated approaches to imputation capture the interdependencies between inputs (Troyanskaya et al., ; Stekhoven and Bühlmann, ), for example by explicitly modeling $p(\bm{X}_{\mathcal{M}}|\bm{X}_{-\mathcal{M}})$ by training a separate model.
At inference time, the point estimate $\bm{x}_{\mathcal{M}}=\operatorname*{arg\,max}_{\bm{X}_{\mathcal{M}}}p(\bm{X}_{
\mathcal{M}}|\bm{X}_{-\mathcal{M}})$ can be used for the missing $\bm{X}_{\mathcal{M}}$.
While properly accounting for interdependencies between inputs, this approach requires fitting a separate model for $p(\bm{X}_{\mathcal{M}}|\bm{X}_{-\mathcal{M}})$.
In multiple imputation, multiple samples from $p(\bm{X}_{\mathcal{M}}|\bm{X}_{-\mathcal{M}})$ are drawn and a Monte Carlo approximation is used to estimate the integral on the RHS of Eq. (1) .
Although this is more accurate than single imputation, it is not effective in high dimensional space.

### 2.2 Knockout

We propose a simple augmentation strategy for neural network training called Knockout that enables estimation of the conditional distribution $p(Y|\bm{X})$ and all desired marginals $p(Y|\bm{X}_{-\mathcal{M}})$ in a single, high capacity, nonlinear model, such as a deep neural network.
During training, features are augmented by randomly “knocking out” and replacing them with constant, “placeholder” values.
At inference time, using the placeholder value corresponds mathematically to estimation with the suitable marginal distribution.

Specifically, let $\bm{M}=[M_{1},M_{2},\dots,M_{N}]\in\{0,1\}^{N}$ denote a binary, induced missingness indicator vector.
Let $\bm{\bar{x}}:=[\bar{x}_{1},\bar{x}_{2},\dots,\bar{x}_{N}]\in\mathbb{R}^{N}$ denote a vector of placeholder values.
Then, define $\bm{X}^{\prime}(\bm{M},\bm{X})=\bm{M}\odot\bm{\bar{x}}+(\bm{1}-\bm{M})\odot\bm
{X}$ as augmented Knockout inputs,
where $\bm{1}$ is a vector of ones and $\odot$ denotes element-wise multiplication.
At every training iteration, a different Knockout input is used corresponding to a different randomly sampled $\bm{M}$. The model weights are trained to minimize the loss function with respect to $Y$, as is done regularly.

Two mild conditions are required to ensure proper training: 1) the placeholder values must be “appropriate,” as we will elaborate below. For our theoretical treatment, we will use out-of range values as appropriate; i.e. $\bm{\bar{x}}_{\mathcal{M}}\not\in\text{Support}(\bm{X}_{\mathcal{M}})$; 2) $\bm{M}$ must be independent of $\bm{X}$ and $Y$, i.e. $\bm{M}\mbox{${}\perp\mkern-11.0mu\perp{}$}\bm{X},Y$.(^1^11Note it is not necessary that $M_{i}\mbox{${}\perp\mkern-11.0mu\perp{}$}M_{j}$ for any $i,j$.)
It follows straightforwardly that these two conditions lead to modeling the desired conditional and marginal distributions simultaneously.
First, since $\bm{\bar{x}}_{\mathcal{M}}$ is not in the support of $\bm{X}_{\mathcal{M}}$,

$$ $\bm{X}_{\mathcal{M}}^{\prime}=\bm{\bar{x}}_{\mathcal{M}}\iff\bm{M}_{\mathcal{M }}=\bm{1},\ \ \ \ \bm{X}_{\mathcal{M}}^{\prime}\neq\bm{\bar{x}}_{\mathcal{M}} \iff\bm{M}_{\mathcal{M}}=\bm{0}\ \text{and}\ \bm{X}_{\mathcal{M}}^{\prime}=\bm {X}_{\mathcal{M}},$ (2) $$

where $\bm{0}$ and $\bm{1}$ are vectors of zeros and ones of appropriate shape.
Second, since $\bm{M}$ is independent of $\bm{X}$ and $Y$, it follows that imputing with the default value $\bm{\bar{x}}_{\mathcal{M}}$ is equivalent to marginalization of the missing variables defined by $\mathcal{M}$:

$$ $\displaystyle p(Y|\bm{X}_{\mathcal{M}}^{\prime}{=}\bm{\bar{x}}_{\mathcal{M}}, \bm{X}_{-\mathcal{M}}^{\prime}{=}\bm{x}_{-\mathcal{M}})=p(Y|\bm{M}_{\mathcal{M }}{=}\bm{0},\bm{M}_{-\mathcal{M}}{=}\bm{1},\bm{X}_{-\mathcal{M}}{=}\bm{x}_{- \mathcal{M}})=p(Y|\bm{X}_{-\mathcal{M}}{=}\bm{x}_{-\mathcal{M}}).$ (3) $$

In particular, at the two extremes, no Knockout ($\bm{M}=\bm{0}$) corresponds to the original conditional distribution, and full Knockout ($\bm{M}=\bm{1}$) corresponds to the full marginal:

$$ $\displaystyle p(Y|\bm{X}^{\prime}{=}\bm{x})=p(Y|\bm{M}{=}\bm{0},\bm{X}{=}\bm{x })=p(Y|\bm{X}{=}\bm{x}),\textrm{when }\bm{x}_{j}\neq\bar{x}_{j},\forall j$ (4) $\displaystyle p(Y|\bm{X}^{\prime}{=}\bm{\bar{x}})=p(Y|\bm{M}{=}\bm{1})=p(Y)$ (5) $$

For a new test input $\bm{x}$, the prediction when $\bm{x}_{\mathcal{M}}$ is missing is simply

$$ $\operatorname*{arg\,max}_{Y}p(Y|\bm{X}_{-\mathcal{M}}{=}\bm{x}_{-\mathcal{M}}) =\operatorname*{arg\,max}_{Y}p(Y|\bm{X}_{\mathcal{M}}^{\prime}{=}\bm{\bar{x}}_ {\mathcal{M}},\bm{X}_{-\mathcal{M}}^{\prime}{=}\bm{x}_{-\mathcal{M}}),$ (6) $$

i.e., the learned estimator with the augmented Knockout input.

#### 2.2.1 Knockout as an Implicit Multi-task Objective

The missingness indicator $\bm{M}$ determines how inputs are replaced with appropriate placeholder values during training.
To satisfy the independence condition of $\bm{M}$ with $\bm{X}$ and $\bm{Y}$, the variables $\bm{M}$ are sampled independently from a distribution $p(\bm{M})$ during training.
We show that this training strategy can be viewed as a multi-task objective (Caruana, ) decomposed as a weighted sum of terms, where each term is a separate marginal weighted by the distribution of $\bm{M}$.
Let $\ell$ denote the loss function to be minimized (e.g. mean-squared-error or cross-entropy loss):

$$ $\displaystyle L(\theta)$ $\displaystyle=\mathbb{E}_{\bm{X}^{\prime},Y}\ \ell(Y;f_{\theta}(\bm{X}^{\prime }(\bm{M},\bm{X}))=\mathbb{E}_{\bm{X},Y}\mathbb{E}_{\bm{M}}\sum_{\bm{m}\in\bm{M }}\mathbb{I}(\bm{M}{=}\bm{m})\ \ell(Y;f_{\theta}(\bm{X}^{\prime}(\bm{M},\bm{X} )))$ (7) $\displaystyle=\mathbb{E}_{\bm{X},Y}\sum_{\bm{m}\in\bm{M}}p(\bm{M}{=}\bm{m})\ \ell(Y;f_{\theta}(\bm{X}^{\prime}(\bm{m},\bm{X})))$ (8) $\displaystyle=\sum_{\bm{m}}p(\bm{M}{=}\bm{m})\ \mathbb{E}_{\bm{X},Y}\ell(Y;f_{ \theta}(\bm{X}^{\prime}(\bm{m},\bm{X}))),$ (9) $$

where $\mathbb{I}$ is the indicator function.

If there is knowledge about the missingness patterns at inference (e.g., some $X_{i}$ and $X_{j}$ exhibit correlated missingness), one can design $p(\bm{M})$ appropriately to cover all the expected missing patterns, i.e. by sampling $\bm{m}$ during training with different weights.
In the absence of such knowledge, the most general distribution for $\bm{M}$ is i.i.d. Bernoulli.
A common way correlated missingness arises in real-world applications is in structured inputs like latent features or images, where the entire feature vector or whole image is missing.
In our experiments, we demonstrate the superiority of structured Knockout, over naive i.i.d. Knockout, when such correlated missingness is known a priori.

### 2.3 Choosing Appropriate Placeholder Values

Our theoretical treatment assumed that the placeholder value $\bar{x}_{i}$ is not in the support of $X_{i}$ (see Appendix A.2 for further analysis).
In the following sections, we relax this assumption and make some recommendations for appropriate placeholder values, based on practical considerations.

#### 2.3.1 Non-structured

In this section, we recommend appropriate placeholder values for non-structured, scalar-valued inputs.

Categorical. If $X_{i}$ is a categorical variable with $N_{X_{i}}$ integer-valued classes from 1 to $N_{X_{i}}$, then $\bar{x}_{i}$ can be $N_{X_{i}}+1$.
If one-hot encoded, $\bar{x}_{i}$ can be a vector of 0s.

Continuous and Non-empty Infeasible Set. If $X_{i}$ is a continuous value within a bounded range, then we can scale the range to $[0,1]$ and choose $\bar{x}_{i}=-1$.
More generally, if $X_{i}$ has unbounded range but a non-empty infeasible set, then $\bar{x}_{i}$ can be set to a value in the infeasible set.
For example, if $X_{i}$ only takes positive values, then we can set $\bar{x}_{i}=-1$.

Continuous and Empty Infeasible Set. When $X_{i}$ has unbounded range and an empty infeasible set, then we suggest applying $Z$-score normalization and choosing $\bar{x}_{i}$ such that it lies in a low probability region of the normalized $X_{i}$, $p(X_{i}{=}\bar{x}_{i})\approx 0$.
As we argue in Appendix A.1, this approach leads to an approximation of the desired marginal.

Although, different distributions have different regions of low probability, we use the behavior of the Gaussian distribution as a guide to choose $\bar{x}_{i}$.
Fig. S1 shows the histogram of the norm of points sampled from standard Gaussian distributions with different dimensionality.
For a univariate standard Gaussian, most of the points lie close to the origin so we should choose $\bar{x}_{i}$ far away from the origin.
However, as the dimension increases, most of the points lie on the hyper-sphere away from the origin so we should choose $\bar{x}_{i}$ to be the point at origin (i.e. a vector of zeros).

Table 1 summarizes the choices of $\bar{x}_{i}$ for different types of random variables $X_{i}$.

**Table 1: List of different types of $X_{i}$ and the recommended $\bar{x}_{i}$**
| Type of $X_{i}$ | Example | Dimension | Support | Normalized? | $\bar{x}_{i}$ |
| --- | --- | --- | --- | --- | --- |
| Categorical | Gender | 1 | $\{1,\dots,N_{X_{i}}\}$ | N/A | $N_{X_{i}}+1$ |
| Continuous | Test scores | 1 | $[a,b]$ | Scale to [0, 1] | -1 |
| Continuous | Temperatures | 1 | $[a,\infty)$ or $(-\infty,b]$ | Scale to $[0,\infty)$ | -1 |
| Continuous | White noise | 1 | $(-\infty,\infty)$ | Z-score | $\pm 10$ |
| Structured | Images | >1000 | $[a,b]$ | Scale to [0, 1] | $\bm{0}$ |
| Structured | Latent vectors | >16 | $(-\infty,\infty)$ | Z-score | $\bm{0}$ |

#### 2.3.2 Structured

For structured inputs like images and feature vectors, we have found that Knockout applied with an out-of-range placeholder like $\bm{-1}$, though theoretically sound, can cause issues like unstable gradients.
Therefore, we recommend an appropriate placeholder to be either the image of all 0s or the mean image.
When $Z$-score normalization is applied, the 0 image and the mean image coincide.
Theoretically, it is well known that the mean of a high dimensional random variable, such as a Gaussian, has very low probability (Vershynin, ) (also see Fig. S1).
We believe this recommendation balances the tension between ensuring an extremely-low probability placeholder with proper convergence and performance.
For an empirical demonstration, see Section A.3.

### 2.4 Observed Missingness during Training

The treatment above assumes complete training data, and inference-time missingness only.
We now consider the situation where training data has observed missing inputs.
Let $\bm{N}$ be the binary mask indicating the observed data missingness.
$\bm{N}$ is different from $\bm{M}$, which denotes the missingness induced by Knockout during training.
Thus, $\bm{N}$ is fixed for a data sample, while $\bm{M}$ is stochastic. Observed missingness generally falls under the following scenarios (Little and Rubin, ).

Missing Completely at Random (MCAR):
This implies that $\bm{N}\mbox{${}\perp\mkern-11.0mu\perp{}$}\bm{X},Y$.
Let $\bm{M}^{\prime}:=\bm{N}\lor\bm{M}$ be the augmented masking indicator, where $\lor$ denotes the logical OR operation.
Since $\bm{N}\mbox{${}\perp\mkern-11.0mu\perp{}$}\bm{X},Y$ and $\bm{M}\mbox{${}\perp\mkern-11.0mu\perp{}$}\bm{X},Y$, so $\bm{M}^{\prime}\mbox{${}\perp\mkern-11.0mu\perp{}$}\bm{X},Y$.
Therefore, we can obtain the same result in Section 2.2 when using $\bm{M}^{\prime}$ instead of $\bm{M}$ as the masking indicator vector.
This implies that Knockout can be applied to MCAR training data simply by masking all the missing values using the same placeholders $\bm{\bar{x}}$.

Missing at Random (MAR) and Missing not at Random (MNAR):
This implies that $\bm{N}\mbox{${}\not\!\perp\mkern-11.0mu\perp{}$}\bm{X},Y$.
Thus, we cannot replace the missing values in training data using the same placeholders.
However, we can substitute these values using placeholders that are different from $\bm{\bar{x}}$ but are also outside the support of the input variables (or very unlikely values).
Let the placeholders for the data missingness be $\bm{\dot{x}}\neq\bm{\bar{x}}$.
During training, Knockout still randomly masks out input variables, including those that are not observed in the data.
Thus, the results in Section 2.2 still hold since $\bm{M}\mbox{${}\perp\mkern-11.0mu\perp{}$}\bm{X},Y$.

During inference, if we know a priori that $x_{i}$ of a sample is missing not at random, then we can use $\dot{x}_{i}$ as the placeholder.
Otherwise, if we know $x_{i}$ is missing at completely random, we use $\bar{x}_{i}$.

## 3 Related Work

Knockout is similar to and inspired by other methods with unrelated aims.
Dropout (Srivastava et al., ; Gal and Ghahramani, ) prevents overfitting by randomly dropping units (hidden and visible) during training and can be viewed as marginalizing over model parameters.
During inference, marginalizing over parameters can be approximated by predicting once without dropout (Srivastava et al., ) or averaging multiple predictions with dropout (Gal and Ghahramani, ).
Blankout (Maaten et al., ) and mDAE (Chen et al., ) learn to marginalize out the effects of corruption over inputs.
In contrast, Knockout learns different marginals to handle different missing input patterns.

Imputation techniques impute missing inputs explicitly, for example by imputing with the mean, median, or mode.
In model-based imputation, a separate model or technique first predicts the missing inputs to impute. These models include k-nearest neighbors (Troyanskaya et al., ), chained equations (Van Buuren and Groothuis-Oudshoorn, ), random forests (Stekhoven and Bühlmann, ), autoencoders (Gondara and Wang, ; Ivanov et al., ; Mattei and Frellsen, ; Lall and Robinson, ), GANs (Yoon et al., ; Li et al., ; Belghazi et al., ), or normalizing flows (Li et al., ).
Although more accurate than simple mean/median imputation, model-based imputation incurs significant additional computation cost, especially when missing inputs are high-dimensional.
In contrast, Knockout makes predictions without having to impute missing inputs explicitly.

Another relevant line of work is causal discovery (Spirtes et al., ), which often involves fitting a model using different subsets of available inputs and multiple distributions simultaneously (Lippe et al., ; James et al., ).
To reduce computational cost, it is common to train a single model that can handle different subset of inputs using dropout (Ke et al., ; Brouillard et al., ; Lippe et al., ).

Techniques like Knockout are often used in practice to train a single neural network that models multiple distributions, but are often justified empirically with little care taken in choosing placeholder values.
Many works use zeros without theoretical justification (Belghazi et al., ; Ke et al., ; Brouillard et al., ; Lippe et al., ).
GAIN (Yoon et al., ) and MisGAN (Li et al., ) impute using out-of-support values similar to Knockout.
However, both are limited in their treatment by assuming that the supports are bounded, and furthermore do not consider categorical variables.
Many self-supervised learning techniques can be interpreted as training to reconstruct the inputs with Knockout.
For example, masked language modeling (Devlin et al., ) randomly maps tokens to an unseen “masked” token.
Masked autoencoders (He et al., ) randomly replace image patches with black patches, which are arguably out of the support of natural images.

## 4 Experiments

In all experiments, unless stated otherwise, we compare Knockout against a common baseline model trained on complete data, which, at inference time, imputes missing variables with mean (if continuous) or mode (if discrete) values. If the training is done on incomplete data with observed missing variables, imputed with mean/mode, we denote this as common baseline*.
For most results we report a variant of Knockout but with sub-optimal placeholders (i.e. mean/mode for continuous/categorical features).
We denote this variant as Knockout*.
In all Knockout implementations, we choose random knockout rates such that, in expectation, half of the mini-batches have no induced missing variable. In batches with induced missingness, variables (or groups of variables in structured Knockout) are independently removed, with a probability equal to the knockout rate.

### 4.1 Simulations

We perform simulations on both regression and classification, where the output $Y$ needs to be predicted from some input $\bm{X}$.
In the regression experiments, we additionally compare against missForest (Stekhoven and Bühlmann, ), a competitive baseline for inference-time imputation.
All methods use the same neural network architecture composed of a 3-layer multi-layer perceptron (MLP) with hidden layers 100 and ReLU activations.
Training is done using Adam (Kingma and Ba, ) with learning rate 3e-3 for 5k steps.
In each run, we sample 30k data points in total and use 10% for training.
We generate training data corresponding to complete training data and missing completely at random data (MCAR).
We restrict our focus in this section to regression results.
For further experimental details and classification experiments and results, see Appendix B.1.

Figure: Figure 1: Test MSE evaluated against Bayes optimal prediction ($\mathbb{E}[Y|\bm{X}]$) from 10 repetitions of the regression simulation. Lower is better. X axis indicates the number of missing variables at inference time. a) Complete training data. b) Missing completely at random (MCAR) data.
Refer to caption: x1.png

We experiment on varying the number of missing features of $\bm{X}\in\mathbb{R}^{9}$ from 0 to 3.
We evaluate the models’ predictions against the MMSE-minimizing Bayes optimal predictions: $\mathbb{E}[Y|\bm{X}]$.
Fig. 1 shows the results of 10 repetitions of this simulation.
Both variants of Knockout outperform baselines.
In particular, Knockout outperforms Knockout* in general; this underscores the importance of choosing an appropriate placeholder value.

### 4.2 Missing Clinical Variables in Alzheimer’s Disease Forecasting

We demonstrate Knockout’s ability to manage observed missingness in a real-world clinical task: predicting the risk of progression from mild cognitive impairment (MCI) to Alzheimer’s Disease (AD) over the next five years, using data from the Alzheimer’s Disease Neuroimaging Initiative (ADNI) database (Mueller et al., ).
Input features $\bm{X}$ include subject demographic variables, genetics, cognitive assessment scores, cerebrospinal fluid (CSF) measurements, and measurements derived from magnetic resonance imaging (MRI) and positron emission tomography (PET) images.
The target $Y$ is a binary vector and indicates AD diagnosis in each of the five follow-up years.
We employ the state-of-the-art model of Karaman et al. (). For Knockout, we use an out-of-range value of 10 for induced missingness and -10 for observed missingness, both during training and testing.
Further details about the dataset and experimental setup are provided in Appendix B.2.

Figure 2 presents the average AUROC (area under the receiver operating characteristic curve) scores obtained when each input feature is missing during inference.
We perform 10 random 80-20 train-test splits and calculate a Composite AUROC by averaging the AUROC scores from the five follow-up years in each split.
We observe that Knockout outperforms the common baseline* in vast majority of cases, suggesting that knocking out input during training enhances the model’s ability to handle missingness at test time.
Furthermore, Knockout is largely better than Knockout*, which underscores the importance of choosing an appopriate placeholder.
We note that we present a similar analysis using the complete portion of the training dataset (i.e., with no observed missingness in training data) in Figure S6 of Appendix B.2, further demonstrating Knockout’s effectiveness.

Figure: Figure 2: Composite AUROC scores obtained for the three model variants when each input feature is missing during inference (x-axis). Displayed are averages of 10 train-test splits. Error bars indicate the standard error across these splits.
Refer to caption: x2.png

### 4.3 Privileged Information for Noisy Label Learning

In this experiment, we show that Knockout can be used for learning with privileged information (PI) that is available in training but absent during testing. Specifically, we evaluate this in a noisy label learning task, where the objective is to use PI, such as annotator ID or annotation time, to enhance model robustness against label noise. Due to the absence of PI in testing, existing methods (Ortiz-Jimenez et al., ; Wang et al., ) require an auxiliary classification head for PI utilization. We demonstrate that Knockout can be directly applied with a method that accepts PI as input and achieve competitive performance. We follow previous experiment setups (Wang et al., ) and evaluated model performance on CIFAR-10H (Peterson et al., ) and CIFAR-10/100N (Wei et al., ). These datasets involve relabeled versions of the original CIFAR. For more details, see Appendix B.6.

As a no-PI baseline, we train a Wide-ResNet-10-28 (Zagoruyko and Komodakis, ) model that ignores PI.
We also compare against recent noisy label learning methods: HET (Collier et al., ) and SOP (Liu et al., ).
We implement Knockout with a similar architecture and training scheme as the no-PI baseline, where we concatenate the PI with the image-derived features and randomly knock them out during training.
As a common baseline, we train the same architecture with complete training, but mean imputation for PI data during inference.
Table 2 lists test accuracy results. For the CIFAR-10H dataset, where we have high quality PI, Knockout outperforms all baselines by a large margin, improving test accuracy by 6%. For CIFAR-10/100N datasets, where we have low quality PI during training, Knockout’s boost is more modest, performing similarly with SOP and slightly better than HET and the no-PI baseline.
We conclude that Knockout can offer competitive results when we have access to high quality PI during training.

**Table 2: Test accuracy of different methods on noisy label datasets with PI. We report mean and standard deviation accuracy over 5 runs. PI quality "High" indicates a sample-wise PI is provided by the dataset. "Low" means only batch average is provided. Best results in bold, second-best underlined.**
| Datasets | PI quality | No-PI | HET | SOP | Common baseline | Knockout |
| --- | --- | --- | --- | --- | --- | --- |
| CIFAR-10H (Worst) | High | 51.1$\pm$2.2 | 50.8$\pm$1.4 | 51.3$\pm$1.9 | 55.2$\pm$0.8 | 57.4$\pm$0.6 |
| CIFAR-10N (Worst) | Low | 80.6$\pm$0.2 | 81.9$\pm$0.4 | 85.0$\pm$0.8 | 82.3$\pm$0.3 | 84.7$\pm$0.7 |
| CIFAR-100N (Fine) | Low | 60.4$\pm$0.5 | 60.8$\pm$0.4 | 61.9$\pm$0.6 | 60.7$\pm$0.6 | 62.1$\pm$0.3 |

### 4.4 Missing Images in Tumor Segmentation

Figure: Figure 3: Dice performance of multi-modal tumor segmentation across varying missingness patterns of modality images. Knockout-trained models have better Dice performance across all missingness patterns than the common baseline. We observe mean placeholders perform better than constant-image placeholders. Error bars depict the 95% confidence interval over test subjects.
Refer to caption: x3.png

Here, we investigate the ability of Knockout to handle missingness in a high-dimensional, 3D dense image segmentation task.
In particular, we experiment on a multi-modal tumor segmentation task (Baid et al., ), where the goal is to delineate adult brain gliomas in 3D brain MRI volumes given 4 modalities per subject: T1, T1Gd, T2, and FLAIR.
We use a 3D UNet as the segmentation model (Ronneberger et al., ).
We minimize a sum of cross-entropy loss and Dice loss with equal weighting and use Adam optimizer with a learning rate of 1e-3.
See Appendix B.3 and  A.3 for further details.

At inference time, we evaluate on all modality missingness patterns.
Fig. 3 shows Dice scores.
We observe that the Knockout-trained model has better Dice performance across all missingness patterns.
When all modalities are available, Knockout and the common baseline achieve the same performance level.

### 4.5 Missing Views in Tree Genus Classification

We demonstrate Knockout’s ability to deal with missing data at latent feature level in a classification task.
The Auto Arborist dataset (Beery et al., ), a multiview (street and aerial) image dataset, is used for this purpose.
In this experiment, we used the top 10 genera for multiclass prediction and reported results from 2 sites.
A frozen ResNet-50 (He et al., ) pretrained with ImageNet-v2 (Recht et al., ) is used as a feature extractor.
The two features from street and aerial images are concatenated and were fed into 3-layer MLP with ReLU activations.
We trained Knockout to randomly replace the whole latent vectors with vectors of 0s as placeholders after normalization.
This variant is denoted as Knockout (Structured).
We additionally trained two baselines for comparison: 1) Knockout (Features) where individual features in the latent vectors are independently replaced with placeholders, and 2) an imputation baseline, substituting latent vectors from missing views with vectors of zeros during inference.
The results in Table 3 and Section S8 shows Knockout (Structured) outperforming Knockout (Features) at various missing rates, suggesting that matching $p(\bm{M})$ with missing patterns that we expect to see at inference can be more effective.

**Table 3: F1-scores of Auto Arborist averaged over 5 random seeds (site: Columbus). Each column represents non-missing modalities at inference time. Best results in bold, second-best underlined.**
|  | Aerial+Street | Aerial | Street |
| --- | --- | --- | --- |
| Common baseline | 0.4834±0.0167 | 0.3129±0.0177 | 0.3565±0.0240 |
| Knockout (Features) | 0.4934±0.0209 | 0.2841±0.0230 | 0.3814±0.0221 |
| Knockout (Structured) | 0.4961±0.0169 | 0.3089±0.0242 | 0.4165±0.0140 |

### 4.6 Missing MR Modalities in Prostate Cancer Detection

We demonstrate structured Knockout in the context of a binary image classification task, where Knockout is applied at the latent level.
The dataset consists of T2-weighted (T2w), diffusion-weighted (DWI) and apparent diffusion coefficient (ADC) MR images per subject (Saha et al., ). A simple “ensemble baseline” approach to address missingness is to train a separate convolutional classifier for each modality, and average the predictions of available modalities at inference time (Kim et al., ; Hu et al., ).

To train a model with latent-level structured Knockout, we use the same 3 feature extractors in the “ensemble baseline”, concatenate the features and pass it through a final fully-connected layer.
The loss function is binary cross entropy loss and we use an Adam optimizer with a learning rate of 1e-3.
We randomly knock out each modality.
In the “common baseline” approach, we trained the same architecture with complete modalities.
At inference time, the latent features from missing modalities are imputed with 0s. See Appendix B.5 for more details.

Knockout generally outperforms the baselines in the majority of scenarios as shown in Table 4 for F1 scores and Table S9 for AUC scores, except for inputs with ADC, where the common baseline achieves the best results.
Notably, the F1 scores from the popular ensemble baseline are significantly lower than Knockout.

**Table 4: F1 scores averaged over 5 random seeds, showing performance of ensemble baseline, common baseline, and Knockout across varying missingness patterns at inference time. Each column represents non-missing modalities. Best results in bold, second-best underlined.**
|  | T2 | ADC | DWI | ADC<br>+DWI | T2<br>+DWI | T2<br>+ADC | All |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Ensemble | 0.212$\pm$0.091 | 0.373$\pm$0.016 | 0.285$\pm$0.032 | 0.327$\pm$0.015 | 0.181$\pm$0.044 | 0.337$\pm$0.033 | 0.305$\pm$0.050 |
| Common | 0.432$\pm$0.014 | 0.687$\pm$0.021 | 0.616$\pm$0.021 | 0.706$\pm$0.009 | 0.510$\pm$0.033 | 0.652$\pm$0.006 | 0.673$\pm$0.016 |
| Knockout | 0.639$\pm$0.023 | 0.601$\pm$0.019 | 0.628$\pm$0.025 | 0.677$\pm$0.016 | 0.667$\pm$0.010 | 0.649$\pm$0.023 | 0.688$\pm$0.014 |

## 5 Conclusion and Limitations

We introduced Knockout, a novel, easy-to-implement strategy designed to handle missing inputs, through a mathematically principled approach.
By simulating the missingness during training via random "knock out" and substitution with appropriate placeholder values, our method enables a single model to learn the conditional distribution and all desired marginals.
Our extensive experimental evaluation underscores the versatility and robustness of Knockout. Across diverse datasets, including both synthetic and real-world scenarios, Knockout consistently achieves competitive performance levels, when compared against conventional imputation and ensemble-based techniques across both low and high-dimensional missing inputs.
We further extend Knockout to handle observed missing values in the training set.
Our results highlight the importance of choosing the appropriate placeholder values for induced and observed missingness in training and during inference.
Furthermore, we present structured Knockout that is more effective for scenarios where entire feature vectors or input modalities might be missing.

There are some future directions for further investigation.
While our paper demonstrates the importance of choosing an appropriate placeholder value, and there seems to be a practical tension between selecting an unlikely/infeasible value versus achieving numerical stability (e.g., avoiding exploding gradients), one can conduct a more detailed study of this to optimize the placeholder value.
In our experiments, we did not compare Knockout with individual strong baseline models that would be trained for specific missingness patterns. We considered this to be out of scope, as it became computationally infeasible for all the scenarios we considered. However, in practice, possible missingness patterns might be limited and this approach might be practical.
It is not clear how Knockout would perform against such a strong baseline and this will need to be evaluated.
Another interesting direction of future research is adapting Knockout to tackle distribution shifts in the presence of missingness.
Finally, Knockout’s theoretical treatment hinges on the use of a very high capacity, non-linear model trained on very large data.
In applications, where low capacity models are used and/or training data are limited, Knockout might not be as effective.

## Acknowledgments and Disclosure of Funding

We thank Rachit Saluja, Leo Milecki, Haomiao Chen, and Benjamin C. Lee for the helpful comments.

Funding for this project was in part provided by the NIH grants R01AG053949, R01AG064027 and R01AG070988, and the NSF CAREER 1748377 grant.

## Appendix A Supplementary Material

### A.1 Proof of Continuous and Unbounded Case

As $p(X_{i}{=}\bar{x}_{i})\approx 0$ implies $p(X_{i}{=}\bar{x}_{i}|\cdot)\approx 0$, as long as the conditioning is on a set of observations that is not extremely improbable. Then:

$$ $\displaystyle p(X^{\prime}_{i}{=}\bar{x}_{i}|\cdot)$ $\displaystyle=p(X^{\prime}_{i}{=}\bar{x}_{i},M_{i}=1|\cdot)+p(X^{\prime}_{i}{= }\bar{x}_{i},M_{i}=0|\cdot)$ (10) $\displaystyle=p(M_{i}=1)p(X^{\prime}_{i}{=}\bar{x}_{i}|M_{i}=1,\cdot)+p(M_{i}= 0)p(X^{\prime}_{i}{=}\bar{x}_{i}|M_{i}=0,\cdot)$ (11) $\displaystyle=p(M_{i}=1)(1)+p(M_{i}=0)p(X_{i}{=}\bar{x}_{i}|\cdot)\approx p(M_ {i}=1)$ (12) $\displaystyle\Rightarrow p(Y|X^{\prime}_{i}{=}\bar{x}_{i},\cdot)$ $\displaystyle=\frac{p(Y|\cdot)p(X^{\prime}_{i}{=}\bar{x}_{i}|Y,\cdot)}{p(X^{ \prime}_{i}{=}\bar{x}_{i}|\cdot)}\approx\frac{p(Y|\cdot)p(M_{i}=1)}{p(M_{i}=1) }\approx p(Y|\cdot)$ (13) $$

A similar analysis can be performed to derive the approximate versions of equations of Eq. (4) to (5), which are special cases.

### A.2 Suboptimal Choice of Placeholder Values

If $\bar{x}_{i}\in\text{support}(X_{i})$, then $X_{i}^{\prime}=\bar{x}_{i}$ is either because of two mutually exclusive cases:

$$ $\begin{cases}M_{i}=0\text{ or}\\ M_{i}=1\text{ and }X_{i}=\bar{x}_{i}\end{cases}$ $$

Let $p(M_{i}=1)=r$ and $\bm{x}_{j}\neq\bar{x}_{j},\forall j\neq i$, then:

$$ $\displaystyle p(Y|X_{i}^{\prime}{=}\bar{x}_{i},\bm{X}_{-i}^{\prime}{=}\bm{x}_{ -i})$ $\displaystyle=p(Y|X_{i}^{\prime}{=}\bar{x}_{i},\bm{X}_{-i}{=}\bm{x}_{-i})= \frac{p(Y,X_{i}^{\prime}{=}\bar{x}_{i}|\bm{X}_{-i}{=}\bm{x}_{-i})}{p(X_{i}^{ \prime}{=}\bar{x}_{i}|\bm{X}_{-i}{=}\bm{x}_{-i})}$ (14) $\displaystyle p(X_{i}^{\prime}{=}\bar{x}_{i}|\bm{X}_{-i}{=}\bm{x}_{-i})$ $\displaystyle=p(M_{i}=0)+p(M_{i}=1,X_{i}{=}\bar{x}_{i}|\bm{X}_{-i}{=}\bm{x}_{- i})$ (15) $\displaystyle=1-r+p(X_{i}{=}\bar{x}_{i}|M_{i}=1,\bm{X}_{-i}{=}\bm{x}_{-i})p(M_ {i}=1|\bm{X}_{-i}{=}\bm{x}_{-i})$ (16) $\displaystyle=1-r+r\times p(X_{i}{=}\bar{x}_{i}|\bm{X}_{-i}{=}\bm{x}_{-i})$ (17) $\displaystyle p(Y,X_{i}^{\prime}{=}\bar{x}_{i}|\bm{X}_{-i}{=}\bm{x}_{-i})$ $\displaystyle=p(Y|\bm{X}_{-i}{=}\bm{x}_{-i})p(X_{i}^{\prime}{=}\bar{x}_{i}|Y, \bm{X}_{-i}{=}\bm{x}_{-i})$ (18) $\displaystyle=p(Y|\bm{X}_{-i}{=}\bm{x}_{-i})(1-r+r\times p(X_{i}{=}\bar{x}_{i} |Y,\bm{X}_{-i}{=}\bm{x}_{-i}))$ (19) $$

From Eq. (17) and (19),

$$ $p(Y|X_{i}^{\prime}{=}\bar{x}_{i},\bm{X}_{-i}^{\prime}{=}\bm{x}_{-i})=p(Y|\bm{X }_{-i}{=}\bm{x}_{-i})\frac{1-r+r\times p(X_{i}{=}\bar{x}_{i}|Y,\bm{X}_{-i}{=} \bm{x}_{-i})}{1-r+r\times p(X_{i}{=}\bar{x}_{i}|\bm{X}_{-i}{=}\bm{x}_{-i})}\\$ (21) $$

Thus, $p(Y|X_{i}^{\prime}{=}\bar{x}_{i},\bm{X}_{-i}^{\prime}{=}\bm{x}_{-i})=p(Y|\bm{X
}_{-i}{=}\bm{x}_{-i})$ is equivalent to

$$ $\displaystyle\Leftrightarrow 1-r+r\times p(X_{i}{=}\bar{x}_{i}|Y,\bm{X}_{-i}{= }\bm{x}_{-i})=1-r+r\times p(X_{i}{=}\bar{x}_{i}|\bm{X}_{-i}{=}\bm{x}_{-i})$ (22) $\displaystyle\Leftrightarrow p(X_{i}{=}\bar{x}_{i}|Y,\bm{X}_{-i}{=}\bm{x}_{-i} )=p(X_{i}{=}\bar{x}_{i}|\bm{X}_{-i}{=}\bm{x}_{-i})$ (23) $\displaystyle\Leftrightarrow Y\mbox{${}\perp\mkern-11.0mu\perp{}$}X_{i}{=}\bar {x}_{i}|\bm{X}_{-i}{=}\bm{x}_{-i}$ (24) $$

It is difficult to find $\bar{x}_{i}$ to satisfy Eq. 24, most likely $p(Y|X_{i}^{\prime}{=}\bar{x}_{i},\bm{X}_{-i}^{\prime}{=}\bm{x}_{-i})\neq p(Y|
\bm{X}_{-i}{=}\bm{x}_{-i})$.

Figure: Figure S1: The region of high density of standard Gaussian shifts away from the origin as the number of dimensions increases. This motivates different choices of placeholder values at different dimensions.
Refer to caption: x4.png

Figure: Figure S2: Test MSE evaluated against observations ($Y$) from 10 repetitions of the regression simulation. Lower is better. X axis indicates the number of missing variables at inference time. a) Complete training data. b) Missing completely at random (MCAR) data.
Refer to caption: x5.png

Figure: Figure S3: Test MSE in the regression simulation decreases as the placeholders move away from 0.
Refer to caption: x6.png

### A.3 Analysis of Placeholders for Structured Knockout

In this subsection, we analyze the effect of different placeholder values on a structured Knockout task, specifically the multi-modal tumor segmentation task from Section 4.4.
We scale image intensity values to $[0,1]$ per image. We train three Knockout models with the following placeholders: a constant image of -1s, a constant image of 0s, and the mean of all images per modality.
At inference time, we evaluate on all modality missingness patterns.
In the event that all images are missing, we randomly select one so that the model sees at least one image.
For Knockout-trained models, the corresponding placeholder is imputed for missing images.

Fig. S4 shows the results.
Interestingly, we observe the mean placeholder (Knockout*) performs better than constant-image placeholders, and the constant image of 0s generally outperforms the constant image of -1s.
We hypothesize that in the context of structured inputs like images in conjunction with limited data and model capacity, placeholders which balance feasibility with practical considerations like causing unstable gradients due to out-of-range inputs is an important consideration.

Figure: Figure S4: Dice performance of multi-modal tumor segmentation across varying missingness patterns of modality images. Knockout-trained models only. We observe mean placeholders perform better than constant-image placeholders. Error bars depict the 95% confidence interval over test subjects.
Refer to caption: x7.png

## Appendix B Experimental Details

All experiments were performed with access to a machine equipped with an AMD EPYC 7513 32-Core processor and an Nvidia A100 GPU. CPU testing was performed on a machine equipped with an Intel Xeon Gold 6330 CPU @ 2.00GHz.
All code is written in PyTorch.

### B.1 Simulations

#### B.1.1 Regression

In each repetition, the data are sampled from a 10-dimensional multivariate Gaussian distribution with mean $\bm{\mu}$ and covariance $\bm{\Sigma}$.
The mean vector $\bm{\mu}$ is sampled uniformly from the interval $[0,1]$, i.e. $\bm{\mu}\sim\mathsf{Uniform}(0,1)\in\mathbb{R}^{10}$.
The covariance matrix is sampled as $\bm{\Sigma}:=\bm{W}^{T}\bm{W}$, whereby $\bm{W}\sim\mathsf{Uniform}(0,1)\in\mathbb{R}^{10\times 10}$.
The first 9th variables of the multivariate Gaussian are assigned as $\bm{X}$ ($\bm{X}\in\mathbb{R}^{9}$) and the 10th variable is assigned as $Y$ ($Y\in\mathbb{R}$).

#### B.1.2 Binary Classification

We evaluate the prediction error rate with full features ($\bm{X}$) and missing feature (only $X_{1}$ or $X_{2}$ as input).
We also evaluate how close the models’ predicted probability distributions with missing feature are against the marginal distributions (Fig. 5(a) and Fig. 5(a) top and left panels) using Jensen-Shannon divergence.
The marginals distributions are estimated empirically using all data.

Figure: (a) Continuous inputs
Refer to caption: x8.png

Continuous Inputs.
All input variables are continuous ($\bm{X}\in\mathbb{R}^{2}$).
Knockout achieves similar error rate compared to standard training but much better performance when a variable is missing (Table S5).
Knockout^∗ performs worse than Knockout due to the sub-optimal choice of placeholders.

**Table S5: Classification simulations. Best results are in bold. Err.: Proportion of test error. JSD: Jensen–Shannon divergence of the estimated and empirical marginal.**
|  | Missing Rate = 0 | Missing Rate = 1/2 |  |  |  |
| --- | --- | --- | --- | --- | --- |
| Method | Err. ($\bm{X}$) | Err. ($X_{1}$) | JSD ($X_{1}$) | Err. ($X_{2}$) | JSD ($X_{2}$) |
| Continuous inputs |  |  |  |  |  |
| Common baseline | 0.0003 | 0.3970 | $\infty$ | 0.4001 | $\infty$ |
| Knockout^∗ (Ours) | 0.0210 | 0.3549 | 0.0179 | 0.3727 | 0.0214 |
| Knockout (Ours) | 0.0007 | 0.2559 | 0.0003 | 0.2563 | 0.0007 |
| Fitted Marginals | N/A | 0.2531 | 0.0006 | 0.2600 | 0.0006 |
| Mixed inputs |  |  |  |  |  |
| Common baseline | 0.0032 | 0.5972 | $\infty$ | 0.5410 | $\infty$ |
| Knockout^∗ (Ours) | 0.0187 | 0.4843 | 0.0038 | 0.3410 | 0.0073 |
| Knockout (Ours) | 0.0031 | 0.4028 | 0.0001 | 0.2844 | 0.0009 |
| Fitted Marginals | N/A | 0.4028 | 0.0000 | 0.2809 | 0.0008 |

Mixed Inputs.
$\bm{X}$ consists of a binary variable and a continuous variable (i.e. $X_{1}\in\{0,1\},\;\;X_{2}\in\mathbb{R}$).
Knockout achieves better results than baselines in all scenarios (Table S5).

### B.2 Alzheimer’s Forecasting

All participants used in this work are from the Alzheimer’s Disease Neuroimaging Initiative (ADNI) database.
ADNI aims to evaluate the structure and function of the brain across different disease states and uses clinical measures and biomarkers to monitor disease progression.
Applications for ADNI data use can be submitted through the ADNI website at [https://adni.loni.usc.edu/data-samples/access-data/](https://adni.loni.usc.edu/data-samples/access-data/).
Others would be able to access the data in the same manner as the authors.
We did not have any special access privileges that others would not have.
The investigators within the ADNI contributed to the design and implementation of ADNI and/or provided data but did not participate in analysis or writing of this report. Michael Weiner (E-mail: Michael.Weiner@ucsf.edu) serves as the principal investigator for ADNI.

We select the participants who have mild cognitive impairment (MCI) at the baseline (screening) visit and had at least one follow-up diagnosis within the next five years.
We excluded participants who were diagnosed as CN in a later follow-up year (n=284) since these subjects might have been diagnosed incorrectly at some point.
After this exclusion, we are left with 789 participants.
Table S6 lists summary statistics for the participants; including sex, age, number of years of education completed, count of Apolipoprotein E4 (APOE4) allele, Clinical Dementia Rating(CDR), and Mini Mental State Examination (MMSE) scores at baseline.

**Table S6: Summary statistics of the participants at baseline. Mean $\pm$ standard deviations are listed. APOE4 row represents the number of alleles.**
| Characteristic | ($n$=789) |
| --- | --- |
| Female/Male | $324/465$ |
| Age $(yr)$ | $73.46\pm 7.39$ |
| Education $(yr)$ | $15.93\pm 2.81$ |
| APOE4 $(0/1/2)$ | $371/313/98$ |
| CDR | $1.55\pm 0.89$ |
| MMSE | $27.52\pm 1.82$ |

As is common in many real-world longitudinal studies, ADNI experiences missing follow-up visits, irregular timings, and high dropout rates before the study’s planned end.
Table S7 shows the number of subjects available in each diagnostic category for annual follow-ups.
In Table S7 and all analyses, any subject who progressed from MCI to AD before withdrawing was considered to remain in the AD state until the fifth year, reflecting AD’s irreversible nature.
We employed the reweighted cross-entropy loss scheme introduced in Karaman et al. [] during training to account for the imbalance in diagnoses.

**Table S7: The number of available subjects in each diagnostic group for annual follow-up visits. The follow-up diagnoses are not actually exactly 12 months apart. They have been rounded to the nearest time horizon in years.**
| Follow-up year | 1 | 2 | 3 | 4 | 5 |
| --- | --- | --- | --- | --- | --- |
| MCI | 674 | 431 | 317 | 202 | 127 |
| AD | 110 | 218 | 261 | 286 | 292 |

Our input features include subject demographics
(age;
and
number of years of education completed, or PTEDUCAT),
genotype (number of APOE4 alleles),
cognitive assessments
(Clinical Dementia Rating, or CDR;
Activities of Daily Living, or FAQ;
Alzheimer’s Disease Assessments 11, 13, and Q4, or ADAS11, ADAS13, ADASQ4, respectively;
Mini-Mental State Exam, or MMSE;
Rey Auditory Verbal Learning Test Trials, or RAVLT immediate, learning, forgetting, and percent forgetting;
Logical Memory Delayed Recall, or LDELTOTAL;
Trail Making Test Part B, Or TRABSCOR;
and
Digit and Trails B versions of Preclinical Alzheimer’s Cognitive Composite score, or mPACCdigit and mPACCtrailsB, respectively).
The biomarkers are
Cerebrospinal Fluid (CSF) measurements
(Amyloid-Beta 1–42, or ABETA;
Total Tau, or TAU;
Phosphorylated Tau, or PTAU),
Magnetic Resonance Imaging (MRI) volume measurements
(Ventricles; Hippocampus; WholeBrain; Entorhinal; Fusiform; MidTemp; Intracranial Volume, or ICV; all computed using the FreeSurfer software [Desikan et al., , Fischl et al., ]),
and Positron Emission Tomography (PET) standardized uptake value ratio (SUVR) score for tracer Fluorine-18-Fluorodeoxyglucose, or FDG.
We note that all of our input features are numerical and we perform z-score normalization using the mean and standard deviation values derived from training data.

We note that we train our models using the hyperparameters stated in Karaman et al. []. Figure S6 shows the Composite AUROC scores obtained using the complete portion of the dataset ($n=256$ subjects). In this experiment, the training data has no observed missing variables. These results are similar to the results included in the main text, where Knockout outperforms the baseline and the choice of the appropriate placeholder has an impact on the performance.

Figure: Figure S6: Composite AUROC scores obtained for the three model variants when each input feature is missing during inference (x-axis) for the complete data case. Displayed are averages of 10 train-test splits. Error bars indicate the standard error across these splits.
Refer to caption: x10.png

### B.3 Multi-modal Tumor Segmentation

The RSNA-ASNR-MICCAI BraTS Baid et al. [] challenge releases a dataset of 1251 subjects with multi-institutional routine clinically-acquired multi-parametric MRI scans of glioma.
Each subject has 4 modalities: native (T1), post-contrast T1-weighted (T1Gd), T2-weighted (T2), and T2 Fluid Attenuated Inversion Recovery (T2-FLAIR).
All the imaging datasets have been annotated manually, by one to four raters, following the same annotation protocol, and their annotations were approved by experienced neuro-radiologists. Annotations comprise the GD-enhancing tumor (ET — label 3), the peritumoral edematous/invaded tissue (ED — label 2), and the necrotic tumor core (NCR — label 1).

The following pre-processing is applied: co-registration to the same anatomical template, interpolation to the same resolution (1 mm3), skull-stripped, and min-max normalized to the range [0, 1].
The ground truth data were created after their pre-processing.
For training, we use 80%/5%/15% data split of the subjects for training/validation/testing.

For the segmentation model, we use a 3D UNet with 4 downsampling layers and 2 convolutional blocks per resolution [Ronneberger et al., ].
We minimize a sum of cross-entropy loss and Dice loss with equal weighting and use an Adam optimizer with a learning rate of 1e-3.

### B.4 Auto Arborist

The Auto Arborist dataset is a comprehensive multiview fine-grained visual categorization dataset designed to advance robust methods for large-scale urban forest monitoring.
It includes over 2 million trees from more than 300 genus-level categories across 23 cities in the US and Canada.
Table S8 shows the F-1 scores.

**Table S8: F1-scores of Auto Arborist averaged over 5 random seeds (site: Buffalo). Best results in bold, second-best underlined.**
|  | Aerial+Street | Aerial | Street |
| --- | --- | --- | --- |
| Common baseline | 0.3684±0.0108 | 0.1871±0.0228 | 0.2969±0.0447 |
| Knockout (Features) | 0.3367±0.0255 | 0.1702±0.0204 | 0.2408±0.0283 |
| Knockout (Structured) | 0.3585±0.0219 | 0.1997±0.0155 | 0.3539±0.0135 |

### B.5 Prostate Cancer Detection

A common clinical workflow for the diagnosis of prostate cancer is to detect and localize abnormalities from 3 MR modalities: T2-weighted (T2w), diffusion-weighted (DWI) and apparent diffusion coefficient (ADC) images [Turkbey et al., ]. T2w images provide anatomical details, while DWI and ADC highlight restricted diffusion, which can be a sign of malignancy.

We divided 1500 biparametric MR image sets provided from Prostate Imaging: Cancer AI (PICAI) challenge [Saha et al., ] "training" dataset into training, validation, test sets in a 0.6/0.2/0.2 ratio. Among the 1500 cases, 425 were confirmed as cancer by biopsy. DWI and ADC images are registered to T2w images and all images are cropped around prostate and resized to $100\times 100\times 40$. For the modality-wise classification tasks, we used 3D CNN with 4 blocks, each with a convolution layer, BatchNorm, leakly ReLU activation and average pooling layer, followed by fully connected layer. We trained the models to predict PCa using binary cross entropy loss and an Adam optimizer with a learning rate of $1e-3$.

**Table S9: AUC performance of ensemble baseline, common baseline, and Knockout across varying missingness patterns at inference time. Each column represents non-missing modalities. Best results in bold, second-best underlined.**
|  | T2 | ADC | DWI | ADC<br>+DWI | T2<br>+DWI | T2<br>+ADC | All |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Ensemble | 0.683$\pm$0.013 | 0.786$\pm$0.010 | 0.718$\pm$0.007 | 0.780$\pm$0.005 | 0.722$\pm$0.005 | 0.766$\pm$0.006 | 0.766$\pm$0.004 |
| Common | 0.687$\pm$0.011 | 0.771$\pm$0.011 | 0.720$\pm$0.007 | 0.784$\pm$0.003 | 0.727$\pm$0.006 | 0.771$\pm$0.008 | 0.774$\pm$0.004 |
| Knockout | 0.694$\pm$0.009 | 0.730$\pm$0.019 | 0.736$\pm$0.009 | 0.789$\pm$0.005 | 0.744$\pm$0.004 | 0.753$\pm$0.011 | 0.774$\pm$0.007 |

### B.6 Privileged information for noisy label learning

We briefly introduce two datasets we used for this experiment: CIFAR-10H [Peterson et al., ] and CIFAR-10/100N [Wei et al., ]. CIFAR-10H relabels the original CIFAR-10 10K test set with multiple annotators and provides high-quality sample-wise annotation information such as annotator ID, reaction time and annotator confidence as PI. Following a previous setup [Wang et al., ], we test on the high-noise version of CIFAR-10H, by selecting incorrect labels when available, denoted as "CIFAR-10H Worst". The estimated noise rate is 64.6%. While we train on the high-noise version, testing is conducted on the original CIFAR-10 50K training set. CIFAR-10/100N provides multiple annotations for CIFAR-10/100 training set. The raw data also includes information about annotation process. But this information is provided as averages over batches of examples rather than sample-wise. The estimated noise rate is 40.2% for CIFAR-10/100N.

For all CIFAR experiments and baselines, we use the Wide-ResNet-10-28 [Zagoruyko and Komodakis, ] architecture. We use SGD optimizer with 0.9 Nesterov momentum, a batch size of 256, 0.1 learning rate and 1e-3 weight decay and minimized the cross-entropy loss with respect to the provided labels. The total training epoch is 90, and the learning rate decayed by a factor of 0.2 after 36, 72 epochs. For the PI features, we use annotator ID and annotation reaction time. In PI features are normalized to [0, 1] for preprocessing. For Knockout, during training, we randomly knock out all PI features at $50\%$ rate and use -1 as placeholder value. All experiments are performed on one A6000.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - The choice of the placeholder value seems crucial for Knockout's performance. While the authors provide recommendations, have you explored techniques to automatically learn or optimize the placeholder values during training?
- In scenarios with limited training data or low-capacity models, how well does Knockout perform compared to simpler baselines like mean/mode imputation? Are there any modifications or variations of Knockout that could improve its effectiveness in such cases?
- Have you explored the potential of Knockout to handle distribution shifts in the presence of missingness? For example, if the missingness patterns or distributions change between training and inference, how would Knockout perform compared to other methods?
- While the paper covers a diverse set of tasks and data modalities, it would be interesting to see Knockout's performance on more complex tasks like language modeling or multimodal learning, where missing modalities could be prevalent.
soundness: 3
strengths: - The theoretical analysis is well-reasoned, and the empirical evaluation is comprehensive, covering a diverse set of tasks and data modalities. The authors have taken care to compare against appropriate baselines and provide ablation studies (e.g., structured vs. unstructured Knockout).
- The paper is well-written and clearly explains the core idea, theoretical justification, and experimental setup. The authors have provided sufficient details to facilitate reproducibility.
summary: The paper introduces "Knockout," a simple yet effective data augmentation strategy for handling missing inputs during inference in machine learning models. Knockout randomly replaces input features with appropriate placeholder values during training. At inference time, using the placeholder value corresponds to marginalization over the missing variables.

The key contributions are:

1) Theoretical justification showing that Knockout implicitly maximizes the likelihood of a weighted sum of the conditional estimators and all desired marginals in a single model.

2) Analysis and recommendations for choosing appropriate placeholder values for different data types (categorical, continuous, structured).

3) Extensive experiments on synthetic and real-world datasets (images, tabular data) demonstrating Knockout's effectiveness in handling missing inputs across various scenarios, outperforming common baselines like mean/mode imputation and ensemble methods.
weaknesses: - The idea of randomly masking/corrupting inputs during training is not entirely new, many papers in related work section essentially use the same approach, eg PartialVAE, VAEAC, ACFlow, 
- While the authors provide theoretical justification for Knockout, the analysis relies on the assumption of using a very high capacity, non-linear model trained on large data. It is unclear how well Knockout would perform in scenarios with limited data or low-capacity models.
- The comparison against strong baselines trained specifically for certain missingness patterns is missing. In practical scenarios where the missingness patterns are known or limited, such specialized models could potentially outperform Knockout.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: I believe this paper is well-written and organized, so I don’t have any direct questions. However, I am curious if the authors have considered applying the Knockout method to every hidden layer of the neural network, rather than just the input layer. This method seems somewhat similar to dropout, which is typically applied across all layers.
soundness: 3
strengths: 1. The manuscript is generally well-written, demonstrating quality and clarity.
2. The author presents a comprehensive review of related work.
3. The new Knockout method is evaluated against multiple strong baselines.
4. The author thoroughly discusses various types of missing data mechanisms and evaluates the performance of Knockout and common baselines on them.
summary: The paper presents Knockout, a strategy for handling missing inputs in complex models. Knockout randomly replaces input features with placeholders during training, enabling a single model to learn conditional and marginal distributions. This method is theoretically sound and intensive simulation and real data application were used to evaluate the performance of Knockout.
weaknesses: 1. Figure 2 shows that selecting an appropriate placeholder value has a strong impact on Knockout. While the author emphasizes the importance of this choice, a general guideline for choosing placeholder values is lacking, leaving it to be determined on a case-by-case basis.
2. The simulation results appear somewhat limited. The input dimension of X is only 9, and the number of missing features ranges from 0 to 3. It would be beneficial to include simulations that better align with real-world dat. Specifically, those with high dimensionality and higher missing rates.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - How does Knockout conceptually and empirically compare to related data augmentation methods like Dropout and DAMP [1], particularly since all these methods involve randomly masking inputs during training?
- How does the method perform with limited training data or when using simpler model architectures?
- What is the computational overhead during training compared to standard training and other baselines?
- How sensitive is the performance to the choice of placeholder values in practice?
- Could the random knockout process affect model calibration or uncertainty estimates?
- When would traditional imputation methods be preferable to this approach?
- Have the authors considered extending this to sequential data or other structured input types?

1. Trinh et al. Improving robustness to corruptions with multiplicative weight perturbations. NeurIPS 2024.
soundness: 3
strengths: - Elegant and practical solution that balances simplicity with theoretical soundness
- Strong theoretical foundation with rigorous mathematical analysis 
- Impressive versatility across different data types and applications
- Practical single-model solution compared to existing multi-model approaches
- Comprehensive empirical evaluation with meaningful baselines
- Clear and actionable implementation guidelines for practitioners
summary: This paper introduces "Knockout," a simple yet theoretically-grounded method for handling missing input features in deep learning models. The key idea is to randomly replace input features with appropriate placeholder values during training, which enables a single model to learn both the conditional distribution using full inputs and the marginal distributions for cases with missing inputs. The authors provide theoretical justification showing that Knockout can be interpreted as an implicit marginalization strategy and demonstrate its effectiveness across diverse scenarios including synthetic data, clinical predictions, noisy label learning, tumor segmentation, and multi-modal classification tasks. Compared to existing approaches like marginalization (computationally expensive), imputation (potentially inaccurate for high-dimensional data), or training multiple models (costly and requires prior knowledge of missing patterns), Knockout offers a more efficient and flexible solution. The authors also carefully analyze how to choose appropriate placeholder values for different types of inputs and show how Knockout can handle both complete and incomplete training data under various missingness mechanisms (MCAR, MAR, MNAR).
weaknesses: - Limited theoretical analysis for finite-capacity models and small datasets, as theory assumes high-capacity models and large data
- Missing comparison against specialized models trained for specific missingness patterns
- No detailed ablation study on optimal placeholder value selection, despite its importance
- Lack of exploration into computational overhead during training compared to simpler approaches
- Limited discussion of failure cases or scenarios where the method might underperform
- No investigation into potential impact on model robustness or calibration

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: I explained my concerns in the weaknesses section.
I may change my rating after rebuttal and discussion with other reviewers.
soundness: 3
strengths: I appreciate the authors efforts to run experiments on various datasets.
The paper is also interesting and practical, proposing a simple and straightforward technique.
summary: This paper proposes a simple technique, Knockout, which randomly replaces input features with appropriate placeholder values during training. The main contribution/novelty of this paper is to introduce specialized placeholder values for different types of features: categorical, bounded vs unbounded continuous, structured vs unstructured.
weaknesses: Although this paper is interesting and practical. But, it is very incremental in terms of research novelty considering the expectations from an ICLR paper. For these types of papers, it is required to have thorough experimental studies and solid comparisons to show the applied contributions. But, I think this paper has lack of comparisons with important baselines or prior works.

* The main missing baseline for comparison is the dropout method. Actually, the comparison between knockout and knockout* shows that most of the advantage comes from random replacement. In fact, since the model has seen more missing data during training, it has become more robust to missingness and performs better on test data. An intuitive baseline is to apply dropout with the same missing rate and compare with knockout. Likewise, it may be possible that by applying the same rate of random missingness and use any other baseline (e.g. mean/media imputation or other methods), their performance also improves. In general, I think the comparisons performed in the paper were not fair since the other methods did not see the same amount of missing values during training

* The paper misses to cite some important related works:
    * Why not to use zero imputation? correcting sparsity bias in training neural networks
    * Debiasing Averaged Stochastic Gradient Descent
    * Learning from data with structured missingness

* Further, the paper cites some important papers such as "Ipsen et al, How to deal with missing data in supervised deep learning" but does not perform any comparison. I understand this paper proposes a complicated model. But, for an ICLR paper, I think stronger baseslines are expected.