Stochastic Doubly Robust Gradient

TL;DR

This paper introduces a stochastic doubly robust gradient method for training machine learning models on incomplete data with missingness dependent on covariates, improving bias correction and variance reduction.

Contribution

It proposes the SDRG algorithm combining inverse propensity weighting and regression adjustment, linking double robustness with variance reduction in SGD.

Findings

Demonstrates convergence in training image classifiers with missing data.

Shows theoretical connection between double robustness and variance reduction.

Empirical results validate the effectiveness of SDRG.

Abstract

When training a machine learning model with observational data, it is often encountered that some values are systemically missing. Learning from the incomplete data in which the missingness depends on some covariates may lead to biased estimation of parameters and even harm the fairness of decision outcome. This paper proposes how to adjust the causal effect of covariates on the missingness when training models using stochastic gradient descent (SGD). Inspired by the design of doubly robust estimator and its theoretical property of double robustness, we introduce stochastic doubly robust gradient (SDRG) consisting of two models: weight-corrected gradients for inverse propensity score weighting and per-covariate control variates for regression adjustment. Also, we identify the connection between double robustness and variance reduction in SGD by demonstrating the SDRG algorithm with a unifying framework for variance reduced SGD. The performance of our approach is empirically tested by showing the convergence in training image classifiers with several examples of missing data.

Equations

$$\underset{\mathbf{\theta }\in {\mathbb{R}}^d}{\min }f(\mathbf{\theta }):=\frac{1}{n}\sum _{i=1}^n{w}_i{f}_i(\mathbf{\theta })$$

$${\mathbf{\theta }}^{t+1}$$

$$\Delta {\mathbf{\theta }}^t$$

$$g_i(\cdot )$$

$$\tilde{\mathbf{\theta }}$$

$$g_i(\tilde{\mathbf{\theta }})$$

$$\Delta {\mathbf{\theta }}^t=\frac{1}{n}\left[\nabla f_{i_t}({\mathbf{\theta }}^t)-g_{i_t}(\tilde{\mathbf{\theta }})\right]+\frac{1}{n}\sum _{i=1}^n{g}_{i_t}(\tilde{\mathbf{\theta }})$$

$${\mathbb{E}}_p[f(\mathbf{\theta })]$$

$$=\frac{1}{n}\sum _{i=1}^n\frac{p({\mathbf{x}}_i)}{q({\mathbf{x}}_i)}{f}_i(\mathbf{\theta }),{\mathbf{x}}_i\sim q$$

$$\Delta {\mathbf{\theta }}^t$$

$$=w_{i_t}\left\{\nabla f_{i_t}({\mathbf{\theta }}^t)-g_{i_t}(\tilde{\mathbf{\theta }})\right\}+g(\tilde{\mathbf{\theta }})$$

$$\Delta {\mathbf{\theta }}^t$$

$${\tilde{\mathbf{\theta }}}_c$$

$$g^t({\tilde{\mathbf{\theta }}}_c)$$

$${\mathbf{\theta }}^{t+1}$$

$$={\mathbf{\theta }}^t-\left(w_{i_t}\nabla f_{i_t}({\mathbf{\theta }}^t)-w_{i_t}{g}_{i_t}(\tilde{\mathbf{\theta }})+g(\tilde{\mathbf{\theta }})\right)$$

$${\mathbf{\theta }}^{t+1}$$

$$:{\mathbf{\theta }}^{t+1}={\mathbf{\theta }}^t+\eta \sum _{j=1}^{t}S(j,i_{1:t})\nabla f_{i_t}({\mathbf{\theta }}^j)$$

$$:{\mathbf{\theta }}^{t+1}={\mathbf{\theta }}^t+\eta \sum _{j=1}^t(\gamma \eta {)}^{t-j}\nabla f_{i_t}({\mathbf{\theta }}^j)$$

$$\displaystyle\Delta\boldsymbol{\theta}^{t}=\frac{1}{10}\sum_{c=0}^{9}\Bigg{[}\frac{1}{|I_{t,c}|}\sum_{i_{t}\in I_{t}}\nabla f_{i_{t,c}}(\boldsymbol{\theta}^{t})-\frac{1}{|I_{t,c}|}\sum_{i_{t}\in I_{t}}g_{i_{t,c}}(\tilde{\boldsymbol{\theta}}_{c})+g^{t}(\tilde{\boldsymbol{\theta}}_{c})\Bigg{]}$$

