Title: M$^3$-Impute: Mask-guided Representation Learning for Missing Value Imputation
Venue: NeurIPS 2024
Accepted: no
Keywords: Missing Value Imputation

Graph Representation Learning

Data Correlations
Abstract: Missing values are a common problem that poses significant challenges to data analysis and machine learning. This problem necessitates the development of an effective imputation method to fill in the missing values accurately, thereby enhancing the overall quality and utility of the datasets. Existing imputation methods, however, fall short of considering the 'missingness' information in the data during initialization and modeling the entangled feature and sample correlations explicitly during the learning process, thus leading to inferior performance. We propose M$^3$-Impute, which aims to leverage the missingness information and such correlations with novel masking schemes. M$^3$-Impute first models the data as a bipartite graph and uses an off-the-shelf graph neural network, equipped with a refined initialization process, to learn node embeddings. They are then optimized through M$^3$-Impute’s novel feature correlation unit (FCU) and sample correlation unit (SCU) that enable explicit consideration of feature and sample correlations for imputation. Experiment results on 15 benchmark datasets under three different missing patterns show the effectiveness of M$^3$-Impute by achieving 13 best and 2 second-best MAE scores on average.

Abstract

Missing values are a common problem that poses significant challenges to data analysis and machine learning. This problem necessitates the development of an effective imputation method to fill in the missing values accurately, thereby enhancing the overall quality and utility of the datasets. Existing imputation methods, however, fall short of considering the 'missingness' information in the data during initialization and modeling the entangled feature and sample correlations explicitly during the learning process, thus leading to inferior performance. We propose M 3 -Impute, which aims to leverage the missingness information and such correlations with novel masking schemes. M 3 -Impute first models the data as a bipartite graph and uses an off-the-shelf graph neural network, equipped with a refined initialization process, to learn node embeddings. They are then optimized through M 3 -Impute's novel feature correlation unit ( FCU ) and sample correlation unit ( SCU ) that enable explicit consideration of feature and sample correlations for imputation. Experiment results on 15 benchmark datasets under three different missing patterns show the effectiveness of M 3 -Impute by achieving 13 best and 2 second-best MAE scores on average.

1 Introduction

Missing values in a dataset are a pervasive issue in real-world data analysis. They arise for various reasons, ranging from the limitations of data collection methods to errors during data transmission and storage. Since many data analysis algorithms cannot directly handle missing values, the most common way to deal with them is to discard the corresponding samples or features with missing values, which would compromise the quality of data analysis. To tackle this problem, missing value imputation algorithms have been proposed to preserve all samples and features by imputing missing values with estimated ones based on the observed values in the dataset, so that the dataset can be analyzed as a complete one without losing any information.

The imputation of missing values usually requires modeling of correlations between different features and samples. Feature-wise correlations help predict missing values from other observed features in the same sample, while sample-wise correlations help predict them in one sample from other similar samples. It is thus important to jointly model the feature-wise and sample-wise correlations in the dataset. In addition, the prediction of missing values also largely depends on the 'missingness' of the data, i.e., whether a certain feature value is observed or not in the dataset. Specifically, the missingness information directly determines which observed feature values can be used for imputation. For example, even if two samples are closely related, it may be less effective to use them for imputation if they have missing values in exactly the same features. It still remains a challenging problem how to jointly model feature-wise and sample-wise correlations with such data missingness.

Among existing methods for missing value imputation, statistical methods [4, 9, 14, 16, 18, 19, 22, 36 28, 30, 31, 37, 43] extract data correlations with statistical models, which are generally not flexible 37

in handling mixed data types and struggles to scale up to large datasets. Learning-based imputation 38 methods [10, 24, 27, 29, 33, 42, 50, 51, 53], instead, take advantage of the strong expressiveness 39 and scalability of machine/deep learning algorithms to model data correlations. However, most of 40 them are still built upon the raw tabular data structure as is, which greatly restricts them from jointly 41 modeling the feature-wise and sample-wise correlations. In light of this, graph-based methods [52, 42 54] have been proposed to model the raw data as a bipartite graph, with samples and features being 43 two different types of nodes. A sample node and a feature node are connected if the feature value 44 is observed in that sample. The missing values are then predicted as the inner product between 45 the embeddings of the corresponding sample and feature nodes. However, this simple prediction 46 does not consider the specific missingness information as mentioned above. For instance, the target 47 feature to impute may have different correlations with features in the samples which have different 48 kinds of missingness; however, the same feature-node embedding is still used for their imputation. 49 A similar issue also arises for sample-node embeddings. 50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

In this work, we address these problems by proposing M 3 -Impute, a mask-guided representation learning method for missing value imputation. The key idea behind M 3 -Impute is to explicitly utilize the data-missingness information as model input with our proposed novel masking schemes so that it can accurately learn feature-wise and sample-wise correlations in the presence of different kinds of data missingness. M 3 -Impute first builds a bipartite graph from the data as used in [52]. In the embedding initialization for graph representation learning, however, we not only use the the relationships between samples and their associated features but also the missingness information so as to initialize the embeddings of samples and features jointly and effectively. We then propose novel feature correlation unit ( FCU ) and sample correlation unit ( SCU ) in M 3 -Impute to explicitly take feature-wise and sample-wise correlations into account for imputation. FCU learns the correlations between the target missing feature and observed features within each sample, which are then further updated via a soft mask on the sample missingness information. SCU then computes the samplewise correlations with another soft mask on the missingness information for each pair of samples that have values to impute. We then integrate the output embeddings of FCU and SCU to estimate the missing values in a dataset. We carry out extensive experiments on 15 open datasets. The results show that M 3 -Impute outperforms state-of-the-art methods in 13 of the 15 datasets on average under three different settings of missing value patterns, achieving up to 11 . 47% improvement in MAE compared to the second-best method.

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

2 Related Work

Statistical methods: These imputation approaches include joint modeling with expectationmaximization (EM) [9, 16, 22], k -nearest neighbors (kNN) [14, 43], and matrix completion [5, 6, 18, 32]. However, joint modeling with EM and matrix completion often lack the flexibility to handle data with mixed modalities, while kNN faces scalability issues due to its high computational complexity. In contrast, M 3 -Impute is scalable and adaptive to different data distributions.

Learning-based methods: Iterative imputation frameworks [1, 2, 15, 20, 23, 24, 35, 41, 44, 45], such as MICE [45] and HyperImpute [23], have been extensively studied. These iterative frameworks apply different imputation methods for each feature and iteratively estimate missing values until convergence. In addition, for deep neural network learners, both generative models [27, 29, 36, 50, 51, 53], such as GAIN [50] and MIWAE [29], and discriminative models [10, 24, 48], such AimNet [48], have also been proposed. However, these methods are built upon raw tabular data structures, which fall short of capturing the complex correlations in features, samples, and their combination [54]. In contrast, M 3 -Impute is based on the bipartite graph modeling of the data, which is more suitable for learning the data correlations for imputation.

Graph neural network-based methods: GNN-based methods [40, 52, 54] are proposed to address the drawbacks mentioned above due to their effectiveness in modeling complex relations between entities. Among them, GRAPE [52] transforms tabular data into a bipartite graph where features are one type of node and samples are the other. A sample node is connected to a feature node only if the corresponding feature value is present. This transformation allows the imputation task to be framed as a link prediction problem, where the inner product of the learned node embeddings is computed as the predicted values. IGRM [54] further enhances the bipartite graph by explicitly introducing linkages between sample nodes to facilitate message propagation between samples. However, these methods do not effectively encode the missingness information of different samples and features into

Γ6ΓDΓPΓSΓOΓHΓıΓZΓLΓVΓHΓΛΓ&ΓRΓUΓUΓHΓOΓDΓWΓLΓRΓQ

Γ6Γ&Γ8

Γ)Γ&Γ8

Γ' ΓD

ΓW

ΓD

Γ)ΓHΓDΓWΓΛΓˇ

Γ)ΓHΓD WΓΛΓ˘

Γ)ΓHΓD WΓΛΓ¯

Γ)ΓHΓD WΓΛΓ˚

Γ´ΓΓ˘Γ˚

Γ1Γ`Γ$

Γ´ΓΓ¸Γø

Γ´ΓΓ˚

Γ1Γ`Γ$

Γ´ΓΓ˘

Γ´ΓΓ¸Γø

Γ´ΓΓˇ Γ´ΓΓ¯ Γ´ΓΓ¯

