ReMasker: Imputing Tabular Data with Masked Autoencoding
Authors: Tianyu Du
Luca Melis
Ting Wang
Venue: ICLR 2024 poster (2024)
Keywords: Masked modeling; Tabular data imputation
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: A simple yet effective imputation method

Abstract: We present ReMasker, a new method of imputing missing values in tabular data by extending the masked autoencoding framework. Compared with prior work, ReMasker is extremely simple -- besides the missing values (i.e., naturally masked), we randomly "re-mask" another set of values, optimize the autoencoder by reconstructing this re-masked set, and apply the trained model to predict the missing values; and yet highly effective -- with extensive evaluation on benchmark datasets, we show that ReMasker performs on par with or outperforms state-of-the-art methods in terms of both imputation fidelity and utility under various missingness settings, while its performance advantage often increases with the ratio of missing data. We further explore theoretical justification for its effectiveness, showing that ReMasker tends to learn missingness-invariant representations of tabular data. Our findings indicate that masked modeling represents a promising direction for further research on tabular data imputation. The code is publicly available.

Published as a conference paper at ICLR 2024
REMASKER: IMPUTING MASKED AUTOENCODING
TABULAR DATA WITH
Tianyu Du Luca Melis²
Ting Wang**
'Zhejiang University Meta
Penn State SStony Brook University
ABSTRACT
We present REMASKER, new method of imputing missing values in tabular data by extending the masked autoencoding framework. Compared with prior work, REMASKER is both simple besides the missing values (i.e., naturally masked), we randomly "re-mask" another set of values, optimize the autoencoder by reconstructing this re-masked set, and apply the trained model predict the missing values; and effective - with extensive evaluation on benchmark datasets, we show that REMASKER performs on par with or outperforms state-of state-of-the-ard -he-rttttt methods in terms of both imputation fidelity and utility under various missingness settings, while its performance advantage often increases with the ratio of missing data. We further explore theoretical justification for its effectiveness, showing that REMASKER tends learn missingness-invarian representations of tabular data. Our findings indicate that masked modeling represents a promising direction for further research on tabular data imputation. (code available: pettps///gitan I ..co/alps-- lab/ http/remaske...
INTRODUCTION
Missing values are ubiquitous in real-world tabular data due to various reasons during data collection, processing, storage, transmission. It is often desirable to know the most likely values of missing data before performing downstream tasks (e.g., classification or synthesis). To this end, intensive research has been dedicated to developing imputation methods ("imputers") that estimate missing
values based on observed data Yoon et al. 2019: Jarrett et al. 2022] Kyono et al.) 2021 Stekhoven & Buhlmann| 2012, Mattei & Frellsen| 2018). Yet, imputing missing values in tabular data with high fidelity and utility remains an open problem, due to challenges including the intricate correlation across different features, the variety of missingness scenarios, and the scarce amount of available data with respect to the number of missing values.
The -fate-of-the-ar imputers can be categorized as either discriminative or generative. The dis- criminative imputers, such as MissForest (Stekhoven & Buhlmann] 2012). MICE (van Buuren & Groothuis-Oudshoomm 2011), and MIRACLE (Kyono et al.) al.]|221]] 2021), impute missing values by model- ing their conditional distributions on the basis of other values. In practice, these methods are often hindered by the requirement of specifying the proper functional forms of conditional distributions and adding the set of appropriate regularizers. The generative imputers, such as GAIN Yoon et al. 2019) MIWAE (Mattei & Frellsen 2018), GAMIN (Yoon & Sull 2020), and II-VAE (Nazabal et al. 2020).
estimate the joint distributions of all the features by leveraging the capacity of deep generative models and impute missing values by querying the trained models. Empirically. GAN- based methods often require a large amount of training data and suffer the difficulties of adversarial training (Goodfellow eet al. 2014), while VAE-based methods often face the limitations of training through variational bounds (Zhao et al.) 2022). Further, some of these methods either require complete data during training or operate on the assumptions of specific missingness patterns.
In this paper, we present REMASKER, novel method that extends the masked autoencoding (MAB) framework (Devlin et al. 2018 He et al. (022) to imputing missing values of tabular data. The idea of REMASKER simple: Besides the missing values in the given dataset (i.e.. naturally masked), we randomly select and "re-mask" another set of values, optimize the autoencoder with the objective of
"The work was done when the first author was postdoc Penn State.

---

Published as a conference paper at ICLR 2024
reconstructing this re- re-masked set, and then apply the trained autoencoder to predict the missing values. Compared with the prior work. REMASKER enjoys the following desiderata: (i) it is instantiated with Transformer (Vaswani et al. 2017) as its backbone, of which the self-attention mechanism is able to capture the intricate inter-feature correlation (Huang et al., 2020); (i)) without specific assumptions about the missingness mechanisms, it is applicable various scenarios even if complete data is unavailable; and (iii) as the re-masking approach naturally accounts for missing values and encourages learning high-level representations beyond low-level statistics, REMASKER works effectively even under high ratio of missing data (e.g., up to 70%).
With extensive evaluation on 12 benchmark datasets under various missingness scenarios, we show that REMASKER performs on par with or outperforms 13 popular methods in terms of both imputation fidelity and utility, while its performance advantage often increases with the ratio of missing data. We further explore the theoretical explanation for its effectiveness. We find that REMASKER encourages learning missingness-innee
representations of tabular data, which are insensitive to missing values. Our findings indicate that, besides its success in the language and vision domains, masked modeling also represents a promising direction for future research on tabular data imputation.
2 RELATED WORK
Tabular data imputation.
The existing imputation methods can be categorized as either discrim-
inative or generative. The discriminative methods (Stekhoven & Buhlmann 2012] van Buuren &
Groothuis- Oudshoomn 2011: Kyono et al. 2021) often specify univariable model for each feature conditional on all others and perform cyclic regression over each target variable until convergence. Recent work has also explored adaptively selecting and configuring multiple discriminative imput- ers (Jrrett et al. 2022). The generative methods either implicitly train imputers as generators within
the GAN framework (Yonn et al. 2019 Yoon & Sull, 2020) or explicitly train deep latent-varis variable models to approximate the joint distributions of all features (Mattei & Frellsen 2018 Nazabal et al. 2020 There are also imputers based on eepresentatiee-vlluee (e.g., mean) substitution (Hawthorne
& Ellioui 2005), EM optimization (García-l
al. (0TO)) matrix completion (Hastie et al., 2015). or optimal transport (Muzellec et al. (022)) The work is also related to that models missing
data by adapting existing model architectures (PrzewigZlikowski al. 2021
Transformer. Transformer has emerged as a dominating design (Vaswani et al. 2017) in the language domain, in which multi-head self-attention and MLP layers are stacked to capture both short- and long-term correlations between words. Recent work has explored the use of Transformer in the vision domain by treating each image as grid of visual words (Dosovitskiy et al. 2020). For instance, has been integrated into image generation models (Jiang et al. 2021] ZZhng et al. 2021 Hudson & Zitnick 2022), achieving performance comparable to CNN-based models. Recent work (Somepalii
al. | Huang et al. 2020 Arik & Plister) 2021) also explores applying Transformer to model labular data, showing its promising expressive power.
Masked autoencoding (MAE). Autoencoding is a classical sel-supervised approach for learning representation (Vincent et al. 2008 Pathak et al. (0066): an encoder maps an input to its representation and decoder reconstructs the original input. Mcanwhile, masked modeling is originally proposed as pre-training method in the language domain: by holding out proportion of a word sequence, it trains the model predict the masked tokens (Devlin al.] [018] Radford & Narasimhan (018)) Recent work has combined autoencoding and masked modeling ling in vision tasks (Dosovitskiy et al. 2020
Bao al. 2022). The seminal MAE (He et al.) (2022) represents the state of the in self-supervised pre-training on the ImageNet-IKK K benchmark. In the tabular data domain, MET (Majmundar et al., 2022) employs the MAE approach to model tabular data. However, given its primary focus on representation learning, MET assumes data completeness and directly applies MAE.
To our best knowledge, this represents the first work to explore an extended MAE approach (with Transformer as the backbone) in the task of tabular data imputation.
REMASKER
Next, we present REMASKER, an extremely simple yet effective method for imputing missing values of tabular data. We begin by formalizing the imputation problem.
2