$$\Delta {\mathbf{\theta }}^t$$

Full text

1 Introduction

The missing data problem is commonly encountered when training a machine learning model with real world data: unlike in the case of clean-cut experimental data, one or more covariates are often missing in recorded observations. Learning from the incomplete data may introduce an undesirable bias, especially when the missingness mechanism is not completely at random. More specifically, if the missingness depends on some covariates (e.g., gender, age, religion, and race) involved in generating the data, the estimation based on these unequally collected observations can be significantly different from the ideal result. This does not only interfere with the consistency in the learning process, but may also have a profound effect on the fairness of learning outcome.

To mitigate this problem, one may want to infer the causal effect of covariates on the missingness mechanism when training models. Doubly robust estimator [25, 27], first introduced in the area of observational study, has been known as an effective method to deal with such causal missingness and still remains popular [3, 15, 26, 10, 37]. It employs two well-known approaches, regression adjustment and inverse propensity score weighting, and by its interesting theoretical property, guarantees that the estimate remains unbiased as long as either of the two is specified correctly. In recent years, the double robustness has emerged in wide range of machine learning areas including covariate shift [24], adversarial training [14], and reinforcement learning [7, 12, 30, 8].

In this paper, we introduce the concept of doubly robust estimator to stochastic gradient descent (SGD) to correct the bias induced by the causal missingness in training data. Our approach, namely stochastic doubly robust gradient (SDRG), consists of per-covariate control variates and weight-corrected gradients that serve as the methods for regression adjustment and inverse propensity score weighting, respectively. To the best of our knowledge, SDRG is the first SGD algorithm that provides the property, double robustness.

Recently, the use of control variate methods have been excessively studied in the literature of variance reduction of SGD for accelerating the convergence [28, 5, 13, 32, 22]. Among these works, stochastic variance reduced gradient (SVRG) [13] and SAGA [4], which is an unbiased estimate of stochastic average gradient (SAG) [28], are closely related to SDRG in that both of them and SDRG can be viewed within a generic framework for control variates of variance reduced SGDs [23], as SDRG involves the use of per-covariate control variates. Finding connections between the ways control variates are used, we will show how pursuing double robustness in gradient estimation is aligned with reducing the variance.

Although SDRG and the variance reduced SGDs look similar to each other, there is a notable difference in the situation they can handle: each gradient estimate should be weighted unequally to reflect the missingness in training data and the weight-corrected gradients in SDRG are devised to address this need, whereas the aforementioned variance reduction methods can only consider the equal contribution of individual gradients. The SDRG algorithm can be straightforwardly applied to practical scenarios such as class imbalance problem, as we demonstrate in this paper considering some contextual information (such as class-label or any kind of tag) of training data as the covariates.

In summary, our contributions are as follows.

• $\bullet$

We propose the first doubly robust SGD algorithm, called SDRG, and demonstrate that SDRG can be devised in much the same way of SAGA and SVRG.

• $\bullet$

We define per-covariate momentum functions as control variates of SDRG, and show that it does not require to either periodically calculate (as SVRG) or store (as SAGA) the full gradients.

• $\bullet$

We provide a relation between SDRG and momentum, which is a much more direct derivation than the previous relationship presented in [28].

• $\bullet$

We experimentally show the performance of SDRG in training image classifiers with class-imbalanced MNIST and fashion-MNIST datasets since they are simple, yet commonly arising form of missing data problems.

To clarify, we remark that our work is not aligned with the approaches that employ non-uniform importance sampling for variance reduction in SGD based [20, 36, 16, 17, 29]. Rather than proposing a sampling criteria, the purpose of our work is to develop a robust learning algorithm when the weights are already determined with regards to the causal missingness.

2 Background and Related Work

A principled optimization problem in modern machine learning is that of the finite-sum form: minimization of an objective function $f(\boldsymbol{\theta})$ that are naturally expressed as a summation over a finite set of data $\mathcal{D}=\{x_{i}\}_{i=1}^{n}$, which is described as,

[TABLE]

