Paper: ReMasker: Imputing Tabular Data with Masked Autoencoding (ICLR-2024)

Reviewer comment (Novelty):
The methodology is not novel, as masking is a common practice for imputation algorithms and masked autoencoders with Transformers are used for other tasks. However, the authors do not claim novelty and this is likely the first work to consider a masked autoencoder with Transformers for imputation specifically.

Author response:
While both Transformer and masked autoencoding (MAE) have been used in prior work, this work adapts the MAE approach (with Transformer as the backbone) in a nuanced yet novel way: besides the missing values in the given dataset, we randomly select and “re-mask” another set of values, optimize the autoencoder with the objective of reconstructing this re-masked set, and then apply the trained autoencoder to predict the missing values. This resulting re-masking approach serves as a strong baseline for developing and evaluating future imputation methods.

Reviewer comment (Theory):
It’s not clear how we could properly learn a missingness-invariant representation in the MNAR setting. Under MNAR, the missing data values fundamentally depend on the missingness mechanism, so recovering a missingness-invariant representation would require external information.

Author response:
We thank the reviewer for the insightful comments. In general, it is impossible to identify the missingness distribution in MNAR without domain-specific assumptions or constraints (Ma & Zhang, 2021). That is why most prior work (e.g., Yong et al., 2019; Yoon and Sull, 2020; Jarrett et al., 2022) focuses on MCAR and MAR. For completeness, following the setting of Jarrett et al., (2022), we also evaluate ReMasker under one specific MNAR setting and show that it also partially generalizes to such settings. We acknowledge that there are MNAR settings (such as the ones suggested by the reviewer) where access to external knowledge is essential. We have revised the statement (**Section 4**) to more accurately reflect our work’s focus on MAR and MCAR and its limitations with respect to MNAR. 

Further, we have added a detailed discussion about how ReMasker’s performance is influenced by the missingness mechanism (**Section 5 Q2**). Intuitively, ReMasker encourages learning representations invariant to re-masked values and then leverages such representations to impute missing values. Thus, while ReMasker takes effect as long as the distributions of re-masked and missing values are correlated, its performance may vary with the disparity between the two distributions. In MACR, the distributions of re-masked and missing values are highly similar; in MAR, the distribution of re-masked values is biased towards the observable values; in MNAR, there is even more disparity between the two distributions. Thus, ReMasker’s performance tends to vary as MCAR > MAR > MNAR.

Reviewer comment (Experiments):
The experimental condition for MNAR (masking via a Bernoulli random variable with fixed mean) is not representative of many MNAR settings. For example, cases where values across possibly many variables are masked due to censoring of one or more variables (e.g. \(X_2\), \(X_5\), \(X_6\) are missing when \(X_5 < 0\)). The missing values often follow a completely different distribution than the observed values, which cannot be recovered using the observed data alone.

Author response:
We thank the reviewer for the insightful comments. In general, it is impossible to identify the missingness distribution in MNAR without domain-specific assumptions or constraints (Ma & Zhang, 2021). That is why most prior work (e.g., Yong et al., 2019; Yoon and Sull, 2020; Jarrett et al., 2022) focuses on MCAR and MAR. For completeness, following the setting of Jarrett et al., (2022), we also evaluate ReMasker under one specific MNAR setting and show that it also partially generalizes to such settings. We acknowledge that there are MNAR settings (such as the ones suggested by the reviewer) where access to external knowledge is essential. We have revised the statement (**Section 4**) to more accurately reflect our work’s focus on MAR and MCAR and its limitations with respect to MNAR. 

Further, we have added a detailed discussion about how ReMasker’s performance is influenced by the missingness mechanism (**Section 5 Q2**). Intuitively, ReMasker encourages learning representations invariant to re-masked values and then leverages such representations to impute missing values. Thus, while ReMasker takes effect as long as the distributions of re-masked and missing values are correlated, its performance may vary with the disparity between the two distributions. In MACR, the distributions of re-masked and missing values are highly similar; in MAR, the distribution of re-masked values is biased towards the observable values; in MNAR, there is even more disparity between the two distributions. Thus, ReMasker’s performance tends to vary as MCAR > MAR > MNAR.

Reviewer comment (Evaluation):
It is not clear that the results are practically different from HyperImpute, and perhaps other methods. The included bar plots are hard to read, making assessment difficult. While the relative simplicity of ReMasker compared to HyperImpute is mentioned as one advantage, the much simpler ICE appears competitive in almost all scenarios.

