RIFLE: Imputation and Robust Inference from Low Order Marginals

Sina Baharlouei 1 Kelechi Ogudu 2 Peng Dai 1 Sze-chuan Suen 1 Meisam Razaviyayn 1

Abstract
The ubiquity of missing values in real-world
datasets poses a challenge for statistical inference
and can prevent similar datasets from being analyzed in the same study, precluding many existing datasets from being used for new analyses.
While an extensive collection of packages andalgorithms have been developed for data imputation, the overwhelming majority perform poorly
if there are many missing values and low sample
sizes, which are unfortunately common characteristics in empirical data. Such low-accuracy
estimations adversely affect the performance ofdownstream statistical models. We develop a statistical inference framework for predicting thetarget variable in the presence of missing data
without imputation. Our framework, RIFLE (Robust InFerence via Low-order moment Estimations), estimates low-order moments of the underlying data distribution with corresponding confidence intervals to learn a distributionally robustmodel. We specialize our framework to ridge
linear regression, where the resulting min-max
problem is efficiently solved by applying the alternating direction method of multipliers (ADMM)
on the dual problem. This framework can also
be adapted to impute missing data. We compare
RIFLE with state-of-the-art approaches (including MICE, Amelia, MissForest, KNN-imputer,
MIDA, and Mean Imputer) in numerical experiments. Our experiments demonstrate that RIFLE
outperforms other benchmark algorithms when
the percentage of missing values is high and/or
when the number of data points is relatively small.
RIFLE is publicly available1 .

1. Introduction
Machine learning algorithms have shown promise whenapplied to various problems, including healthcare, finance,social data analysis, image processing, and speech recognition. However, this success mainly relied on the availability of large-scale, high-quality datasets, which may bescarce in many practical problems, especially in medical
and health applications (Pedersen et al.; Sterne et al., 2009;
Beaulieu-Jones et al., 2018). Moreover, many experiments
and datasets suffer from the small sample size in such applications. Despite the availability of a small number ofdata points in each study, an increasingly large number ofdatasets are publicly available. To fully and effectively utilize information on related research questions from diverse
datasets, information across various datasets (e.g., different questionnaires from multiple hospitals with overlapping
questions) must be combined in a reliable fashion. Afterappending these datasets together, the obtained dataset cancontain large blocks of missing values, as they may notshare the same features (Figure 1).

Industrial and Systems Engineering Department, Universityof Southern California 2 Computer Science Department, University of Southern California. Correspondence to: Sina Baharlouei<baharlou@usc.edu>.

Figure 1. Consider the problem of predicting the trait y fromfeature vector (x1 , . . . , x100 ). Suppose that we have accessto three data sets: The first dataset includes the measurements of (x1 , x2 , . . . , x40 , y) for n1 individuals. The second
dataset collects data from another n2 individuals by measuring(x30 , . . . , x80 ) with no measurements of the target variable y init; and the third dataset contains the measurements from the variables (x70 , . . . , x100 , y) for n3 number of individuals. How one
should learn the predictor ŷ = h(x1 , . . . , x100 ) from these threedatasets?

ICML 2023 Workshop on Duality for Modern Machine Learning,
Honolulu, Hawaii, USA. Copyright 2023 by the author(s).
1
https://github.com/
optimization-for-data-driven-science/RIFLE

Related Works: There are three general approaches forhandling missing values in classification and regression
tasks. A Naive method is to remove the rows containing
missing entries. However, such an approach is not an option

1

RIFLE: Imputation and Robust Inference from Low Order Marginals

when the percentage of missingness in a dataset is high.
For instance, as demonstrated in Figure 1, the entire dataset
will be discarded if we eliminate the rows with at least one
missing entry.
The most common approach for handling missing values ina learning task is to impute them in a pre-processing stage.
The general idea behind data imputation approaches is thatthe missing values can be predicted using the other available
data points and correlated features. Imputation algorithms
cover a wide range of methods, including imputing missing entries with the columns means Little & Rubin (2019,
Chapter 3) (or median), least-square and linear regressionbased methods (Raghunathan et al., 2001; Kim et al., 2005;
Zhang et al., 2008; Cai et al., 2006; Buuren & GroothuisOudshoorn, 2010), matrix completion and expectation maximization approaches Dempster et al. (1977); Ghahramani& Jordan (1994); Honaker et al. (2011), KNN based (Troyanskaya et al., 2001), Tree based methods (Stekhoven &Bühlmann, 2012; Xia et al., 2017), and methods using different neural network structures. Appendix A presents acomprehensive review of these methods.
The imputation of data allows practitioners to run standard
statistical algorithms requiring complete data. However, theprediction model’s performance can be highly reliant on theaccuracy of the imputer. High error rates in the predictionof missing values by the imputer can lead to the catastrophic
performance of the downstream statistical methods executedon the imputed data.
Another class of methods for inference in the presence of
missing values relies on robust optimization over the uncertainty sets on missing entries. Shivaswamy et al. (2006)
and Xu et al. (2009) adopt robust optimization to learn theparameters of a support vector machine model. They consider uncertainty sets for the missing entries in the dataset
and solve a min-max problem over those sets. The obtainedclassifiers are robust to the uncertainty of missing entries
within the uncertainty regions. In contrast to the imputationbased approaches, the robust classification formulation does
not carry the imputation error to the classification phase.
However, finding appropriate intervals for each missing entry is challenging, and it is unclear how to determine theuncertainty range in many real datasets. Moreover, their
proposed algorithms are limited to the SVM classifier.
In this paper, we propose RIFLE (Robust InFerence via
Low-order moment Estimations) for the direct inferenceof a target variable based on a set of features containing
missing values. The proposed framework does not require
the data to be imputed in a pre-processing stage. However,it can also be used as a pre-processing tool for imputing
data. The main idea of the proposed framework is to estimate the first and second-order moments of the data andtheir confidence intervals by bootstrapping on the available

Figure 2. Prediction of the target variable without imputation. RIFLE estimates confidence intervals for low-order (first and secondorder) marginals from the input data containing missing values.
Then, it solves a distributionally robust problem over the set of all
distributions whose low-order marginals are within the estimatedconfidence intervals.

data matrix entries. Then, RIFLE finds the optimal parameters of the statistical model for the worst-case distribution
with the low-order moments (mean and variance) within theestimated confidence intervals (See Figure 2). Comparedto Shivaswamy et al. (2006); Xu et al. (2009) we estimate
uncertainty regions for the low-order marginals using the
Bootstrap technique. Furthermore, our framework is not
restricted to any particular machine learning model, such assupport vector machines (Xu et al., 2009).

2. Robust Inference via Estimating Low-order
Moments
RIFLE is based on a distributionally robust optimization(DRO) framework over low-order marginals. Assume that(x, y) ∈ Rd × R follows a joint probability distribution P ∗ .
A standard approach for predicting the target variable ygiven the input vector x is to find the parameter θ thatminimizes the population risk with respect to a given loss
function `:h 
imin E(x,y)∼P ∗ ` x, y; θ .
(1)
θSince the underlying distribution of data is rarely available
in practice, the above problem cannot be directly solved.
The most common approach for approximating (1) is tominimize the empirical risk with respect to n given i.i.d
samples (x1 , y1 ), . . . , (xn , yn ) drawn from the joint distribution P ∗ :n1X
min`(xi , yi ; θ).
θn i=1
The above empirical risk formulation assumes that all entriesof xi and yi are available. Thus, to utilize the empirical risk
minimization (ERM) framework in the presence of missing
values, one can either remove or impute the missing data
points in a pre-processing stage. Training via robust optimization is a natural alternative in the presence of missing
data. Shivaswamy et al. (2006); Xu et al. (2009) suggest
the following optimization problem that minimizes the loss
function for the worst-case scenario over the defined uncertainty sets per data points:
n

min
θ

maxn

{δi ∈Ni }i=1

1X`(xi − δi , yi ; θ),
n i=1

(2)

RIFLE: Imputation and Robust Inference from Low Order Marginals

where Ni represents the uncertainty region of data point i.
Shivaswamy et al. (2006) obtains the uncertainty sets byassuming a known distribution on the missing entries ofdatasets. The main issue in their approach is that the constraints defined on data points are totally uncorrelated. Xu
et al. (2009) on the other hand defines Ni as a “box” constraint around the data point i such that they can be linearlycorrelated. For this specific case, they show that solving thecorresponding robust optimization problem is equivalent tominimizing a regularized reformulation of the original loss
function. Such an approach has several limitations: First, it
can only handle a few special cases (SVM loss with linearlycorrelated perturbations on data points). Furthermore, Xu
et al. (2009) is primarily designed for handling outliers andcontaminated data. Thus, they do not offer any mechanism
for the initial estimation of xi when several vector entries
are missing. In this work, we instead take a distributionallyrobust approach by consiering uncertainty on the distribution of the data instead of defining uncertainty set for eachdata point. In particular, we aim to fit the best parametersof a statistical learning model for the worst distribution in agiven uncertainty set by solving:
min maxP ∈Pθ