where $\boldsymbol{\theta}$ is a parameter to be optimized over, each term $f_{i}(\boldsymbol{\theta})$ contributes with the weight $w_{i}$, and $w_{i}=1$ for typical setup. Such objective in Eqn. (1) commonly appears in the empirical risk minimization framework where the objective is the average of losses computed over the data in $\mathcal{D}$, that is, $f_{i}(\boldsymbol{\theta}):=L(\boldsymbol{\theta};x_{i})$

Stochastic gradient descent (SGD) is a method of choice to deal with such optimization problem. It iteratively updates the design parameter as follows: for each training iteration $t=1,2,...,T$,

[TABLE]

where $\eta>0$ is a learning rate, $i_{t}\in\{1,...n\}$ and $f_{i_{t}}(\cdot)$ is the loss computed with $x_{i_{t}}$ which is drawn iteratively from a training set $\mathcal{D}$.

In recent years, a class of algorithms to improve the convergence of SGD by reducing the variance of the estimates has been proposed [28, 13, 32, 5]. Especially, Reddi et al. [23] provides a formal unifying framework as Alg. 1 for stochastic variance reduction methods proposed in the literature, including SVRG [13], SAGA [5], and SAG [28]. The basic idea behind the variance reduction methods is to augment the gradient with a control variate and its expectation as,

[TABLE]

where $\tilde{\boldsymbol{\theta}}$ is an approximation of $\boldsymbol{\theta}$. The resulted estimate is unbiased, and has smaller variance if $g_{i_{t}}(\tilde{\boldsymbol{\theta}})$ has a high correlation with the target estimate $\nabla f_{i_{t}}(\boldsymbol{\theta}^{t})$.

As studied in [23], the mechanisms of updating control variates, $\{g_{i}(\tilde{\boldsymbol{\theta}})\}_{i=1}^{n}$ can be arranged within the unifying framework (see line 5 of Alg. 1) for the well-known variance reduction methods:

SVRG

Parameter $\tilde{\boldsymbol{\theta}}$ is updated after every $m$ iterations as

[TABLE]

SAGA

The gradients of all functions $\{\nabla f_{i}(\boldsymbol{\theta})\}_{i=1}^{n}$ are kept in memory, and one of them corresponding to the training instance is updated at every iteration as

[TABLE]

For SAG, the only difference with SAGA is that the line 3 of Alg. 1 is changed into,

[TABLE]

One may notice that SAG update rule does not exactly fit in the formulation of Eqn. (2) since the last term in Eqn. (5) does not become an expectation of the control variate by the scale of $1/n$. However, we categorize SAG as control variate-based variance reduction methods along with other methods, since they are all similar in the sense of incorporating an additional parameter to reduce the variance of estimates.

The aforementioned approaches are originally under strong convexity assumptions and has been extended to non-convex optimization problems [2, 22]. Asynchronous [23, 19, 11], proximal [35, 1] and accelerated variants have also been proposed.

3 Stochastic Doubly Robust Gradient

Before to introduce our main algorithm, we begin by formalizing the notion of weighted finite-sum problem that we are interested in and introduce notation that we use throughout this paper.

3.1 Problem Setting

We consider the cases where the individual loss term, $f_{i}$, in Eqn. (1) contributes unequally to the optimization, and thus, should be weighted differently according to certain criteria (i.e., $w_{i}\neq 1$). The weight $w_{i}$ is defined by the generic importance sampling literature [33, 31, 6]: when training and testing data come from different distributions, it can be specified to correct the difference. To elaborate, we are only given a collected set of data from a distribution $q$ (which we call the sampling distribution). We originally intend to compute the expected loss over some proper distribution $p$ (which we refer to as the target distribution), that is, $\mathbb{E}_{p}[f(\boldsymbol{\theta};\boldsymbol{x})]$, $\boldsymbol{x}\sim q$. Since we may not have direct access to $p$, however, we want to do a finite-sum approximate to the expectation over samples $\{\boldsymbol{x}_{i}\}_{i=1}^{n},~{}\boldsymbol{x}_{i}\in\mathcal{\mathbb{R}}^{d}$ drawn from the distribution at hand:

[TABLE]

where $p(\boldsymbol{x}_{i})/q(\boldsymbol{x}_{i})$, the ratio between target distribution and sampling distribution is regarded as the weighting factor $w_{i}$.

3.2 Weight-Corrected Gradient with Variance Reduction

