Consistency Regularization for Domain Generalization
with Logit Attribution Matching

Let us first consider the case where we have a set of unlabeled SS pairs $\{({x}_{i},\tilde{{x}}_{i})\}_{i=1}^{N}$. The distinction between labeled and unlabeled SS pairs is if the semantic information is invariant for just one particular class or all classes, not whether the original examples $x_{i}$ is labeled. Unlabeled SS pairs contain stronger information than labeled SS pairs: two examples ${x}_{i}$ and $\tilde{{x}}_{i}$ contain the same semantic information for all classes implies that they contain the same semantic information for every class.

With unlabeled SS pairs, the first two conditions of Theorem 1 can be approximately satisfied by solving the following constrained optimization problem:

	$\displaystyle\min_{\theta}\hskip 8.00003pt$	$\displaystyle{\mathbb{E}}_{({x},y)\sim P^{\mathrm{s}}}[-\log\hat{P}_{\theta}(\hat{Y}=y|{x})]$
	$\displaystyle\mathrm{subject\ to}\hskip 8.00003pt$	$\displaystyle\hat{P}_{\theta}(\hat{Y}|{X}={x}_{i})=\hat{P}_{\theta}(\hat{Y}|{X}=\tilde{{x}}_{i}),\hskip 5.69046pti\in[N].$

Of course, how well the two conditions are actually satisfied depends on how representative the unlabeled SS pairs we have are of all possible SS pairs.

If we turn the equality constraints into a consistency regularization (CR) term, the problem becomes:

$$\min_{\theta}\hskip 3.99994pt\mathbb{E}_{({x},y)\sim P^{\mathrm{s}}}[-\log\hat{P}_{\theta}(\hat{Y}=y|{x})]+\lambda\mathbb{E}_{i}[r_{\theta}({x}_{i},\tilde{{x}}_{i})],$$

where $\lambda$ is a balancing parameter and the summation over $r_{\theta}({x}_{i},\tilde{{x}}_{i})$ is a regularization term that relaxes the corresponding equality constraints.

Some notations are needed in order to discuss specific choices for $r_{\theta}$. Suppose $\hat{P}_{\theta}$ consists of a feature extractor $f_{\phi}$ with parameters $\phi$ and a linear classification head $g_{{w}}$ with parameters ${w}$. Hence, $\theta=(\phi,{w})$. For an input ${x}$, let $f_{\phi}^{u}({x})$ be the component of the feature vector $f_{\phi}({x})$ for a feature unit $u$. Let $w_{uy}$ be the weight between a feature unit $u$ and the output unit for a class $y$. The logit for class $y$ is

$$z_{\theta}^{y}({x})=\sum_{u}f_{\phi}^{u}({x})w_{uy},$$

where the summation is over all feature units $u$ and the bias is omitted.

For each unlabeled SS pair $({x}_{i},\tilde{{x}}_{i})$, the CR term $r_{\theta}({x}_{i},\tilde{{x}}_{i})$ can be defined in several ways:

	$\displaystyle r_{\theta}^{\texttt{KL}}({x}_{i},\tilde{{x}}_{i})$	$\displaystyle={D}_{\texttt{KL}}\big{[}\hat{P}_{\theta}(\hat{Y}|X={x}_{i})\;\|\;\hat{P}_{\theta}(\hat{Y}|X=\tilde{x}_{i})\big{]},$
	$\displaystyle r_{\theta}^{\texttt{JS}}({x}_{i},\tilde{{x}}_{i})$	$\displaystyle={D}_{\texttt{JS}}\big{[}\hat{P}_{\theta}(\hat{Y}|X={x}_{i})\;\|\;\hat{P}_{\theta}(\hat{Y}|X=\tilde{x}_{i})\big{]},$
	$\displaystyle r_{\theta}^{\texttt{LM}}({x}_{i},\tilde{{x}}_{i})$	$\displaystyle=\sum\nolimits_{y}\big{[}z_{\theta}^{y}({x}_{i})-z_{\theta}^{y}(\tilde{x}_{i})\big{]}^{2},$
	$\displaystyle r_{\theta}^{\texttt{FM}}({x}_{i},\tilde{{x}}_{i})$	$\displaystyle=\sum\nolimits_{u}\big{[}f_{\phi}^{u}({x}_{i})-f_{\phi}^{u}(\tilde{x}_{i})\big{]}^{2}.$