E(x,y)∼P [`(x, y; θ)],(3)

where P is an uncertainty set over the underlying distribution of data. A key observation is that defining the uncertainty set P in (3) is easier and computationally more
efficient than defining the uncertainty sets {Ni }ni=1 in (2).
In particular, the uncertainty set P can be obtained naturally
by estimating low-order moments of data distribution using
only available entries. To explain this idea and to simplify
the notations, let z = (x, y), µ̄z , E[z], and C̄z , E[zzT ].
While µ̄z and C̄z are typically not known exactly, one canestimate them (within certain confidence intervals) from theavailable data by simply ignoring missing entries. Moreover, we can estimate the confidence intervals via bootstraping. Particularly, we can estimate µzmin , µzmax , Czmin ,and Czmax from data such that µzmin ≤ µ̄z ≤ µzmax and
Czmin ≤ C̄z ≤ Czmax with high probability (where the inequalities for matrices and vectors denote component-wise
relations). In Appendix B, we show how a bootstrapping
strategy can be used to obtain the confidence intervals described above. Given these estimated confidence intervals
from data, (3) can be reformulated asminθ

Let us specialize our framework to the ridge linear regression model. In the absence of missing data, ridge regression
finds optimal regressor parameter θ by solving
min kXθ − yk22 + λkθk22 ,θor equivalently by minimizing:θ T XT Xθ − 2θ T XT y + λkθk22 .

(5)Thus, having the second-order moments of the data C =XT X and b = XT y is sufficient for finding the optimal
solution. In other words, it suffices to compute the innerproduct of any two column vectors ai , aj of X, and theinner product of any column ai of X with vector y. Since
the matrix X and vector y are not fully observed due to theexistence of missing values, one can use the same approachas (12) to compute the point estimators C0 and b0 . These
point estimators can be highly inaccurate, especially when
the number of non-missing rows for two given columns issmall. In addition, if the pattern of missing entries does notfollow the MCAR assumption, the point estimators are notunbiased estimators of C and b.
3.1. A Distributionally Robust Formulation of Linear
Regression
As we mentioned above, to solve the linear regression problem, we only need to estimate the second-order moments ofthe data (XT X and XT y). Thus, the distributionally robust
formulation described in (4) is equivalent to the following
optimization problem for the linear regression model:
θPµzmin ≤ EP [z] ≤ µzmax ,3. Robust Linear Regression in the Presenceof Missing Values

min max θ T Cθ − 2bT θ + λkθk22

max EP [`(z; θ)]s.t.

low order moments of the data, we use `∞ constraints thatare computationally more natural in the presence of missing entries when combined with bootstrapping. Furthermore, while it can be applied to general convex losses, their
method relies on the ellipsoid and the existence of oracles
for performing the steps of the ellipsoid method, which is
not applicable in modern high-dimensional problems. Moreover, they assume concavity in data (the existence of someoracle to return the worst-case data points) that is practicallyunavailable even in convex loss functions (including linear
regression in our work).

(4)

Czmin ≤ EP [zzT ] ≤ Czmax .
Gao & Kleywegt (2017) utilize the distributionally robust
optimization as (3) over the set of positive semi-definite(PSD) cones for robust inference under uncertainty. While
their formulation considers `2 balls for the constraints on

s.t.

C,b

C0 − c∆ ≤ C ≤ C0 + c∆,(6)b0 − cδ ≤ b ≤ b0 + cδ,C  0,
where the last constraint guarantees that the covariance matrix is positive and semi-definite. We dicuss the procedureof estimating the confidence intervals (b0 , C0 , δ, and ∆) inAppendix B.

RIFLE: Imputation and Robust Inference from Low Order Marginals

3.2. RIFLE for Ridge Linear Regression
Since the objective function in 6 is convex in θ (ridge regression) and concave in b and C (linear), the minimizationand maximization sub-problems are interchangeable (Sionet al., 1958). Thus, we can equivalently rewrite Problem (6)as:max g(C, b)C,bs.t. C0 − c∆ ≤ C ≤ C0 + c∆,(7)b0 − cδ0 ≤ b ≤ b0 + cδ,C  0,
where g(b, C) = minθ θ T Cθ−2bT θ+λkθk2 . Function gcan be computed in closed-form given any pair of (C, b) byθ = (C + λI)−1 b. Based on that, we can apply projectedgradient ascent to function g to find an optimal solutionof (7) as described in Algorithm 1. At each iteration of thealgorithm, we first perform one step of projected gradientascent on matrix C and vector b; then we update θ in closedform for the obtained C and b. We use C0 and b0 obtained
from Equation (12) (see Appendix B) as the initialization
points for C and b respectively. The projection of b to the
Algorithm 1 Projection Gradient Ascent on functiong(C, b) in the Presence of Missing Values
1: Input: C0 , b0 , ∆, δ, T
2: Initialize: C = C0 , b = b0 .
3: for i = 1, . . . , T do 
4:Update C = Π∆+ C + αθθ T
5:Update b = Πδ (b − 2αθ)6:Set θ = (C + λI)−1 b
7: end forbox contraint b0 − cδ ≤ b ≤ b0 + cδ can be done entriwise
and has the following closed-form
if b0i − cδi ≤ bi ≤ b0i + cδi ,biΠδ (bi ) = b0i − cδi if bi ≤ b0i − cδi ,bi + cδi if b0i + cδi ≤ bi .
Theorem 1. Assume that D = kC0 − C∗ k2F + kb0 − b∗ k22 .
Then Algorithm 1 computes ansolution of the -stationaryLDobjective function in (7) in O  iterations, where L =2
λ.

Proof. The proof is relegated to Appendix F.
In Appendix C, we show how using the acceleration methodof Nesterovcanthe convergence rate of Algoq improveLD
rithm 1 to O
. A technical issue of Algorithm 1 and
its accelerated version presented in Appendix C is that applying the simultaneous projection of C to both box constraints
and the set of positive semidefinite matrices (Π∆+ [C]) is

challenging and cannot be done in closed-form. A more resilient approach to handle both constraints simultaneously is
to write the dual problem of the maximization problem andhandle the resulting constrained minimization problem with
the Alternating Direction Method of Multipliers (ADMM).

4. Solving the Dual Problem of the Robust
Ridge Linear Regression via ADMM
The Alternating Direction Method of Multipliers (ADMM)
is a popular algorithm for efficiently solving linearly constrained optimization problems (Gabay & Mercier, 1976;
Hong et al., 2016). It has been extensively applied to largescale optimization problems in machine learning and statistical inference in recent years (Assländer et al., 2018;
Zhang et al., 2018). Consider the following optimization
problem consisting of two blocks of variables x and y thatare linearly coupled:
min f (w) + g(z)
w,z

(8)

s.t. Aw + Bz = c,
The augmented Lagrangian of the above problem can be
written as:
min f (w) + g(z) +hAw + Bz − c, λiw,z+ ρ2 kAw + Bz − ck2 ,(9)ADMM schema updates the primal and dual variables iteratively as presented in Algorithm 2.
Algorithm 2 Schema of ADMM algorithm Applied to theAugmented Lagrangian Problem
1: for t = 1, . . . , T do2:
wt+1 = argminw f (w) + hAw + Bzt − c, λi +ρt2
2 kAw + Bz − ck3:

zt+1 = argminz f (wt+1 )+hAwt+1 +Bz−c, λi+ρt+1+ Bz − ck2
2 kAw

4:
λt+1 = λt + ρ(Awt+1 + Bzt+1 − c)5: end for

As we mentioned earlier, simultaneous projection of C to
the set of positive semi-definite matrices and the box constraint Cmin ≤ C ≤ Cmax in Algorithm 1 is computationally expensive. Moreover, careful step-size tuning isnecessary to avoid inconsistency and guarantee convergencein that algorithm.
An alternative approach for solving Problem (6) that avoids
removing the PSD constraint in the implementation of Algorithm 1 is to solve the dual of the inner maximization
problem. Since the maximization problem is concave with

RIFLE: Imputation and Robust Inference from Low Order Marginals

respect to C and b, and the relative interior of the feasible set of constraints is non-empty, the duality gap is zero.
Hence, instead of solving the inner maximization problem,
we can solve its dual which is a minimization problem. Theorem 2 describes the dual problem of the inner maximization
problem in (6). Thus, Problem (6) can be alternatively formulated as a minimization problem rather than a min-max
problem. We can solve such a constrained minimization
problem efficiently via the ADMM algorithm. As we willshow, the ADMM algorithm applied to the dual problem
does not need tuning of step size or applying simultaneous
projections to the box constraints and positive semi-definite(PSD) constraints.

Problem (10) in order, we can write the corresponding augmented Lagrangian function as:
0

θ,A,B,d,e,H−hCmin , Ai + hCmax , Bi + λkθks.t.