To solve the weighted finite-sum problem as Eqn. (3.1), in the standard SGD algorithms, the $t$-th iteration involves picking an instance from sampling distribution $q$ over all instances and updates parameters as

[TABLE]

where $w_{i_{t}}:=p(\boldsymbol{x}_{i_{t}})/q(\boldsymbol{x}_{i_{t}})$ is per-sample importance weight. This weight-corrected gradient can be thought of as an inverse propensity score estimate, and we call it importance weighted SGD.

Within this literature of doubly robust estimator, our goal is to reduce the variance of stochastic gradient algorithm by introducing a control variate method to accelerate SGD. Given Eqn. (2) and Eqn. (7), an intuitive approach to employ a control variate to stochastic weighted gradient descent is as follows,

[TABLE]

where $\tilde{\boldsymbol{\theta}}$ is required to be highly correlated with $\boldsymbol{\theta}$ and $g(\tilde{\boldsymbol{\theta}}):=\mathbb{E}_{p}[g_{i}(\tilde{\boldsymbol{\theta}})]$. The resulted estimate is unbiased as the stochastic weighted gradient in Eqn. (7) with reduced variance.

The constructed gradient estimator in Eqn. (8) involves two variables: per-sample importance weight $w_{i_{t}}$ and control variate $g_{i_{t}}$. In other words, the estimation accuracy of our approach relies on how correctly the two could be specified. From this perspective, we see an advantage of our formulation by observing that either one of two models needs to be correctly specified to obtain an unbiased estimator:

Seen in a broader context, such distinguishing property of our formulation in Eqn. (8) arises in many areas. Indeed, in the area of observational study, the property called double robustness has been well studied [25, 3, 15, 26]: doubly robust (DR) estimators involves models for both the propensity score and the conditional mean of the outcome, and remain consistent even if one of those models (but not both) is misspecified. By observing that the constitution and the property of DR estimator are similar with that of our constructed gradient estimator, we see an opportunity to bring the insights of DR estimation into stochastic gradient optimization.

Theorem 1**.**

The weight-corrected gradient with variance reduction in Eqn. (8) satisfies the double robustness. Thus, Eqn. (8) is a doubly robust gradient estimator.

Proof.

First, we rewrite Eqn. (8) as follows:

[TABLE]

Since $\mathbb{E}[w_{i_{t}}g_{i_{t}}(\tilde{\boldsymbol{\theta}})]$ is equal to $g(\tilde{\boldsymbol{\theta}})$ and $\mathbb{E}[\nabla f_{i_{t}}(\boldsymbol{\theta}^{t})-g_{i_{t}}(\tilde{\boldsymbol{\theta}})]=0$, we can observe that $\Delta\boldsymbol{\theta}^{t}$ is unbiased if either $w_{i_{t}}$ is accurately estimated or the model $g_{i_{t}}(\tilde{\boldsymbol{\theta}})$ is correct. Therefore, we can conclude that Eqn. (8) is a doubly robust estimator.

∎

In the statistics community, and particularly in causal inference settings, DR estimators provide an estimation on average causal effect from observational data, adjusting appropriately for confounders. The ability of DR estimators to taking account to the causal effect of confounders can be also useful in the general machine learning literature. For instance, in supervised machine learning, the goal is to seek a function $h:X\rightarrow Y$, given $n$ pairs of inputs and corresponding target outputs $\{\boldsymbol{x}_{i},y_{i}\}_{i=1}^{n}$. Meanwhile, there often exist contextual information associated with instances $\boldsymbol{x}_{i}$, and it is not directly used to compute the objective (loss) value but may indirectly influence the process of learning the relation between $\boldsymbol{x}_{i}$ and $y_{i}$. In that case, one may want to address causal effect of contextual information in the process of learning the relation between $\boldsymbol{x}_{i}$ and $y_{i}$. By regarding the contextual information as confounding factors in the approximation of the expected loss over training set, we propose a method to adjust the causal effect of contextual information that may confound the gradient estimation.

3.3 Confounded Mini-Batch Gradient

