BlockEcho: Retaining Long-Range Dependencies for Imputing Block-Wise Missing Data

1 Introduction

The issue of missing data is prevalent in real-world datasets, stemming from factors such as users' reluctance to share data, variations in feature structures, privacy concerns, and unintended data corruption over time [Scheffer, 2002]. In scenarios involving large-scale and highly correlated missing data, matrix completion becomes essential to impute the unknown entries accurately and efficiently. Accurately and efficiently estimating the unknown elements in a matrix is pivotal for effective dataset analysis and supports subsequent machine learning and operations research tasks. Matrix completion finds wide application in recommender systems, trajectory recovery, phase retrieval, computer vision, genotype imputation and other fields of research [Santos et al., 2019]. Missing data isn't limited to randomly scattered elements; it can also occur in blocks, referred to as block-wise missing data. For example, discrepancies in feature sets provided by different agents can result in non-overlapping blocks of NA values during data aggregation. Moreover, certain regions may experience skipped clinical trials due to technical and financial constraints, while continuous data corruption can arise from malfunctioning monitoring devices over time.

A diverse array of methods have been developed over the years to address the matrix completion problem, spanning traditional techniques like linear interpolation, K-nearest neighbors imputation, Multiple Imputation by Chained Equations (MICE), and ensemble methods such as MissForest [Williams, 2015]. In the new data regime, Generative Adversarial Networks (GANs) have also demonstrated immense potential in accurately estimating high-dimensional data distributions [Goodfellow et al., 2020]. Consequently, recent works have explored GAN architectures for data recovery in the missing data context. However, most prevailing completion techniques display suboptimal effectiveness on blockmissing matrices, especially with increasing missing proportions. They predominantly leverage neighboring available elements to predict unknown entries, making them vulnerable to systematic feature gaps. An exception is Matrix Factorization (MF) which implicitly retains inter-sample dependencies across longer ranges [Hastie et al., 2015]. But its linearity assumptions limit capturing complex nonlinear relationships. To unite their complementary strengths and offset limitations, we propose BlockEcho - an integrated approach tailored for block-missing data that capitalizes on the strengths of both GAN and MF. It explicitly retains long-range inter-element associations via MF while modeling intrinsic nonlinear patterns through GANs. The key contributions are:

1. We have conducted an in-depth analysis and formally defined the concept of "Block-wise" missing data. Subsequently, we have innovatively introduced a solution to address this challenge, named BlockEcho.
2. We extensively benchmark BlockEcho against SOTA baselines on diverse public datasets for traffic, epidemi-

ology, and recommendations. Results demonstrate significant gains in imputation accuracy under both blockmissing and high-rate scattered missing regimes. Downstream forecasting tasks further showcase advantages.

1. We provide theoretical analysis into global optimality and convergence for the proposed integrated objective. We particularly justify the synergistic effects of MF encoding long-range dependencies and GANs locally adapting complex distributions. Rigorous proofs supplement empirical observations.

The rest of this paper is structured as follows. In Section 2 we provide a concise literature review in missing data imputation studies. Section 3 defines the problem. We elaborate our model BlockEcho in Section 4. Section 5 provides a theoretical discussion. Section 6 details our experiments and results. We conclude the paper in Section 7.

2 RelatedWork

2.1 Data Imputation Approaches

Broadly, prevailing techniques for missing data imputation fall under two categories - discriminative or generative. Discriminative methods directly estimate the conditional distribution of missing values given observed entries. This includes MICE [Buuren and Oudshoorn, 2000], MissForest [Stekhoven and B¨uhlmann, 2011] and Matrix Factorization [Hastie et al., 2015] [Yu et al., 2016]. Generative approaches instead model the complete joint data distribution and draw imputes from conditionals on observed parts. Traditionally, expectation maximization algorithms iteratively refining estimates were common [Gold and Bentler, 2000]. Recently, advances in deep generative modeling have spawned powerful imputation paradigms. Notably, J. Yoon et al. [Yoon et al., 2018] pioneered Generative Adversarial Networks (GANs) for missing data imputation through the GAIN framework in 2018. Instead of assuming explicit parametric forms, GANs can flexibly approximate arbitrary data distributions. Consequently, GAN-based solutions have risen as SOTA, significantly improving over pre-GAN techniques.

2.2 GAN-BasedMethods