The first two terms aim to match the output probability distributions of ${x}_{i}$ and $\tilde{{x}}_{i}$ by minimizing either the KL or JS divergence between them. The third term aims to match their logit vectors, and the fourth term aims to match their feature vectors. They are used in previous methods ReLIC [Mitrovic et al., 2021], AugMix [Hendrycks et al., 2020], CoRE [Heinze-Deml and Meinshausen, 2021], and MatchDG [Mahajan et al., 2021] respectively. Note that while we focus on pairs for simplicity, logit and feature matching can also be extended to the case of multiple examples that share the same semantic contents. To achieve this, we can simply replace the sum of squared differences with the sum of variances. This is done in CoRE and MatchDG.

Now consider the case where we have a set of labeled SS pairs $\{({x}_{i},\tilde{{x}}_{i};y_{i})\}_{i=1}^{N}$. Here, each pair ${x}_{i}$ and $\tilde{{x}}_{i}$ share the same semantic information only for the class $y_{i}$. It is no longer justifiable to match all the features, logits or probabilities of all classes. In the following, we propose three methods for leveraging labeled SS pairs.

First, we can match the probabilities or logits of the target class $y_{i}$ only, leading to what we call target probability matching (TPM) and target logit matching (TLM):

	$\displaystyle r_{\theta}^{\texttt{TPM}}({x}_{i},\tilde{x}_{i};y_{i})$	$\displaystyle=\big{[}\hat{P}_{\theta}(\hat{Y}=y_{i}|{x}_{i})-\hat{P}_{\theta}(\hat{Y}=y_{i}|\tilde{x}_{i})\big{]}^{2},$
	$\displaystyle r_{\theta}^{\texttt{TLM}}({x}_{i},\tilde{x}_{i};y_{i})$	$\displaystyle=\big{[}z_{\theta}^{y_{i}}({x}_{i})-z_{\theta}^{y_{i}}(\tilde{x}_{i})\big{]}^{2}.$

To introduce the third method, note that $f^{u}_{\phi}({x}_{i})w_{uy_{i}}$ is the contribution to the logit $z_{\theta}^{y_{i}}({x})$ of $y_{i}$ from the feature unit $u$. We can match the logit contributions $f^{u}_{\phi}({x}_{i})w_{uy_{i}}$ and $f^{u}_{\phi}(\tilde{x}_{i})w_{uy_{i}}$ from all feature units $u$ to $y_{i}$. This gives rise to logit attribution matching (LAM):

$$r_{\theta}^{\texttt{LAM}}({x}_{i},\tilde{{x}}_{i};y_{i})=\sum_{u}\big{[}f_{\phi}^{u}({x}_{i})w_{uy_{i}}-f_{\phi}^{u}(\tilde{x}_{i})w_{uy_{i}}\big{]}^{2}.$$

LAM is of finer grain than TLM. Small $r_{\theta}^{\texttt{LAM}}$ implies small $r_{\theta}^{\texttt{TLM}}$, but not vice versa:

$$\begin{split}r_{\theta}^{\texttt{LAM}}({x}_{i},\tilde{{x}}_{i};y_{i})&\geq\frac{1}{m}\Big{[}\sum_{u}f^{u}_{\phi}(x_{i}){{w_{uy_{i}}}}-\sum_{u}f^{u}_{\phi}(\tilde{x}_{i}){{w_{uy_{i}}}}\Big{]}^{2}\\ &=\frac{1}{m}r_{\theta}^{\texttt{TLM}}({x}_{i},\tilde{x}_{i};y_{i}),\end{split}$$