We maintain different models that estimate two key parameters of Eqn. (8) – importance weights and control variates, conditioned on each configuration of contextual information. In practice, the contextual information could represent a class-label or any tag associated with $x_{i}$ [9]. In this paper, we decide to use class-label as contextual information as they are almost always available. Given a finite set of class-labels as the observed contextual information, we confine our interest to class imbalance problem, one of the most common scenario we may encounter in classification tasks. By taking a mini-batch variant of our estimator, we view the class imbalance problem in the perspective of importance weighting: if the data set collected in the mini-batch is sampled from highly skewed distribution, we want to correct the difference between the skewed distribution and the target distribution which is assumed to be uniformly balanced in terms of classes. Let $I_{t}$ be the set of indices of mini-batch instances at training iteration $t$, $I_{t,c}$ is a disjoint subset of $I$, whose instances belong to class $c$, $w_{c}(\cdot)$ is per-class model for estimating importance weights and $g_{t,c}(\cdot)$ is per-class control variates. Then, our proposed algorithm for class imbalance problem is described as,

[TABLE]

We call this method stochastic doubly robust gradient (SDRG). In practical implementations, an intuitive way of setting importance weights in the above setting is to compute the proportion of the number of instances that belong to each class-label over the mini-batch size.

We show that the existing variance reduction methods are related to SDRG in the sense their mechanism of updating control variates and computing the expectation of them (line 5 of Alg. 1 can be directly used in SDRG update rule) as follows:

[TABLE]

where $m$ is the parameter update frequency and $\tilde{\boldsymbol{\theta}}_{c}$ is initialized by $\boldsymbol{\theta}^{0}$ for all $c$, and

[TABLE]

where $g^{0}(\tilde{\boldsymbol{\theta}}_{c})$ is initialized by 0 for all $c$.

4 Relation to Momentum

In practical implementations, it is natural to take $\tilde{\boldsymbol{\theta}}$ as the average or a snapshot from the past iterations [13]. However, we propose to set $\tilde{\boldsymbol{\theta}}$ as a geometric weighting of previous gradients and by doing so, show an simple analysis on the relation between the above formulation and momentum optimizer [21]: the proposed control variate method which is described as,

[TABLE]

reduces to the formulation of momentum under a special setting where $g_{i_{t}}(\tilde{\boldsymbol{\theta}})=g(\tilde{\boldsymbol{\theta}})=\frac{\gamma}{1-\eta}\Delta\boldsymbol{\theta}^{t-1}$ and $w_{i_{t}}=\eta~{}(\neq 1)$, as follows,

[TABLE]

where $\gamma$ is a momentum coefficient. In other words, the control variate method with constant importance weight has the exact same formulation with momentum update rule.

Proposition 1**.**

The weight-corrected gradient with variance reduction in Eqn. (8) generalizes the momentum update rule in Eqn. (9) into the cases of importance weighting.

From the perspective of the classic control variate schemes that involves an additional parameter and its expectation, momentum method can be regarded as a biased estimator since the expectation of weighted control variates $w_{i_{t}}g_{i_{t}}$ does not correspond to $g(\tilde{\boldsymbol{\theta}})$. The interpretation on the momentum as a biased estimator can be also find in SAG of Eqn. (5) and from this observation, we see a connection between momentum and SAG using our control variate formulation. It is noteworthy that there has been an attempt to find a relation between SAG and momentum optimizer:

in the work of [28], Eqn. (5) and Eqn. (9) methods can be expressed in the following formulation,

[TABLE]

where $S(i,i_{1:t})$ is the selection function and equal to $1/n$ if $j$ corresponds to the last iteration where $j=i_{t}$ and is set to [math] otherwise. Roux et al. [28] finds a connection of SAG on momentum method by showing that they can be written in a similar formation. However, we provide an simpler but stronger analysis by showing a direct connection of SAG on momentum method by proving that they reduce to the exactly same form of equation (Prop. 1).

5 Experiments

In this section, we experiment with SDRG in comparison with importance weighted gradient descent for classification task with MNIST [18] and Fashion-MNIST [34] datasets. Both of MNIST and Fashion-MNIST are well balanced for 10 classes, but as we want to test the convergence rate in the setting of class imbalance problem, we modify the sampling distribution to make the sampled instances in mini-batches to be highly unbalanced in terms of class-label. Assuming that we want the training instances are uniformly distributed for all classes, we correct the gradient computed over the mini-batch samples drawn from the skewed distribution by employing importance weighting. In such cases where importance weighting is employed, we demonstrate that our SDRG which augments the importance weighted stochastic gradient with control variates whose expectation is replaced by a momentum, shows empirical improvements on the convergence rate.