---

Published as a conference paper at ICLR 2024
Fitting
Imputation
Poddied Embedding
Paddedd Embedking
Impul
m
Recorstructed Value
Input
Embedding
m
Impuled Value
Embedkiing
23
Frnc ood
Decoder
Z6
[+] Mizsing Value
Masked Value
[m] Mask Token
Figure I: Overall framework of REMASKER. During the fitting stage, for each input. in addition to its missing values, another subset of values tre-masked values) is randomly selected and masked out. The encoder is applied the remaining values generate embedding, which padded with mask tokens and processed by the decoder re-construct the re-masked values. During imputation stage, the optimized model is applied to predict the missing values.
3.1 PROBLEM FORMALIZATION
Incomplete data. To model tabular data with d features, consider d-dimensional random variable _ (x1. Xx) X1 X Xa, where X; is either continuous categorical for i {1,.... d). The observational access is mediated by an mask variable - (m.,..., ma) {0, 13d, which indicates the missing values of x, such that x; is accessible only m; 1. In other words, we observe x in its incomplete form & A (K1..... Xa) with
X if m, =
if m; = (IE
where denotes the unobserved value.
Missingness mechanisms. Missing values occur due to various reasons. To simulate different scenarios, following the prior work (Yoon et al.) [2019] Parrett et al.) (022). consider three missingness mechanisms: MCAR ("missing completely at random") the missingness does not depend on the data, which indicates that Vm, x, x', p(mmmm====-))) = ppmmmmm'' x'); MAR ("missing random") the missingness depends on the observed values, which indicates that Vm, x, x', if the observed values of x and x' are the same, then p(m[x) p(m] (inxxx)), and MNAR ("missing not at random") the missingness depends on the missing values as well, which is the case if the definitions of MCAR and MAR not hold. general, impossible identify the missingness distribution of MNAR without Nomain-speciffe assumptions or constraints (Ma & Zhang 2021).
Imputation task. In this task, we are given an incomplete dataset D (x((.) m )
which
consists of i..d. realizations of x and m. The goal is recover the missing values of each input by generating an imputed version & (81...., &g) such that
m,
(1 (1,..,d))))
m, = =0
where X; is the imputed value. Note that the imputation task is not concerned about optimizing data for concrete downstream tasks (e.g., training regression generative models); such settings motivate concerns fundamentally entangled with each downstream task and often require optimizing the end objectives. Here, we focus solely the imputation task itself.
3.2 DESIGN REMASKER
The REMASKER imputer extends the MAE framework (Dosovitskiy et al.) 2020 Bao et al. 2022 He et al. 2022) that reconstructs masked components based on observed components. As illustrated in Figure REMASKER comprises encoder that maps the observed values to their representations and a decoder that reconstructs the masked values from the latent representations. However, unlike conventional MAE, as the data in the imputation task is inherently incomplete (i.e., naturally masked),
Without ambiguity, omit the superscript in the following notations.
3

---

Published as a conference paper at ICLR 2024
employ "re-masking" approach that explicitly accounts for this incompleteness in applying masking and reconstruction. At a high level, REMASKER works in two phases: fitting it optimizes the model with respect to the given dataset, and imputation - it applies the trained model to predict the missing values of the dataset.
Re-masking. In the fitting phase, for each input X, in addition its missing values, we also randomly select and mask out another subset (e.g.. 25%) of X's 's values. Formally, letting m be x's mask, we define another mask vector m' E {0, 1] which is randomly sampled without replacement, following a uniform distribution. Apparently, m and m' entail three subsets:
Imaak {ilm;- = =}, Tremask = {ilm; = LA A = 0}. Tunmask = { = A m = 1} (3)
Let Xm. X *mmm", and XmAm respectively be the masked, re- me-maasedd and unmasked values. With a sufficient number of re-masked values, in addition to the missing values, we create challenging task that encourages the model to learn missingness-invane -invarrann representations (more details in 5 In the imputation phase, we do not apply re-masking.
Encoder. The encoder embeds each value using an encoding function and processes the resulting embeddings through a sequence of Transformer blocks. In implementation, we apply linear encoding function to each value x: enc(x) wa + b. where w and b are learnable parameters. We also add positional encoding to x' embedding to force the model to memorize x position in the input (e.g.. the k-th feature): pe(A: 21) sin(k/ ///000///0///))) where kt and i respectively denote x's position in the input and the dimension of the embedding, and d is the embedding width. The encoder is only applied to the observed values: in the fitting phase, it operates on the observed values after re-masking (i.e., the unmasked set Zummask); in the imputation phase, it operates the non-missing values (i.e.,
the union of re- re-masked and unmasked sets Iunmask U Iremask), as illustrated in Figure
Decoder. The REMASKER decoder is instantiated as sequence of Transformer blocks followed by an MLP layer. Different from the encoder, the decoder operates on the embeddings of both observed and masked values. Following (Devlin al., 2018 He et al. 2022), token as the initial embedding of each masked value. The decoder first adds positional encoding to the embeddings of all the values (observed and masked), processes the embeddings through sequence of Transformer blocks, and finally applies linear projection to map the embeddings to scalar values as the predictions. Similar to (He al. 2022), we use an asymmetric design with deep encoder and shallow decoder (e.g., 8 versus 4 blocks), which often suffices to re- re-construct the masked values. Conventional MAE focuses on representation learning and uses the decoder only in the training phase. In REMASKER the decoder is required to re-construct the missing values and is thus used both fitting and imputation phases.
shared, learnable mask
Reconstruction loss. Recall that the REMASKER decoder predicts the value for each input feature. We define the reconstruction loss functions as the mean square error (MSE) between the reconstructed and original values on (i) the re-masked set Tremask and (ii) unmasked set Zunmask. We empirically experiment with different reconstruction loss functions (e.g., only the re-met re-masked set or both re-masked and unmasked sets).
Putting everything together, Algorithm| I sketches the implementation of REMASKER.
Algorithm 1 REMASKER
Inpot: (( mn ()))) incomplete dataset; remask: re-masking function: fo. da: encoder and decoder; max_cpoch: training
epochs; E reconstruction loss
Oulput: D- = (())))))2r imputed dataset
reached do 2: for (x. m) E D do
1: while max_ mex_h
8 fitting phese
3: * * &mmmm t- remask(E,, m):
b remasking b If encoding unmasked values
4: z fo (Rmmmm))
5:
pad with mask lokens;
6: end for
7: update 9, by VE(do (12)), {8 (mmam))
minimizing reconstruction
8:
end while
9:
for (x. m) D do
// imputation phase
10:
2 - (8m))
b- If caceding observed valuos
11:
pad z musk tokens
12: X.. <- (x))
b // predicting missing values
13:
Km. Nuni
14:
end for
15:
return D = (x).
4

