Distributionally Robust Learning with Stable Adversarial Training

Traditional machine learning algorithms which optimize the average empirical loss often suffer from the poor generalization performance under distributional shifts induced by latent heterogeneity, unobserved confounders or selection biases in training data[1, 2, 3]. However, in high-stake applications such as medical diagnosis[4], criminal justice[5, 6] and autonomous driving [7], it is critical for the learning algorithms to ensure the robustness against potential unseen data. Therefore, robust learning methods have recently aroused much attention due to its favorable property of robustness guarantee[8, 9, 10].

Instead of optimizing the empirical cost on training data, robust learning methods seek to optimize the worst-case cost over an uncertainty set and can be further separated into two main branches named adversarially and distributionally robust learning. In adversarially robust learning, the uncertainty set is constructed point-wisely[9, 11, 10, 12]. Specifically, adversarial attack is performed independently on each data point within a $L_{2}$ or $L_{\infty}$ norm ball around itself to maximize the loss of current classification model. In distributionally robust learning, on the other hand, the uncertainty set is characterized on a distributional level[13, 14, 15]. A joint perturbation, typically measured by Wasserstein distance or $f$-divergence between distributions, is applied to the entire distribution entailed by training data. These methods can provide robustness guarantees under distributional shifts when testing distribution is captured in the built uncertainty set. However, in real scenarios, to contain the possible true testing distribution, the uncertainty set is often overwhelmingly large, and results in learned models with fairly low confidence, which is also referred to as the over pessimism or the low confidence problem[16, 17]. That is, with an overwhelmingly large set, the learner optimizes for implausible worst-case scenarios, resulting in meaningless results (e.g. the classifier assigns equal probability to all classes). Such a problem greatly hurts the generalization ability of robust learning methods in practice.

The essential problem of the above methods lies in the construction of the uncertainty set. To address the over pessimism of the learning algorithm, one should form a more practical uncertainty set which is likely to contain the potential distributional shifts in the future and meanwhile is as small as possible. More specifically, in real applications, we observe that different covariates may be perturbed in a non-uniform way, which should be considered when building a practical uncertainty set. Taking the problem of waterbirds and landbirds classification as an example[18]. There exist two types of covariates where the stable covariates (e.g. representing the bird itself) preserve immutable correlations with the target across different environments, while those unstable ones (e.g. representing the background) are pretty likely to change (e.g. waterbirds on land). Therefore, for the example above, the construction of the uncertainty set should be anisotropic which mainly focuses on the perturbation of those unstable covariates (e.g. background) to generate more practical and meaningful samples.

Further, we illustrate the anisotropic uncertainty set in figure 1, where blue points denote the observed training distribution ($\mathcal{N}(0,I_{2})$). And we sample data points from all distributions in the uncertainty set captured by an isotropic Wasserstein ball around the observed distribution, which are colored orange. We can see from figure 1 that the original distribution is perturbed equally along both the stable and unstable direction. With the intuition above, we propose that the ideal uncertainty should be like green points, which only perturb the training distribution along unstable directions. Following this, there are several work[19, 20] based on the adversarial attack which focus on perturbing the color or background of images to improve the adversarial robustness. However, these methods mainly follow a step by step routine where the segmentation is conducted first to separate the background from the foreground and cannot theoretically provide robustness guarantees under unknown distributional shifts, which greatly limits their applications on more general settings.

In this paper, we propose the Stable Adversarial Learning (SAL) algorithm to address this problem in a more principled and unified way, which leverages heterogeneous data sources to construct a more practical uncertainty set. Specifically, we adopt the framework of Wasserstein distributionally robust learning(WDRL) and further characterize the uncertainty set to be anisotropic according to the stability of covariates across the multiple environments, which induces stronger adversarial perturbations on unstable covariates than those stable ones. A synergistic algorithm is designed to jointly optimize the covariates differentiating process as well as the adversarial training process of model’s parameters. Compared with traditional robust learning techniques, the proposed method is able to provide robustness under strong distributional shifts while not hurting much confidence of the learner. Theoretically, we prove that our method constructs a more compact uncertainty set, which as far as we know is the first analysis of the compactness of adversarial sets in WDRL literature. Empirically, the advantages of our SAL algorithm are demonstrated on both synthetic and real-world datasets in terms of uniformly good performance across distributional shifts.

In this section, we investigate several strands of related literature more thoroughly, including domain adaptation, domain generalization, stable learning and distributionally robust learning.