−θθ T − A + B − H = 0,

H  0.
Proof. The proof is relegated to Appendix F.
To apply the ADMM method to the dual problem, we needto divide the optimization variables into two blocks as in (8)to solve both sub-problems in Algorithm 2 efficiently. To do0000so, first, we introduce the auxiliary variables d , e , θ , A
0
0 0
and B to the dual problem. Also, let G = H + θ θ T .
Therefore, Problem (2) is equivalent to:−hbmin , di + hbmax , ei
−hCmin , Ai + hCmax , Bi + λkθk2s.t.

B − A = G,
2θ − d + e = 0,
0

(10)

0

A = A ,B = B ,000d = d ,e = e ,θ = θ ,
0

0

0

0

A , B , d , e ≥ 0,
0

0

0

0

+hB − B , MB i + ρ2 kB − B k2F00+hd − d , µd i + ρ2 kd − d k200

+he − e , µe i + ρ2 ke − e k200+hθ − θ , µθ i + ρ2 kθ − θ k2
+h2θ − d + e, ηi + ρ2 k2θ − d + ek2+hB − A − G, Γi + ρ2 kB − A − Gk2F
0

0

0

0

A , B , d , e ≥ 0,
0A, B, d, e ≥ 0,

min

0

0

G  θ θ T,22θ − d + e = 0,

θ,A,B,d,e,H−hbmin , di + hbmax , ei+hA − A , MA i + ρ2 kA − A k2Fs.t.

−hbmin , di + hbmax , ei

0−hCmin , Ai + hCmax , Bi + λkθk2

Theorem 2. (Dual Problem) Problem (6) can be equivalently formulated as:
min

min0

θ,θ ,A,A ,B,B
00d,d ,e,e ,G0G  θ θ T.
Since handling both constraints on θ in Problem (2) is diffi0
cult, we interchange θ with θ in the first constraint. Moreover, the non-negativity constraints on A, B, d and e are000exchanged with non-negativity constraints on A , B , d0and e . Now, we show how to apply ADMM schema to
Problem (10) to obtain Algorithm 5. As we discussed earlier, we consider two separate blocks of variables w =
0
0
0
0
0
(θ, d, e, G, B , A ) and z = (d , e , θ , B, A). Assigning Γ, η, MA , MB , µd , µe , and µθ to the constraints of(11)
At each iteration of the ADMM algorithm, the parametersof one block are fixed, and the optimization problem issolved with respect to the parameters of the other block.
The details of the aforementioned ADMM procedure for
optimizing (11) and its convergence analysis are presented
in Appendix D.
Remark 3. The optimal solution obtained from the ADMM
algorithm can be different from the one given by Algorithm 1
because we remove the positive semi-definite constraint on
C in the latter. We investigate the difference between solutions of two algorithms in three cases: First, we generatea small positive semi-definite matrix C∗ and the matrix ofconfidence intervals (∆) as follows:
97 40 92
0.2 0.3 0.2
C∗ = 40 17 38 , ∆ = 0.3 0.1 0.2 .
92 38 88
0.1 0.3 0.1
Moreover, let b∗ and δ are generated as follows:


 
6.65
0.1
b∗ = 8.97 , δ = 0.2 .
5.40
0.2
Initializing both algorithms with a random matrix within
Cmin = C∗ − ∆ and Cmax = C∗ + ∆, and a random
vector within bmin = b∗ − δ and bmax = b∗ + δ, ADMM
algorithm returns a different solution from Algorithm 1.
Besides, the difference in the performance of algorithms
during the test phase can be observed in the experiments on
synthetic datasets depicted in Figure 5 as well, especiallywhen the number of samples is smaller.

5. Experiments
In this section, we evaluate the performance of RIFLE on
a diverse set of inference tasks in the presence of missing values. We compare RIFLE’s performance to several

RIFLE: Imputation and Robust Inference from Low Order Marginals

state-of-the-art approaches for data imputation on synthetic
and real-world datasets. The experiments are designed ina manner that the sensitivity of the model to factors suchas the number of samples, data dimension, types, and proportion of missing values can be evaluated. The descriptionof all datasets used in the experiments can be found in Appendix G.
5.1. Generating MCAR and MNAR Missing Values
To evaluate RIFLE and other state-of-the-art imputation
approaches, we need to have access to the ground-truth
values of the missing entries. Hence, we artificially mask aproportion of available data entries and predict them withdifferent imputation methods. A method performs better
than others if the predicted missing entries are closer to theground-truth values. To measure the performance of RIFLE
and the existing approaches on a regression task for a given
test dataset consisting of N data points, we use normalized
root mean squared error (NRMSE), defined as:
q P
N
1
2
i=1 (yi − ŷi )N
NRMSE = q PN12i=1 (yi − ȳ)Nwhere yi , ŷi , and ȳ represent the true value of the i-th data
point, the predicted value of the i-th data point, and theaverage of true values of data points, respectively. In all experiments, generated missing entries follow either a missing
completely at random (MCAR) or a missing not at random(MNAR) pattern. A discussion on the procedure of generating these patterns can be found in Appendix E.

