CACTI: Leveraging Copy Masking and Contextual Information to Improve Tabular Data Imputation

We would like to thank the reviewer for the helpful feedback and questions.

Comments Addressed

We've added CMAE (CACTI without context) as an additional benchmark in the main results table (qrHN rebuttal Table. R1). We’ve also added a comparison between CACTI, CMAE to quantify the statistical significance of the improvements driven by context awareness (qrHN rebuttal Table. R2). We hope these results further clarify that context yields statistically significant (p<0.05) improvement over masked learning alone and is a useful source of inductive bias to improve imputation. We really appreciate the reviewer's attention to detail and highlighting the typos. They have now been addressed, and we apologize for the oversight.

Next, our theoretical results in Appendix G show that copy masking can be more effective than random masking under MNAR. For MAR, consider a structured MAR setting [1] where the missing variables are dependent on some observed variable but also have additional dependencies, e.g., feature i is missing if feature j is missing. Such settings are quite common in questionnaire and other datasets. Under this setting, copy masking is still able to maintain the structured missingness between the variables whereas random masking completely diverges from the missingness process. The crux of our argument in Appendix G is that copy masking is better suited to modeling MAR and MNAR than random masking. For example, unlike random masking, applying copy masking to MAR results in masking only the missing features. This encourages the MAE to learn a conditional model that maximizes the probability of the masked feature given the observed features.

Contributions

While MAEs with random masking (ReMasker) and MLPs with naive copy masking (AutoComplete) have both been used to tackle imputation in prior works, we believe our work makes novel, concrete contributions.

1. MT-CM, an upgrade to copy masking, is a novel, key contribution. As highlighted by our results in Table A7 the application of naive copy masking to transformer-based MAE architecture leads to suboptimal performance and reduced learning efficiency. This is a significant gap in the field, as transformers are arguably the most popular architecture in current literature. MT-CM enables the application of copy masking effectively to such architectures and creates an upper bound on the worst-case scenario in batch-based training, i.e. the proportion of null tokens in any batch is upper bounded by 50%. Our MT-CM strategy unlocks the effective use of copy masking for any transformer-based masked learning model and can have applications beyond just tabular imputation.
2. CACTI is the first approach (to the best of our knowledge) to effectively leverage context as a useful inductive bias and show statistically significant improvements in various imputation settings. Our results highlight the importance of investigating new inductive biases as a promising direction for future tabular imputation research, particularly in domains with MNAR missingness.
3. Our results establish CACTI as a strong baseline for evaluating and developing future approaches in tabular imputation.
4. Although our theoretical section (Appendix G) is not highlighted in the main text (due to space restrictions), we believe we do provide a new empirical risk framework for evaluating MAE training without fully observed data (previous works don’t consider the probabilistic nature of the missingness process and only view the objective as representation loss minimization), provide the first formal motivation for copy masking and show how random masking is suboptimal compared to copy masking in non-MCAR settings. If given the chance we would like to introduce this theoretical framework in the main text in the final draft.

We hope this addresses the reviewer's concerns and find this paper to be a valuable contribution to the current literature. We would appreciate their reconsideration of their overall recommendation. And as always we're happy to address any additional questions or concerns!

[1] Jackson, J., et. al. A complete characterisation of structured missingness, 2023.

We thank the reviewer for their thoughtful review and suggestions. We would like to highlight that all the individual per dataset/missingness condition results for all analyses performed in this paper were previously reported with 95% confidence intervals in the Appendix. We choose to exclude standard errors (SEs) in the main text to ensure visual parsimony and due to space constraints. Please see the silhouette for the main benchmarking results table with SEs reported in Table. R1 which we will include in the final version. Finally, we have a new set of results in Table. R2, which shows that CACTI demonstrates, via a paired t-test, statistically significant (p < 0.05) improvement in R2 and RMSE over all missingness conditions when compared to the next best method (ReMasker). This relative ranking improvement holds even after we exclude the 4 datasets used for the sensitivity and ablation analysis (Table. R3). This provides further evidence to show that leveraging inductive biases from dataset-specific missingness patterns and context information can help improve imputation accuracy. Finally, if the reviewer believes we have sufficiently addressed their concerns and finds this paper to be a valuable contribution to the current literature, we would appreciate their consideration in increasing their score. As always, we're happy to address any additional questions or concerns!

[Please note that we’re only reporting the top few methods since we’re limited to 5000 characters per rebuttals, final version of the paper will contain all methods when appropriate]

