AdvST: Revisiting Data Augmentations for Single Domain Generalization

We define a semantics transformation as a composition of several standard data augmentation functions that manipulate certain kinds of semantics of a sample. For example, we can perturb both the hue and brightness of an image $x$ with $\tau(x;\omega)=o_{h}(o_{b}(x;\omega_{b});\omega_{h})$, where $o_{h}$ is the function that changes the hue of $x$, $o_{b}$ changes the brightness of $x$, and $\omega=\omega_{b}\cup\omega_{h}$ denotes the set of parameters for $\tau$. We construct a set of $M$ semantics transformations $\mathcal{T}=\{\tau_{i}(\cdot;\omega_{i}),i=1,\ldots,M\}$ by randomly composing $L(1\leq L\leq L_{\max})$ unique standard data augmentation functions (see Appendix).

Intuitively, a semantics transformation with a large $L$ can produce more diverse samples. However, depending on the target domain data, the semantics transformations that produce more diverse samples are not necessarily better than those producing less diverse ones. Since we have no knowledge about target domains in SDG, we first uniformly choose the length for semantics transformations and then uniformly choose a semantics transformation with the selected length. Thus, we derive the distribution over $M$ semantics transformations as $G(\tau^{L})=\frac{1}{M_{L}L_{\max}}$, where $\tau^{L}$ denotes a semantics transformation with $L$ standard augmentations, $L_{\max}$ is the maximum number of standard augmentations in $\tau^{L}$, and $M_{L}$ is the total number of $\tau^{L}$ and satisfies $M=\sum_{L=1}^{L_{\max}}M_{L}$.

SDG aims to train a model that is robust to unseen domain shifts with the training samples from a single source domain. The robustness of the trained model to unseen domain shifts depends on how much the training data covers target domains. Therefore, with semantics transformations, we aim to generate new data samples that have large domain shifts from the source domain, increasing the chance of covering data samples from unseen target domains.

A key property of the samples from target domains is that they often yield a high average prediction loss because of their large domain shifts from the source. This motivates AdvST, an adversarial learning framework that learns semantics transformations to generate challenging samples with significant semantics variations for model training.

Given a model $f_{\theta}$ with parameters $\theta$, a set of source domain samples $\mathcal{D}_{S}=\{(x_{n},y_{n})\}_{n=1}^{N}$ with $N$ pairs of training sample $x_{n}$ and its label $y_{n}$, and a distribution $G$ over a set of $M$ semantics transformations $\{\tau_{i}(\cdot;\omega_{i})\}_{i=1}^{M}$, we express the learning objective of AdvST as:

where $\xi=(x,y)$ denotes a tuple of a sample $x$ and its label $y$, $\xi^{\prime}=(\tau(x;\omega),y)$ is the tuple of the same label $y$ and a new sample obtained by applying the semantics transformation $\tau$ to $x$, $\ell(\theta;\xi)$ is the prediction loss for $\xi=(x,y)$, $\Theta$ denotes the set of all possible values of $\theta$, $\psi=\cup_{i=1}^{M}\omega_{i}$ denotes the union of the parameters of $M$ semantics transformations, $\Psi$ is the set of all possible values of $\psi$, $\lambda$ is a nonnegative regularization parameter, and $d_{\theta}$ is the squared Euclidean distance function between $\xi$ and $\xi^{\prime}$ in the deep feature space of the model $f_{\theta}$, i.e., $d_{\theta}(\xi,\xi^{\prime})=\|v-v^{\prime}\|_{2}^{2}$ with the embeddings $v$ and $v^{\prime}$ of $x$ and $x^{\prime}$, respectively.

The objective in (1) aims to train a robust model with the challenging samples generated from the samples in the source domain while maintaining the core features of the original data. The novel part of (1) is that instead of generating images in the pixel space, we adversarially learn the parameters of semantics transformations, exploiting the domain knowledge in standard data augmentations to generate diverse images.

We adopt an iterative optimization algorithm (Volpi et al. 2018; Zhao et al. 2020) to solve (1). Specifically, the algorithm consists of a minimization and a maximization optimization procedures.