Γ1Γ`Γ$

Γ*ΓUΓDΓSΓKΓΛΓ1ΓHΓXΓUΓDΓO

Γ1ΓHΓWΓZΓRΓUΓN

ΓffΓHΓΓJΓΓfflΓΛΓ*ΓUΓDΓSΓKΓ6Γ$Γ*Γ(Γfi

Γ0Γ/Γ3

Γ,Γ¯

ΓΓΓ

ΓΓΓ

ΓΓΓ

ΓΓΓ

ΓΓΓ

Γ,ΓQΓLΓWΓDΓOΓLΓ]ΓDΓWΓLΓRΓQ

Γ/Γ(Γ*Γ(Γ1Γ'

Γ(Γ[ΓLΓVΓWΓLΓQΓJΓΛΓ(ΓGΓJΓH

Γ'ΓDΓWΓDΓΛΓ)ΓOΓRΓZ

Γ,ΓQΓLΓWΓLΓDΓOΓLΓ]ΓHΓGΓΛΓ1ΓRΓGΓH

Γ6ΓDΓPΓSΓOΓHΓΛΓ1ΓRΓGΓH

Γ)ΓHΓDΓWΓXΓUΓHΓΛΓ1ΓRΓGΓH

Γ.ΓQΓRΓZΓQΓΛΓ0ΓDΓVΓN

Γ)ΓHΓDΓWΓXΓUΓHΓıΓZΓLΓVΓHΓΛΓ&ΓRΓUΓUΓHΓOΓDΓWΓLΓRΓQ

the imputation process, which can impair their imputation accuracy. In contrast, M 3 -Impute enables 93 explicit modeling of missingness information through novel masking schemes so that feature-wise 94 and sample-wise correlations can be accurately captured in the imputation process. 95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

3 M 3 -Impute

3.1 Overview

We here provide an overview of M 3 -Impute to impute the missing value of feature f for a given sample s , as depicted in Figure 1. Initially, the data matrix with missing values is modeled as an undirected bipartite graph, and the missing value is imputed by predicting the edge weight ˆ e sf of its corresponding missing edge (Section 3.2). M 3 -Impute next employs a GNN model, such as GraphSAGE [17], on the bipartite graph to learn the embeddings of samples and features. These embeddings, along with the known masks of the data matrix (used to indicate which feature values are available in each sample), are then input into our novel feature correlation unit ( FCU ) and sample correlation unit ( SCU ), which shall be explained in Section 3.3 and Section 3.4, to obtain feature-wise and sample-wise correlations, respectively. Finally, M 3 -Impute takes the feature-wise and sample-wise correlations into a multi-layer perceptron (MLP) to predict the missing feature value ˆ e sf (Section 3.5). The whole process, including the embedding generation, is trained in an end-to-end manner.

3.2 Initialization Unit

Let A ∈ R n × m be an n × m matrix that consists of n data samples and m features, where A ij denotes the j -th feature value of the i -th data sample. We introduce an n × m mask matrix M ∈ { 0 , 1 } n × m for A to indicate that the value of A ij is observed when M ij = 1 . In other words, the goal of imputation here is to predict the missing feature values A ij for i and j such that M ij = 0 . We define the masked data matrix D to be D = A ⊙ M , where ⊙ is the Hadamard product, i.e., the element-wise multiplication of two matrices.

̸

As used in recent studies [52, 54], we model the masked data matrix D as a bipartite graph and tackle the missing value imputation problem as a link prediction task on the bipartite graph. Specifically, D is modeled as an undirected bipartite graph G = ( S ∪ F , E ) , where S = { s 1 , s 2 , . . . , s n } is the set of 'sample' nodes and F = { f 1 , f 2 , . . . , f m } is the set of 'feature' nodes. Also, E is the set of edges that only exist between sample node s and feature node f when D sf = 0 , and each edge ( s, f ) ∈ E is associated with edge weight e sf , which is given by e sf = D sf . Then, the missing value imputation problem becomes, for any missing entries in D (where D sf = 0), to predict their corresponding edge weights by developing a learnable mapping F ( · ) , i.e.,

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

The recent studies that use the bipartite graph modeling [52, 54] initialize all sample node embeddings as all-one vectors and feature node embeddings as one-hot vectors, which have a value 1 in the positions representing their respective features and 0's elsewhere. We observe, however, that such an initialization does not effectively utilize the information from the masked data matrix, which leads to inferior imputation accuracy, as shall be demonstrated in Section 4.3. Thus, in M 3 -Impute, we propose to initialize each sample node embedding based on its associated (initial) feature embeddings instead of initializing them separately. While the feature embeddings are randomly initialized, the sample node embeddings are initialized in a way that reflects the embeddings of the features whose values are available in their corresponding samples.

Let h 0 f be the initial embedding of feature f , which is a randomly initialized d -dimensional vector, and define H 0 F = [ h 0 f 1 h 0 f 2 . . . h 0 f m ] ∈ R d × m . Also, let d s ∈ R m be the s -th column vector of D ⊤ , which is a vector of the feature values of sample s , and let m s ∈ R m be its corresponding mask vector, i.e., m s = col s ( M ⊤ ) , where col s ( · ) denotes the s -th column vector of the matrix. We then initialize the embedding h 0 s of each sample node s as follows:

where 1 ∈ R m is an all-one vector, and ϕ ( · ) is an MLP. Note that the term d s + ϵ ( 1 -m s ) indicates a vector that consists of observable feature values of s and some small positive values ϵ in the places where the feature values are unavailable (masked out).

3.3 Feature Correlation Unit

To improve the accuracy of missing value imputation, we aim to fully exploit feature correlations which often appear in the datasets. While the feature correlations are naturally captured by GNNs, we observe that there is still room for improvement. We propose FCU as an integral component of M 3 -Impute to fully exploit the feature correlations.

To impute the missing value of feature f for a given sample s , FCU begins by computing the feature 'context' vector of sample s in the embedding space that reflects the correlations between the target missing feature f and observed features. Let h f ∈ R d be the learned embedding vector of feature f from the GNN, and let H F be the d × m matrix that consists of all the learned feature embedding vectors. We first obtain dot-product similarities between feature f and all the features in the embedding space, i.e., H ⊤ F h f . We then mask out the similarity values with respect to nonobserved features in sample s . Here, instead of applying the mask vector m s of sample s directly, we use a learnable 'soft' mask vector, denoted by m ′ s , which is defined to be m ′ s = σ 1 ( m s ) ∈ R m , where σ 1 ( · ) is an MLP with the GELU activation function [21]. In other words, we obtain featurewise similarities with respect to sample s , denoted by r f s , as follows:

where σ 2 ( · ) denotes another MLP with the GELU activation function. FCU next obtains the Hadamard product between the learned embedding vector of sample s , h s , and the feature-wise similarities with respect to sample s , r f s , to learn their joint representations in a multiplicative manner. Specifically, FCU obtains the feature context vector of sample s , denoted by c f s , as follows:

where σ 3 ( · ) is also an MLP with the GELU activation function. That is, FCU fuses the representation vector of s and the vector that has embedding similarity values between the target feature f and the available features in s through the effective use of the soft mask m ′ s . From (3) and (4), the operations of FCU can be written as

3.4 Sample Correlation Unit

To measure similarities between s and other samples, a common approach would be to use the 166 dot product or cosine similarity between their embedding vectors. This approach, however, fails 167 to take into account the observability or availability of each feature in a sample. It also does 168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

ϕ p ( h p , r f p ) = σ 5 ( h p ⊙ r f p ) ∈ R d , (8) where σ 5 ( · ) is also an MLP with the GELU activation function. Once we have the masked embedding vectors of samples (excluding s ) in P , we finally compute the sample context vector of sample s , denoted by z f s , which is a weighted sum of the masked embedding vectors with weights being the pairwise similarity values, i.e., where σ 6 ( · ) is again an MLP with the GELU activation function. From (6)-(9), the operations of 211 SCU can be written as 212

not capture the fact that different observed features are of different importance to the target feature to impute when it comes to measuring the similarities. We introduce SCU as another integral component of M 3 -Impute to compute the sample 'context' vector of sample s by incorporating the embedding vectors of its similar samples as well as different weights of observed features. SCU works based on the two novel masking schemes, which shall be explained shortly.

Suppose we are to impute the missing value of feature f for a given sample s . SCU aims to leverage the information from the samples that are similar to s . As a first step to this end, we create a subset of samples P ⊂S that are similar to s . Specifically, we randomly choose and put a sample into P with probability that is proportional to the cosine similarity between s and the sample. This operation is repeated without replacement until P reaches a given size.

Mutual Sample Masking: Given a subset of samples P that include s , we first compute the pairwise similarities between s and other samples in the subset P . While they are computed in a similar way to FCU , we only consider the commonly observed features (or the common ones that have feature values) in both s and its peer p ∈ P \ { s } , to calculate their pairwise similarity in the sense that the missing value of feature f is inferred. Specifically, we compute the pairwise similarity between s and p ∈P\{ s } , which is denoted by sim ( s, p | f ) , as follows:

SCU

FCU

FCU

dot product

weighted sum

MLP

masking

where h s and h p are the learned embedding vectors of samples s and p from the GNN, respectively, and m s and m p are their respective mask vectors. Note that the multiplication in the RHS of (6) is the dot product.

Irrelevant Feature Masking: After we obtain the pairwise similarities between s and other samples in P , it would be natural to consider a weighted sum of their corresponding embedding vectors, i.e., ∑ p ∈P\{ s } sim ( s, p | f ) h p , in imputing the value of the target feature f . However, we observe that h p contains the information from the features whose values are available in p as well as possibly other features as it is learned via the so-called neighborhood aggregation mechanism that is central to GNNs, but some of the features may be irrelevant in inferring the value of feature f . Thus, instead of using { h p } directly, we introduce a d -dimensional mask vector r f p for h p , which is to mask out potentially irrelevant feature information in h p , when it comes to imputing the value of feature f . Specifically, it is defined by

where m f is an m -dimensional one-hot vector that has a value 1 in the place of feature f and 0's elsewhere, [ · ; · ] denotes the vector concatenation operation, and σ 4 ( · ) is an MLP with the GELU activation function. Note that the rationale behind the design of r f p is to embed the information on the features whose values are present in p as well as the information on the target feature f to impute. The mask r f p is then applied to h p to obtain the masked embedding vector of p as follows:

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

Algorithm 1 Forward computation of M 3 -Impute to impute the value of feature f for sample s .

1: Input: Bipartite graph G , initial feature node embeddings H 0 F , GNN model (e.g., GraphSAGE) GNN ( · ) , known mask matrix M , and a subset of samples P ⊂ S .

3: Obtain initial sample node embeddings H 0 S according to Equation (2).

▷ Perform graph representation learning

2: Output: Predicted missing feature value ˆ e sf .

4: H S , H F = GNN ( H 0 S , H 0 F , G ) .

5: c f s = FCU ( h s , m s , H F ) .

6: z f s = SCU ( H P , M P , H F ) .

7: Predict the missing feature value ˆ e sf using Equation (11).

3.5 Imputation

For a given sample s , to impute the missing value of feature f , M 3 -Impute obtains its feature context vector c f s and sample context vector z f s through FCU and SCU , respectively, which are then used for imputation. Specifically, it is done by predicting the corresponding edge weight ˆ e sf as follows:

where ϕ α ( · ) denotes an MLP with a non-linear activation function (i.e., ReLU for continuous values and softmax for discrete ones), and α is a learnable scalar parameter. This scalar parameter α is introduced to strike a balance between leveraging feature-wise correlation and sample-wise correlation. It is necessary because the quality of z f s relies on the quality of the samples chosen in P , so overly relying on z f s would backfire if their quality is not as desired. To address this problem, instead of employing a fixed weight α , we make α learnable and adaptive in determining the weights for c f s and z f s . Note that this kind of learnable parameter approach has been widely adopted in natural language processing [26, 34, 38, 46] and computer vision [8, 55, 56], showing superior performance to its fixed counterpart. In M 3 -Impute, the scalar parameter α is learned based on the similarity values between s and its peer samples p ∈ P \ { s } as follows:

where ∥ represents the concatenation operation, and ϕ γ ( · ) is an MLP with the activation function γ ( x ) = 1 -1 / e | x | . The overall operation of M 3 -Impute is summarized in Algorithm 1. To learn network parameters, we use cross-entropy loss and mean square error loss for imputing discrete and continuous feature values, respectively.

4 Experiments

4.1 Experiment Setup

Datasets: We conduct experiments on 15 open datasets. These real-world datasets consist of mixed data types with both continuous and discrete values and cover different domains including civil engineering (CONCRETE, ENERGY), physics and chemistry (YACHT), thermal dynamics (NAVAL), etc. Since the datasets are fully observed, we introduce missing values by applying a randomly generated mask to the data matrix. Specifically, as used in prior studies [23, 24], we apply three masking generation schemes, namely missing completely at random (MCAR), missing at random (MAR), and missing not at random (MNAR). 1 We use MCAR with a missing ratio of 30% , unless otherwise specified. We follow the preprocessing steps adopted in [52, 54] to scale feature values to [0, 1] with a MinMax scaler [25]. Due to the space limit, we below present the results of eight datasets that are used in Grape [52] and report the other results in Appendix.

Baseline models: M 3 -Impute is compared against popular and state-of-the-art imputation methods, including statistical methods, deep generative methods, and graph-based methods listed as follows: MEAN : It imputes the missing value ˆ e sf as the mean of observed values in feature f from all the samples. K-nearest neighbors ( kNN ) [43]: It imputes the missing value ˆ e sf using the kNNs that have observed values in feature f with weights that are based on the Euclidean distance to sample s . Multivariate imputation by chained equations ( Mice ) [45]: This method runs multiple regressions where each missing value is modeled upon the observed non-missing values. Iterative

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

SVD ( Svd ) [18]: It imputes missing values by solving a matrix completion problem with iterative low-rank singular value decomposition. Spectral regularization algorithm ( Spectral ) [30]: This matrix completion algorithm uses the nuclear norm as a regularizer and imputes missing values with iterative soft-thresholded SVD. Miwae [29]: It works based on an autoencoder generative model trained to maximize a potentially tight lower bound of the log-likelihood of the observed data and Monte Carlo techniques for imputation. Miracle [24]: It uses the imputation results from naive methods such as MEAN and refines them iteratively by learning a missingness graph (m-graph) and regularizing an imputation function. Gain [50]: This method trains a data imputation generator with a generalized generative adversarial network in which the discriminator aims to distinguish between real and imputed values. Grape [52]: It models the data as a bipartite graph and imputes missing values by predicting the weights of the missing edges, each of which is done based on the inner product between the embeddings of its corresponding sample and feature nodes. HyperImpute [23]: HyperImpute is a framework that conducts an extensive search among a set of imputation methods, selecting the optimal imputation method with fine-tuned parameters for each feature in the dataset.

Model configurations: Parameters of M 3 -Impute are updated by the Adam optimizer with a learning rate of 0.001 for 40,000 epochs. For graph representation learning, we use a variant of GraphSAGE [17], which not only learns node embeddings but also edge embeddings via the neighborhood aggregation mechanism, as similarly used in [52]. We consider its three-layer GNN model. We employ mean-pooling as the aggregation function and use ReLU as the activation function for the GNN layers. We set the embedding dimension d to 128. It is known that randomly dropping out a subset of observable edges during training improves the model's generalization ability. We also leverage the observation and randomly drop 50% of observable edges during training. For each experiment, we conduct five runs with different random seeds and report the average results.

4.2 Overall Performance

We first compare the feature imputation performance of M 3 -Impute with popular and state-of-theart imputation methods. As shown in Table 1, M 3 -Impute achieves the lowest imputation MAE for six out of the eight examined datasets and the second-best MAE scores in the other two, which validates the effectiveness of M 3 -Impute. For KIN8NM dataset, M 3 -Impute underperforms Miracle. It is mainly because each feature in KIN8NM is independent of the others, so none of the observed features can help impute missing feature values. For NAVAL dataset, the only model that outperforms M 3 -Impute is HyperImpute [23]. In the NAVAL dataset, nearly every feature exhibits a strong linear correlation with the other features, i.e., every pair of features has correlation coefficient close to one. This allows HyperImpute to readily select a linear model from its model pool for each feature to impute. Nonetheless, M 3 -Impute exhibits overall superior performance to the baselines as it can be well adapted to each dataset that possesses different amounts of correlations over features and samples. In other words, M 3 -Impute benefits from explicitly incorporating feature-wise and sample-wise correlations together with our carefully designed mask schemes. Furthermore, we evaluate the performance of M 3 -Impute under MAR and MNAR settings. We observe that M 3 -Impute consistently outperforms all the baselines under all datasets and achieves a larger margin in the improvement compared to the case with MCAR setting. This implies that M 3 -Impute is also effective in handling different patterns of missing values in the input data. Comprehensive results are provided in Appendix.

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

4.3 Ablation Study

To study the effectiveness of three integral components of M 3 -Impute, we consider three variants of M 3 -Impute, each with a subset of the components, namely initialization only (Init Only), initialization + FCU (Init + FCU ), and initialization + SCU (Init + SCU ). The performance of these variants are evaluated against the top-performing imputation baselines such as Grape and HyperImpute. As shown in Table 2, the three variants derived from M 3 -Impute achieve lower MAE values than both baselines in most datasets, demonstrating the effectiveness of our novel components in M 3 -Impute.

Specifically, for initialization only, the key difference between M 3 -Impute and Grape lies in our refined initialization process of feature-node and sample-node embeddings. The reduced MAE values observed by the Init Only variant demonstrate that our proposed initialization process is more effective in utilizing information between samples and their associated features, including missing ones, as compared to the basic initialization used in [52]. In addition, we observe that when FCU or SCU is incorporated, MAE values are further reduced for most datasets. This validates that explicitly modeling feature-wise or sample-wise correlations through our novel masking schemes can improve imputation accuracy. When all the three components are combined together as in M 3 -Impute, they work synergistically to lower MAE values, validating the efficacy of explicit consideration of both sample-wise and feature-wise correlations (in addition to the refined initialization process) for missing data imputation.

4.4 Robustness

Missing ratio: In practice, datasets may possess different missing ratios. To validate the model's robustness under such circumstances, we evaluate the performance of M 3 -Impute and other baseline models with varying missing ratios, i.e., 0.1, 0.3, 0.5, and 0.7. Figure 3 shows their performance. We use the MAE of HyperImpute ( HI ) as the reference performance and offset the performance of each model by MAE x -MAE HI , where x represents the considered model. For clarity, we here only report the results of four top-performing models. As shown in Figure 3, M 3 -Impute outperforms other baseline models for almost all the cases, especially under YACHT, CONCRETE, ENERGY, and HOUSING datasets. It is worth noting that modeling feature correlations in these datasets is particularly challenging due to the presence of considerable amounts of weakly correlated features, along with a few strongly correlated ones. Nonetheless, FCU and SCU in M 3 -Impute were able to better capture such correlations with our efficient masking schemes, thereby resulting in a large improvement in imputation accuracy. In addition, for KIN8NM dataset, M 3 -Impute ties with the second-best model, Grape. As mentioned in Section 4.2, each feature in KIN8NM is independent of the others, so none of the observed features can help impute missing feature values. For NAVAL dataset, where each feature strongly correlates with the others, M 3 -Impute surpasses Grape but falls short of HyperImpute, due to the same reason as discussed above. Overall, M 3 -Impute is robust to various missing ratios. Comprehensive results for all the baseline models can be found in Appendix.

Sampling strategy in SCU: While SCU uses a sampling strategy based on pairwise cosine similarities to construct a subset of samples P , the simplest sampling strategy to build P would be to choose samples uniformly at random without replacement (M 3 -Uniform). Intuitively, this approach cannot identify similar peer samples accurately and thus would lead to inferior performance. Nonetheless, as shown in Table 2, even with this naive uniform sampling strategy, M 3 -Uniform still outperforms the two leading imputation baselines.

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

Wine

Yacht

Naval

Power

Energy

Kin8nm

Housing

Concrete

0.5

0.0

0.5

Missing Ratio 0.1

Wine

Yacht

Naval

Power

Energy

Kin8nm

Housing

Concrete

0.5

0.0

0.5

Missing Ratio 0.3

Wine

Yacht

Naval

Power

Energy

Kin8nm

Housing

Concrete

0.5

0.0

0.5

Missing Ratio 0.5

Wine

Yacht

Naval

Power

Energy

Kin8nm

Housing

Concrete

0.5

0.0

0.5

Missing Ratio 0.7

MAE Offset by HyperImpute

MICE

KNN

GRAPE

M

3 -Impute

Size of P in SCU: Intuitively, neither an excessively small nor overly large size of the sample subset P is optimal. Too few peer samples leave SCU with insufficient information to learn sample-wise correlations, while too many peer samples may include quite a few dissimilar ones, which may introduce significant noise to the computation of SCU and thus degrade the performance. Table 3 shows the performance of M 3 -Impute with varying numbers of peer samples. In general, the trends agree with our intuition. Although the optimal size varies across different datasets, we observe that having the number of peer samples to be 5 to 10 achieves the overall best imputation accuracy.

Initialization parameter ϵ : We also evaluate whether a non-zero value of ϵ in the initialization process of M 3 -Impute indeed lead to an improvement in imputation accuracy. As shown in Table 3, for YACHT and WINE datasets, the introduction of a non-zero value of ϵ results in lower MAE scores. Another insight that we have from Table 3 is that ϵ should not be set too large, as a large value of ϵ might impose incorrect weights to the features with missing values. We observe that it is an overall good choice to set ϵ to 1 × 10 -5 or 1 × 10 -4 .

5 Conclusion

We have presented M 3 -Impute, a mask-guided representation learning for missing data imputation. M 3 -Impute improved the initialization process by considering the relationships between samples and their associated features (including missing ones) even in initializing the embeddings. In addition, for more effective representation learning, we introduced two novel components in M 3 -Impute FCU and SCU , which learn feature-wise and sample-wise correlations, respectively, to capture data correlations explicitly and leverage them for imputation. Extensive experiment results demonstrate the effectiveness of M 3 -Impute. M 3 -Impute achieves overall superior performance to popular and state-of-the-art methods on 15 open datasets, with 13 best and two second-best MAE scores on average under three different settings of missing value patterns.

358

References

Melissa J Azur, Elizabeth A Stuart, Constantine Frangakis, and Philip J Leaf. Multiple imputa359 tion by chained equations: what is it and how does it work? International journal of methods 360 in psychiatric research , 20(1):40-49, 2011. 361

Jaap Brand. Development, implementation and evaluation of multiple imputation strategies for 362 the statistical analysis of incomplete data sets . 1999. 363

Thomas Brooks, D. Pope, and Michael Marcolini. Airfoil self-noise. https://doi.org/10. 364 24432/C5VW2C , March 2014. 365

Lane F Burgette and Jerome P Reiter. Multiple imputation for missing data via sequential 366 regression trees. American journal of epidemiology , 172(9):1070-1076, 2010. 367

Jian-Feng Cai, Emmanuel J Candès, and Zuowei Shen. A singular value thresholding algorithm 368 for matrix completion. SIAM Journal on optimization , 20(4):1956-1982, 2010. 369

Emmanuel Candes and Benjamin Recht. Exact matrix completion via convex optimization. 370 Communications of the ACM , 55(6):111-119, 2012. 371

Paulo Cortez, Antonio Cerdeira, Fernando Almeida, Telmo Matos, and José Reis. Wine quality. 372 https://doi.org/10.24432/C56S3T , October 2009. 373

Jifeng Dai, Haozhi Qi, Yuwen Xiong, Yi Li, Guodong Zhang, Han Hu, and Yichen Wei. De374 formable convolutional networks. In 2017 IEEE International Conference on Computer Vision 375 (ICCV) , pages 764-773, 2017. 376

Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete 377 data via the em algorithm. Journal of the royal statistical society: series B (methodological) , 378 39(1):1-22, 1977. 379

Tianyu Du, Luca Melis, and Ting Wang. Remasker: Imputing tabular data with masked au380 toencoding. In The Twelfth International Conference on Learning Representations , 2024. 381

Dheeru Dua and Casey Graff. Uci machine learning repository, 2017. 382

Bradley Efron, Trevor Hastie, Iain Johnstone, and Robert Tibshirani. Least angle regression. 383 Ann. Statist. , 32(1):407-499, 2004. 384

Ronald A Fisher. The use of multiple measurements in taxonomic problems. Annals of eugen385 ics , 7(2):179-188, 1936. 386

Pedro J García-Laencina, José-Luis Sancho-Gómez, and Aníbal R Figueiras-Vidal. Pattern 387 classification with missing data: a review. Neural Computing and Applications , 19:263-282, 388 2010. 389

Andrew Gelman. Parameterization and bayesian modeling. Journal of the American Statistical 390 Association , 99(466):537-545, 2004. 391

Zoubin Ghahramani and Michael Jordan. Supervised learning from incomplete data via an em 392 approach. Advances in neural information processing systems , 6, 1993. 393

William L. Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on 394 large graphs. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob 395 Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information 396 Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 397 December 4-9, 2017, Long Beach, CA, USA , pages 1024-1034, 2017. 398

Trevor Hastie, Rahul Mazumder, Jason D. Lee, and Reza Zadeh. Matrix completion and low399 rank svd via fast alternating least squares. J. Mach. Learn. Res. , 16(1):33673402, jan 2015. 400

Trevor Hastie, Rahul Mazumder, Jason D Lee, and Reza Zadeh. Matrix completion and 401 low-rank svd via fast alternating least squares. The Journal of Machine Learning Research , 402 16(1):3367-3402, 2015. 403

David Heckerman, David Maxwell Chickering, Christopher Meek, Robert Rounthwaite, and 404 Carl Kadie. Dependency networks for inference, collaborative filtering, and data visualization. 405 Journal of Machine Learning Research , 1(Oct):49-75, 2000. 406

407

408

Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR , abs/1606.08415, 2016.

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

James Honaker, Gary King, and Matthew Blackwell. Amelia ii: A program for missing data. Journal of statistical software , 45:1-47, 2011.

Daniel Jarrett, Bogdan Cebere, Tennison Liu, Alicia Curth, and Mihaela van der Schaar. Hyperimpute: Generalized iterative imputation with automatic model selection. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA , volume 162 of Proceedings of Machine Learning Research , pages 9916-9937. PMLR, 2022.

Trent Kyono, Yao Zhang, Alexis Bellot, and Mihaela van der Schaar. MIRACLE: causallyaware imputation via learning missing data mechanisms. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 2380623817, 2021.

Jure Leskovec, Anand Rajaraman, and Jeffrey D. Ullman. Mining of Massive Datasets, 2nd Ed . Cambridge University Press, 2014.

Mingzhe Li, Xiuying Chen, Shen Gao, Zhangming Chan, Dongyan Zhao, and Rui Yan. VMSMO: Learning to generate multimodal summary for video-based news articles. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pages 9360-9369, Online, November 2020. Association for Computational Linguistics.

Steven Cheng-Xian Li, Bo Jiang, and Benjamin M. Marlin. Misgan: Learning from incomplete data with generative adversarial networks. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019 . OpenReview.net, 2019.

Roderick JA Little and Donald B Rubin. Statistical analysis with missing data , volume 793. John Wiley & Sons, 2019.

Pierre-Alexandre Mattei and Jes Frellsen. MIWAE: deep generative modelling and imputation of incomplete data sets. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA , volume 97 of Proceedings of Machine Learning Research , pages 4413-4423. PMLR, 2019.

Rahul Mazumder, Trevor Hastie, and Robert Tibshirani. Spectral regularization algorithms for learning large incomplete matrices. J. Mach. Learn. Res. , 11:22872322, aug 2010.

Rahul Mazumder, Trevor Hastie, and Robert Tibshirani. Spectral regularization algorithms for learning large incomplete matrices. The Journal of Machine Learning Research , 11:22872322, 2010.

Rahul Mazumder, Trevor Hastie, and Robert Tibshirani. Spectral regularization algorithms for learning large incomplete matrices. The Journal of Machine Learning Research , 11:22872322, 2010.

Boris Muzellec, Julie Josse, Claire Boyer, and Marco Cuturi. Missing data imputation using optimal transport. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 7130-7140. PMLR, 13-18 Jul 2020.

Romain Paulus, Caiming Xiong, and Richard Socher. A deep reinforced model for abstractive summarization. In International Conference on Learning Representations , 2018.

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

491

492

493

494

Trivellore E Raghunathan, James M Lepkowski, John Van Hoewyk, Peter Solenberger, et al. A multivariate technique for multiply imputing missing values using a sequence of regression models. Survey methodology , 27(1):85-96, 2001.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022 , pages 10674-10685. IEEE, 2022.

Joseph L Schafer. Analysis of incomplete multivariate data . CRC press, 1997.

Abigail See, Peter J. Liu, and Christopher D. Manning. Get to the point: Summarization with pointer-generator networks. In Regina Barzilay and Min-Yen Kan, editors, Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 1073-1083, Vancouver, Canada, July 2017. Association for Computational Linguistics.

V. Sigillito, S. Wing, L. Hutton, and K. Baker. Ionosphere. https://doi.org/10.24432/ C5W01B , December 1988.

Indro Spinelli, Simone Scardapane, and Aurelio Uncini. Missing data imputation with adversarially-trained graph convolutional networks. Neural Networks , 129:249-260, 2020.

Daniel J Stekhoven and Peter Bühlmann. Missforestnon-parametric missing value imputation for mixed-type data. Bioinformatics , 28(1):112-118, 2012.

Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. CSDI: conditional scorebased diffusion models for probabilistic time series imputation. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 2480424816, 2021.

Olga Troyanskaya, Michael Cantor, Gavin Sherlock, Pat Brown, Trevor Hastie, Robert Tibshirani, David Botstein, and Russ B. Altman. Missing value estimation methods for dna microarrays. Bioinformatics , 17(6):520-525, 2001.

Stef Van Buuren, Jaap PL Brand, Catharina GM Groothuis-Oudshoorn, and Donald B Rubin. Fully conditional specification in multivariate imputation. Journal of statistical computation and simulation , 76(12):1049-1064, 2006.

Stef van Buuren and Karin Groothuis-Oudshoorn. mice: Multivariate imputation by chained equations in r. Journal of Statistical Software , 45(3):167, 2011.

Wenbo Wang, Yang Gao, Heyan Huang, and Yuxiang Zhou. Concept pointer network for abstractive summarization. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan, editors, Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLPIJCNLP) , pages 3076-3085, Hong Kong, China, November 2019. Association for Computational Linguistics.

William H. Wolberg, Olvi L. Mangasarian, and W. Nick Street. Breast cancer wisconsin (diagnostic). https://doi.org/10.24432/C5DW2B , October 1995.

Richard Wu, Aoqian Zhang, Ihab Ilyas, and Theodoros Rekatsinas. Attention-based learning 495 for missing data imputation in holoclean. Proceedings of Machine Learning and Systems , 496 2:307-325, 2020. 497

I-Cheng Yeh. Blood transfusion service center. https://doi.org/10.24432/C5GS39 , Oc498 tober 2008. 499

500

501

502

503

504

505

506

507

508

509

510

511

512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

Jinsung Yoon, James Jordon, and Mihaela van der Schaar. GAIN: missing data imputation using generative adversarial nets. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018 , volume 80 of Proceedings of Machine Learning Research , pages 5675-5684. PMLR, 2018.

Seongwook Yoon and Sanghoon Sull. GAMIN: generative adversarial multiple imputation network for highly missing data. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020 , pages 8453-8461. Computer Vision Foundation / IEEE, 2020.

Jiaxuan You, Xiaobai Ma, Daisy Yi Ding, Mykel J. Kochenderfer, and Jure Leskovec. Handling missing data with graph representation learning. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

Shuhan Zheng and Nontawat Charoenphakdee. Diffusion models for missing value imputation in tabular data. CoRR , abs/2210.17128, 2022.

Jiajun Zhong, Ning Gui, and Weiwei Ye. Data imputation with iterative graph reconstruction. In Brian Williams, Yiling Chen, and Jennifer Neville, editors, Thirty-Seventh AAAI Conference on Artificial Intelligence, AAAI 2023, Thirty-Fifth Conference on Innovative Applications of Artificial Intelligence, IAAI 2023, Thirteenth Symposium on Educational Advances in Artificial Intelligence, EAAI 2023, Washington, DC, USA, February 7-14, 2023 , pages 11399-11407. AAAI Press, 2023.

Xizhou Zhu, Dazhi Cheng, Zheng Zhang, Stephen Lin, and Jifeng Dai. An empirical study of spatial attention mechanisms in deep networks. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) , October 2019.

Xizhou Zhu, Han Hu, Stephen Lin, and Jifeng Dai. Deformable convnets v2: More deformable, better results. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 9300-9308, 2019.

A Appendix 528

529

530

531

532

533

534

535

536

537

538

539

540

541

542

543

544

545

546

547

548

549

550

551

0

0

1

2

3

4

5

6

7

0

1

2

3

4

5

6

7

Concrete

0

1

2

3

4

5

6

7

0

1

2

3

4

5

6

7

Energy

0

1

2

3

4

5

6

7

8

9

10

11

12

1

2

3

4

5

6

7

8

9

10

11

12

Housing

0

1

2

3

4

5

6

7

0

1

2

3

4

5

6

7

Kim8nm

0

2

4

6

8

10 12 14 16

0

2

4

6

8

16 14 12 10

Naval

0

1

2

3

0

1

2

3

Power

0

1

2

3

4

5

6

7

8

9 10

0

1

2

3

4

5

6

7

8

10 9

Wine

0

1

2

3

4

5

0

1

2

3

4

5

Yacht

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

In this section, we discuss further experimental details. We first give an overview of the dataset details in Section A.1, followed by the implementation of different missing types and present corresponding imputation performance under MAR and MNAR settings (Section A.2). We then provide the comprehensive results of the robustness experiments (Section A.3). Finally, we extend our evaluation of M 3 -Impute to seven additional datasets (Section A.4) and elaborate on the computational resources in Section A.5.

A.1 Dataset Details

Table 4 presents the statistics of the eight UCI datasets [11] used throughout Section 4. Figure 4 illustrates the Pearson correlation coefficients among the features. In the Kin8nm dataset, all features are linearly independent, whereas the Naval dataset exhibits strong correlations among its features. Under the MCAR setting, M 3 -Impute performs comparably to the baseline imputation methods on these two datasets (shown in Table 1). However, in real-world scenarios, features are not always entirely independent or strongly correlated. In the other six datasets, we observe a mix of weakly correlated features along with a few that are strongly correlated. In these cases, M 3 -Impute consistently outperforms all baseline methods.

A.2 Detailed Results of Different Missing Types

We adopt the same procedure outlined in [52, 54] to generate missing values under different settings.

MCAR : A n × m matrix is sampled from a uniform distribution. Positions with values no greater than the ratio of missingness are viewed as missing and the remaining positions are observable.

MAR : First, a subset of features is randomly selected to be fully observed. Then, these remaining features have values removed according to a logistic model with random weights, using the fully observed feature values as input. The desired rate of missingness is achieved by adjusting the bias term.

MNAR : This is done by first apply the MAR mechanism above. Then, the remaining feature 552 values are masked out by the MCAR mechanism. 553

In addition to the results for MCAR setting presented in Table 4.2, Table 5 and Table 6 present the 554 MAEscores under MAR and MNAR settings, respectively. M 3 -Impute consistently outperforms all 555 baseline methods in both scenarios. 556

557

558

559

560

A.3 Robustness against Various Ratios of Missingness

Table 8 presents the performance of various imputation methods across different ratios of missingness. M 3 -Impute achieves the lowest MAE scores in most cases and the second-best MAE scores in the remaining ones.

A.4 Further Evaluation on Seven Additional Datasets 561

In this experiment, we further evaluate M 3 -Impute on seven datasets: Airfoil [3], Blood [49], Wine562 White [7], Ionosphere [39], Breast Cancer [47], Iris [13], and Diabetes [12]. An overview of dataset 563 details is provided in Table 7, and feature correlations are illustrated in Figure 5. We simulate 564 missingness in data under MCAR, MAR, and MNAR conditions, each with a missing ratio of 0.3. 565 Results are demonstrated in Table 9. Across all three types of missingness, M 3 -Impute achieves five 566 best and two second-best MAE scores on average. 567

0

0

1

2

3

4

5

0

1

2

3

4

5

Airfoil

0

1

2

3

4

0

1

2

3

4

Blood

0

3

6

9 12 15 18 21 24 27 30

2

4

6

8

10

12

14

16

18

20

22

24

26

28

30

Breast

0

1

2

3

4

5

6

7

8

9 10

0

1

2

3

4

5

6

7

8

10 9

Diabetes

0

3

6

9 1215182124273033

0

3

6

3330272421181512 9

Ionosphere

0

1

2

3

4

0

1

2

3

4

0

1

2

3

4

5

6

7

8

9 1011

0

1

2

3

4

5

6

7

8

1110 9

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

A.5 Computational Resources 568

All our experiments are conducted on a GPU server running Ubuntu 22.04, with PyTorch 2.1.0 569 and CUDA 12.1. We train and test M 3 -Impute using a single NVIDIA A100 80G GPU. With the 570 experimental setup described in Section 4.1, the total runtime (including both training and testing) 571

572

for each of the five repeated runs ranged from 1 to 5 hours, depending on the scale of the datasets.

573

574

575

576

577

578

579

580

581

582

583

584

585

586

587

588

589

590

591

592

593

594

595

596

597

598

599

600

601

602

603

604

605

606

607

608

609

610

611

612

613

614

615

616

617

618

619

620

621

622

623

624

625

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: In the abstract and introduction sections, we clearly define the scope of this paper, focusing on missing value imputation. We propose M 3 -Impute, a mask-guided imputation method designed to compute feature-wise and sample-wise correlations based on missing data patterns. A concise summary of the experimental results is provided at the end of both sections.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: In Section 4.2, we discussed two cases of MAE degradation for the KIN8NM and NAVAL datasets. It is mainly because 1. Each feature in KIN8NM is independent of the others, so none of the observed features can help impute missing feature values. 2. In the NAVAL dataset, nearly every feature exhibits a strong linear correlation with the other features. While it is true that M 3 -Impute does not achieve the best MAE on these two datasets, our model has outperformed all the other baselines on the majority of datasets. This demonstrates the unique strengths of graph modeling in M 3 -Impute over tabular data modeling in baselines like Hyperimpute. In real-world scenarios, the correlation structure of datasets is often unpredictable, and such extreme cases are relatively rare. Thus, we design a scheme to handle general cases for data imputation tasks. The empirical evidence suggests that our approach has been quite successful and exhibits overall superior performance to the baselines as it can be well adapted to each dataset that possesses different levels of correlations over features and samples.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: This paper does not present theoretical results. We do not assume that the data is missing under MCAR, MAR, or MNAR conditions for M 3 -Impute to be effective. Instead, M 3 -Impute demonstrates robust performance across all three settings.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: In Section 3, we explain the computational pipeline of the proposed model in detail and provide a pseudo-code to better outline the methodology. The experimental setup is comprehensively described in Section 4.1. In addition, supplementary materials include our complete codebase to reproduce the results presented in this paper, including the model implementation, training and testing pipeline, configuration files, and execution scripts.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via

detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: In the supplementary material, we provide the complete code for our model, including the scripts for experiments and evaluations, as well as the execution scripts used in the experiments. We have also included the data preprocessed by us, along with the download links for publicly available datasets. We will release the formatted codebase on GitHub following the conclusion of the anonymity period.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so No is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Experimental setup is detailed in Section 4.1 and Appendix A.2. We also explore the hyperparameters utilized in M 3 -Impute. Results are presented in Table 3.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We conduct all the experiments over five runs and report the mean MAE scores, along with the standard deviations.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We train and test M 3 -Impute on a single Nvidia A100 80G GPU (Detailed setup described in A.5). With the experimental setup described in Section 4.1, the total running time (including training and testing) for one of the five repeated runs varies from 1 to 5 hours, depending on the scale of the datasets.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our paper adheres to the NeurIPS Code of Ethics in every respect. 1. We ensured fair wages for all human participants involved in our study, abiding by regional minimum hourly rates. 2. Our research methodology adhered to institutional protocols for human subjects and data privacy. 3. We obtained informed consent from all participants and minimized exposure of personally identifiable information. 4. The datasets used are publicly available and have not been deprecated, with all copyrights respected. 5. We have transparently communicated the societal impact of our research, considering potential misuse and its effects on discrimination, surveillance, and environmental impact. 6. We have also reflected on the biases in our models and datasets and taken steps to mitigate them. 7. Our data and models are documented and released with appropriate licenses, and we've employed secure data storage and distribution practices. 8. We ensured legal compliance and provided all necessary elements for the reproducibility of our research.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: The method proposed in this work is only applicable for missing value imputation and is unlikely to have a negative social impact.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

838

839

840

841

842

843

844

845

846

847

848

849

850

851

852

853

854

855

856

857

858

859

860

861

862

863

864

865

866

867

868

869

870

871

872

873

874

875

876

877

878

879

880

881

882

883

884

885

886

887

888

889

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The model we propose does not carry the risk of misuse; the datasets were selected under fair use conditions, from publicly available sources with undisputed licenses. Therefore, our work does not require additional safeguard protections.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: In our research, we have carefully credited all the code and data used, providing explicit citations for each. The licenses for this code and data are notably permissive, including MIT, BSD 3-clause, and CC BY 4.0. In accordance with these licenses, we have properly acknowledged the contributions of the original authors.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We have detailed the new datasets employed in this research in Appendix A.4. We commit to making these datasets publicly accessible following the anonymity period to foster transparency and reproducibility.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our work does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our work does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces M3-Impute, a mask-guided representation learning method for missing value imputation. The core idea of M3-Impute is to leverage missingness information as an explicit input to the model through masking schemes. This approach allows M3-Impute to effectively learn both feature-wise and sample-wise correlations, accommodating various types of data missingness. The model employs a variant of GraphSAGE for graph representation learning, incorporating edge embeddings via neighborhood aggregation. It outperforms traditional tabular data models in various benchmark datasets.\n\nSTRENGTHS:\n1. The paper presents a novel approach to missing value imputation through the introduction of a mask-guided representation learning method (M3-Impute). The originality of the work lies in its utilization of missingness information as a model input, employing innovative masking schemes. This allows M3-Impute to accurately capture feature-wise and sample-wise correlations despite varying types of missing data (MCAR, MAR, MNAR). The use of GraphSAGE for graph representation learning, combined with edge embeddings via neighborhood aggregation, further distinguishes this work from traditional tabular data models.\n\n2. The quality of the research is demonstrated through comprehensive experiments across multiple datasets and missing data mechanisms. The empirical results show that M3-Impute consistently outperforms baseline methods. The authors include a code package and datasets with the submission.\n\nWEAKNESSES:\n1. The paper evaluates the sensitivity of the M3-Impute model to the initialization parameter ϵ (Table 3), demonstrating that a non-zero value of ϵ improves imputation accuracy. However, the lack of detailed sensitivity analysis for other critical hyperparameters, such as the learning rate, batch size, number of GNN layers, and the dropout rate, represents a weakness.\n\n2. The paper also has notable limitations in its contextualization relative to prior work. While it effectively presents M3-Impute and compares it against several baseline models, it lacks a deeper analysis of how these baseline models have evolved and the specific innovations they have introduced over time. For instance, the paper mentions GRAPE and IGRM as key prior graph-based imputation methods but does not adequately explore their strengths and weaknesses or how M3-Impute directly addresses the limitations of these methods. This omission makes it challenging to understand the novelty and improvements offered by M3-Impute.\n\n3. Relying solely on MAE to evaluate the performance of imputation models has several limitations. MAE measures the average magnitude of errors but does not account for the variance or distribution of those errors, making it insensitive to outliers and providing no insight into model bias. This can result in an incomplete understanding of a model's performance, particularly in contexts where large errors or systematic biases are important considerations. To address these limitations, incorporating RMSE alongside MAE would be beneficial. RMSE penalizes larger errors more heavily, offering additional insight into the presence and impact of significant errors in the model's predictions.\n\nQUESTIONS:\n- The statement that most learning-based methods are``built upon the raw tabular data structure as is, which greatly restricts them from jointly modeling the feature-wise and sample-wise correlations” (line 41) is not entirely accurate for two prominent tabular generative models, MIDA and GAIN. MIDA transforms raw tabular data into a higher-dimensional space through its encoder-decoder architecture. This transformation allows MIDA to capture more complex, nonlinear relationships that are not immediately apparent in the raw data. The adversarial process of GAIN allows it to model the joint distribution of the data, thus capturing complex correlations between features and samples. \n\n      - Gondara, L., Wang, K. (2018). MIDA: Multiple Imputation Using Denoising Autoencoders. In: Phung, D., Tseng, V., Webb, G., Ho, B., Ganji, M., Rashidi, L. (eds) Advances in Knowledge Discovery and Data Mining. PAKDD 2018. Lecture Notes in Computer Science(), vol 10939. Springer, Cham. https://doi.org/10.1007/978-3-319-93040-4_21\n\n      - Wang, Zhenhua, et al. \"Are deep learning models superior for missing data imputation in surveys? Evidence from an empirical comparison.\" Survey Methodology 48 (2022): 375-399.\n\n- When discussing statistical methods (line 70), the authors should mention that FCS approaches such as MICE are flexible in imputing different types of variables. \n\n- Related to the previous point, MICE is generally considered a statistical method rather than a learning-based method, although a learning algorithm such as CART can be used as the imputer. See the paper by Wang et. al. (2022).\n\n- In Sec. 4.2, it should be noted that the M3-Impute model tends to perform slightly better under MAR and MNAR settings for most datasets, indicating its effectiveness in handling missingness that depends on the observed data.\n\n- How can interpretability techniques be incorporated into M3-Impute to help users understand the imputation decisions?\n\nLIMITATIONS:\nThe authors should consider summarizing the limitations of their method in the conclusion to provide a comprehensive overview of their work. Specifically, in Section 4.2, the authors discuss the cases of MAE degradation for the Kin8nm and Naval datasets. They attribute this to the independence of features in Kin8nm, which prevents observed features from aiding in the imputation of missing values, and the strong linear correlations between nearly all features in the Naval dataset. Summarizing these points in the conclusion would give readers a clear understanding of the method's limitations and the contexts in which it performs best.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-30"}

comments: SUMMARY:
This study proposes a missing value imputation method. The proposed method tackles the missing value imputation problem as a link prediction task on the bipartite graph. It represents a data matrix with missing values as a bipartite graph, then uses a graph neural network on the bipartite graph to learn the embeddings of samples and features. Next, a feature correlation unit and a sample correlation unit are employed to obtain feature-wise and sample-wise correlations, which are then fed into an MLP to obtain imputed values.

STRENGTHS:
- An advanced and novel approach to missing value imputation. The idea is intuitive.
- Experimental comparison is done with various existing methods, including recent ones.
- Significant performance improvements achieved.

WEAKNESSES:
- The size of the bipartite graph may drastically increase with data size and dimensionality.
- No investigation on various missingness scenarios and missing rates.

QUESTIONS:
- Computational and space complexity analysis is needed. Can the proposed method be used for a very large training dataset with high dimensionality?
- How does the proposed method perform under lower and higher missing rates?
- The inference phase is unclear. Please elaborate on how the proposed method makes imputations for a query instance containing missing values?
- While the authors reviewed other recent methods that leveraged graph neural networks ([40], [52], [54]), only the method in [52] was compared in the experiments. Why were the methods in [40] and [54] not compared?

LIMITATIONS:
-
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the challenge of missing values in data analysis and machine learning by proposing M3-Impute, a novel imputation method. Traditional imputation techniques often neglect 'missingness' information and fail to explicitly model feature and sample correlations, leading to suboptimal results. M3-Impute innovatively incorporates missingness information and correlations through advanced masking schemes. It represents data as a bipartite graph and utilizes a graph neural network with a refined initialization process to learn node embeddings. These embeddings are further optimized using feature correlation and sample correlation units, which explicitly consider the correlations during imputation. The method's effectiveness is demonstrated through experiments on 15 benchmark datasets with three different missing patterns.\n\nSTRENGTHS:\nThe feature correlation unit (FCU) and sample correlation unit (SCU) are particularly compelling. The FCU learns correlations between the target missing feature and observed features within each sample, refined by a soft mask on missingness information. Similarly, the SCU computes sample-wise correlations, enhanced by another soft mask on missingness information for pairs of samples.\n\nIntegrating FCU and SCU outputs to estimate missing values is methodologically sound. Extensive experiments on 15 open datasets show M3-Impute's superior performance in 13 out of 15 cases under various missing value patterns. The reported improvements in mean absolute error (MAE), up to 11.47% over the second-best method, underscore M3-Impute's practical relevance and robustness.\n\nWEAKNESSES:\nM3-Impute demonstrates strong performance across many datasets but shows limitations in handling datasets with highly independent features or strong linear correlations, such as the KIN8NM and NAVAL datasets.\n\nThe model's robustness in general scenarios may degrade when confronted with datasets containing extreme correlation structures.\n\nFuture improvements could concentrate on enhancing M3-Impute's adaptability to these challenging cases to broaden its applicability and robustness across diverse datasets.\n\nQUESTIONS:\nIn ablation study (4.3), Could you provide more details on the refined initialization process of feature-node and sample-node embeddings? How does it differ specifically from the initialization used in Grape?\n\n\nCan you provide more insights into the computational complexity and runtime performance of M3-Impute compared to the baselines, especially for large datasets?\n\nIn missing ratio test, why does M3-Impute perform similarly to Grape on the KIN8NM dataset, and what characteristics of KIN8NM contribute to this result?\n\nLIMITATIONS:\nAlthough M3-Impute outperforms other baselines on most datasets, two cases (KIN8NM, NAVAL dataset) highlight its limitations in handling datasets with either highly independent features or strongly linear correlations. This suggests that while M3-Impute is robust in general scenarios, its performance may degrade in datasets with extreme correlation structures.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposed a new imputations method called M3-impute. M3-impute follows the basic structure of some recent imputation methods: a undirected bipartite graph is constructed with nodes for features and samples, where edge weights correspond to observed data at the given feature-sample pair. Previous approaches use Graph Neural Networks (GNNs) to impute missing values via edge weight prediction. M3-impute improves these approaches by adding two new components on top of an initial GNN to model feature-wise and sample-wise correlations respectively. Empirical results show that M3-impute achieves competitive performance in terms of MAE for imputation across several tabular datasets.\n\nSTRENGTHS:\n- The paper is generally well written.\n- Empirical results are extensive. Many other imputations methods are included for comparison, providing a good representation for the state-of-the-art for tabular data imputation. Ablation studies and robustness studies also further strengthen the credibility of the methodology.\n\nWEAKNESSES:\n- The paper does not support categorical features. This is a big weakness compared to other imputation methods that can handle categorical features such as iterative approaches like hyperimpute. \n- The paper does not discuss the impact of missing value imputation on downstream tasks. Imputation is usually a preprocessing step, and thus assessing the impact on possible downstream tasks is paramount. For example, in supervised learning, some recent evidence suggests that mean/zero imputation is as good as more complex imputations [1, 2]. \n\n[1] Le Morvan, Marine, et al. \"What’s a good imputation to predict with missing values?.\" Advances in Neural Information Processing Systems 34 (2021): 11530-11540.\n[2] Van Ness, Mike, and Madeleine Udell. \"In defense of zero imputation for tabular deep learning.\" NeurIPS 2023 Second Table Representation Learning Workshop. 2023.\n\nQUESTIONS:\n- What is the current approach for handling categorical features? Do none of the datasets in the experiments have any categorical features, or are these feature simply being one-hot encoded?\n- How does the runtime of M3-impute compare to the other methods? I know this is briefly discussed in appendix 5, but more details runtimes would be appreciated. \n- Why are standard deviations not included for Table 1?\n\nLIMITATIONS:\nSee weaknesses. In particular, not handling categorical features is not mentioned in the paper anywhere as a limitation of the method.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes M^3-Impute, a missing value imputation method that utilizes GNNs to learn embeddings of samples and features. By incorporating feature correlation unit and sample correlation unit, M^3-Impute effectively captures correlations between features and samples for accurate imputation.\n\nSTRENGTHS:\nS1. This paper introduces a novel masking scheme that effectively utilizes missing information for modeling.\n\nS2. This paper proposes the feature correlation unit (FCU) and sample correlation unit (SCU), which help to consider feature and sample correlations during imputation.\n\nS3. Experimental evaluations on various datasets compare the proposed method with state-of-the-art approaches, demonstrating good imputation performance.\n\nWEAKNESSES:\nW1. SCU takes into account the pairwise similarity of \\mathcal{P} during its construction, which subsequently determines the scalar parameter \\alpha during imputation. The initialization of \\mathcal{P} seems to directly impact the model's performance and remains unchanged once set. It would be beneficial for the authors to discuss this aspect and, if possible, provide some experimental evidence to support their approach.\n\nW2. In Section 4.4, the paper mentions different sampling strategies for SCU and uses a new strategy in the ablation study (Table 2), which is different from the strategy mentioned in Section 3.4. The authors claim that this strategy leads to inferior performance compared to previous strategies, thereby highlighting the superiority of the ablation study results. This lacks experimental evidence and results in inconsistency in the experimental setup.\n\nW3. As a crucial parameter, the size of \\mathcal{P} directly affects the construction of SCU. Table 3 presents experimental results with different sizes, but the differences are not significant, which is somewhat counterintuitive. Although the authors discuss the experimental results in Section 4.4, the performance fluctuation of 0.01 to 0.02 does not clearly reflect the \"decrease then increase\" trend mentioned by the authors. Exploring larger peer values and providing more detailed analysis and guidance on parameter selection might be beneficial.\n\nW4. The authors emphasize the importance of specific missingness information throughout the paper. Intuitively, different types of missing data (MAR, MNAR, MCAR) might offer varying types of missingness information, potentially impacting the model's performance. While the paper experiments with data of different missingness types, a more thorough discussion of the results could enhance the motivation and clarity of the paper.\n\nW5. In Section 4.3, the caption of Table 2 references a concept, the uniform sampling strategy, which is introduced for the first time in a later subsection. The authors might consider adjusting the structure of the paper for better clarity.\n\nW6. In Figure 3, in the subfigure with the missing ratio of 0.7, two bars exceed the upper boundary and need adjustment.\n\nW7. It would be helpful to add independent labels to all subfigures, such as (a), (b), etc., to facilitate referencing.\n\nQUESTIONS:\nQ1. In Section 3.4, the authors mention that \\mathcal{P} is constructed by randomly choosing with a certain probability. Why not select directly based on the degree of similarity rather than introducing randomness? How might these different approaches impact the model's performance?\n\nQ2. To what extent does applying different GNN models to learn embeddings impact the performance of the model proposed in this paper?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
This is a novel approach for imputing missing data using mask-guided representation learning. The main contributions include the development of an imputation model that leverages both feature and sample correlations. This model improves imputation accuracy and robustness compared to existing methods. The paper also provides comprehensive experiments and ablation studies to validate the effectiveness of the proposed approach across various datasets and missing data scenarios.

STRENGTHS:
- Novelty: A unique mask-guided representation learning method that effectively combines feature-wise and sample-wise correlations.
- Comprehensive experiments
- Strong empirical performance

WEAKNESSES:
- Computational complexity 

## Minor Points

l4: "Existing imputation methods, however, fall short of considering the ‘missingness’ information in the data  during initialization and modeling the entangled feature and sample correlations  explicitly during the learning process,"
-> This is not true. Many existing methods consider missingness patterns.

The distinction between "statistical" and "learning based" methods seems off. Certainly most learning based methods are statistical, and vice versa.

l. 38 "struggles" -> struggle

QUESTIONS:
- (How) do you do HPO for competing methods?

LIMITATIONS:
- Gives little insight into why it works better on some datasets than others.
- Would be interesting to understand better how robust results are under systematic changes in datasets, e.g., different types of missingness.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a novel imputation method, based on a bipartite graph constructed from the data and the missing-data patterns, and two components which allow to measure similarities between the features and samples.\n\nThe method shows very good results in terms of MAE on several datasets for MCAR, MAR and MNAR data.\n\nSTRENGTHS:\n- The paper is well written. \n\n- Experiments are well conducted, on several datasets, with different missing-data ratios and considering MCAR, MAR or MNAR data. The authors have made an effort to compare themselves with many other imputation methods. \n\n- There is a true discussion on the parameters to choose in the experiments. The authors are honest about the performance of their method, and give explanations when another method is better.\n\nWEAKNESSES:\n- Although well presented, the method is complicated to understand. \n\n- The methods uses 8 MLPs and one GNN. The authors discuss in Appendix the computational resources, but do not compare other methods on this point.\n\nQUESTIONS:\nGeneral remarks: \n- How does this methodology relate to the simple concatenation of the mask to the data matrix, and the execution of an imputation method on the augmented matrix? (see Josse, Julie, et al. \"On the consistency of supervised learning with missing values.\")\n- Is M3-Impute supposed to work well for MNAR? This should be discussed more in details, as the authors claim that the method utilizes the data-missingness information. A remark: there exists for MIWAE an extension specifically designed for MNAR data, called not-MIWAE. Ipsen, Niels Bruun, Pierre-Alexandre Mattei, and Jes Frellsen. \"not-MIWAE: Deep generative modelling with missing not at random data.\"\nIt can be interesting to have a comparison of M3-Impute with this one in a final version. \n\nAlgorithm: \n- Figure 1: maybe the authors should add numbers in the graphics to refer to them when describing the method in the text (especially in 3.1)\n- In Algorithm 1, one of the input is the GNN model. How are hyperparameters of the GNN managed in practice? \n\nNumerical experiments: \n- In the final version, the authors should add a comparison with the missForest algorithm, which is one of the most widely used imputation methods. \n- for other methods, such as MIWAE, which hyperparameters did the authors choose? \n- in the MAR setting, how many features are selected to be observed? Did the authors take the best subset for the results?\n\nMinor comments:\n- l.15 mechanisms instead of patterns\n- l.56 \"the the\"\n- l.137 notation col_s: harmonise d and m\n- l.548 \"these remaining\" <- \"the remaining\"\n- l.291 mechanisms instead of patterns\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}