Author response:
To better visualize the improvement of ReMasker over other baselines, we have summarized the performance of ReMasker and baselines (measured and ranked in terms of RMSE) on 12 benchmark datasets under MAR with 0.3 missingness ratio in Table 10 (**Appendix B3**). The key observation is that ReMasker consistently outperforms other baselines across most datasets (ranked top 1 or 2). Compared with HyperImpute, ReMasker is a much simpler and more modular imputer; therefore, it may serve as a strong baseline to develop or evaluate future imputation methods or it can be integrated into an ensemble imputer such as HyperImpute. In addition, we further compare the running time of ReMasker (under the default setting of Table 7) and HyperImpute on datasets with over 1,000 records, with results summarized as follows (**Appendix B.2**). Observe that compared with HyperImpute, ReMasker is not only more scalable but also less sensitive to data size. 

| | compression |  wine |  spam | credit | bike | obesity | california |
| :----: | :----: | :----: | :----: | :----: | :----: | :----: | :----: |
| **HyperImpute** | 154.9s | 297.7s | 185.8s | 39.9s | 259.9s | 196.1s | 195.8s |
| **ReMasker** | 72.9s | 73.8s | 74.1s | 74.0s | 66.9s | 66.8s |66.9s |

Reviewer comment (Theory):
While the theoretical justification is appropriately included in the Discussion section, the justification is a conjecture with minimal validation. In practice, the missingness mechanism is different for \(m^+\) and \(m^-\), so it is not clear that representation will be invariant in general. Further experimentation is needed.

Author response:
To further validate the masking-invariant properties of ReMasker, we measure the CKA similarity (Kornblith et al., 2019) between the latent representations (i.e., the output of ReMasker’s encoder) of inputs under the masks of $m^+$ and $m^-$, with results shown in Figure 6 (**Appendix B4**). Observe that the CKA measures between the inputs under the two masks steadily increase with the training epochs, which empirically corroborates the analysis in Section 5 Q1 that ReMasker tends to learn representations that are invariant to various missingness.

Reviewer comment (Experiments):
While it is clearly stated that imputation for downstream tasks is not part of the scope of the paper, can any more discussion be provided? Considering imputation is almost always followed by downstream tasks, more justification for this decision is needed.

Author response:
We thank the reviewer for the question. Following the line of work on imputing tabular data (e.g., Mattei & Frellsen, 2018; Yong et al., 2019; Nazabal et al., 2020; Yoon and Sull, 2020; Jarrett et al., 2022), we do not consider imputing missing values as a means to obtain data for known downstream tasks. Under those settings, the imputation is fundamentally entangled with the downstream task, typically requiring joint training to optimize the downstream task and the imputation task simultaneously. Due to the potential conflict between the objectives of the two tasks (i.e., data fidelity versus downstream-task performance), it is not straightforward to evaluate the performance of imputation methods directly. We thus focus on the imputation task with data fidelity as the main metric.

Again, we thank the reviewer for the valuable feedback. Please let us know if there are any other questions or suggestions.

Best,

Authors

Reviewer comment (Novelty):
The authors fail to mention and discuss multiple previous works where transformers have been applied to tabular data, specifically Somepalli et al. 2021 (SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training), Arık et al. 2020 (TabNet: Attentive Interpretable Tabular Learning), and Huang et al. 2020 (TabTransformer: Tabular Data Modeling Using Contextual Embeddings).

Author response:
We thank the reviewer for pointing out the missing references. We have included a discussion of these works (**Section 2**). While SAINT, TabNet, and TabTransformer also use Transformer as the backbone to model tabular data, the main idea of ReMasker is to adapt the masked autoencoding (MAE) approach to imputing missing values of tabular data: besides the missing values in the given dataset, randomly select and “re-mask” another set of values, optimize the autoencoder with the objective of reconstructing this re-masked set, and then apply the trained autoencoder to predict the missing values. Meanwhile, although MET also employs the MAE approach, its primary focus is on the representation learning of tabular data, operating under the assumption of data completeness and applying MAE in a relatively direct manner. In contrast, ReMasker, with its emphasis on imputing missing values, inherently assumes incomplete data and adapts the MAE technique in a nuanced yet novel way.

Reviewer comment (Novelty):
Specific to Huang et al. 2020, the authors cite this paper for supporting their model choice but fail to include it in the related works section.