The learning objective of SDG can be expressed via DRO (Gao and Kleywegt 2022) since it does not rely on the notion of a known target distribution. Specifically, DRO chooses a set of probability distributions $\mathcal{U}$ called uncertainty set, and then finds a decision $\theta$ from $\Theta$ that provides the best hedge against $\mathcal{U}$ by solving the following mini-max problem:

where $\ell(\theta;\xi)$ is the prediction loss with the data-label pair $\xi=(x,y)$, $\Theta$ denotes the set of all possible model parameters, and $\mathcal{U}$ contains distributions that are at most $\delta$-distance away from the source distribution $P$. The uncertainty set, $\mathcal{U}=\{Q|D(P,Q)<\delta\}$, depends on a distance metric $D(\cdot,\cdot)$ and a predefined threshold $\delta>0$. The objective in (5) finds an optimized model under the worst-case distribution $Q^{*}$ found in $\mathcal{U}$ that maximizes the prediction loss.

Given a set of $M$ semantics transformations, a semantics-induced distribution $Q_{\psi}(\xi^{\prime})$ is defined as follows

where $\xi^{\prime}=(x^{\prime},y^{\prime})$, $\xi=(x,y)$ is a sample from the source distribution $P$, $\psi=\cup_{i=1}^{M}\omega_{i}$ denotes the parameters of $M$ semantics transformations, and $p(\xi^{\prime}|\tau_{i},\xi,\omega_{i})$ is the probability of obtaining $\xi^{\prime}$ from $\xi$ and the $i$th semantics transformation $\tau_{i}$ with parameters $\omega_{i}$. We require that transformed samples are still assigned with their original labels. Therefore, we have $p(\xi^{\prime}|\tau_{i},\xi,\omega_{i})=0$ if $y^{\prime}\neq y$. Moreover, if $\tau_{i}$ is a deterministic transformation, then $p(\xi^{\prime}|\tau_{i},\xi,\omega_{i})=1$ when $\tau_{i}(x;\omega_{i})=x^{\prime}$ and $y^{\prime}=y$ and $p(\xi^{\prime}|\tau_{i},\xi,\omega_{i})=0$ otherwise. If $\tau_{i}$ is a stochastic transformation, then $p(\xi^{\prime}|\tau_{i},\xi,\omega_{i})$ follows the distribution of $\tau_{i}(x;\omega_{i})$. A sample $\xi^{\prime}$ from $Q_{\psi}$ can be obtained by first sampling $\xi$ from $P$ with $y=y^{\prime}$ and $\tau_{i}$ from $G$, and then obtaining $x^{\prime}=\tau_{i}(x;\omega_{i})$. Given $G$ and $P$, $Q_{\psi}$ fully depends on $\psi$. We denote the set of all semantics-induced distributions as $\mathcal{Q}_{\Psi}=\{Q_{\psi}|\psi\in\Psi\}$, where $\Psi$ is the set of all possible parameters $\psi$.

The uncertainty set of AdvST consists of semantics-induced distributions $Q_{\psi}$ around the source distribution $P$ to simulate unseen target distributions. These distributions should not deviate too much from the source to avoid hedging against noisy distributions that are not learnable. Hence, we need a proper distance metric $D(\cdot,\cdot)$ to control the distribution shifts. Since semantics transformations create new data samples, we use Wasserstein distances (Definition 1) as the metric $D$ to allow a data distribution $Q_{\psi}$ to have a different support from that of $P$.

We define the transportation cost function $c$ in the deep feature space (Zhao et al. 2020; Volpi et al. 2018) to include distributions whose samples have large style variations since these samples may still be close to the samples from the source distribution in the deep feature space. To exclude noisy distributions whose data samples change their original labels in the source domain after transformations, we design the cost of moving a source distribution sample to such a sample as infinity. Specifically, the cost function of moving $\xi=(x,y)\sim P$ to $\xi^{\prime}=(x^{\prime},y^{\prime})\sim Q_{\psi}$ given the model $\theta$ is defined as follows

where $v$ and $v^{\prime}$ are the model-dependent embeddings for $x$ and $x^{\prime}$, respectively. Therefore, the uncertainty set that we consider in AdvST is