MNIST

MNIST is a large database of handwritten digits from 0 to 9 that is commonly used to evaluate various machine learning algorithms. The images are grayscale with size of 28 x 28 pixel and they are uniformly balanced for all 10 classes.

We devise two different mechanisms to generate settings for the class imbalance problem: (a) first, the sampling distribution is skewed consistently for a single class during the entire training process. For instance, the probability for the instances from class 0 to be sampled is forced as 0.8, where the instances from the rest of 9 classes are sampled uniformly. (b) Otherwise, the class to be skewed with the sampling probability of 0.8 is selected in turn from class 0, 1, 2, … to 9.

We test the performance of SDRG update rule which can be specifically written as the follows for 10 class classification task: Let $I_{t}$ be a set of indices for mini-batch samples at $t$ th training iteration, and $I_{t,c}$ be the disjoint subset of $I_{t}$ that is the set of indices of instances that belong to class-label $c$.

[TABLE]

We compare the performance of SDRG with two algorithms: importance weighted gradient descent which is described as,

[TABLE]

and vanilla SGD.

To compare SDRG and other algorithms, we train a neural network (with one fully-connected hidden layer of 100 nodes and ten softmax output nodes) using cross-entropy loss with mini-batches of size 20 and learning rate of 0.01. We evaluate the performance of SDRG with different frequency of updating $\tilde{\boldsymbol{\theta}}$: $m=50,100$. For momentum parameters, $\gamma=0.9$ and $\eta=0.1$ are used. And we add two extra parameters $\alpha$ ($0.5$ for (a) and $1.5$ for (b)) and $\beta$ ($1.5$ for (a) and $0.5$ for (b)) for additional weighting sample gradients and control variate function, respectively. These parameters are corresponding to weights of $w_{i}=\eta$ and $g_{i_{t}}(\tilde{\boldsymbol{\theta}})=g(\tilde{\boldsymbol{\theta}})=\frac{\gamma}{1-\eta}\Delta\boldsymbol{\theta}^{t-1}$ in Eqn. (9). The results are all generated by taking the average from 20 runs of experiments. The confidence intervals are too insignificant to be noted and we decided not to include them in the figures.

In both Fig. 1 (a) and (b), SDRG empirically shows a faster convergence rate than importance weighted SGD and vanilla SGD. However, in Fig 1. (b) where the class to be skewed is changed after $m$ iterations, which might be a harder case than (b), the overall convergance rates flucture more than Fig, 1 (a), but SDRG shows more robust convergence than the other algorithms.

Fashion-MNIST

Fashion-MNIST is an MNIST-like database of clothes. The images are grayscale, with size of 28 x 28 and associated with labels from 10 classes. We evaluate the performance of SDRG in comparison with importance weighted SGD and vanilla SGD, under experimental settings which are exactly same with the MNIST experiment above.

In Fig 2. we can observe the same tendency in the convergence rate with Fig. 1, where SDRG shows less variant converging pattern for both of cases (a) and (b).

6 Conclusion and Future Work

In this paper, we proposed a SGD algorithm that addresses the caual effects of covariates on the missingness of incomplete data. Along with the previous studies that extended the use of doubly robust estimators to a variety of machine learning areas [24, 14, 7, 12, 30, 8], this paper has been the first approach to apply the idea of doubly robust estimator to stochastic optimization. In SDRG, employing control variates for regression adjustment allowed us to view the proposed method in the framework for variance reduced SGDs that also utilize control variate schemes, except that ours include the additional weight correction term. For the efficiency in computing and storing control variates, we suggested to choose momentum function as control variates and by doing so, the direct derivation from SDRG to the momentum has been found as a byproduct. In addition to these notable findings, empirical results have demonstrated that the proposed SDRG shows faster convergence than vanilla SGD and importance weighted SGD.

Future work includes further empirical studies to evaluate the performance of SDRG in various cases of missing data problem: for instance, the setting where multiple covariates are involved in the missingness mechanism. Also, theoretical analysis on convergence of SDRG should be provided. It would be essential to do deeper investigation on how the property of double robustness affects the convergence of SGD, compared to the existing variance reduced SGDs such as SAGA and SVRG.