Many works have since built upon GAIN, either enhancing its convergence like SGAIN (slim Wasserstein GAIN) [Neves et al., 2021], integrating auxiliary models as in GAMIN (Generative adversarial multiple imputation network) [Yoon and Sull, 2020], or exploiting complementary data modalities like images via CollaGAN (Collaborative GAN for missing image data imputation) [Lee et al., 2019]. Some explore sophisticated relationships within data beyond distributions, including: 1) Dual generative modeling of values and highdimensional feature abstractions in MisGAN [Li et al., 2019], 2) time-series self-attention mechanisms in [Zhang et al., 2021] to capture sensor inter-dependencies, 3) Pretraining classifier-guided generative models as PC-GAIN [Wang et al., 2021], 4) Distribution-centered losses like STGAN [Yuan et al., 2022], and 5) Multimodal spatio-temporal modeling via GAN+RNN+GCN [Kong et al., 2023]. ANODE-GAN [Chang et al., 2023] also augments GANs with auxiliary variational autoencoders.

Nonetheless, prevailing GAN imputation approaches overlook systematic data deficiencies like block-missing data. Mostly relying on local neighborhoods, they remain less effective as gaps exacerbate, especially at higher missing rates. Our proposed BlockEcho framework aims to address this limitation by uniquely blending GANs with matrix factorization.

3 ProblemDefinition

Consider an original data matrix X ∈ Rm∗n . M ∈ {0, 1}m∗n is set as a binary mask matrix codifying missingness, where Mij = 0 denotes element Xij as unobserved. Imposing M on X gives the incomplete matrix:

We refer to X as the original matrix, M as the mask matrix, and ˜X as the missing data matrix. The data imputation problem aims to obtain an estimation ˆX of X given ˜X and M, with the objective of minimizing the discrepancy between the imputed data and the original data. Our final output is the imputed matrix ¯X defined as

where ⊙ denotes element-wise multiplication.

For any given coordinates (il, jl), ∃ iu ≥ il + 3 and ju ≥ jl + 3, ∀ il ≤ i ≤ iu and jl ≤ j ≤ ju, Mij = 0, then we define this area as exhibiting a local block-wise missing pattern. Following the definition, the missing data in datasets can manifest as uni-block-wise, multiple-separateblock-wise, scattered, or as a mixed combination thereof.

We conducted a straightforward experiment to understand the harm of missing data. We employed a random forest algorithm to forecast the next timestamp's traffic flow for dataset PE-BAY (outlined in Section 7), without any feature engineering but normalization. Each category of missing data exhibited a 60% missing rate. Our experimental findings, as presented in Table 1, highlighted the notably adverse impact of a substantial, uniformly absent data block on the accuracy of predictions. Henceforth, unless explicitly specified, "block-wise missing" denotes the uni-block.

4 The Architecture of "BlockEcho"

In this section, we will elaborate on our integrated model, elucidating how the MF and GAN components synergize to enhance overall performance. The overarching model architecture of BlockEcho is depicted in Figure 2.

4.1 Matrix Factorization (MF) and Pre-training

Within the MF component, the missing data matrix is decomposed into two smaller matrices, denoted as the imputed row embedding matrix U and the imputed column embedding matrix V . Their dimensions are m ∗ h (adjustable to mini batch size∗h during training) and h ∗ n, where h is a hyperparameter. In essence, h serves as an indicator of the original matrix's level of freedom, with smaller values reflecting a higher degree of interdependency among the data in the original matrix.

As illustrated in Figure 2, there are two MF components in BlockEcho. One is deployed during the pre-training stage, breaking down the matrix into Up and Vp to extract longrange relationship features, which are then used to supervise model training in a transformed space through a discriminator DI . The other is directly involved in neural network training and result generation: the imputed row embedding matrix U is produced by the generator, while the imputed column embedding matrix V is directly calculated through backpropagation. Subsequently, these matrices undergo a matrix completion layer (MCL) to match the known parts of the missing data matrix. Rather than employing direct multiplication, MCL incorporates a fully connected neural network layer after multiplication. This technique is employed to introduce nonlinear relationships while upholding the inherent advantage of low constraints on the degree of freedom in the matrix factorization method. Both MF components are oriented towards minimizing the Kullback-Leibler divergence of known elements, thereby establishing the MF objective as

Noted that KL divergence is used as the loss function and it demonstrates the best performance experimentally, since all three of our datasets are non-negative and closer to a Poisson distribution. Alternatively, for datasets containing negative values, we can chose the MSE loss function as a substitute.

4.2 Generative Adversarial Nets (GANs)

The GAN framework consists of a generator G and two discriminators DI , DII as denoted in Figure 2. The generator G