where $m$ is the number of feature units. Also, note that

$$r_{\theta}^{\texttt{LAM}}({x}_{i},\tilde{{x}}_{i};y_{i})=\sum_{u}\big{[}f_{\phi}^{u}({x}_{i})-f_{\phi}^{u}(\tilde{x}_{i})\big{]}^{2}w^{2}_{uy_{i}}.$$

Hence, LAM exerts two complementary regularization forces, one on $g_{{w}}$ and the other on $f_{\phi}$:

• 1).

  It encourages the classification head $g_{{w}}$ to put large weights $|w_{uy_{i}}|$ on the feature units $u$ where the values of ${x}_{i}$ and $\tilde{x}_{i}$ are similar, i.e., $f_{\phi}^{u}({x}_{i})\approx f_{\phi}^{u}(\tilde{x}_{i})$. In other words, it makes $g_{{w}}$ rely on the feature units that reflect the common information contents of ${x}_{i}$ and $\tilde{x}_{i}$.

• 2).

  It encourages the feature extractor $f_{\phi}$ to make $f_{\phi}^{u}({x}_{i})\approx f_{\phi}^{u}(\tilde{x}_{i})$ for those feature units $u$ that $g_{{w}}$ relies on heavily, i.e., with large weights $|w_{uy_{i}}|$. In other words, it encourages $f_{\phi}$ to channel the common information contents of ${x}_{i}$ and $\tilde{x}_{i}$ toward the units that $g_{{w}}$ considers important.

As ${x}_{i}$ and $\tilde{x}_{i}$ share the causal factors for class $y_{i}$ but not the non-causal factors, those forces help a model focus more on the causal factors.

Our experiments involve five DG datasets, three with background shifts and two with style shifts.

iWildCam2020 (iWildCam) [Beery et al., 2020, Koh et al., 2021] consists of camera trap photos of animals taken at different locations for wildlife classification. The training domain comprises images from 200 locations, while the test and validation domains contain images from some other locations. Targeted DA is performed by Copy-Paste the animals in a training image to another image (with no animal) taken at a different location where the same animals sometimes appear [Gao et al., 2023].

ImageNet-9 [Xiao et al., 2020] includes images of nine coarse-grain classes from ImageNet [Deng et al., 2009]. Several synthetic variations are created by segmenting the foreground of each image and place it onto a different background. In our experiments, the synthetic images with a black background are used as target-augmented examples. For the test domain, we use the samples where the foreground of an original image is placed onto the background of a random image.

NICO [He et al., 2020] includes around 25,000 images across 19 classes of animals or vehicles in different contexts such as “at home” or “on the beach”. As there is no predefined train-test split, we randomly select one context per class for testing and use the remaining contexts for training. Target-augmented training examples and test domains are created in a way similar to ImageNet-9.

Camelyon17 (Camelyon) [Tellez et al., 2018, Koh et al., 2021] contains histopathology images from multiple hospitals for binary tumor classification. Images from three hospitals are used for training, while images from two other hospitals are used for testing and validation respectively. There are stylistic variations among images from different hospitals. One key stylistic difference often observed is the stain color. Therefore, the stain color jitter is applied to training images to create target-augmented examples [Gao et al., 2023].

PACS  [Li et al., 2017] contains images of objects and creatures in four different styles: photo, art, cartoon and sketch. Following common practice [Li et al., 2017, Gulrajani and Lopez-Paz, 2021], we train a model using three of the domains and test the model on the fourth domain. For Targeted DA, we apply Stable Diffusion [Rombach et al., 2022] to images in the photo domain to create target-augmented examples in the other three domains. The photo domain is therefore not used as the test domain, while the other three domains are used as the test domain in turn. See Appendix C for details.