Domain adaptation methods[21] leverage the data from target domain to assist the model training on source domain. Therefore the resulted model could capture the possible distributional shift in testing. Shimodaira[22] proposes to assign each training data a new weight equal to the density ratio between source and target distribution, and therefore guarantee the optimality of learned model on test distribution. Then several techniques have been proposed to estimate the ratio more accurately, such as discriminative estimation[23], kernel mean matching[24] and maximum entropy[25]. Apart from reweighting methods, deep learning based methods[26, 27] learn a transformation in feature space to characterize both source and target domain. However, the deployment of domain adaptation methods in real applications, where one can hardly access data from target domain, is quite limited.

Compared with domain adaptation, domain generalization techniques do not require the availability of target domain data and become more and more popular these years due to its practicability. Different from domain adaptation, domain generalization methods propose to learn a domain-invariant classifier with multiple training domains. Muandet et al. [28] propose a kernel-based optimization algorithm to learn an invariant latent space of data across training domains. Through the lens of causality[29, 30], M. Arjovsky et al. [31] propose Invariant Risk Minimization to learn invariant representation with theoretical guarantee of the optimality of out-of-distribution generalization, which gains the most attention recently. Also, stable learning methods [32, 3, 33] propose to decorrelate the covariates via sample reweighting to estimate the real causal effects, which enhances the stability under distributional shifts. However, they only deal with the covariate shift problem and do not apply to other kinds of distributional shifts (e.g. concept shifts brought by anti-causal variables).

Distributionally robust learning (DRL), from the optimization literature, proposes to optimize for the worst-case cost over an uncertainty distribution set, so as to protect the model against the potential distributional shifts in the uncertainty set, which is constrained by moment or support conditions [34, 35], or $f$-divergence [36, 17]. As the uncertainty set formulated by Wasserstein ball is much more flexible, Wasserstein Distributionally Robust Learning (WDRL) has been widely studied [37, 13, 14]. WDRL for logistic regression was proposed by Abadeh et al. [37]. Sinha et al. [13] achieved moderate levels of robustness with little computational cost relative to empirical risk minimization with a Lagrangian penalty formulation of WDRL. Esfahani and Kuhn [14] reformulated the distributionally robust optimization problems over Wasserstein balls as finite convex programs under mild assumptions. Although DRL offers an alternative to empirical risk minimization for robust performance under distributional perturbations, there has been work questioning its real effect in practice. Hu et al. [38] proved that when the DRL is applied to classification tasks, the obtained classifier ends up being optimal for the observed training distribution, and the core of the proof lies in the over-flexibility of the built uncertainty set. And Fronger et al. [16] also pointed out the problem of overwhelmingly-large decision set, and they used large number of unlabeled examples to further constrain the distribution set.

As mentioned above, the uncertainty set built in WDRL is often overwhelmingly large in wild high-dimensional scenarios. To demonstrate this over pessimism problem of WDRL, we design a toy example in 6.1.1 to show the necessity to construct a more practical uncertainty set. Indeed, without any prior knowledge or structural assumptions, it is quite difficult to design a practical set for robustness under distributional shifts.

Therefore, in this work, we consider a dataset $D=\left\{D^{e}\right\}_{e\in\mathcal{E}_{tr}}$, which is a mixture of data $D^{e}=\left\{(x_{i}^{e},y_{i}^{e})\right\}_{i=1}^{n_{e}}$ collected from multiple training environments $e\in\mathcal{E}_{tr}$, $x_{i}^{e}\in\mathcal{X}$ and $y_{i}^{e}\in\mathcal{Y}$ are the $i$-th data and label from environment $e$ respectively. Specifically, each dataset $D^{e}$ contains examples identically and independently distributed according to some joint distribution $P_{XY}^{e}$ on $\mathcal{X}\times\mathcal{Y}$. Given the observations that in real scenarios, different covariates have different extents of stability, we propose assumption 1.

Intuitively, assumption 1 indicates that the correlation between stable covariates $S$ and the target $Y$ stays invariant across environments, which is quite similar to the invariance in [39, 31, 40]. Moreover, assumption 1 also demonstrates that the influence of $V$ on the target $Y$ can be wiped out as long as whole information of $S$ is accessible. Under the assumption 1, the disparity among covariates revealed in the heterogeneous datasets can be leveraged for better construction of the uncertainty set. And based on assumption 1, we propose our problem:

In this work, we propose the Stable Adversarial Learning (SAL) algorithm, which leverages heterogeneous data to build a more practical uncertainty set with covariates differentiated according to their stability.