takes ˜X with 0 imputed, M, and a noise matrix Z with the same size as U as the inputs, and U as the output. Next, MCL takes U as input and outputs the estimation ˆX. With ˆX, ˜X and M, we build the imputed matrix ¯X following Equation (2). Given a vector of length h, the discriminator DI will determine whether it originates from pre-trained matrix factorization or from the generator. To do this, we randomly generate Y ∈ {0, 1}m∗1 as the index vector and build a mixed-row matrix UD combining Up and U:

where I ∈ {1}1∗h . The discriminator DI will output a result of the same size as Y and the objective is

The discriminator DII is designed to discern whether the elements in ¯X is from ˜X (real) or ˆX (fake). This differs from a standard GAN model which aims to differentiate inputs ranging from entirely true to entirely false. In addition to ¯X, DII incorporates the hint matrix H as input. Hint matrix, proposed in GAIN method [Yoon et al., 2018], modifies the elements of M to 1/2 at a specified rate. This modification facilitates DII to converge more swiftly and accurately by furnishing "enough" information regarding M. The objective of this component is

In particular, the MF component serves as a structured constraint and guidance for the GAN component (analogous to the relationship between CNN and MLP), enabling the entire network to explicitly preserve long-range relationships (similar to CNN's explicit preservation of proximity relationships). This fundamental characteristic underpins the superior performance of our model compared to baselines on datasets with missing blocks or high missing rates.

4.3 Objective

According to (3), (5) and (6), the objective of the whole multiloss model is

Similar to a standard GAN, we alternate between training the generator and discriminators separately. The MF loss (3) is only connected to the generator and will be utilized for multiloss training alongside the generator's objective.

5 Theoretical Discussion

5.1 Global Optimality

Without loss of generality, we assume the elements of X as Xij = f(i, j) + sij. Here, f represents the bias effect of coordinates on X. The greater the internal correlation of X's elements, the stronger regularity f shows. {sij : i ∈ [0,m - 1] j ∈ 0, n-1} are a series of random variables which follow independent and identical distribution pd. We can define ˆX in same way, i.e., ˆXij = ˆ f (i, j)+ˆ sij , where ˆ s's probability distribution is pg. In real-world data, whether an element is masked or not may be affected by the value of the element itself. It also has different effects on the corresponding elements of ˆX in generative model. So we need to discuss the probability distribution under different mask conditions. We assume that the distributions of s corresponding to the unmasked elements in X and ˆX are pd|m=1 and pg|m=1, and to other masked elements are pd|m=0 and pg|m=0. The objective of this problem can be written as

Throughout this section, we assume that the binary mask matrix M is independent of the original matrix X. Then we can obtain pd|m=0 = pd|m=1 = pd, and (8) can be transformed into

Obviously (8), (9) and (10) are equivalent.

Below we will prove that the theoretical global optimal of (6) is the same as that in (9). The loss function for standard GAN is

For this question, it can be re-written as:

For any a, b > 0, the function alogx+blog(1-x) reaches its minimum if and only if x = a a+b . So if G is fixed, the optimal discriminator D is pd|m=1(x-f ) pd|m=1(x-f )+pg|m=0(x- ˆ f) . Substituting the optimal discriminator into (12), we can obtain

We recognize that (12) is essentially optimizing KullbackLeibler divergence. It will achieve its global minimum -log(4) if and only if ˆ f = f & pg|m=0 = pd|m=1, which is the same as (9).

Next, we will calculate the theoretical global optimal of (3). To do this, we initially introduce the entropy of the variable, Ent(A), a measure of the disorder of A. According to the definitions of entropy and mutual information, we can find that

where I denotes mutual information. Noted that I( ˆX ⊙ M, ˆX⊙(1-M)) ≤ Ent( ˆX⊙(1-M)) and the equality holds when pg|m=0 = pg|m=1, which means that if pg|m=1 is fixed, the global optimal of min(Ent( ˆX)) is pg|m=0 = pg|m=1. Therefore, (10) is equivalent to a multi-objective optimization problem

Where D( ˆX ⊙M||X ⊙M) is the primary objective, which can be Euclidean distance or Kullback-Leibler divergence. Its theoretical optimal solution is ˆ f = f & pg|m=1 = pd|m=1 under the assumption that ˆX could be fitted by any functions (actually not). Ent( ˆX) is the second objective. We assume the global optimality of (15) is pg|m=0 = pg|m=1 = pd|m=1 = p∗ , in which case Ent( ˆX) = Ent∗ ( ˆX). Thus, we have shown that problem (10) is equivalent to the multiobjective optimization problem (15).

Actually, it is very difficult to quantify Ent( ˆX) in (15) in practice. Matrix factorization is essentially a constraint on entropy in the form of the network structure. From this premise, we target at solving (15) approximately, and the optimization problem for the matrix factorization can be written as:

Where ent is a hyperparameter decided by the network structure (mainly h mentioned in section 4.1). We proceed with the following two cases: 1. ent > Ent∗ ( ˆX). In this case, the optimal solution of ˆ f = f & pg|m=1 → pd|m=1 can be reached theoretically, but it lacks a sufficient solution for pg|m=0 → pg|m=1. It corresponds to overfitting in general machine learning models. 2. ent < Ent∗ ( ˆX). In this case, ˆx cannot be fully represented by the model due to its low degree of freedom limited by the tighter constraint on entropy, which hinders ˆx from reaching the optimal of both ˆ f = f & pg|m=1 → pd|m=1 and pg|m=0 → pg|m=1. It corresponds to underfitting in general machine learning models. In either case, the theoretical global optimal of MF, i.e., problem (16), is the approximate solution of (15).

5.2 Convergence

In the previous subsection, we demonstrated that loss (3) has a similar but relatively inferior theoretical optimal solution compared to loss (6), even when we introduce non-linear transformations to MF through neural networks. However, it is anticipated intuitively (also be verified by ablation experiments in subsection 6.3), that the long-range dependencies introduced by MF will exert a significant influence on the outcomes of block-missing data imputation. This is due to the inherent challenge faced by models with a higher degree of freedom in effectively converging to fit long-range dependencies, particularly in the case of GAN models. The convergence capacity of GANs has posed a substantial challenge since its inception. In 2017, LMescheder et al. [Mescheder et al., 2017] postulated that poor convergence is attributed to the eigenvalues of the Jacobian matrix of the gradient vector field having a zero real part and an excessively large imaginary part. Consequently, strategies such as gradient penalty [Gao et al., 2020] and spectral normalization [Miyato et al., 2018] have been adopted to alleviate this limitation. In tackling the specific matrix completion problem at hand, we leverage (3), the MF loss function, and (5), which enables the generator to more accurately emulate MF, to effectively steer the convergence of (6).

At last of the section, we will prove the convergence of loss (3), indicating that D( ˆX⊙M||X⊙M) is nonincreasing under certain update rules.

Let F denote this expression. Let u denote the set of parameters of the model, and ut the set of parameters upon update at step t. We introduce an auxiliary function G(u, ut ) for F(u) satisfying G(u, u) = F(u) and G(u, ut ) ≥ F(u).Consequently, we can obtain F(ut+1 ) ≤ F(ut ) when G(ut+1 , ut ) ≤ G(ut , ut ), then F is nonincreasing under the update ut+1 = argminuG(u, ut ).

In our paper, we choose Kullback-Leibler divergence as the loss (3) because the datasets are closer to Poisson distributions. In 2001, Daniel et al.[Lee and Seung, 2000] proposed the general auxiliary function for divergence. For this problem, we can make F(u) equal to loss (3) as ˆXi,j := ∑a uiavaj and ˆXt i,j := ∑a utia vaj, and define G(u, ut ) as

It is straightforward to verify that G satisfies G(u, u) = F(u) and G(u, ut ) >= F(u) , and that F(ut+1 ) ≤ G(ut+1 , ut ) ≤ G(ut , ut ) ≤ F(ut ). Hence, F is nonincreasing under certain update rules and the convergence of loss (3) is proved.

6 Experiments and Evaluation

In this section, we rigorously test the performance of our data imputation method, BlockEcho, through a series of four carefully designed experiments using various real-world datasets available in the public domain. The initial set of experiments compare BlockEcho with state-of-the-art baselines, focusing on a fixed missing rate as high as 60%. Following this, further experiments are conducted to assess the stability of these models as the missing rate is incrementally increased. Additionally, an ablation study is performed to examine the impact of specific components of the method. Finally, we designed a downstream prediction task using the imputed data as input, to evaluate how the performance of the data imputation method influences the end-to-end performance of a realworld prediction task. The experiments in this project are based on the following open-source real-world datasets from three distinct fields:

1. Traffic Datasets (Time-Street matrix). We employ two traffic flow datasets, ME-LA and PE-BAY, which collected in the highway of Los Angeles County and from California Transportation Agencies Performance Measurement System[Li et al., 2017]. The traffic data are characterised by strong periodicity such as peak hours.
2. COVID-19 Dataset (Date-City matrix). This government dataset records daily COVID-19 cases (Cov-ca) and deaths (Cov-de) in major cities around the world
3. since 2020. Infectious disease datasets are characterized by local continuity and sudden bursts.
4. Movie Dataset (User-Movie matrix). This dataset (Movie) records movie ratings collected from the MovieLens website[Harper and Konstan, 2015], which contains a large amount of missing data.

For "scattered missing data," we generate a random matrix of the same size as the original data and rank the random numbers within it. We compare the rankings with missrate∗m∗n to determine which original elements should be masked. As for the "block-wise" missing data, we randomly select the top-left and bottom-right corners of a "block" within the feasible region limited by the missing rate, and then mask the original elements within that region.

All models were trained on an Nvidia Tesla V100S PCIE GPU and each experiment is repeated for ten times with different random seeds, and the results are averaged.

6.1 Performance of BlockEcho

In the first series of experiments, we synthetically mask 60% of the data in each dataset to generate a missing data matrix, and we compare the performance of BlockEcho with 5 baseline matrix completion models - MissForest[Stekhoven and B¨uhlmann, 2011], Matrix Factorization[Hastie et al., 2015], GAIN[Yoon et al., 2018], PC-GAIN [Wang et al., 2021] and STGAN[Yuan et al., 2022] - in terms of imputation accuracy. Data are masked in two ways: block-wise mask and scattered mask. To measure the imputation accuracy, we use RMSE to calculate the error between the modelimputed data and the original masked data, which is defined as RMSE = ||( ¯X-X)⊙(1-M)||F ∑i,j (1-M)ij . In order to have a relatively uniform measure for different datasets, we normalize the data before computing RMSE. Table 2 reports the RMSE for BlockEcho and 5 other imputation models. Results show that BlockEcho performs significantly better than baselines, especially for block-wise missing data.

6.2 Models Performance in Different Missing Rates

In the second series of experiments, we take the traffic dataset PE-BAY as a representative, dynamically adjust the data missing rate from 20% to 80%, and explore the variation trend of each model performance with the data missing rate. Figure 3 quantitatively shows this trend: the SOTA machine learning algorithms represented by GAIN perform better on datasets with a low data missing rate, but will deteriorate rapidly as the missing rate increases; the Matrix Factorization method performs more stable but has a lower ceiling for accuracy. BlockEcho absorbs the advantages of both, and thus gives the best performance at high missing rates. Although at lowmissing rates it is not as accurate as some SOTAmodels, BlockEcho significantly outperforms them at high missing rates when the datasets become more incomplete. This superiority is more obvious in block-wise missing datasets.

6.3 Ablation Study

In this subsection we design ablation experiments to verify the contribution of each component of BlockEcho to the results. The table 3 shows that when we remove any part or replace it with a more conventional alternative, it will negatively affect the results. It is worth mentioning that when the matrix factorization loss function 3 is removed and only the GAN framework is used for matrix completion, the accuracy of the imputed data will be significantly reduced, which is why most generative models (including ours) use direct loss functions to guide the convergence of GAN.

6.4 Case Study: Traffic Forecasting

In the final series of experiments, we design a downstream prediction task to illustrate how the performance of the data imputation method affects the end-to-end performance of a real-world prediction task. We take PE-BAY dataset as input to forecast the traffic conditions at the next timestamp. For the input datasets, we take the original data, the imputed data with our model BlockEcho, and the imputed data from each baseline model for comparison. For traffic forecasting, we use Random Forests (RF) with the same set of hyperparameter settings with the same feature engineering for all input datasets. We use Weighted Mean Absolute Percentage Error (WMAPE) as the error metric to measure the prediction performance. Table 4 summarizes the prediction results. Comparing Table 2 and Table 4, we find that models with higher data imputation accuracy tend to give better results in subsequent prediction tasks, and BlockEcho, the best-performing data imputation model, also gives the lowest prediction errors in subsequent traffic forecasting tasks.

7 Conclusion and FutureWork

In this paper, we excavate and mathematically define the issue of "block-wise" missing data and innovatively propose the solution, BlockEcho. Experiments on various data sets, especially with block missing data, show that our method outperforms other SOTA methods. Our future work will extend to federated learning where block-wise missing data widely appear.

Acknowledgments

This work was supported by National Key R&D Program of China (2022YFB4501500,2022YFB4501504).