---

Published as a conference paper at ICLR 2024
4 EVALUATION
We empirically evaluate REMASKER in various scenarios using benchmark datasets. Our experiments are designed to answer the following key questions: (i) Does REMASKER work? We compare REMASKER with a variety of state-of-the-art imputers in terms of imputation quality. (ii) How does it work? We conduct an ablation study to assess the contribution of each component of REMASKER to its performance. (iti) What is the best way of using REMASKER? We explore the standalone imputer well one component of an ensemble imputer to
of REMASKER understand its best practice.
Datasets. For reproducibility and comparability, similar the prior work (Yoon al. et al. 2022), we use 12 real-world datasets from the UCI Machine Learning repository (Dua & Graff
2019 Jarrett
2017 with their characteristics deferred to S A..
Missing mechanisms. Following prior work (Yoon et al. 2019 Jarrett et al. 2022, Mattei & Frellsen
2018 Hastie et al.] 2015), we mainly focus on the MCAR and MAR settings. In MCAR, the mask vector of each input is realized following a Bernoulli random variable with a fixed mean. In MAR, with random subset of features fixed to be observable, the remaining features are masked using logistic model. For completeness, we also simulate one particular MNAR setting in experiments: the input features of MAR are further masked following a Bernoulli random variable with fixed mean. We use Hyperlmpute (Jarrett et al.) 2022) to simulate the above missingness mechanisms.
Baselines. We compare REMASKER with 13 state-of- sfate-of-the-art the-art imputation methods: Hyperlmpute (Jr rett et al. 2022). MIWAE (Mattei & Frellsen 2018), EM (García-Laancci) et al. 2010). GAIN (Yoon
et al. 2019), ICE, MICE, MIRACLE (Kyono et al. 2021), MissForest (Stekhoven & Buhlmann 2012), Mean (Hawthome & Ellioti 2005). Median, Frequent, Sinkhorn (Muzellec et al. 2020)), and SoftImpute (Hastic et al. 2015), with details deferred to A/A.2
Metrics. For each imputation method, we evaluate the imputation fidelity and utility by comparing its imputed data with the ground-truth data. In terms fidelity, we mainly use two metries: root mean square error (RMSE) to measure how the individual imputed values match the ground-truth data, and the Wasserstein distance (WD) to measure how the imputed distribution matches the ground-truth distribution. In terms of utility, we use area under the receiver operating characteristic curve (AUROC) as the metric on applicable datasets (i.e., ones associated with classification tasks). In the case of multi-class classification, we use the one versus rest (OvR) setting. To be fair, we use logistic regression as the predictive model across all the cases.
ReMesker
Hyperfmqut
GAIN
ICF
MICE
MissForest
Sinchoon
MINAE
MIRACLE
Mean
Median
Frequentt
Softimpuue
0.6
0.2
0.0
bike
califoenia
elimate
compression
disbetes
abesity
spam
0.4
0.2
1
0.0
bike
california
climate
compression
credit
disbetes
letter
abesity
raisin
spam
wine
yocct
Ideeal AURROC
0.9
0.8
abesity
rasin
spam
wine
Benchmark Dataser
Figure Overall performance of REM EEMA and baseline imputers on 12 benchmark datasets under MAR with 0.3 missingness ratio. The results are shown as the mean and standard deviation of RMSE, WD, and AUROC scores (AUROC is only applicable to datasets with classification tasks). Note that REMASKER outperforms all the baseline imputers under at least one metric across all the datasets.
5

---

Published as a conference paper at ICLR 2024
-+- ReNasker
Repedmpuue
ICE
Misfores
Sirkhorn
MWWA
GAN
0.1
0.2
RADSE
0.15
0.15
0.2
0.1
0.65
0.1
1000
13000
0.125
0.3
20.075
0.1
0.7
0.05
0.05
0.025
1000
7000
13000
20000
0.5
0.90
0.79
0.7
AUROC
0.6
0.5
1000
7000
13000
20000
10
1.2
16
0.1
0.5 0.7
Misiingess
Figure 3: Sensitivity analysis of REMASKER on the letter dataset under the MAR setting. The scores measured with respect to (a) the dataset size, (b) the number features, and (c) missingness ratio. The default setting is as follows: dataset size 20,000, number of features = 16, and missingness ratio = 0.3.
4.1 OVERALL PERFORMANCE
We evaluate REMASKER and baseline imputers on the benchmark datasets under the MAR setting with 0.3 missingness ratio, with results summarized in Figure Z Observe that REMASKER consis- tently outperforms all the baselines in terms of both fidelity (measured by RMSE and WD) and utility (measured by AUROC) across all the datasets. Recall that the benchmark datasets are collected from a variety domains with highly varying characteristics (ef. Ta) the dataset size varies from 308 to 20,060, while the number of features ranges from 7 to 57. Its superior performance across all the datasets demonstrates that REMASKER effectively models the intricate correlation among different features, even if the amount of available data is scarce. The only imputer with performance close to REMASKER is Hyperlmpute (Jarrett et al. 2022). which is an ensemble method that integrates multiple imputation models and automatically selects the most fitting model for each column of the given dataset. This highlights that the modeling capacity of REMASKER'S 's masked autoencoder comparable with ensemble models. In $B.44, we conduct more comprehensive evaluation by simulating all three missingness scenarios (MCAR, MAR, and MNAR) with different missingness ratios. The results show that REMASKER consistently performs better across a range of settings.
4.2 SENSITIVITY ANALYSTS
To assess the factors influencing REMASKER'S performance, we conduct sensitivity analysis by varying the dataset size, the number of features in the dataset, and the missingness ratio under the MAR setting. Figure 3 shows the performance of REMASKER within these experiments against the six closest competitors (Hyperlmpute, ICE, MissForest, GAIN, MIWAE, and Sinkhorn) on the letter dataset. We have the following observations. (a) The performance of REMASKER improves with the size of available data, while its advantage over other imputers (with the exception of Hyperlmpute) grows with the dataset size. (b) The number of features has significant impact on the performance of REMA with its advantage over other imputers increasing steadily with the number of features. This may be explained by that REMASKER relies lcarning the holistic representations of inputs, while including more features contributes to better representation learning. (c) REMASKER is fairly insensitive the missingness ratio. For instance, even with 0.7 missingness ratio, it achieves RMSE below 0.1, suggesting that it effectively fits sparse datasets. $ B.5 we also conduct an evaluation on other datasets with similar observations.
4.3 ABLATION STUDY
We conduct an ablation study of REMASKER to understand the contribution of different components to its performance using the letter dataset. Results other datasets are deferred B

---