Firstly, we introduce the Wasserstein Distributionally Robust Learning (WDRL) framework which attempts to learn a model with minimal risk against the worst-case distribution in the uncertainty set characterized by Wasserstein distance:

Following the intuition above that the uncertainty should not be isotropic along stable and unstable directions, we propose to learn an anisotropic uncertainty set with the help of heterogeneous environments. The objective function of our SAL algorithm is:

where $P_{0}$ denotes the training distribution, $W_{c_{w}}$ denotes the Wasserstein distance with transportation cost function $c_{w}$ defined as

and $\mathcal{W}=\left\{w:w\in[1,+\infty)^{m+1}\ \ \text{and}\ \min(w)=1\right\}$ denotes the covariate weight space ($\min(w)$ denotes the minimal element of $w$), and $\mathcal{L}^{e}$ denotes the average loss in environment $e\in\mathcal{E}_{tr}$, $\alpha$ is a hyper-parameter to adjust the tradeoff between average performance and the stability.

The core of our SAL is the covariate weight learning procedure in equation 3. In our algorithm, the uncertainty set is built to achieve stable performance across heterogeneous multiple environments. Intuitively, $w$ controls the perturbation level of each covariate and formulates an anisotropic uncertainty set compared with the conventional WDRL methods. The objective function of $w$ (equation 3) contains two parts: the average loss in training environments as well as the maximum margin, which aims at learning such $w$ that the resulting uncertainty set leads to a learner with uniformly good performance across environments. Equation 2 is the objective function of model’s parameters via distributionally robust learning with the learnable covariate weight $w$. During training, the covariate weight $w$ and model’s parameters $\theta$ are iteratively optimized.

Details of the algorithm are delineated below. We first will introduce the optimization of model’s parameter in section 4.1, then the transportation cost function learning procedure in section 4.2. The pseudo-code of the whole Stable Adversarial Learning (SAL) algorithm is shown in Algorithm 1.

Here we first provide the robustness guarantee for our method, and then we analyze the rationality of our uncertainty set, which also demonstrates the uncertainty set built in our SAL is more practical. And we finally derive the generalization bounds for our method.

In this section, we validate the effectiveness of our method on simulation data and real-world data.

Baselines    We compare our proposed SAL with the following methods.

• •

  Empirical Risk Minimization(ERM):

  $$\min\limits_{\theta}\mathbb{E}_{P_{0}}\left[\ell(\theta;X,Y)\right]$$

  (41)

• •

  Wasserstein Distributionally Robust Learning(WDRL):

  $$\min\limits_{\theta}\sup\limits_{Q\in W(Q,P_{0})\leq\rho}\mathbb{E}_{Q}\left[\ell(\theta;X,Y)\right]$$

  (42)

• •

  Invariant Risk Minimization(IRM[31]):

  $$\min\limits_{\theta}\sum_{e\in\mathcal{E}}\mathcal{L}^{e}+\lambda\|\nabla_{w|w=1.0}\mathcal{L}^{e}(w\cdot\theta)\|^{2}$$

  (43)

For completeness, we also compare with LASSO [42], and Ridge regression [43].

For ERM and WDRL, we simply pool the multiple environments data for training. For fairness, we search the hyper-parameter $\lambda$ in $\{0.01,0.1,\dots,1e0,1e1,\dots,1e4\}$ for IRM and the hyper-parameter $\rho$ in $\{1,5,10,20,50,80,100\}$ for WDRL. And we search the hyper-parameters $\lambda$ for LASSO and Ridge in $\{1e-3,1e-2,\dots,1e-1,\dots,1e1\}$. The best hyper-parameter is selected according to the validation set, which is sampled i.i.d from the training environments.

Kinds of Distributional Shifts  To demonstrate the superiority of our methods, we design two typical kinds of distributional shifts, including selection bias[32, 33] and anti-causal effects[31]. In our simulation data, we introduce strong distributional shifts, where the spurious correlation between training and testing data varies a lot.

Evaluation Metrics  We use $\mathrm{Mean\_Error}$ defined as $\mathrm{Mean\_Error}=\frac{1}{|\mathcal{E}_{te}|}\sum_{e\in\mathcal{E}_{te}}\mathcal{L}^{e}$ and $\mathrm{Std\_Error}$ defined as $\mathrm{Std\_Error}=\sqrt{\frac{1}{|\mathcal{E}_{te}|-1}\sum_{e\in\mathcal{E}_{te}}\left(\mathcal{L}^{e}-\mathrm{Mean\_Error}\right)^{2}}$ which are the mean and standard deviation error across testing environments $e\in\mathcal{E}_{te}$.