For all datasets, RandAugment [Cubuk et al., 2020] is performed on all training examples. Targeted DA [Gao et al., 2023] is also performed on all training examples in iWildCam and Camelyon. However, it is performed on only about 5% of the training data in ImageNet-9 and NICO, and about 10% of the training data for PACS.

All CR methods have a balancing parameter $\lambda$, which is tuned on the validation domain for iWildCam and Camelyon, and on a test set from the training domain for the other three datasets. For CR and single-source methods, multiple training domains are simply combined into one. More details on how the training data are organized for different types of methods can be found in Table 5 (Appendix E).

Following Gao et al. [2023], we use a variety of models for different datasets. Specifically, we use an ImageNet pretrained ResNet-50 model [He et al., 2016] for iWildCam, and a randomly initialized DenseNet-121 model [Huang et al., 2017] for Camelyon. We use a CLIP-pretrained ViT-B/16 model [Radford et al., 2021] for ImageNet-9 and NICO, and a CLIP-pretrained ResNet-50 model for PACS.

To showcase the combined use of LAM with advanced CLIP model fine-tuning method can yield SOTA-level performance on iWildCam, we also employ CLIP-pretrained ViT-L/14 and ViT-L/14@336 model for iWildCam.

The use of various model architectures and weight initializations allows us to assess the relative merits of DG algorithms on a mixture of datasets and models. Implementation details about hyperparameters for each dataset and method can also be found in Appendix E.

Table 1 shows the results for LAM, ERM+DA, and ERM. We see that simply adding augmented data to the training set (ERM+DA) increases the average OOD score from 71.4% to 76.5% with RandAugment [Cubuk et al., 2020], and to 78.7% with Targeted DA [Gao et al., 2023]. Applying LAM on the resulting SS pairs further increases the scores to 78.2% and 81.1% in the two cases respectively. In the case of Targeted DA, the average OOD score on those five benchmarks is improved by 81.1-71.4 = 9.7%, with 78.7-71.4 = 7.3% due to data augmentation and 81.1-78.7 = 2.4% due to LAM. The improvements are especially pronounced on the iWildCam and Camelyon datasets, where Targeted DA increases the OOD scores drastically. This is consistent with what was reported in Gao et al. [2023]. LAM further improves the scores by 4.7% and 3.0% respectively.

While trying to gain some insights, we find that LAM makes a model focus on much fewer feature units (see Figure 11 in Appendix F) as compared with ERM+DA. We also use an XAI method called Grad-CAM [Selvaraju et al., 2017] to explain the outputs of the model trained on ImageNet-9 by LAM and they some other methods. Examples are shown in Figure 4 (and Figure 12 in Appendix F). We see that, in all those examples, the LAM model focuses on the foreground objects and gives the correct predictions. Those corroborate with the analysis we make at the end of Section 4.2. In contrast, the ERM+DA model is more inclined to focus on the wrong part of an input image and predict incorrectly.

In addition to comparing LAM over the traditional ERM which is based on the standard cross-entropy loss, it has been shown in Goyal et al. [2023] that when fine-tuning CLIP models, the use of CLIP contrastive loss with utilizing the CLIP text encoder is more effective. The proposed method is colloquially known as “finetune like you pretrain” (FLYP). In Table 3, we show that the use of LAM can also yield improved OOD performance over FLYP+DA.

Both LAM and ERM+DA achieve better results with Targeted DA than with RandAugment. We believe this is because Targeted DA generally yields higher quality augmentations than the latter. To further support the claim, we perform additional experiments with ImageNet-9 in the Targeted DA setting. Specifically, we test three different ways to create augmented examples: 1). use a segmentation method called GrabCut [Rother et al., 2004], 2). use another less effective segmentation method called FCN [Long et al., 2015], and 3). simply use bounding boxes that come with ImageNet-9 (Box). The resulting OOD scores are as follows:

We see that, as expected, the results with GrabCut are the best, followed by those with FCN and Box, in that order.

We also perform additional experiments with ImageNet-9 to investigate how the quantity of augmented examples influences LAM. Specifically, GrabCut is applied to different percentages of the training examples and LAM is run on the resulting SS pairs. To make a comparison, we do the same thing for the ERM+DA. The resulting OOD scores are as follows:

It is clear that the increase in the quantity of SS pairs benefits LAM, and the availability of SS pairs for a small fraction of training examples can significantly improve OOD performance already. While providing more SS pairs can also improve the performance of ERM+DA, it is obvious that the improvement is smaller than that of LAM.

Table 2 shows the results for LAM and other CR methods. Let us first compare LAM and two other CR methods we propose in this paper, namely target probability matching (TPM) and target logit matching (TLM). We see that LAM achieves higher OOD scores than the other two methods on average, and it outperforms ERM+DA in all cases while the other two methods do not. Those show that when making use of labeled SS pairs, it is more effective to apply consistency regularization to the logit contributions of the target classes (LAM) rather than the logits themselves (TLM) or the probabilities of the target classes (TPM).

Next, we compare LAM with previous strong CR methods, namely probability matching with KL or JS, logit matching (LM) and feature matching (FM). LAM achieves higher OOD scores than those methods on average. Moreover, it achieves the highest score in all cases except for Camelyon with Targeted DA. Moreover, it outperforms ERM+DA in all cases, while the other methods do not. Those show that it is generally beneficial to regard SS pairs created using both Targeted DA and RandAugment as labeled and apply LAM on them, rather than considering them unlabeled and applying any of the previous CR methods on them.

In LAM, a labeled SS pair $({x}_{i},\tilde{x}_{i};y_{i})$ is used only to regularize the contributions from feature units to the logit of the ground-truth class $y_{i}$. It does not impact the other classes. In the previous CR methods, on the other hand, the pair is used to regularize the entire feature, logit, or probability vector for ${x}_{i}$. It affects other classes as well as $y_{i}$. This is problematic when a training example ${x}_{i}$ contains multiple objects of interest. Some objects that appear in the background of the main object in ${x}_{i}$ might be removed during data augmentation. In such a case, the features of those minor objects would be suppressed. To further demonstrate the adverse consequences, we created a variant of the iWildCam dataset [Beery et al., 2020, Koh et al., 2021] by adding a small segmented image of another animal to the background of each image. The new dataset is named iWildCam-N (examples of this dataset are given in Appendix D). On this dataset, LAM still improves over ERM+DA. However, the performances of all four previous methods are substantially worse than that of ERM+DA.

Camelyon is a binary classification problem. There is no issue of suppressing features of other classes. This is probably why probability matching with JS is superior to LAM on Camelyon in the case of Targeted DA.

Table 4 shows the OOD performances of LAM with six representative single-source and multi-source DG methods reviewed in Section 2. Here only Targeted DA [Gao et al., 2023] is considered. On the first four datasets, LAM outperforms all the six DG methods on average. In particular, it outperforms them by large margins on iWildCam. While LAM improves over ERM+DA on all the first four datasets, the other methods are inferior to ERM+DA in the majority of the cases. On the binary classification dataset Camelyon, however, LAM is on par with SD [Pezeshki et al., 2021], but it still outperforms ERM+DA.

Recall that augmented examples are simply added to the training set for the single-source methods (RSC [Huang et al., 2020] and SD), and they are treated as an additional training domain for the multi-sources methods (DANN [Ganin et al., 2016], GDRO [Sagawa et al., 2020], IRM [Arjovsky et al., 2019] and V-REx [Krueger et al., 2020]). In contrast, LAM applies consistency regularization on the resulting SS pairs. The results in Table 4 show that consistency regularization with LAM is a more effective way to use augmented examples than representative previous single-source and multi-source DG methods.