where $\delta$ ($\delta>0$) denotes the predefined distance threshold between the source $P$ and the semantics-induced distributions $\mathcal{Q}_{\Psi}$.

Directly solving Eq. (5) with $\mathcal{U}=\mathcal{U}_{\Psi}$ is intractable since it requires searching over the infinite dimension space of distribution functions. We consider the following Lagrangian relaxation with the penalty parameter $\lambda$:

However, Eq. (10) is still hard to compute. For the inner maximization term of Eq. (10), Proposition 1 provides a tractable form which only requires the source distribution $P$ and the distribution over semantics transformations $G$.

The proof of Proposition 1 (see Appendix) includes taking the dual reformulation of Eq. (10) and considering a semantics-induced distribution $Q$ as a mixture of $M$ distributions. We observe that the objective in (1) actually minimizes the empirical version of (11) with $P$ and $c_{\theta}$ being replaced by $\mathcal{D}_{S}$ and $d_{\theta}$, respectively.

We ran our experiments on Nvidia Quadro RTX 8000 GPUs. We ran our experiments 5 times with different random seeds and reported the average accuracy with standard deviation.

We conducted ablation studies on AdvST-ME using the Digits dataset. We evaluated how semantics transformations (Semantics), the contrastive regularizer (Contrastive), and the entropy regularizer (Entropy) affect the average SDG performance. We observe from Table 1 that semantics transformations can significantly boost the average classification accuracy by 17.7% (Configurations 1 and 2). The contrastive and entropy regularizers can further boost the performance of Configuration 1 by 0.8% (Configuration 3) and 1.8% (Configuration 4), respectively. Our method (Configuration 5) achieves the highest average classification accuracy with all three components.

We further compared the coverage of generated samples on target domain data between our methods, AdvST and AdvST-ME, and their pixel-level counterparts, ADA (Volpi et al. 2018) and ME-ADA (Zhao et al. 2020), which directly generate images in the pixel space. We visualized how the samples generated by ADA, ME-ADA, and our methods distribute in the embedding space in Figure 1. We color the samples from the source domain MNIST orange and the samples from the four target domains gray. We give details for obtaining the figure in Appendix. From Figure 1(a) and (c), we observe that most of the synthetic samples distribute very close to the source domain data and have little coverage on the target domains. In contrast, the synthetic samples in Figure 1(b) and (d) deviate from the source domain and have broad coverage on the target domains.

We provide analyses on the sensitivity of $\lambda$ and the effect of different semantics transformations in Appendix.

Baselines. We included ADA, ME-ADA, and the following methods for comparison: ERM, which trains a model only using the standard cross-entropy loss; CCSA (Motiian et al. 2017), which aligns samples from different domains to improve generalization; d-SNE (Xu et al. 2019), which minimizes the maximum distance between sample pairs of the same class and maximizes the minimum distance among sample pairs of different categories; JiGen (Carlucci et al. 2019), which is a multi-task learning method that combines the target recognition task and the Jigsaw classification task; M-ADA (Qiao, Zhao, and Peng 2020), which uses generative models and meta-learning (Finn, Abbeel, and Levine 2017; Chen and Zhang 2021, 2022) to improve ADA; AutoAug (Cubuk et al. 2018) and RandAug (Cubuk et al. 2020b), which augment data based on the searched augmentation policies; RSDA (Volpi and Murino 2019), which randomly searches image transformations to train a robust model; and PDEN (Li et al. 2021) and L2D (Wang et al. 2021), which use generative models for data augmentation.

We further demonstrated the utility of our methods by evaluating the average classification accuracy of our methods on target domains with limited training data. We used the Art dataset from PACS as the source domain and the remaining three datasets in PACS as the target domains. We used partial training data of the Art domain and reported the average classification accuracy over the three target domains in Figure 3. We observe that under different ratios of available training data, our methods, AdvST and AdvST-ME, consistently outperform ADA and ME-ADA, respectively. The gains are significant when the ratio is small, demonstrating the effectiveness of our method when there is a lack of available training data.