Published as a conference paper at ICLR 2024
depth
RMSF.
WD
AUROC
width
RMSF
WD
AUROC
depth
RMSF
WD
AUROC
2
0,0729
0.0263
0.7898
16
0,0902
0.0379
0.7902
2
0.0637
0.0239
0.7887
0.0636
0.0228
0.7903
32
0,0714
0.0289
0.7885
0.0625
0.0236
0.7877
0.0616
0.0219
0.7909
64
0.0616
0.0219
0.7909
0.0644
0.0239
0.7889
0.0611
0.0217
0.7892
128
0.0795
0.0305
0.7845
0.0616
0.0219
0.7909
00.0633
0.0245
0.7879
256
0.1040
0.0403
0.7868
10
0.0637
0.0227
0.7878
(a) Decoder depth
(b) Einhedding width
(c) Encoder depth
Table 1. Ablation study of REMASKER the letter dataset. The default setting as follows: encoder depth
8, decoder depth 6, embedding width = 64, masking ratio 50%, and training epochs 600.
backbone
letter
california
loss
letter
california
RMSE
WD
AUROC
RMSE WD
RMSE WD AUROC
RMSE WD
Transformer
0.0611
0,0217
0.7892
0.0663
0.0172
Tmankt U U.uamank
0,0616 0.0219 0.7909 0.0663 0.0172
Linear
0.1732
0,1604
0.7821
0.1786
0.1329
Imaak+
0.0629 0.0237 0.7890 0.0000 0.0311
Convolutional
0.1694
0.1582
0.7836
0.1715 0.1286
Lynmeok
0.2079 0.1129 0.7901 00.1932 0.1906
Table 2. Performance with different backbones. (note:
Table 3. Performance of REMASKER with reconstruc-
AUROC is inapplicable to the california dataset)
tion loss w/ or w/o unmasked values.
Model design. The encoder and decoder of REMASKER can be flexibly designed. Here, we study the impact of three key parameters, the encoder depth (the number of Transformer blocks), the embedding width (the dimensionality of latent representations), and the decoder depth, with results summarized in Table Ia] Table Tb and Table [c. Observe that the performance with a proper model configuration (encoder depth = 8, decoder depth 8, and embedding width 64). This observation suggests that the model complexity needs to fit the given dataset: it needs to be sufficiently complex to effectively learn the holistic representations of inputs but not overly complex to overfit the dataset. We also compare the performance of REMASKER with different backbone models (i.e., Transformer, linear, and convolutional) with the number of layers and the size of each layer fixed as the default setting. As shown in Table Transformer based REMASKER largely outperforms the other variants, which may be explained by that the self-attention mechanism can effectively capture the intricate inter-feature correlation under limited data Huang et al. (2020).
REMASKER reaches its peak
Reconstruction loss. We define the reconstruction loss as the error between the reconstructed and original values on the re-masked values Imaek+ and the unmasked values Tunmask. We measure the performance of REMASKER under three different settings of the construction loss: (i) Imaak+ U Zummask. (ii) Imaak+ only, and (iii) Zunmask only on the letter and call fornia datasets, with results shown in Table Observe that using the reconstruction of unmasked values only insufficient and yet including the reconstruction loss of unmasked values improves the performance, which is especially the case on the cal fornia dataset. This finding is different from the vision domain in which computing the loss on unmasked image patches reduces accuracy (He et al.) 2022). We hypothesize that this difference is explained as follows. Unlike conventional MAE, due to the naturally missing values in tabular data, relying on re-masked values provides limited supervisory signals. Moreover, while images are signals with heavy spatial redundancy (i.e., a missing patch can be recovered from its neighboring patches), tabular data tends to be highly semantic and information- dense. Thus, including the construction loss of unmasked values improves the model training.
4.4 PRACTICE OF REMASKER
Training regime. The ablation study above by default uses 600 training epochs. Figure a) shows the impact of training epochs, in which we vary the training epochs from 100 to 1,600 and measure the performance of REMASKER on the letter dataset. Observe that the imputation performance
x10#
8.D XT0*
x10=1
8.0
2,6
RMSE
WD
AUROC
Reconstruction Less
3.0
7.0
2.4
7.9
2.5
2.0
2.0
6.0
2.2
7.8
10 20
1.5
1.0
5.0
100
200
400
800
1600
100
200
400
800
1600
100
200
400
800
1600
200
400
600
(al Max Training Epochs
(b) Training Epochs
Figure 4: (a) REMASKER performance with respect to the maximum number of training epachs; (b) Convergence of REMASKER'S 's reconstruction loss. Performed on letter under MAR with 0.3 missingness ratio.
7

---

Published as a conference paper at ICLR 2024
letter
1 fornia
latter
fornia
masking ratio
base imputer
RMSE
WD
AUROC
RMSE
WD
RMSE
WD
AUROC
RMSE
0.0668
0.0215
0.0789
0.0888
0.0230
default 0.0564 0.0215
0.7899
0.0722 0.0134
0.0562
0.0207
0.7897
0.0654
0.0151
REMASKER 0.0554 0.0212 0.7935
0.0702 0.0115
0.5
0.0554
0.0212
0.7935
0.0663
0.0172
Table S. REMASKER as the base imputer within Hy-
0.7
0.0906
0.0366
0.7878
0,1320
0.0650
perlmpute. The results evaluated letter and california under 0.3 MAR.
Table 4. Performance with varying masking ratio. The
results are evaluated on letter california under MAR with 0.3 missingness ratio.
improves (as RMSE and WD decrease and AUROC increases) steadily with longer training and does not fully saturate even at 1,600 epochs. However, for efficient training, it is often acceptable to terminate earlier (e.g., 600 epochs) with sufficient imputation performance. To further validate the trainability of REMASKER, with the maximum number of training epochs fixed at 600 (which affects the learning rate scheduler), we measure the reconstruction loss as a function the training epochs. As shown in Figure bb), the loss quickly converges plateau within about 100 epochs and steadily decreases after that, demonstrating the trainability of REMASKER.
Masking ratio. The masking ratio controls the number of re-masked values (after excluding missing values). Table shows its impact on REMASKER. Observe that the optimal ratio differs across different datasets, which may be explained by the varying number of features of different datasets (16 versus in letter and california) a)) Intuitively, a larger number of features affords higher masking ratio to balance (i) encouraging the model to learn iiissingnessinvariann representations and (ii) having sufficient supervisory signals to facilitate the training.
Standalone vs. ensemble. Besides using REMASKER as a standalone imputer, we explore its use as a base imputer within the ensemble imputation framework of Hyperlmpute, with results summarized in Ta It is observed that compared with the default setting (with mean substitution as the base imputer), using REMASKER as the base imputer improves the imputation performance, suggesting another effective way of operating REMASKER.
5 DISCUSSION
Q1: What is REMASKER learning? The empirical evaluation above shows REMASKER'S 's superior performance in imputing missing values of tabular data. Next, we provide theoretical justification for its effectiveness. By extending the siamese form of MAE (Kong& & Zhang, 2022). we show that REMASKER encourages learning missinguess--nvarioon representations of input data, which requires a holistic understanding of the data even in the presence of missing values.
Let fa(+) and (-)-) respectively be the encoder and decoder. For given input x, mask m, and re-mask m', the reconstruction loss of REMASKER training is given by (here we focus on the reconstruction of re- masked values):
f(x, m. m') [[d(( O m m')) о (1 m')
x
xx
(1 m ml|
where @ denotes element multiplication. Let m + & m ② m' and m - m ② (1 m') Eq
can be simplified as: m(x,m*, m .=) [f @ m o m x O m 12 As the embedding dimensionality is typically much larger than the number of features, it is possible autoencoder lossless. In other words, for a given encoder fo(.), there exists a decoder der (-), such as:
make the
that day (fe (x C m-)) @ m x m We can further re-write Eq
E(x, m", mn = =I[do(Ja(x0) m @m m da- (fo(x 1))m 10ml
(5)
s.t. 0- mig Bx (folx' @ mmmi Am -x ② m
We define new distance metric A (2, [dd)(z) - dgr @ m 11. Then, Eq 4 is reformulated as:
f(x m = (folx fo m
(6)
0- arg min Exe ||der (Jo(x' m m m
Note that optimizing Eq (6) essentially minimizes the difference between x's representations under m* and m (with respect to the decoder). As m and m mask out different values, this formulation promotes learning representations insensitive to missing values.
8

---

Published as a conference paper at ICLR 2024
To validate the analysis above, we empirically mea- sure the CKA similarity (Komblith et al.) (019) between the latent representations (i.e, the output of REM ASKER'S 's encoder) of complete inputs and inputs with missing values, with results shown in Figure 5 Observe that the CKA under different missingness ratios all steadily increase with the training length, indicating that REMASKER tends to invariant representations, which may explain its imputation effectiveness.
1.0
0.9
Smillge
0.8
Misingnos
0.5
learn missingness ienssigges--nvaariant
0.5
Training lpochs
400
600
Figure CKA similarity between the representa- tions of complete and incomplete inputs (with the number of missing values controlled by the missing- ness ratio). The tested model is trained on letter under the MAR setting with 0.5 masking ratio.
Note that while the idea of masking observable val- ues is also explored in prior work (Condara & Wang 2018), its adaptation in the context of MAE admits more elegant imputation solution: it does not re- quire initializing missing or re-masked values in the input while representing them with the same trainable tokens in the latent space encourages learning the distribution of re re-masked values close to the distribution of missing values.
Q2: How is REMASKER'S performance influenced by the missingness mechanism? As revealed above, REMASKER encourages learning representations invariant to re-masked values and then leverages such representations impute missing values. Thus, while REMASKER takes effect long as the distributions of re-masked and missing values are correlated, its performance may vary with the disparity between the two distributions. In MCAR, both re-masked and missing values follow similar distributions (i.e., evenly distributed across different features). In MAR, as the observable features are fixed, the re-masked values are more likely to be selected from the observable features, which biases the representation learning toward the observable features. This bias is reflected in our experimental results, in which REMARRER tends to perform better under MCAR. MNAR, as the missingness depends on the missing values themselves, it causes even more disparity between the distributions of re- masked and missing values, which also explains REMASKER'S ' better performance under MAR and MCAR compared with MNAR.
Q3: Why is Transformer effective for tabular data imputation? While the theoretical understand- ing of the Transformer architecture is still limited, 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 the latent representations of different tokens, which makes Transformer robust against severe occlusions perturbations (Naseer et al. 202T). In the same vein, in the context of tabular data imputation, MSA benefits learning representations that are invariant to missing values.
Q4: What are REMASKER'S 's limitations? One kind of bias we observe is that in some cases (e.g., the cl ilmatte dataset in Figure REMASKER outperforms alternative methods in terms of RMSE but underperforms terms of WD. One explanation is that REMASKER is trained to optimize the reconstruction loss measured by MSE. Thus, it is biased towards re- re-constructing individual missing values rather than the distributions of missing values. Also, as REMASKER solely focuses on the imputation task itself, if the downstream tasks are known, its performance may not be optimal. In such contexts, it is essential to optimize REMASKER with respect to concrete tasks such as uncertainty estimation (van Buuren & Groothuis-Oudsshoom 2011 Gondara & Wang 2018).
6 CONCLUSION
In this paper, we conduct a pilot study exploring the masked autoencoding approach for tabular data imputation. We present REMASKER, a novel imputation method that learns missingness-invant intiaaaannt representations of tabular data and effectively imputes missing values under various scenarios. With extensive evaluation on benchmark datasets, we show that REMASKER outperforms state-of-the-an art methods in terms of both imputation utility and fidelity. Our findings indicate that masked tabular modeling represents a promising direction for future research on tabular data imputation.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The authors Ting Wang and Tianyu Du are partially supported by the National Science Foundation under Grant No. 1951729, 1953813, 2119331, and 2212323. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.
REFERENCES
Sercan O. Arik and Tomas Pfister. TabNet: Attentive Interpretable Tabular Learning. In Proceedings of AAAI Conference on Artificial Intelligence (AAAI), 2021.
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. BEiT: BERT Pre- Trainiri of Image Transform- ers. In Proceedings of International Conference on Learning Representations (ICLR), 2022.
Jacob Devlin, Ming-Was Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of Annual Conference of the North American Chapter of the Association for Computational Linguistics (NACCL), 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv e-prints, 2020.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URL http:://rcchi// ics. uci. edu/ml
Pedro J. Gareía-Laencinn José-Luis Sancho-Gómez, and Aníbal R. Figueiras Vidal. Pattern Clas- sification with Missing Data: A Review. Neural Computing and Applications, 19(2):263-282, 2010.
Lovedeep Gondara and Ke Wang. MIDA: Multiple Imputation using Denoising Autoencoders. In Proceedings of the Pacific-Asia Conference on Knowledge Discovery and Data Mining (PAKDD), 2018.
Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative Adversarial Networks. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2014.
Trevor Hastie, Rahul Mazumder, Jason D. Lee, and Reza Zadeh. Matrix Completion and Low-Rank
SVD via Fast Alternating Least Squares. J. Mach. Learn. Res., (6(1):336-3342, 2015.
G. Hawthorne and Elliott. Imputing sectionattional missing data: comparison of common tech- niques. Aust N ZJ Psychiatry. 39(7):583 2005.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked Autoencoders Are Scalable Vision Learners. In Proceedings of IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Xin Huang, Ashish Khetan, Milan Cvitkovic, and Zohar Karnin. TabTransformer: Tabular Data Modeling Using Contextual Embeddings. ArXiv e-prints, 2020.
Drew A. Hudson and C. Lawrence Zitnick. Generative Adversarial Transformers. In Proceedings of IEEE Conference on Machine Learning (ICML), 2021.
Daniel Jarrett, Bogdan Cebere, Tennison Liu, Alicia Curth, and Mihaela van der Schaar. Hyperlmpute: Generalized Iterative Imputation with Automatic Model Selection. In Proceedings of IEEE Conference on Machine Learning (ICML), 2022.
Yifan Jiang, Shiyu Chang, and Zhangyang Wang. TransGAN: Two Pure Transformers Can Make One Strong GAN, and That Can Scale Up. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2021.
10

---

Published as a conference paper at ICLR 2024
Xiangwen Kong and Xiangyu Zhang. Understanding Masked Image Modeling via Learning Occlusion Invariant Feature. arXiv e-prints. 2022.
Simon Korblith,, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of Neural Network Representations Revisited. In Proceedings of IEEE Conference Machine Learning (ICML), 2019.
Trent Kyono, Yao Zhang, Alexis Bellot, and Mihaela van der Schaar. MIRACLE: Causally-AAey Awwaee Imputation via Learning Missing Data Mechanisms. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2021.
Chao Ma and Cheng Zhang. Identifiable Generative Models for Missing Not at Random Data Imputation. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2021.
Kushal Majmundar, Sachin Goyal, Praneeth Netrapalli, and Prateek Jain. MET: Masked Encoding for Tabular Data. ArXiv e-prints, 2022.
Fferre-Alexandr Mattei and Jes Frellsen. MIWAE: Deep Generative Modelling and Imputation of Incomplete Data. In Proceedings of IEEE Conference on Machine Learning (ICML), 2018.
Boris Muzellee, Julie Josse, Claire Boyer, and Marco Cuturi. Missing Data Imputation Using Optimal Transport. In Proceedings of IEEE Conference Machine Learning (ICML), 2020.
Muzammal Naseer, Kanchana Ranasinghe, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Yang. Intriguing Properties of Vision Transformers. In Proceedings of Advances in Neural Information Processing Systems (NeurlPS), 2021.
Alfredo Nazabal, Pablo M. Olmos, Zoubin Ghahramani, and Isabel Valera. Handling Incomplete Heterogeneous Data Using VAEs. Pattern Recognition, 107:107501, 2020.
Namuk Park and Songkuk Kim. How Do Vision Transformers Work? In Proceedings of International Conference on Learning Representations (ICLR), 2022.
Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A. Efros. Context Encoders: Feature Learning by Inpainting. Proceedings of IEEE Conference on Computer
Vision and Puttern Recognition (CVPR), 2016.
Marein Przewięžlikowski, Marek Smieja, Lukasz Struski, and Jacek Tabor. MisConv: Convolutional Neural Networks for Missing Data. ArXiv e-prints, 2021.
Alec Radford and Karthik Narasimhan. Improving Language Understanding by Generative Pre- Training. ArXiv e-prints, 2018
Gowthami Somepalli, Micah Goldblum, Avi Schwarzschild, C. Bayan Bruss, and Tom Goldstein. SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre- Training. ArXiv e-prints, 2021.
D. J. Stekhoven and P. Buhlmann. MissFForest-onn-aaameett missing value imputation for mixed- type data. Bioinformaticx, 28(1):112-8, 2012.
Stef van Buuren and Karin Soatbhui--OUdhhoom.. mice: Multivariate Imputation by Chained Equations in R. Journal of Statistical Software, 45(3):1-67, 2011.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2017.
Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol. Extracting and Composing Robust Features with Denoising Autoencoders. In Proceedings of IEEE Conference Machine Learning (ICML), 2008
Jinsung Yoon, James Jordon, and Mihaela der Schaar. GAIN: Missing Data Imputation using Generative Adversarial Nets. In Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2019.
11

---

Published as a conference paper at ICLR 2024
Seongwook Yoon and Sanghoon Sull. GAMIN: Generative Adversarial Multiple Imputation Network for Highly Missing Data. Proceedings of IEEE Conference on Computer Vision and Pattern Recognition 2020.
Bowen Zhang, Shuyang Gu, Bo Zhang, Jianmin Bao, Dong Chen, Fang Wen, Yong Wang, and Baining Guo. StyleSwin: Transformer-basee GAN for High-resolution Image Generation. In Proceedings of IEEE Conference Computer Vision and Pattern Recognition (CVPR), 2021.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Towards Deeper Understanding of Variational Autoencoding Models. In Proceedings of ACM SIGIR International Conference on Theory of Information Retrieval (ICTIR), 2022.
A EXPERIMENTAL DETAILS
A.1 DATASETS
Table 6 summarizes the characteristics of the datasets used in our experiments.
Experiment Dataset
Dataset Size
# Features
(California) Housing
20,640
(Climate) Model Simulation Crashes
540
18
Concrete (Compressive) Strength
1,030
(Diabetes)
442
Estimation of (Obesity) Levels
2,111
(Credit) Approval
690
15
(Wine) Quality
1,599
12
(Raisin)
900
(Spam) Base
4,601
57
(Bike) Sharing Demand
8,760
14
(Letter) Recognition
20,000
16
(Yacht) Hydrodynamies
308
Table 6. Characteristics of the datasets used in the experiments.
A.2 BASELINES
We compare REMASKER with 13 state-of-the-arr imputation methods: Hyperlmpute (Jarrett et al. 2022), hybrid imputer that performs iterative imputation with automatic model selection; MI- WAE Mattei & Frellsen 2018), an autoencoder model that fits missing data by optimizing a variational bound; EM (García-Laencina al. 2010), an iterative imputer based on expectation- maximization optimization; GAIN (Yoon et al. 2019). generative adversarial imputation network that trains the discriminator classify the generator's output in an element-wiss manner; ICE, an iterative imputer based on regularized linear regression; MICE, an ICE-like, iterative imputer based on Bayesian ridge regression; MIRACLE (Kyono et al. 2021). an iterative imputer that refines imputation of a baseline by simultaneously modeling the missingness generating mechanism; MissFor- est (Stekhoven & Buhlmann] 2012). an iterative imputer based on random forests; Mean (Hawthorne & Ellioti 2005). Median, and Frequent, which impute missing values using column- wine uncondi- tional mean, median, and the most frequent values, respectively; Sinkhorn (Muzellec et al. 2020), an imputer trained through the optimal transport metrics of Sinkhorn divergences; and Softlmpute (Hastie et al. 2013), which performs imputation through soft-thresholded singular value decomposition.
A.3 PARAMETER SETTING
The default parameter setting of REMASKER is listed in Table
12

---

Published as a conference paper at ICLR 2024
model
parameter
setting
optimizer
Adam
initial learning
le-3
LR scheduler
cosine annealing
global
gradient clipping threshold
5.0
training epochs
600
batch size
64
masking ratio
Transformer block
encoder
embedding width
numher of heads
Transformer block
decoder
embedding width 64
numher of heads 4
Table 7. Default parameter setting of REMASKER.
B ADDITIONAL RESULTS
B. EXECUTION EFFICIENCY
One design objective of MAE (e.g., masking out over 75% of an image) to scale up the training on large-scale image datasets without sacrificing its performance. This benefit is naturally inherited by REMASKER. We measure the running time (including both training and imputation) of REMASKER (under the default setting of T)ble datasets with over 1,000 records. For comparison, also measure the running time of Hyperlmpute under the same setting, with results summarized in Table
compression
wine
spam
credit
bike
obesity
california
Hyperlmpute 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
Table 8. Total running time of Hyperlmpute and REMASKER under the default setting.
Observe that compared with Hyperlmpute, REMASKER is not only more scalable but also less sensitive to data size.
B.2 GENERALIZATION TO NEW DATA
To evaluate REMASKER'S 's ability to generate new data, we run the following evaluation. We halve each dataset into two subsets 'D and D'' (both with 30% missing value missing under MAR). We train REMASKER on D and then apply REMASKER to impute the missing values in D'. The results are shown in Table Note that compared with the setting where REMASKER is trained and applied on the same dataset (Figure REMASKER performs comparably well under the setting that applied on new dataset, indicating its generalizability.
compression
wine
spam
credit
bike
obesity
california
RMSE WD 1
0.117
0.120
0.052
0,251
0.108
0.197
0.068
0.047
0.033
0.009
0.131
0.035
0.059
0.020
AUROC
N/A
0.784
0.098
0.921
0.968
0.937
N/A
Table 9. Performance of REMASKER in imputing new datasets.
B.3 CKA MEASURES OF INPUTS UNDER MASKS m AND m
To further validate the analysis in $ we empirically 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 Observe that while m and m represent different missing mechanisms, the CKA measures between the inputs under m and m still steadily increase with the training epochs, which empirically corroborates the analysis $ that REMASKER tends to learn masking-invarian representations.
13

---

Published as a conference paper at ICLR 2024
0.90
0.85
CM Similara
0.70
0.65
0.60
0.40
Training Epechs
Figure 6: CKA similarity between the representations of inputs under masks m² and (with the number of missing values controlled by the missing ratio). The tested model is trained letter under the MAR setting with 0.5 masking ratio.
depth
RMSE
WD
width
RMSE
WD
depth
RMSE
WD
2
0.0783
0.0230
16
0.0678
0.0213
2
0.0886
0.0334
0.0663
0.0172
32
0.0663
0.0172
0.0738
0.0203
0.0821
0.0225
64
0.0974
0.0322
0.0663
0.0172
8
0.0834
0.0244
128
0.1125
0.0388
0.0878
0.0322
10
0.0726
0.0196
256
0.0877
0.0324
10
0.0776
0.0239
(a) Decoder depth
(b) Embedding width
(c) Encoder depth
Table 10. Ablation study of REMASKER the
dataset. The default setting follows: encoder depth =6, decoder depth = 4, embedding width = 32, masking ratio 50%, training epochs 600.
B.4 OVERALL PERFORMANCE
and 12 respectively show the imputation performance of REMASKER and 8 basclines on 12 benchmark datasets under the MAR, MCAR, and MNAR scenarios with the missingness ratio varying from 0.1 to 0.7. Observed that REMASKER performs on with or outperforms almost all the baselines across wide range of settings. Note that the MIRACLE imputer does not work on the Compress i on dataset and the Raisin dataset under some settings, of which the results are not reported. Given that both Compress inn and Rais iisinn are relatively small datasets, one possible explanation is that MIRACLE requires sufficient amount of data to train the model.
Figure 7 g
B.5 SENSITIVITY ANALYSIS
Figure 13 and 14 show the sensitivity analysis of REMASKER and other 6 baselines the folifornia dataset under the MAR, MCAR, and MNAR settings. The observed trends are generally similar that in Figure which further demonstrates the observations we made in $ about how different factors may impact REMASKER'S 's imputation performance.
B.6 ABLATION STUDY
The ablation study of REMASKER on the ca fornia a dataset is shown in Table TO Observed that the performance of REMASKER reaches its peak with encoder depth = 6, decoder depth 4, and embedding width 32.
B.7 TRAINING REGIME
Figure 15 shows the imputation performance of REMASKER the cali california dataset when the training Length varies from 100 to 1,600 epochs. Figure 16 plots the convergence of reconstruction loss in REMASKER, showing trend similar to Figure b).
14

---

Published as a conference paper at ICLR 2024
0.4
Farsed
LMA
(b) MAR-0.1
Reliyder
Hever
0.6
2-
1.
ALDOC
(d) MAR-0.3
Figure 7: Overall performance of REMASKER and baselines on benchmark datasets under MAR scenario with 0.1 and 0.3 missingness ratio. The results are shown as the mean and standard deviation of RMSE, WD, and AUROC scores (AUROC only applicable to datasets with classification tasks).
B.8 OPTIMAL MASKING RATIO VS PAIRWISE MUTUAL INFORMATION
We report below the average pairwise mutual information (PMI) between different features and the performance of ReMasker under varying masking ratios over each dataset. It is observed that in general the PMI value is positively correlated with the optimal masking ratio. This may be explained by that with higher feature correlation (or more information redundancy), larger masking ratio often leads to more effective representation learning, which corroborates the existing studies on MAE He et al. (2022).
15

---

Published as a conference paper at ICLR 2024
gan
(a) MAR-0.5
Partenst
river
2-
(b) MAR-0.7
Figure 8: Overall performance of REMASKER and 8 baselines on 12 benchmark datasets under MAR scenario with 0.5 and 0.7 missingness ratio. The results shown the mean and standard deviation of RMSE, WD. and AUROC scores (AUROC only applicable datasets with classification tasks).
16

---

Published as a conference paper at ICLR 2024
(a) MCAR-0.1
sindor
2-
A300C
(b) MCAR-0.3
Figure 9: Overall performance of REMASKER and 8 baselines on 12 benchmark datasets under MCAR scenario with 0.1 and 0.3 missingness ratio. The results are shown as the mean and standard deviation of RMSE, WD. and AUROC scores (AUROC is only applicable datasets with classification tasks).
17

---

Published as a conference paper at ICLR 2024
aindorn
PINCC
(a) MCAR-0.5
Partent
PiNCe
g
(b) MCAR-0.7
Figure 10: Overall performance of REMASKER and baselines on 12 benchmark datascts under MCAR scenario with 0.5 and 0.7 missingness ratio. The results shown as the mean and standard deviation of RMSE, WD, and AUROC scores (AUROC only applicable datasets with classification tasks).
18

---

Published as a conference paper at ICLR 2024
gal
(a) MNAR-0.1
2-
(b) MNAR-0.3
Figure 11: Overall performance of REMASKER and basclines on 12 benchmark datasets under MNAR with 0.1 and 0.3 missingness ratio. The results are shown as the mean and standard deviation of RMSE, WD, and AUROC scores (AUROC only applicable datasets with classification tasks).
19

---

Published as a conference paper at ICLR 2024
aindorn
PiNCC
(a) MN
Partent
Pince
(b) MNAR-0.7
Figure 12: Overall performance of REMASKER and basclines on 12 benchmark datasets under MNAR with 0.5 and 0.7 missingness ratio. The results are shown as the mean and standard deviation of RMSE, WD, and AUROC scores (AUROC is only applicable datasets with classification tasks).
20

---

Published as a conference paper at ICLR 2024
ReMasker
hyperimpute
Ice
missforest
snihtom
miebe
gain
0:225
0.35
0.225
0.200
0.30
0.200
0.175
0.25
0.150
RM0E
0.150
0.20
0.325
0.100
0.300
0.10
0.075
0.05
0.050
0.050
1000
7000
13000
20000
0.00
0.025
0.1
0.3
0.5
0.30
0.32
0.10
0.08
0.08
0.06
0.06
0.04
0.02
0.02
0.00
0.00
0.00
1000
7000
13300
20000
0.3
0.5
Observed dara $22
Feature count
Missigness rate
(a) MAR
hyperimpute
misforres
sinchom
+ miwae
gain
0.275
0.30
0.20
0.18
0.25
R055
0.36
0.200
0.20
0.14
0.175
0 150
0.32
0.30
0.100
1000
7000
13000
20000
0.1
0.20
0.18
1.0
0.16
0.8
0.14
0.14
0.6
2 0.12
0.10
0.30
0.08
0.08
0.2
0.06
0.04
0.34
0.02
Observed size
Feature
Miss
(b) MNAR
Figure 13: Sensitivity analysis of REMASKER the cali fornia dataset under MAR and MNAR scenarios. The results are shown terms RMSE and WD, with the scores measured with respect to (a) the dataset size, (b) the number of features, and (c) the missingness ratio. The default setting is as follows: dataset size = 20,000, number of features = 9, and missingness ratio = 0.3.
21

---

Published as a conference paper at ICLR 2024
typerimpute
missforest
sinchom
miwae
gain
0.20
0.275
0.250
0,250
0.18
0 225
0.225
2.0
0.200
0.175
0.150
0.150
0.225
0.125
0.10
0.100
0.075
1000
7000
13000
20000
0.1
0.3
0.5
0.7
0.18
0.14
0.6
0 0.12
0.10
0.08
0.4
0.08
0.06
gness
(a) MCAR
Figure 14: Sensitivity analysis of REMASKER on california dataset under the MCAR setting. The
results are shown in terms of RMSE and WD, with the scores measured with respect to (a) the dataset size, (b) the number of features, and (c) the missingness ratio. The default setting is as follows: dataset size 20,000, number of features = 9, and missingness ratio 0.3.
0.0700
0.0675
RMSE
0.0630
0.020
0.0625
0.015
200
400
1600
0.030
100
200
400
1500
epoch
Max epoch
(a) RMSE
(b)
Figure REMASKER performance with respect the number of training epochs on the california dataset under MAR with 0.3 missingness ratio.
0.40
0.35
0.30
Loss
3.25
0.10
100
200
300
100
500
600
Figure 16: Convergence of REMASKER's fitting california under MAR with 0.3 missingness ratio.
22

---

Published as a conference paper at ICLR 2024
climate
compressive ve
diabetes
obesity
masking ratio
RMSE
WD
AUROC
RMSE
WD
RMSE
WD
RMSE
WD
AUROC
0.1
0.2825
0,2315
0.9544
0.1739
0.0544
0.1827
0.0691
0.2078
0.0588
0,9411
0.3
0.2832
0,2336
0.953
0.1129
0.0439
0.1432
0.0614
0.1943
0.0617
0.9422
0.2823
0.2481
0.9543
0.1046
0.0306
0.1496
0.0694
0.193
0.0743
0.9417
0.7
0.2837
0.2544
0.9533
0.1635
0.0825
0.1542
0.0784
0.2129
0.1093
0.9408
PMI
0.023
0.7829
0.1106
0.147
masking ratio
credit
wine
raisin
spam
RMSE
WD
AUROC
RMSE
AUROC
RMSE
WD
AUROC
RMSE
WD
AUROC
0.2298
0.1075
0.9255
0.1171
0.0255
0.7954
0.1497
0.0727
0.9931
0.0451
0.0097
0.9998
0.3
0.2277
0.0993
0.9279
0.1122
0.0303
0,7949
0.0917
0.0431
0.9929
0.0434
0.0084
0.9998
0.5
0.245
0.1243
0.9247
0,1151
0.0395
0.7955
0.0904
0.0347
0.9943
0.0435
0.0099
0.9998
0.7
0.3175
0.2176
0.914
0.1401
0.0757
0.7934
0.1172
0.0727
0.9946
0.0507
0.0197
0.9998
PMI
0.069
0.1612
0.5246
0.0771
bike
yacht
letter
californis
masking ratio
RMSE
WD
AUROC
RMSE
WD
RMSE
WD
AUROC
RMSE
WD
0.1
0.1068
0.0244
0.97
0.2882
0.1339
0.0668
0.0215
0.0789
0.0888
0.023
0.3
0.096
0.0226
0.9736
0.2185
0.0812
0.0562
0.0207
0.7897
0.0654
0.0151
0.5
0.1015
0.0277
0.9707
0.2174
0.0996
0.0554
0.0212
0.7935
0.0663
0.0172
0.7
0.1306
0.0561
0.9687
0.3188
0.1948
0.0906
0,0366
0.7878
0.132
0.065
PMI
0.1614
0.8229
0.1421
0.1576
Table ASKER performance with respect to masking ratio. The results are evaluated under MAR with 0.3 missingness ratio.
23


Reviews:
[{"review_id": "0t3yfhKnWD", "reviewer": "Reviewer_w2Hd", "summary": "The paper Remasker: Imputing Tabular Data with Masked Autoencoding presents a novel algorithm for missing data imputation in tabular datasets using the masked autoencoding framework. The algorithm operates in two stages: a first stage where extra variables are masked and the autoencoder is optimized to reconstruct the re-masked variables, and a second stage where the trained model is used to predict the missing values. Authors evaluate the method on a wide range of state-of-the-art methods in several tabular datasets, showing competitive performance.", "questions": "Some questions I have are the following:\n\n1. 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?\n2. 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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written and grammatically correct with no evident typos. Every section is very well outlined with clear paragraphs describing the content of the manuscript. \n\nThe proposed method, namely the Remasker method, is easy to understand for experts and non-experts in the field. The idea is intuitive yet effective, showing promising results on how we can leverage the promises of transformer-based architecture for missing data imputation. \n\nThe evaluation of the method is quite thorough, where standard tabular datasets are tested under different state-of-the-art methods. This reveals that the authors are aware of the most recent advances in missing data imputation. Besides, different metrics and ablation studies are performed to motivate the performance of the method further. Another precious point is that the authors do comment on the limitations of their proposed method, which is something that is sometimes overlooked.", "weaknesses": "I have a few concerns and questions that I'd like the authors to address:\n\n1. While I understand the primary motivation behind using a transformer-based architecture in this work, I would appreciate a more in-depth discussion of why the authors chose to use masked autoencoders over other approaches, like deep generative models. It's crucial to outline the advantages and disadvantages of this choice. Additionally, it would be beneficial to address questions like the computational cost of training the Remasker method and its ability to generate new data, which is a common concern in the missing data imputation field.\n2. 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.\n3. The manuscript appears to be densely packed in terms of space. Figures and tables are positioned very closely to the text, which can make the paper seem cluttered and condensed. Consider the possibility of making some tables smaller or relocating them to the Appendix (e.g., Table 1), and possibly reducing the size of certain figures (e.g., Figure 2 with fewer datasets). This issue can be addressed in the camera-ready version if the paper is accepted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s89cvPTsU7", "reviewer": "Reviewer_Fz3c", "summary": "The authors propose ReMasker, a method for imputing missing values in tabular data based on masked auto encoders (MAE). Specifically, given a dataset, they mask some of the features and train a model to impute these features. The approach does not require all features to be present in training, so a dataset with missing values can be used for training. The authors provide extensive empirical study and a theoretical justification.", "questions": "I have one minor question:\n- On page 1, you state that discriminative imputers are hindered by the requirement of specifying a functional form. It seems to me that your architecture is also a choice of \"functional form\". Am I wrong?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I think the authors propose an elegant solution to a problem many practitioners encounter. The paper is well written and the empirical results are convincing. \nIn general, for an approach such as the one presented, it is difficult to provide a theoretical analysis. I appreciate the theoretical justification in Section 5, but I would have liked it to have a more prominent position in the paper.", "weaknesses": "The authors fail to mention and discuss multiple previous works. Specifically, transformers have been previously been applied to tabular data (Somepalli et al. 2021, Arık et al. 2020,  Huang et al 2020). Specific to Huang et al, the authors cite this paper for supporting their model choice, but fail to mention it in the related works. The MAE work by Majmundar et al. (2022) is also not discussed.\n\n\nSomepalli et al. 2021: SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training\nArık et al. 2020: TabNet: Attentive Interpretable Tabular Learning\nHuang et al 2020: TabTransformer: Tabular Data Modeling Using Contextual Embeddings\nMajmundar et al. 2022: MET: Masked Encoding for Tabular Data", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vR4UfGFsFT", "reviewer": "Reviewer_YQFA", "summary": "The paper presents an imputation method, ReMasker, that leverages a remasking autoencoder with a Transformer-based architecture. A significant amount of experimentation is included to validate the proposed procedure.", "questions": "Is there a better way to represent the results visually? It is very difficult to read the plots (e.g. Figure 2), and in particular to the confidence intervals.\n\nWhile 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, I believe more justification for this decision is needed.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The proposed method provides a simpler heuristic than competing SOTA methods such as HyperImpute, while displaying similar or slightly improved accuracy. \n\nThe experiments are relatively thorough and show strong results for the MCAR setting in particular.", "weaknesses": "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 I agree this is likely the first work to consider a masked autoencoder with Transformers for imputation specifically. \n\nThe claims that ReMasker generalizes to MNAR settings is not well justified:\n\n* 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.\n\n* 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.\n\nIt is not clear that the results are practically different from HyperImpute, and perhaps other methods. The included bar plots are hard to read, so any assessment is difficult. The relative simplicity of ReMasker compared to HyperImpute is mentioned as one advantage, although the much simpler ICE appears competitive in almost all scenarios.\n\nWhile 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)