We keep the criterion as “MSE” since our performance evaluation measure is NRMSE. Moreover, we tune the numberof iterations and number of estimations (number of trees)by checking values from {5, 10, 20} and {50, 100, 200}, respectively. We do not change the structure of the neural networks for MIDA (Gondara & Wang, 2018) and GAIN (Yoonet al., 2018), and the default versions are performed for imputing datasets.
5.3. RIFLE Convergence
We presented three algorithms for solving the robust linear
regression problem formulated in (6): Projected gradientascent (Algorithm 1, Nesterov acceleration method (Algorithm 4), and Alternating Direction Method of Multipliers(ADMM) (Algorithm 5) applied on the dual problem. We
established the convergence rate of the projected gradientascent and Nesterov acceleration methods in Theorem ,1and Theorem ,5 respectively. To investigate the convergence of the ADMM algorithm and its dependence on ρ,
we perform Algorithm 5 on the Super Conductivity dataset(Description in Appendix G) with 30% MCAR missing values. Figure 3 demonstrates the convergence of the ADMM
algorithm for multiple values of ρ applied to the Super Conductivity dataset as described above. As can be observed,
decreasing the value of ρ accelerates the ADMM convergence to the optimal value. Note that for ρ = 0.2, theobjective function is smaller than the final value in the firstfew iterations. The reason is that for those iterations, thesolution is not feasible (as observed in the right figure). The
final solution is the optimal feasible solution.

5.2. Tuning Hyper-parameters of RIFLE
The hyper-parameter c in (6) controls the robustness of themodel by adjusting the size of confidence intervals. This parameter is tuned by performing a cross-validation procedure
over the set {0.1, 0.25, 0.5, 1, 2, 5, 10, 20, 50, 100}, and theone with the lowest NMRSE is chosen. The default value inthe implementation is c = 1 since it consistently performs
well over different experiments. Furthermore, λ, the hyperparameter for the ridge regression regularizer, is tuned bychoosing 20% of the data as the validation set from the set{0.01, 0.1, 0.5, 1, 2, 5, 10, 20, 50}. To tune K, the numberof bootstrap samples for estimating the confidence intervals,
we tried 10, 20, 50, and 100. No significant difference isobserved in terms of the test performance for the above
values.

Figure 3. Convergence of ADMM algorithm to the optimal solution of Problem (2) for different values of ρ. The left plot measures
the objective function of Problem (2) per iteration (without considering the constraints), while the right plot demonstrates theconstraint violation of the algorithm per iteration. The constraint
violation can be measured by adding all regularization terms in theaugmented Lagrangian function formulated in Problem (11).

Furthermore, we tune the hyper-parameters of the competing packages as follows. For KNN-Imputer (Troyanskaya
et al., 2001), we try {2, 10, 20, 50} for the number of neighbors (K) and pick the one with the highest performance.
For MICE (Buuren & Groothuis-Oudshoorn, 2010) andAmelia (Honaker et al., 2011), we generate 5 different imputed data and pick the one with the highest performanceon the test data. MissForest has multiple hyper-parameters.

In the next experiment, we compare the three proposed algorithms in terms of the number of iterations required to reacha certain level of test accuracy on the Super Conductivity
dataset. The number of training samples is 1000, containing
40% of MCAR missing values on both input features andthe target variable. The test dataset contains 2000 samples.
As depicted in Figure 4, ADMM and Nesterov’s algorithms

RIFLE: Imputation and Robust Inference from Low Order Marginals

require less number of iterations to reach the -optimal solution compared to Algorithm 1. However, the cost periteration of the ADMM algorithm (Algorithm 5) is higher
than the Nesterov acceleration and Algorithm 1.

Figure 5. Comparing the consistency of RIFLE, MissForest, KNN
Imputer, MICE, Amelia, and Expectation Maximization methods
on a synthetic dataset containing 40% of missing values.

Figure 4. The performance of the Nesterov acceleration method,
projected gradient ascent, and ADMM on the Super Conductivity
dataset vs. the number of iterations.

wherethe proportion of missing values are relatively highn(1−p)≈ O(1) . RIFLE outperforms these methods indalmost all cases and performs slightly better than MissForest, which uses a highly non-linear model (random forest)
to impute missing values.

5.4. RIFLE Consistency
In Figure 5, we investigate the consistency of RIFLE on
synthetic datasets with different proportions of missing values. The synthetic data has 50 input features following a
jointly normal distribution with the mean whose entries are
randomly chosen from the interval (−100, 100). Moreover,the covariance matrix equals Σ = SS T where S elements
are randomly picked from (−1, 1). The dimension of S is
50 × 20. The target variable is a linear function of input
features added to a mean zero normal noise with a standard
deviation of 0.01. As depicted in Figure 5, RIFLE requires
fewer samples to recover the ground-truth parameters of themodel compared to MissForest, KNN Imputer, Expectation
Maximization (Dempster et al., 1977), and MICE. Amelia’s
performance is significantly good since the predictors havea joint normal distribution and the linear underlying model.
Note that by increasing the number of samples, the NRMSE
of our framework converges to 0.01, which is the standard
deviation of the zero-mean Gaussian noise added to eachtarget value (the dashed line).
5.5. Data Imputation via RIFLE
As explained in Section 3, while the primary goal of RIFLE is to learn a robust regression model in the presenceof missing values, it can also be used as an imputation
tool. We run RIFLE and several state-of-the-art approaches
on five datasets from the UCI repository (Dua & Graff,2017) (Spam, Housing, Clouds, Breast Cancer, and Parkinson datasets) with different proportions of MCAR missing values (the description of the datasets can be found in
Appendix G). Then, we compute the NMRSE of imputed
entries. Table 1 (see Appendix I) shows the performanceof RIFLE compared to other approaches for the datasets

Figure 6. Performance Comparison of RIFLE, MICE, and MissForest on four UCI datasets: Parkinson, Spam, Wave Energy Converter, and Breast Cancer. For each dataset, we count the numberof features that each method outperforms the others.

5.6. Sensitivity of RIFLE to the Number of Samples
and Proportion of Missing Values
In this section, we analyze the sensitivity of RIFLE and otherstate-of-the-art approaches to the number of samples and the
proportion of missing values. In the experiment in Figure 6,
we create 5 datasets containing 40%, 50%, 60%, 70%, and
80% of MCAR missing values, respectively, for four real
datasets (Spam, Parkinson, Wave Energy Converter, andBreast Cancer) from UCI Repository (Dua & Graff, 2017)(the description of the datasets can be found in Appendix G).

RIFLE: Imputation and Robust Inference from Low Order Marginals

Given a feature in a dataset containing missing values, we
say an imputer wins that feature if the imputation error interms of NRMSE for that imputer is less than the error of theother imputers. Figure 6 reports the number of features won
by each imputer on the created datasets described above. As
we observe, the number of wins for RIFLE increases as we
increase the proportion of missing values. This observationshows that the sensitivity of RIFLE as an imputer to the
proportion of missing values is less than MissForest andMICE in general.
Figure 6 does not show how the NRMSE of one imputeris changed when the proportion of missing values is increased. Next, we analyze the sensitivity of RIFLE andseveral imputers to change in missing value proportions.
Fixing the proportion of missing values, we generate 10
random datasets containing missing values in random locations on the Drive dataset (the description of datasets isavailable in Appendix G). We impute the missing values
for each dataset with RIFLE, MissForest, Mean Imputation,and MICE. Figure 7 shows the average and the standard

Figure 7. Sensitivity of RIFLE, MissForest, Amelia, KNN Imputer,
MIDA, and Mean Imputer to the percentage of missing values on
the Drive dataset. Increasing the percentage of missing value entries degrades the benchmarks’ performance compared to RIFLE.
KNN-imputer implementation cannot be executed on datasets containing 80% (or more) missing entries. Moreover, Amelia and
MIDA do not converge to a solution when the percentage of missing value entries is higher than 70%.

deviation of these 4 imputers’ performances for different
proportions of missing values (10% to 90%). Figure 7 depicts the sensitivity of MissForest and RIFLE to the proportion of missing values in the Drive dataset. We select 400
data points for each experiment with different proportionsof missing values (from 10% to 90%) and report the average
NRMSE of imputed entries. Finally, in Figure 8, we haveevaluated RIFLE and other methods on the BlogFeedback
dataset (see Appendix G) containing 40% missing values.
The results show that RIFLE’s performance is less sensitive
to decreasing the number of samples.

Figure 8. Sensitivity of RIFLE, MissForest, MICE, Amelia, Mean
Imputer, KNN Imputer, and MIDA to the number of samples
for the imputations of Blog Feedback dataset containing 40% of
MCAR missing values. When the number of samples is limited,
RIFLE outperforms other methods, and its performance is very
close to the non-linear imputer MissForest for larger samples.

5.7. Performance Comparison on Real Datasets
In this section, we compare the performance of RIFLE toseveral state-of-the-art approaches, including MICE (Buuren & Groothuis-Oudshoorn, 2010), Amelia (Honaker
et al., 2011), MissForest (Stekhoven & Bühlmann,2012), KNN Imputer (Raghunathan et al., 2001), andMIDA (Gondara & Wang, 2018). There are two primaryways to do this. One method to predict a continuous target
variable in a dataset with many missing values is first to impute the missing data with a state-of-the-art package, thenrun a linear regression. An alternative approach is to directly
learn the target variable, as we discussed in Section 3.
Table 2 compares the performance of mean imputation,
MICE, MIDA, MissForest, and KNN to that of RIFLE
on three datasets: NHANES, Blog Feedback, and superconductivity. Both Blog Feedback and Superconductivity
datasets contain 30% of MNAR missing values generated
by Algorithm 6, with 10000 and 20000 training samples,
respectively. The description of the NHANES data and its
distribution of missing values can be found in Appendix G.
Conclusion: In this paper, we proposed a distributionallyrobust optimization framework over the distributions withthe low-order marginals within the estimated confidence
intervals for inference and imputation of datasets in thepresence of missing values. We developed algorithms forrobust ridge linear regression with convergence guarantees.
In particular, we apply the alternating direction methodof multipliers (ADMM) on the dual of the original minmax problem. The resulting algorithm is convergent to
the original problem’s optimal solution. The performanceof the method is evaluated on synthetic and real datasetswith different numbers of samples, dimensions, missing
value proportions, and types of missing values. In mostexperiments, RIFLE consistently outperforms other existing
methods.

RIFLE: Imputation and Robust Inference from Low Order Marginals

References
Assländer, J., Cloos, M. A., Knoll, F., Sodickson, D. K.,
Hennig, J., and Lattanzi, R. Low rank alternating direction method of multipliers reconstruction for mr fingerprinting. Magnetic resonance in medicine, 79(1):83–96,
2018.
Barazandeh, B. and Razaviyayn, M. Solving non-convex
non-differentiable min-max games using proximal gradient method. arXiv preprint arXiv:2003.08093, 2020.
Beaulieu-Jones, B. K., Lavage, D. R., Snyder, J. W., Moore,
J. H., Pendergrass, S. A., and Bauer, C. R. Characterizing
and managing missing structured data in electronic healthrecords: data analysis. JMIR medical informatics, 6(1):e11, 2018.
Bertsimas, D., Pawlowski, C., and Zhuo, Y. D. From predictive methods to missing data imputation: an optimization
approach. The Journal of Machine Learning Research,
18(1):7133–7171, 2017.
Bubeck, S. Convex optimization: Algorithms and complexity. arXiv preprint arXiv:1405.4980, 2014.
Buuren, S. v. and Groothuis-Oudshoorn, K. mice: Multivariate imputation by chained equations in r. Journal of
statistical software, pp. 1–68, 2010.
Cai, Z., Heydari, M., and Lin, G. Iterated local least squares
microarray missing value imputation. Journal of bioinformatics and computational biology, 4(05):935–957, 2006.
Dempster, A. P., Laird, N. M., and Rubin, D. B. Maximum likelihood from incomplete data via the em algorithm. Journal of the Royal Statistical Society: Series B(Methodological), 39(1):1–22, 1977.
Dua, D. and Graff, C. UCI machine learning repository,
2017. URL http://archive.ics.uci.edu/ml.
Gabay, D. and Mercier, B. A dual algorithm for the solution of nonlinear variational problems via finite element
approximation. Computers & mathematics with applications, 2(1):17–40, 1976.
Gao, R. and Kleywegt, A. J. Distributionally robust
stochastic optimization with dependence structure. arXiv
preprint arXiv:1701.04200, 2017.

He, B. and Yuan, X. On non-ergodic convergence rate ofdouglas–rachford alternating direction method of multipliers. Numerische Mathematik, 130(3):567–577, 2015.
Honaker, J., King, G., and Blackwell, M. Amelia ii: A
program for missing data. Journal of statistical software,45(7):1–47, 2011.
Hong, M., Luo, Z.-Q., and Razaviyayn, M. Convergence
analysis of alternating direction method of multipliers
for a family of nonconvex problems. SIAM Journal on
Optimization, 26(1):337–364, 2016.
Kim, H., Golub, G. H., and Park, H. Missing value estimation for dna microarray gene expression data: local
least squares imputation. Bioinformatics, 21(2):187–198,
2005.
Little, R. J. and Rubin, D. B. Statistical analysis withmissing data, volume 793. John Wiley & Sons, 2019.
Nesterov, Y. E. A method for solving the convex programming problem with convergence rate o (1/kˆ 2). In Dokl.
akad. nauk Sssr, volume 269, pp. 543–547, 1983.
Pedersen, A., Mikkelsen, E., Cronin-Fenton, D., Kristensen,N., Pham, T., Pedersen, L., et al. Missing data and multiple imputation in clinical epidemiological research. clin
epidemiol. 2017; 9: 157-66.
Raghunathan, T. E., Lepkowski, J. M., Van Hoewyk, J.,
Solenberger, P., et al. A multivariate technique for multiply imputing missing values using a sequence of regression models. Survey methodology, 27(1):85–96, 2001.
Schafer, J. L. Analysis of incomplete multivariate data.
Chapman and Hall/CRC, 1997.
Shivaswamy, P. K., Bhattacharyya, C., and Smola, A. J.
Second order cone programming approaches for handling
missing and uncertain data. Journal of Machine Learning
Research, 7(Jul):1283–1314, 2006.
Sion, M. et al. On general minimax theorems. Pacific
Journal of mathematics, 8(1):171–176, 1958.
Stekhoven, D. J. and Bühlmann, P. Missforest—nonparametric missing value imputation for mixed-type data.
Bioinformatics, 28(1):112–118, 2012.

Ghahramani, Z. and Jordan, M. I. Supervised learningfrom incomplete data via an em approach. In Advancesin neural information processing systems, pp. 120–127,
1994.

Sterne, J. A., White, I. R., Carlin, J. B., Spratt, M., Royston,P., Kenward, M. G., Wood, A. M., and Carpenter, J. R.
Multiple imputation for missing data in epidemiologicaland clinical research: potential and pitfalls. Bmj, 338:
b2393, 2009.

Gondara, L. and Wang, K. MIDA: Multiple imputation
using denoising autoencoders. In Pacific-Asia Conferenceon Knowledge Discovery and Data Mining, pp. 260–272.
Springer, 2018.

Troyanskaya, O., Cantor, M., Sherlock, G., Brown, P.,
Hastie, T., Tibshirani, R., Botstein, D., and Altman, R. B.
Missing value estimation methods for dna microarrays.
Bioinformatics, 17(6):520–525, 2001.

RIFLE: Imputation and Robust Inference from Low Order Marginals

Xia, J., Zhang, S., Cai, G., Li, L., Pan, Q., Yan, J., andNing, G. Adjusted weight voting algorithm for random
forests in handling missing values. Pattern Recognition,
69:52–60, 2017.
Xu, H., Caramanis, C., and Mannor, S. Robustness andregularization of support vector machines. Journal ofmachine learning research, 10(7), 2009.
Yoon, J., Jordon, J., and Van Der Schaar, M. Gain: Missing
data imputation using generative adversarial nets. arXiv
preprint arXiv:1806.02920, 2018.
Zhang, T., Ye, S., Zhang, K., Tang, J., Wen, W., Fardad, M.,
and Wang, Y. A systematic dnn weight pruning framework using alternating direction method of multipliers.
In Proceedings of the European Conference on Computer
Vision (ECCV), pp. 184–199, 2018.
Zhang, X., Song, X., Wang, H., and Zhang, H. Sequential
local least squares imputation estimating missing valueof microarray data. Computers in biology and medicine,38(10):1112–1120, 2008.

RIFLE: Imputation and Robust Inference from Low Order Marginals

A. A Review of Missing Value Imputation Methods in the Literature
The fundamental idea behind many data imputation approaches is that the missing values can be predicted based on theavailable data of other data points and correlated features. One of the most straightforward imputation techniques is toreplace missing values by the mean (or median) of that feature calculated from what data is available see Little & Rubin(2019, Chapter 3). However, this naı̈ve approach ignores the correlation between features and does not preserve the varianceof features. Another class of imputers has been developed based on the least-square methods (Raghunathan et al., 2001;
Kim et al., 2005; Zhang et al., 2008; Cai et al., 2006). Raghunathan et al. (2001) learns a linear model with multivariate
Gaussian noise for the feature with the least missing entries. It repeats the same procedure on the updated data to impute thenext feature with the least missing entries until all features are completely imputed. One drawback of this approach is thatthe error from the imputation of previous features can be propagated to subsequent features. To impute entries of a givenfeature in a dataset, Kim et al. (2005) learns several univariate regression models that consider that feature as the response.
Then it takes the average of these predictions as the final value of imputation. This approach fails to learn the correlations
involving more than two features.
Many more complex algorithms have been developed for imputation, although many are sensitive to initial assumptions andmay not converge. For instance, KNN-Imputer imputes a missing feature of a data point by taking the mean value of the K
closest complete data points (Troyanskaya et al., 2001). MissForest, on the other hand, imputes the missing values of eachfeature by learning a random forest classifier using other training data features (Stekhoven & Bühlmann, 2012). MissForest
does not need to assume that all features are continuous (Honaker et al., 2011) or categorical (Schafer, 1997). However, both
KNN-imputer and MissForest do not guarantee statistical or computational convergence for their algorithms. Moreover,when the proportion of missing values is high, both are likely to have a severe drop in performance, as demonstrated in
Section 5. The Expectation Maximization (EM) algorithm is another popular approach that learns the parameters of a priordistribution on the data using available values based on the EM algorithm of Dempster et al. (1977); see also Ghahramani &Jordan (1994) and Honaker et al. (2011). The EM algorithm is also used in Amelia, which fits a jointly normal distribution
to the data using EM and the bootstrap technique (Honaker et al., 2011). While Amelia demonstrates a superior performanceon datasets following a normal distribution, it is highly sensitive to the violation of the normality assumption (as discussed
in Bertsimas et al. (2017)). Ghahramani & Jordan (1994) adopt the EM algorithm to learn a joint Bernoulli distribution forthe categorical data and a joint Gaussian distribution for the continuous variables independently. While those algorithms
can be viewed as inference methods based on low-order estimates of moments, they do not consider uncertainty in such
low-order moments estimates. By contrast, our framework utilizes robust optimization to consider the uncertainty around theestimated moments. Moreover, our optimization procedure for imputation and prediction is guaranteed to converge despitesome of the algorithms mentioned above.
Another popular method for data imputation is multiple imputations by chained equations (MICE). MICE learns a parametric
distribution for each feature conditional on the remaining features. For instance, it assumes that the current target variable
is a linear function of other features with a zero-mean Gaussian noise. Each feature can have its distinct distribution andparameters (e.g., Poisson regression, logistic regression). Based on the learned parameters of conditional distributions,
MICE can generate one or more imputed datasets (Buuren & Groothuis-Oudshoorn, 2010). More recently, several neuralnetwork-based imputers have been proposed. GAIN (Generative Adversarial Imputation Network) learns a generative
adversarial network based on the available data and then imputes the missing values using the trained generator (Yoon et al.,2018). One advantage of GAIN over other existing GAN imputers is that it does not need a complete dataset during thetraining phase. MIDA (Multiple Imputation using Denoising Autoencoders) is an auto-encoder-based approach that trains adenoising auto-encoder on the available data considering the missing entries as noise. Similar to other neural network-based
methods, these algorithms suffer from their black-box nature. They are challenging to interpret/explain, making them
unpopular in mission-critical healthcare approaches. In addition, no statistical or computational guarantees are provided forthese algorithms.
Bertsimas et al. (2017) formulates the imputation task as a constrained optimization problem where the constraints are
determined by the underlying classification model such as KNN (k-nearest neighbors), SVM (Support Vector Machine),and Decision Trees. Their general framework is non-convex, and the authors relax the optimization for each choice of thecost function using first-order methods. The block coordinate descent algorithm then optimizes the relaxed problem. Theyshow the convergence and accuracy of their proposed algorithm numerically, while a theoretical analysis that guarantees thealgorithm’s convergence is absent in their work.

RIFLE: Imputation and Robust Inference from Low Order Marginals

B. Estimating Confidence Intervals of Low-order Moments
In this section, we explain the methodology of estimating confidence intervals for E[zi ] and E[zi zj ]. Let Xn×d and y be
the data matrix and target variables for n given data points respectively whose entries are in R̃ = R ∪ {∗}, where ∗ symbol
represents a missing entry. Moreover, assume that ai represents the i-th column (feature) of matrix X. We define:(ai (k) if ai (k) 6= *ãi (k) =0
if ai (k) = *Thus, ã is obtained by replacing the missing values with 0. We estimate the confidence intervals for the mean and covarianceof features using multiple bootstrap samples on the available data. Let C0 [i][j] and ∆0 [i][j] be the center and the radius of
the confidence interval for C[i][j], respectively. We compute the center of the confidence interval for C[i][j] as follows:
C0 [i][j] =1 Tã ãjmij i(12)where mi = |{k : ai (k) 6= ∗}| and mij = |{k : ai (k) 6= ∗, aj (k) 6= ∗}|. This estimator is obtained from the rows where
both features are available. More precisely, let M be the mask of the input data matrix X defined as:(0, if Xij is missing,
Mij =1, otherwise.
Assume that mij = (MT M)ij , which is the number of rows in the dataset where both features i and j are available. To
estimate the confidence intervals for Cij , we use Algorithm 3. First, we select multiple (K) samples of size N = mij fromthe rows where both features are available. Each one of these samples with size mij is obtained by applying a bootstrap
sampler (sampling with replacement) on the mij rows where both features are available. Then, we compute the second-ordermoment of two features for each sample.
To find the radius of confidence intervals for each given pair (i, j) of features, we choose k different bootstrap samples
with length n on the rows where both features i and j are available. Then, we compute C0 [i][j] of two features in eachbootstrap sample. The standard deviation of these estimations determines the radius of the corresponding confidence interval.
Algorithm 3 summarizes the required steps for computing the confidence interval radius for the ij-th entry of covariancematrix ∆. Note that the confidence intervals for µ can be computed similarly. Having C0 and ∆, the confidence interval
Algorithm 3 Estimating Confidence Interval Length ∆ij for Feature i and Feature j.
1: Input: K : Number of bootstrap estimations
2: for t = 1, . . . , K do
3:
Pick n samples with replacement from the rows where both i-th and j-th are available.
Let (X̂i1 , X̂j1 ), . . . , (X̂in , X̂jn ) be the i-th and j-th features of the selected samples
Pn
4:
Ct = n1 r=1 X̂ir X̂jr
5: end for
6: ∆ij = std(C1 , C2 , . . . , CK )for the matrix C is computed as follows:
Cmin = C0 − c∆Cmax = C0 + c∆,
Computing bmin and bmax can be done in the same manner. The hyper-parameter c is defined to control the robustness ofthe model by tuning the length of confidence intervals. A larger c corresponds to bigger confidence intervals and, thus, amore robust estimator. On the other hand, large values for c lead to very large confidence intervals that can adversely affectthe performance of the trained model.
Remark 4. Since the computation of confidence intervals for different entries of the covariance matrix are independent ofeach other, they can be computed in parallel. In particular, if γ cores are available, d/γ features (columns of the covariancematrix) can be assigned to each one of the available cores.

RIFLE: Imputation and Robust Inference from Low Order Marginals

C. Solving Robust Ridge Regression with the Optimal Convergence Rate
The convergence rate of Algorithm 1 to the optimal solution of Problem (6) can be slow in practice since the algorithm
requires to do a matrix inversion for updating θ and applying the box constraint to C and b at each iteration. While we
update the minimization problem in closed-form with respect to θ, we can speed up the convergence rate of the maximization
problem by applying Nesterov’s acceleration method to function g(b, C) in (7). Since function g is the minimum of convex
functions, its gradient with respect to C and b can be computed using Danskin’s theorem. Algorithm 4 describes the steps
to optimize Problem (7) using Nesterov’s acceleration method.
Algorithm 4 Applying the Nesterov’s Acceleration Method to Robust Linear Regression
1: C0 , b0 , ∆, δ, T
2: Initialize: C1 = C0 , b1 = b0 , γ0 = 0, γ1 = 1.
3: for i = 1, . . . √, T do
4:
5:
6:
7:

1+

1+4γ 2i
γi+1 =
2
γi −1
YCi = Ci + γi+1 (Ci − Ci−1 )
Ci+1 = Π∆+ YCi + L1 θθ T

−1(bi − bi−1 )Ybi = bi + γγii+1

8:
bi+1 = Πδ (Ybi − 2θL)9:Set θ = (Ci+1 + λI)−1 bi+1
10: end for

Theorem 5. Assume that D =qkC0 − C∗ k2F + kb0 − b∗ k22 . Then Algorithm 4 computes an -stationary solution of theLD
objective function in (7) in O
iterations, where L = λ2 .
Proof. The proof is relegated to Appendix F.

D. ADMM Algorithm For Solving the Dual Problem
00

For the simplicity of presentation, assume that ct1 = bmin − µtd + ρd t + η t , ct2 = −bmax − µte + ρe t − η t , ct3 = −µθ +000ρθ t − 2η t , Dt1 = ρA t − ρGt + Γt − MtA + Cmin , and Dt2 = ρB t + ρGt − Γt − MtB − Cmax . Algorithm 5 describes theADMM algorithm applied to Problem (10). At each iteration of the ADMM algorithm, the parameters of one block are fixed,
and the optimization problem is solved with respect to the parameters of the other block. For the simplicity of presentation,0000
let ct1 = ρθ t −µtθ −2η t , ct2 = ρd t −µtd −bmin +η t , ct3 = ρe t −µte +bmax −η t , Dt1 = ρA t −ρGt +Γt −MtA +Cmin ,0
and Dt2 = ρB t + ρGt − Γt − MtB − Cmax .
Corollary 6. If the feasible set of Problem (6) has non-empty interior, then Algorithm 5 converges to an -optimal solutionof Problem (10) in O( 1 ) iterations.
Proof. Since the inner maximization problem, in (6) is convex, and its feasible interior set is not empty, the duality gap iszero by Slater’s condition. Thus, according to Theorem 6.1 in He & Yuan (2015), Algorithm 5 converges to an optimal
solution of the primal-dual problem with a linear rate. Moreover, the sequence of constraint residuals converges to zero witha linear rate as well.
We have two non-trivial problems containing positive semi-definite constraints. The sub-problem with respect to G can be
written as:ρmin hBt − At − G, Γt i + kBt − At − Gk2F
G
2
(13)
0
0
s.t. G  θ t θ tT ,000

By completing the square, and changing the variable G = G − θ t θ tT , equivalently we require to solve the following

RIFLE: Imputation and Robust Inference from Low Order Marginals

Algorithm 5 Applying ADMM to the Dual Problem of Robust Linear Regression
1: Given: bmin , bmax , Cmin , Cmax , λ, ρ
2: Initialize: C1 = C0 , b1 = b0 , γ0 = 0, γ1 = 1.
3: for t = 0, . . . , T do
1
(2ct1 − 2ct2 − 3ct3 )4:
θ t+1 = 6λ+7ρ5:1( 6ρ+4λct1 + 4ρ+4λ
ct2 + 2ct3 )dt+1 = 6λ+7ρρρ

6:3ρ+2λ t2t
( ρ+2λc2 − ct3 )et+1 = 6λ+7ρρ c1 +
ρ

7:
8:
9:
10:
11:
12:
13:
14:

MtA
ρ , 0)
0
MtB
t+1
t
B
= max(B + ρ , 0)
t
0
0
0
0
t+1
G
= [Bt − At + Γρ − θ t θ tT ]+ + θ t θ tT0µt
d t+1 = max(dt + ρd , 0)0µt
e t+1 = max(et + ρe , 0)000
θ t+1 = argminθ0 kθ t+1 − θ k2 + hµtθ , θ t+1 − θ i
1
t+1
t
t
A
= 3ρ (2D1 + D2 )1
t+1
(Dt1 + 2Dt2 )
B
= 3ρ
0
t+1
MA = MtA + ρ(At+1 − A t+1 )0
t+1t− B t+1 )Mt+1
B = MB + ρ(B
0
t+1
t
t+1
µd = µd + ρ(d− d t+1 )0= µte + ρ(et+1 − e t+1 )
µt+1
e
0
t+1
µθ = µtθ + ρ(θ t+1 − θ t+1 )
t+1
t
t+1
t+1
t+1
0

A t+1 = max(At +0s.t. Gt+1  θ T θ

0

15:
16:
17:
18:
19:
20:
η
= η + ρ(2θ
−d
+e )
21:
Γt+1 = Γt + ρ(Bt+1 − At+1 − Gt+1 )22: end forproblem:
0
0
0
ρ
Γt 2
kG − (Bt − At − θ t θ tT +
)k
ρ F
G 2
0
s.t. G  0,

min
0

t

0

0

(14)

0

Thus, G ∗ = [Bt − At + Γρ − θ t θ tT ]+ , where [A]+ is the projection to the set of PSD matrices, which can be done bysetting the negative eigenvalues of A in its singular value decomposition to zero.
0

The other non-trivial sub-problem in Algorithm (5) is the minimization with respect to θ (Line 10). By completing thesquare, it can be equivalently formulated as:0minkθ − (θ t+1 +
0
θ

0

µtθ 2
)k
ρ 2

0

(15)

s.t. Gt+1  θ θ T ,Let G = U ΛU T be the singular value decomposition of the matrix G where Λ is a diagonal matrix containing theµteigenvalues of the matrix G. Set α = θ t+1 + 2θ . Since U T U = I, we have:kU T θ − U T αk2 = kθ − αk220

Set β = U T θ , then Problem (15) can be reformulated as:
min kβ − U T αk22
βs.t. ββ T  Λ.

(16)

RIFLE: Imputation and Robust Inference from Low Order Marginals

Note that the constraint of the above optimization problem is equivalent to the following:


n
X
βi2
1 βT
T
T −1
ββ  Λ ⇔0⇔β Λ β≤1⇔
≤ 1,
β Λ
λi
i=1
where λi = Λii Since the block matrix is symmetric, using Schur Complement, it is positive semi-definite if and only if Λ
is positive semi-definite and 1 − β t Λ−1 β ≥ 0 (the third inequality above).
Set γ = U T α, then we can write Problem (16) as:
min kβ − γk22βPn βi2s.t.
i=1 λi ≤ 1,
It can be easily shown that the optimal solution has the form βi∗ =

(17)

γi
∗
∗ , where µ is the optimal Lagrangian multiplier
1+ µ
λ
i

corresponding to the constraint of Problem (17). The bisection algorithm can efficiently obtain the optimal Lagrangian
multiplier. Having β ∗ , the optimal θ can be computed by solving the linear equation U T θ ∗ = β ∗ .

E. Generating Missing Values Patterns in Numerical Experiments
In this appendix, we define MCAR and MNAR patterns and discuss how to generate them in a given dataset. Formally, thedistribution of missing values in a dataset follows a missing completely at random (MCAR) pattern if the probability of
having a missing value for a given entry is constant, independent of other available and missing entries. On the other hand, adataset follows a Missing At Random (MAR) pattern if the missingness of each entry only depends on the available data ofother features. Finally, if the distribution of missing values does not follow an MCAR or MAR pattern, we call it missing
not at random (MNAR).
To generate the MCAR pattern on a given dataset, we fix a constant probability 0 < p < 1 and make each data entry
unavailable with the probability of p. On the other hand, the generation of the MNAR pattern is based on the idea that if thevalue of an entry is farther from the mean of its corresponding feature, then the probability of missingness for that entry islarger.
The generation of the MNAR pattern is based on the idea that if the value of an entry is farther from the mean of its
corresponding feature, then the probability of missingness for that entry is larger. Algorithm 6 describes the procedure ofgenerating MNAR missing values for a given column of a dataset:
Algorithm 6 Generating MNAR Pattern for a Given Column of a Dataset
1: Input: x1 , x2 , . . . , xn : The entries of the current column in the dataset, a, b: Hyper-parameters controlling thepercentage of missingP
valuesPnn1
22
2
2: Initialize: Set µ = n1i=1 xi and σ = ni=1 xi − µ .
3: for i = 1, . . . , n do
04:xi = xiσ−µ05:pi = F (a|xi | + b)
6:
Set xi = ∗ with probability of pi
7: end forNote that F in the above algorithm is the cumulative distribution function of a standard Gaussian random variable. a and b
control the percentage of missing values in the given column. As a and b increase, the probability of having more missing
values is higher. Since the availability of each data entry depends on its value, the generated missing pattern is missing not
at random (MNAR).

F. Proof of Lemmas and Theorems
In this appendix, we prove all lemmas and theorems presented in the article.
First, we prove the following lemma that is useful in several convergence proofs:

RIFLE: Imputation and Robust Inference from Low Order Marginals

Lemma 7. The Lipschitz constant of the gradient of the function g defined in Problem (7) is equal to L = λ2 .
Proof. Since, the problem is convex in θ and concave in C and b, we have:minθθ T Cθ − 2bT θ + λkθk2 = − min

max
C,b

C,b

max
θ

−θ T Cθ + 2bT θ − λkθk2

The new objective function is convex in C and b and strongly concave with respect to θ. According to Lemma 1
0in Barazandeh & Razaviyayn (2020), g = −g = maxθ −θ T Cθ + 2bT θ − λkθk2 is Lipschitz continuous with the
Lipschitz constant equal to:
L12
Lg = Lg0 = L11 +,σwhere σ is the strong-concavity modulus of −θ T Cθ +2bT θ −λkθk2 . Assume that h(θ, C, b) , −θ T Cθ +2bT θ −λkθk2
and L11 , L12 are defined as follows:k∇b,C h(θ, b1 , C1 ) − ∇b,C h(θ, b2 , C2 )k ≤ L11 k(C1 , b1 ) − (C2 , b2 )kk∇θ h(θ, b1 , C1 ) − ∇θ h(θ, b2 , C2 )k ≤ L12 k(C1 , b1 ) − (C2 , b2 )k
It is easy to see that L11 = 0 and L12 = 2 max kθk2 ≤ 2, which completes the proof.
Proof of Theorem 1: Since the set of feasible solutions for b and C defines a compact set, and function g is a concave
function with respect to b and C, the projected gradient ascent algorithm converges to the global maximizer of g in
2
∗ 2
∗ 2
T = O( LD ) iterations (Bubeck, 2014), where D = kC0 − C kF + kb0 − b k2 and L is the Lipschitz constant of2function g, which is equal to λ according to Lemma 7.
Proof of Theorem 2: First, note that if we multiply the objective function by −1, Problem (6) can be equivalently formulatedas:max min −θ T Cθ + 2bT θ − λkθk22
θ

C,b

s.t.

−C + Cmin ≤ 0,
C − Cmax ≤ 0,

(18)

−b + bmin ≤ 0,b − bmax ≤ 0,
−C  0If we assign A, B, d, e, H to the constraints respectively, then the Lagrangian function can be written as:
L(C, b, A, B, d, e, H) = −θ T Cθ + 2bT θ + hA, −C + Cmin i+hB, C − Cmax i + hd, −b + bmin i(19)+he, b − bmax i − hC, Hi − λkθk22 ,The dual problem is defined as:
max

min L(C, b, A, B, d, e, H)

A,B,d,e,H C,b(20)

The minimization of L takes the following form:
min hC, −θθ T − A + B − Hi + hb, 2θ − d + ei − λkθk22
C,b−hB, Cmax i + hA, Cmin i − he, bmax i + hd, bmin i,(21)

RIFLE: Imputation and Robust Inference from Low Order Marginals

To avoid −∞ value for the above minimization problem, it is required to set −θθ T − A + B − H and 2θ − d + e to zero.
Thus the dual problem of (18) is formulated as:
max

A,B,d,e,Hs.t.

bTmin d − bTmax e + hCmin , Ai − hCmax , Bi − λkθk22−θθ T − A + B − H = 0,(22)2θ − d + e = 0,
A, B, d, e ≥ 0,H0
Since the duality gap is zero, Problem (6) can be equivalently formulated as:
max

θ,A,B,d,e,Hs.t.

bTmin d − bTmax e + hCmin , Ai − hCmax , Bi − λkθk2−θθ T − A + B − H = 0,(23)2θ − d + e = 0,
A, B, d, e ≥ 0,
H  0.

We can multiply the objective function by −1 and change the maximization to minimization, which gives the dual problem
described in (2).
Proof of Theorem 5 Algorithm 4 applies the projected Nesterov acceleration method on the concave function g. As proved
in Nesterov (1983), the rate of convergence of this method conforms to q
the lower bound of first-order oracles for the generalconvex minimization (concave maximization) problems, which is O(
appeared in the iteration complexity bound by Lemma 7.

LD 2 ). We compute the Lipschitz constant L that

G. Dataset Descriptions
In this section, we introduce the datasets used in Section 5 to evaluate the performance of RIFLE. Except for the NHANES,
all other datasets contain no missing values. For those datasets, we generate MCAR and MNAR missing values artificially(for MNAR patterns, we apply Algorithm 6 to the datasets).
• NHANES: The percentage of missing values varies for different features of the NHANES dataset. There are twosources of missing values in NHANES data: Missing entries during data collection and missing entries resulting frommerging different datasets in the NHANES collection. On average, approximately 20% of data is missing.
• Super Conductivity2 : Super Conductivity datasets contains 21263 samples describing superconductors and theirrelevant features (81 attributes). All features are continuous, and the assigned task is to predict the critical temperature
based on the given 81 features. We have used this dataset in experiments summarized in Figure 3, Figure 4, and Table 2.
• BlogFeedback3 : BlogFeedback data is a collection of 280 features extracted from HTML-documents of the blog posts.
The assigned task is to predict the number of comments in the upcoming 24 hours based on the features of more than
60K data training data points. The test dataset is fixed and is originally separated from the training data. The dataset isused in experiments described in Table 2.
• Breast Cancer(Prognostic)4 : The dataset consists of 34 features and 198 instances. Each record represents follow-up
data for one breast cancer case collected in 1984. We have done several experiments to impute the MCAR missing
values generated artificially with different proportions. The results are depicted in Table 1 and Figure 6.
• Parkinson5 : The dataset describes a range of biomedical voice recording from 31 people, 23 with Parkinson’s disease(PD). The assigned task is to discriminate healthy people from those with PD. There are 193 records and 23 features inthe dataset. The dataset is processed similarly to the Breast Cancer dataset and used in the same experiments.
2

https://archive.ics.uci.edu/ml/datasets/Superconductivty+Data
https://archive.ics.uci.edu/ml/datasets/BlogFeedback4
https://archive.ics.uci.edu/ml/datasets/Breast+Cancer+Wisconsin+(Prognostic)
5
https://archive.ics.uci.edu/ml/datasets/parkinsons
3

RIFLE: Imputation and Robust Inference from Low Order Marginals

Figure 9. Consistency of ADMM (Algorithm 5) and Projected Gradient Ascent on function g (Algorithm 1) on the synthetic datasets with
40%, 60% and 80% missing values.

• Spam Base6 : The dataset consists of 4601 instances and 57 attributes. The assigned classification task is to predictwhether the email is spam. To evaluate different imputation methods, we randomly mask a proportion of data entries
and impute them with different approaches. The results are depicted in Table 1 and Figure 6.
• Boston Housing7 : Boston Housing dataset contains 506 instances and 14 columns. We generate random missing
entries with different proportions and impute them with RIFLE and several state-of-the-art approaches. The results aredemonstrated in Table 1 and Figure 6.
• Cloud8 : The dataset has 1024 instances and 10 features extracted from clouds images. We use this dataset inexperiments depicted in Table 1 with 70% artificial MCAR missing values.
• Wave Energy Converters9 : We sample a subset of 3000 instances with 49 features from the original Wave Energy
Converter dataset. We have executed several imputation methods on the dataset, and the results are shown in Figure 6.
• Sensorless Drive Diagnosis10 : The 49 continuous features in this dataset are extracted from electric current drive
signals, and the associated classification task is to determine the condition of device’s motor. We choose different
random samples with size 400 to run experiments (imputation) in Figure 7.

H. Further Discussion on the Consistency of RIFLE
The three developed algorithms in Section 3 for solving robust ridge regression are all consistent. To show this, we havegenerated a synthetic dataset with 50 input features following a jointly normal distribution. As observed in Figure 9, by
increasing the number of samples, the NRMSE of all three algorithms converges to 0.01, which is the standard deviationof the zero-mean Gaussian noise added to each target value (the dashed line). The pattern can be observed for different
percentages of missing values.
6

https://archive.ics.uci.edu/ml/datasets/spambase
https://www.kaggle.com/c/boston-housing8
https://archive.ics.uci.edu/ml/datasets/Cloud9
https://archive.ics.uci.edu/ml/datasets/Wave+Energy+Converters
10
https://archive.ics.uci.edu/ml/datasets/dataset+for+sensorless+drive+diagnosis
7

RIFLE: Imputation and Robust Inference from Low Order Marginals

I. Imputation Table
Dataset Name
Spam (30%)
Spam (50%)
Spam (70%)
Housing (30%)
Housing (50%)
Housing (70%)
Clouds (30%)
Clouds (50%)
Clouds (70%)
Breast Cancer (30%)
Breast Cancer (50%)
Breast Cancer (70%)
Parkinson (30%)
Parkinson (50%)
Parkinson (70%)

RIFLE
0.87 ±0.009
0.90 ±0.013
0.92 ±0.017
0.86 ±0.015
0.88 ±0.021
0.92 ±0.026
0.81 ±0.018
0.84 ±0.026
0.87 ±0.029
0.52 ±0.023
0.56 ±0.026
0.59 ±0.031
0.57 ±0.016
0.62 ±0.022
0.67 ±0.027

QRIFLE
0.82 ±0.009
0.86 ±0.014
0.91 ±0.019
0.89 ±0.018
0.90 ±0.024
0.95 ±0.028
0.79 ±0.019
0.84 ±0.028
0.90 ±0.033
0.54 ±0.027
0.59 ±0.027
0.65 ±0.034
0.55 ±0.016
0.64 ±0.025
0.74 ±0.033

MICE
1.23 ±0.012
1.29 ±0.018
1.32 ±0.028
1.03 ±0.024
1.14 ±0.029
1.22 ±0.036
0.98 ±0.024
1.10 ±0.041
1.16 ±0.044
0.74 ±0.031
0.79 ±0.029
0.86 ±0.042
0.71 ±0.019
0.77 ±0.029
0.85 ±0.038

Amelia1.26 ±0.007
1.33 ±0.024
1.37 ±0.032
1.02 ±0.016
1.09 ±0.027
1.18 ±0.038
1.04 ±0.027
1.13 ±0.046
1.19 ±0.048
0.81 ±0.032
0.85 ±0.033
0.92 ±0.044
0.67 ±0.021
0.74 ±0.034
0.82 ±0.037

GAIN
0.91 ±0.005
0.93 ±0.015
0.97 ±0.014
0.82 ±0.015
0.88 ±0.019
0.95 ±0.027
0.76 ±0.021
0.82 ±0.027
0.89 ±0.035
0.58 ±0.024
0.64 ±0.025
0.70 ±0.037
0.53 ±0.015
0.61 ±0.022
0.69 ±0.031

MissForest
0.90 ±0.013
0.92 ±0.011
0.95 ±0.016
0.84 ±0.018
0.88 ±0.018
0.93 ±0.024
0.71 ±0.011
0.75 ±0.023
0.81 ±0.031
0.55 ±0.016
0.59 ±0.022
0.63 ±0.028
0.54 ±0.010
0.65 ±0.014
0.73 ±0.022

MIDA
0.97 ±0.008
0.99 ±0.011
0.99 ±0.018
0.93 ±0.025
0.98 ±0.029
1.02 ±0.037
0.83 ±0.022
0.88 ±0.033
0.93 ±0.044
0.70 ±0.026
0.76 ±0.035
0.82 ±0.035
0.62 ±0.017
0.71 ±0.027
0.78 ±0.038

EM
0.94 ± 0.004
0.97 ± 0.008
0.98 ± 0.017
0.95 ± 0.011
0.96 ± 0.016
0.98 ± 0.017
0.86 ± 0.013
0.89 ± 0.018
0.92 ± 0.023
0.67 ± 0.014
0.69 ± 0.022
0.67 ± 0.014
0.64 ± 0.011
0.69 ± 0.022
0.75 ± 0.029

Table 1. Performance comparison of RIFLE, QRIFLE (Quadratic RIFLE), and state-of-the-art methods on several UCI datasets. We
applied to impute methods on three different missing-value proportions for each dataset. The best imputer is highlighted with bold font,
and the second-best imputer is underlined. Each experiment is done 5 times, and the average and the standard deviation of performances
are reported.

Methods
Regression on Complete Data
RIFLE
Mean Imputer + Regression
MICE + Regression
EM + Regression
MIDA Imputer + Regression
MissForest
KNN Imputer

Super Conductivity
0.4601
0.4873 ± 0.0036
0.6114 ± 0.0006
0.5078 ± 0.0124
0.5172 ± 0.0162
0.5213 ± 0.0274
0.4925 ± 0.0073
0.5438 ± 0.0193

Datasets
Blog Feedback
0.74320.8326 ± 0.0085
0.9235 ± 0.0003
0.8507 ± 0.0325
0.8631 ± 0.0117
0.8394 ± 0.0342
0.8191 ± 0.0083
0.8828 ± 0.0124

NHANES
0.62870.6304 ± 0.0027
0.6329 ± 0.0008
0.6612 ± 0.0282
0.6392 ± 0.0122
0.6542 ± 0.0164
0.6365 ± 0.0094
0.6427 ± 0.0135

Table 2. Normalized RMSE of RIFLE and several state-of-the-art Methods on Superconductivity, blog feedback, and NHANES datasets.
The first two datasets contain 30% Missing Not At Random (MNAR) missing values in the training phase generated by Algorithm 6.
Each method applied 5 times to each dataset, and the result is reported as the average performance ± standard deviation of experiments interms of NRMSE.

Since MICE and MIDA cannot predict values during the test phase without data imputation, we use them in a pre-processing
stage to impute the data. Then we apply the linear regression to the imputed dataset. On the other hand, RIFLE, KNN
imputer, and MissForest can predict the target variable without imputing the training dataset. Table 2 shows that RIFLE
outperforms all other state-of-the-art approaches executed on the three mentioned datasets. In particular, RIFLE outperforms
MissForest, while the underlying model used by RIFLE is simpler (linear) compared to the nonlinear random forest modelutilized by Missforest.