Table. R1: Main benchmarking results with standard errors in parentheses and two sets of results indicate train|test.

Table. R2: T-test to quantify the statistical significance of the contributions of MT-CM and context awareness (CMAE = CACTI w/o context).

Table. R3: Results excluding the 4 datasets used for sensitivity analysis.

We thank the reviewer for their valuable feedback.

Does context help?

Our ablation analysis focused on the relative contributions of median truncated copy masking (MT-CM) and context awareness with fixed architecture settings and hyperparameters (Appendix C.2). Random masking autoencoder (RMAE) and ReMasker perform nearly identically when hyperparameters match (internally verified). Performance differences primarily stem from different masking rates: 30% for ReMasker (optimal per original paper) versus 90% for CACTI/RMAE.

To address potential ambiguity, we extended our analysis of CMAE (CACTI without context but with MT-CM) from 4 to all 10 datasets (qrHN rebuttal Table. R1). CMAE consistently improves or matches ReMasker across R2, RMSE, and WD metrics, while CACTI strictly dominates CMAE in all conditions. This confirms that context can improve imputation accuracy. We also performed a comparison using one-sided paired t-tests (qrHN rebuttal Table. R2). CACTI shows statistically significant (p<0.05) improvement over CMAE for R2 (and RMSE) in most settings, reinforcing that context enhances imputation. MT-CM (via CMAE) delivers statistically significant R2 improvements over random masking (via ReMasker) across all settings. CACTI, combining both approaches, shows significant improvement across nearly all metrics and conditions. While MT-CM provides larger improvements, context contributes meaningfully, demonstrating that both MT-CM and context awareness provide useful inductive biases.

Datasets

We selected 10 open datasets from those used by Hyperimpute, ReMasker, and Diffputer authors for fair comparison. Our selection criteria were: 1) 6 with mixed features and 4 with continuous-only features datasets to ensure type diversity, and 2) half the datasets having feature descriptions. Table A9 will be expanded with information about which datasets have column descriptions. Our anonymous repo provides code to obtain all datasets and generate colinfo.txt files (see "Generating UCI ML Datasets" in README).

The 4 datasets used in ablation and sensitivity analysis were randomly chosen from the 10 used in main benchmarking. While this approach is common in prior works, we acknowledge the potential risk of inflating overall results. To maintain fairness, we used similar/same datasets and optimal parameters from previous works. To directly address this concern, we re-analyzed our main benchmarking results excluding these 4 datasets (qrHN rebuttal Table. R3). Results confirm CACTI still strictly outperforms existing state-of-the-art methods in R2 and RMSE, addressing concerns about fairness and result inflation.

Writing

We appreciate the reviewer highlighting these typos. While our notation was chosen for precision, we will move significant portions to the appendix, replacing them with more intuitive text descriptions and figures. ReMasker and [3] are indeed key previous works and we’ll better highlight ReMasker and move the discussion on copy masking from methods to related works. The [MASK] token is indeed fixed and identical for all features, following standard practice in MAEs and ViTs. The hyperparameters of baseline method will be specified in an appendix table. Finally, we apologize for the ambiguity - we mean row-wise shuffling (permuting rows of the observed missingness mask matrix). All these will be fixed in the final draft.

Other Comments

Regarding the comment about including “the mask as additional input features for the imputation model”, we included notMIWAE, which leverages this idea. It maximizes the ELBO of the joint likelihood of the observed features and missing pattern via IWVAE, this ensured we compared against methods that directly factor in missingness structure. As seen in Table 2, we outperform notMIWAE.

Next, while higher copy-masking rates sometimes improve performance, keeping this as a tunable hyperparameter is justified because: 1) performance decreases at rates >95% in MNAR settings for R² and WD; 2) our conservative 90% rate maximizes average performance across sensitivity analysis datasets but isn't universally optimal; 3) practitioners often prioritize specific features based on domain needs, making flexible masking valuable; and 4) optimal masking rates vary significantly across datasets. Next, the embedding size and depths chosen here perform consistently well with reasonable resource usage (<200MB GPU RAM). Since imputation is rarely the end application, we maintain these as tunable parameters for users to optimize their specific use cases.

Finally, we refer the reviewer to point 4 in the contributions section of rebuttal to reviewer eZeh, which addresses concerns about the theoretical contributions.

We hope our responses have addressed the reviewer's primary concerns and demonstrated the paper's contribution to current literature. If so, we would appreciate a reconsideration of the overall recommendation. We're happy to address any additional questions!