Author response:
We thank the reviewer for pointing out the missing references. We have included a discussion of these works (**Section 2**). While SAINT, TabNet, and TabTransformer also use Transformer as the backbone to model tabular data, the main idea of ReMasker is to adapt the masked autoencoding (MAE) approach to imputing missing values of tabular data: besides the missing values in the given dataset, randomly select and “re-mask” another set of values, optimize the autoencoder with the objective of reconstructing this re-masked set, and then apply the trained autoencoder to predict the missing values. Meanwhile, although MET also employs the MAE approach, its primary focus is on the representation learning of tabular data, operating under the assumption of data completeness and applying MAE in a relatively direct manner. In contrast, ReMasker, with its emphasis on imputing missing values, inherently assumes incomplete data and adapts the MAE technique in a nuanced yet novel way.

Reviewer comment (Theory):
On page 1, the paper states that discriminative imputers are hindered by the requirement of specifying a functional form, but the reviewer questions whether the authors' architecture is also effectively a choice of "functional form" and asks if this is wrong.

Author response:
We thank the reviewer for the question. We would like to clarify that the functional form refers to the conditional distributions of missing values. Specifically, the discriminative methods often impute missing values by modeling their conditional distributions on the basis of observable values, therefore these methods are hindered by the requirement of specifying the proper functional forms about the conditional distributions.

Reviewer comment (Theory):
The argument presented in Section 4.1 about why Remasker generalizes different missing data assumptions could be strengthened. While transformer-based architectures have proven their superior performance in various tasks, a more comprehensive analysis of why transformers excel in missing data imputation would enhance the paper's quality and its contribution to both the missing data imputation and transformer-based research fields.

Author response:
Following the reviewer’s suggestion, we have added a detailed discussion about how ReMasker’s performance is influenced by the missingness mechanism (**Section 5 Q2**). Intuitively, ReMasker encourages learning representations invariant to re-masked values and then leverages such representations to impute missing values. Thus, while ReMasker takes effect as long as the distributions of re-masked and missing values are correlated, its performance may vary with the disparity between the two distributions. In MACR, the distributions of re-masked and missing values are highly similar; in MAR, the distribution of re-masked values is biased towards the observable values; in MNAR, there is even more disparity between the two distributions. While ReMasker works across different missingness patterns, its performance varies as MCAR > MAR > MNAR.

Following the reviewer’s suggestion, we have added a discussion about why Transformer is effective for tabular data imputation (**Section 5 Q3**). While the theoretical understanding about the Transformer architecture is still limited, we borrow insights from recent work to explain the effectiveness of Transformer for tabular data imputation. Specifically, it is shown in (Park &
Kim, 2022) that the multi-head self-attention (MSA) mechanism in Transformer essentially performs (optimizable) smoothing over feature maps, which makes Transformer highly robust against severe occlusions or perturbations (Naseer et al., 2021). In the same vein, in the context of tabular data imputation, MSA benefits learning representations that are invariant to missing values.

Reviewer comment (Theory):
I'd like to understand the rationale behind using a deep encoder and a shallow decoder when describing the decoder. While the Ablation Study in Section 4.3 provides some insight, could you provide a more abstract explanation for this choice? Is there a specific reason why you believe the first fitting stage of the Remasker benefits from a more complex encoder?

Author response:
As shown in our experiments, the optimal encoder/decoder depth varies from 4 to 10. While a relatively shallow decoder often suffices to reconstruct the masked values, the encoder needs to be sufficiently complex to effectively learn the holistic representations of inputs (but not too complex to overfit). Such asymmetric encoder-decoder leads to high-fidelity imputation while minimizing the computational cost.

Reviewer comment (Evaluation):
Some figures display results for different baselines (e.g., Figure 3 RMSE), and while they may appear distinct, the values on the y-axis are actually quite close. Could you elaborate on how these slight differences influence the final results? For instance, in Figure 3, Remasker outperforms other methods by only a narrow margin in terms of RMSE and AUROC. How does this marginal improvement make Remasker a more beneficial choice compared to other methods that utilize deep generative models to learn the data distribution?

Author response:
To better visualize the improvement of ReMasker over other baselines, we have summarized the performance of ReMasker and baselines (measured and ranked in terms of RMSE) on 12 benchmark datasets in Table 10 (**Appendix B3**). The key observation is that ReMasker (ranked top 1 or 2) consistently outperforms other baselines across different datasets. Also note that the only method with performance close to ReMasker is HyperImpute, which is an ensemble method that integrates multiple imputation models and automatically selects the most fitting model for each column of the given dataset. As ReMasker is a much simpler and more modular imputer, it may serve as a strong baseline for developing and evaluating future imputation methods or it can be integrated into an ensemble imputer such as HyperImpute.