Seeking Flat Minima with Mean Teacher on Semi- and Weakly-Supervised Domain Generalization for Object Detection

Many methods have been proposed for domain generalization on image classification tasks as summarized in recent survey papers [42, 30]. Among a variety of domain generalization methods, finding flat minima is one of the most common approaches [9, 4, 3, 14, 2, 34, 15, 39]. Those studies empirically and theoretically showed that finding flat minima in parameter space results in a better generalization ability. Ismailov et al. [14] and Cha et al. [3] demonstrated that empirical risk minimization (ERM) with stochastic gradient descent (SGD) converges to the vicinity of a flat minimum, and averaging the parameter weights over a certain number of training steps/epochs results in reaching the flat minimum. Cha et al. [3] also derived the theoretical relationship between flat minima and generalization gap. Inspired by these findings, we reveal that the Mean Teacher learning framework leads to flat minima, and thus can obtain good generalization ability.

Domain generalization for object detection has not been widely explored, compared with image classification. Lin et al. [18] proposed a method for disentangling domain-specific and domain-invariant features by adversarial learning on both image-level and instance-level features for DGOD. Liu et al. [20] investigated DGOD in underwater object detection and proposed DG-YOLO. For Single-DGOD, Wang et al. [35] proposed a self-training method that uses the temporal consistency of objects in videos. Wu et al. [36] proposed a method for disentangling domain-invariant features by contrastive learning and self-distillation. Fan et al. [8] proposed perturbing the channel statistics of feature maps, which can be interpreted as data augmentation of image styles to a variety of domains. Wang et al. [32] proposed a disentangle method on frequency space for object detection from unmanned aerial vehicles. Vidit et al. [29] proposed an augmentation method using a pre-trained vision-language model (CLIP) with textual prompts.

Unlike the above methods, as discussed in Sects. 1 and 2, we tackle SS-DGOD (Semi-Supervised Domain Generalization for Object Detection) and a new problem setting called WS-DGOD (Weakly-Supervised Domain Generalization for Object Detection). SS-DGOD and WS-DGOD are more practical than DGOD because labeled data is required only from one domain. The most closely related to our work is Malakouti and Kovashka’s work [21]. They tackled SS-DGOD and proposed a language-guided alignment method. However, the limitation of their method is that it requires a backbone network that was pre-trained on vision-and-language tasks. Our experiments show that the object detectors trained with the Mean Teacher learning framework and our regularization outperform their method.

There are a few methods that use both labeled and unlabeled data for domain generalization (SSDG) on image classification [41, 43]. Zhang et al. [41] proposed an unsupervised pre-training method called DARLING, which performs contrastive learning on unlabeled images to obtain domain-irrelevant feature representation. Zhou et al. [43] extended a semi-supervised learning method called FixMatch [26] to SSDG.

In contrast to those studies, we tackle SSDG for object detection. We also tackle the “weakly-labeled” setting (i.e., WS-DGOD), which has not been explored even for image classification.

Mean Teacher learning framework was originally proposed for semi-supervised image classification [28]. Several studies have investigated the use of the Mean Teacher learning framework for a variety of tasks such as domain generalization on image classification [37], (in-domain) weakly-supervised object detection [33], (in-domain) semi-supervised object detection [22], UDA-OD [6, 17], and UDA for semantic segmentaion [1, 31, 12, 38]. Lee et al. [16] provided a theoretical analysis of the Mean Teacher learning framework on masked image modeling pretext tasks for semi-supervised image classification. We show that the Mean Teacher learning framework also works well on different settings (SS-DGOD and WS-DGOD), provide their interpretations, and propose incorporating a regularization method.

On both SS-DGOD and WS-DGOD, our goal is to obtain object detectors that work well on the unseen target domain data $\mathcal{D}_{t}=\{X_{t}\}$. Gulrajani and Lopez-Paz [10] reported that if carefully implemented, empirical risk minimization (i.e., the image classifier simply trained with supervised learning on multiple domains) outperformed state-of-the-art domain generalization methods on several benchmark datasets for image classification. Following this important finding, we expect similar behavior on object detection and aim to train an object detector on multiple domains $\mathcal{D}_{s_{i}}(i=1,\cdots,N_{D})$. However, we have no ground-truth labels (or have only weak labels) for $\mathcal{D}_{s_{i}}(i=2,\cdots,N_{D})$ although ground-truth labels are available for $\mathcal{D}_{s_{1}}$. Therefore, the question is how to train a detector on those domains. Our solution is to use the Mean Teacher learning framework for object detection [17, 5] shown in Fig. 1, where we have two networks (teacher and student) with the same structure and train the student network using the pseudo-labels generated by the teacher network. Note that this Mean Teacher learning framework can be applied to any object detector, but we hereafter describe the loss functions of FasterRCNN [23] as an example for ease of explanation.

If we start the Mean Teacher learning from randomly initialized parameters, the teacher network cannot output reliable pseudo labels. Therefore, we first perform supervised learning with the labeled data of one source domain $\mathcal{D}_{s_{1}}=\{(X_{s_{1}},B_{s_{1}},C_{s_{1}})\}$. \linenomathAMS

where $\mathcal{L}^{\mathrm{cls}}_{\mathrm{RPN}}$ and $\mathcal{L}^{\mathrm{reg}}_{\mathrm{RPN}}$ are the classification and regression losses for region proposal networks (RPN), respectively. $\mathcal{L}^{\mathrm{cls}}_{\mathrm{RoI}}$ and $\mathcal{L}^{\mathrm{reg}}_{\mathrm{RoI}}$ are those for RoIhead. We initialize both the teacher and student networks with the parameters $\theta^{*}$ obtained from this pre-training.

We define an empirical risk as $\mathcal{E}_{\mathrm{ER}}(\theta)\vcentcolon=\sum_{i=1}^{N_{D}}\mathcal{L}^{\mathrm{sup}}_{s_{i}}(\theta)$ when we assume that ground-truth labels are available on all the training domains. A risk at the target domain is defined as $\mathcal{E}_{t}(\theta)\vcentcolon=\mathcal{L}^{sup}_{t}(\theta)$. The goal is to minimize the test risk $\mathcal{E}_{t}(\theta)$ by only solving the empirical risk minimization (ERM), i.e., $\min_{\theta}\mathcal{E}_{\mathrm{ER}}(\theta)$. Hereafter, we use the terms risk and loss interchangeably.

Previous studies for domain generalization demonstrated both theoretically and empirically that neural networks with flatter minima in parameter space exhibit superior generalization ability to unseen domains [9, 4, 3, 14, 2, 34, 15, 39]. Fig. 3 shows its intuitive interpretation. We can see that when there is a domain shift between training and testing, the flat minimum of the training loss results in a lower test loss than the sharp minimum.

Cha et al. [3] theoretically revealed the relationship between the flat minima and generalization gap (i.e., performance drop by domain shift). We briefly describe the theorem for the subsequent explanation. We consider the worst-case loss within neighbor regions in parameter space, which is defined as a robust risk $\mathcal{E}_{RR}^{\gamma}(\theta)\vcentcolon=\max_{{\|\Delta}\|\leq\gamma}\mathcal{E}_{ER}(\theta+\Delta)$. Here, $\gamma$ is the radius of the neighbor region. As shown in Fig. 3, when $\gamma$ is sufficiently large, sharp minima of the empirical risk are not minima of the robust risk. In contrast, the minima of the robust risk (i.e., $\operatorname*{arg\,min}_{\theta}\mathcal{E}_{RR}^{\gamma}(\theta)$) are also minima in the flat regions of the empirical risk. The following theorem shows the relationship between the optimal solution of robust risk minimization (RRM):

For its proof, see [3]. From the theorem, we can interpret that the gap between the RRM and ERM (i.e., $\mathcal{E}_{RR}^{\gamma}(\theta^{\gamma})-\min_{\theta^{\prime\prime}}\mathcal{E}_{ER}(\theta^{\prime\prime})$) upper bounds the generalization gap in the test domain (i.e., $\mathcal{E}_{t}(\theta^{\gamma})-\min_{\theta^{\prime}}\mathcal{E}_{t}(\theta^{\prime})$). Intuitively, as shown in Fig 3, the gap between the RRM and ERM narrows at flat regions of ERM. Therefore, we can interpret that lowering the gap leads to flat minima of ERM and results in better generalization performance on the target domain.

We explain why the EMA update in the Mean Teacher learning framework leads to flat minima. Mandt et al. [27] showed that optimizing with constant SGD (i.e., SGD with a fixed learning rate) converges to a Gaussian distribution centered on the optimum. On the basis of this finding, Izmailov et al. [14] and Cha et al. [3] showed that the ERM with SGD converges to the marginal of a flat minimum, and averaging the weights of the parameters over some training steps/epochs leads to the flat minima. Although they [14, 3] proposed sophisticated algorithms for averaging the weights to avoid overfitting, we found that a simple EMA leads to flat minima. The experiments presented in Sec. 7 show that the teacher network with only the EMA update of the student (i.e., without pseudo labeling) as shown in Eqs. (15-16) can reach flatter minima and perform better than the student.

We explain why learning from pseudo-labels in the Mean Teacher learning framework leads to flat minima. Assuming that the pseudo-labels from the teacher are accurate enough (i.e., similar enough to ground truth), $\mathcal{L}_{s_{i}}^{unsup}$ in Eq. (9) can be approximated by $\mathcal{L}_{s_{i}}^{sup}$, and we can regard the student network as the ERM in Sec. 5.1. On the other hand, as explained in Sec. 5.3 and shown in the experiments, because the teacher network updated with EMA has a better ability to reach flat minima than the student, we can regard the teacher network as the RRM. Therefore, from Eq. (14), the smaller the difference between the losses of the teacher and student, the smaller the generalization gap in the target domain is. Fig. 5 shows its intuitive interpretations. At the flat region, the trajectory of the student over the training steps and their mean (teacher) have similar loss values. In contrast, there is a large difference between the loss values of the trajectory of the student and their mean at the sharp valley.

Next, we show that learning from pseudo-labels in the Mean Teacher learning framework makes the losses of the student and teacher similar. Because the student is trained with the output from the teacher as pseudo-ground truth, the training promotes the outputs from the student similar to those from the teacher. When we use monotonically increasing/decreasing functions with respect to the outputs as loss functions $\mathcal{E}$ (e.g., cross-entropy loss $\mathcal{E}(p)=p_{gt}\log(p)$), the more similar the outputs are, the more similar the loss values are, as shown below:

Let us consider $p_{3}$ as the teacher’s output, and $p_{2}$ and $p_{1}$ as the outputs of the student. Since $p_{2}$ is closer to $p_{3}$ than $p_{1}$, the loss of $p_{2}$ becomes more similar to the loss of $p_{3}$ than that of $p_{1}$. Therefore, we can interpret that learning from pseudo-labels align the outputs from the student to be similar to those from the teacher, thereby aligning the loss values, consequently leading to flat minima.

As discussed in Sec. 5, when the output from the student and teacher are similar, the networks tend to reach flat minima. To this end, we propose incorporating a simple regularization method to make the two networks’ outputs more similar by training the student using raw outputs from the teacher.

Fig. 5 shows an overview of the method. The concept is to apply regularization so that the outputs from the two networks are similar for the same input image. Specifically, we perform weak data augmentations to the unlabeled (or weakly labeled) image $x_{s_{i}}^{j}$ and input the image into the teacher network. We then use the output from the teacher $\{(\hat{b}_{s_{i}}^{jr},\hat{p}_{s_{i}}^{jr})\}_{r=1}^{N_{R}}$ directly as pseudo-ground truth without post-processing in Eq. (5). To update the student, we input the same weakly augmented image $x_{s_{i}}^{j}$ into the student and calculate the loss as follows: \linenomathAMS

where $\hat{B}_{s_{i}}=\{\hat{b}_{s_{i}}^{j}\}_{j=1}^{N_{s_{i}}}$ and $\hat{C}_{s_{i}}=\{\hat{c}_{s_{i}}^{j}\}_{j=1}^{N_{s_{i}}}$ are the raw pseudo-labels from the teacher, and $\beta$ is a hyperparameter to tune the strength of the regularization.

We can regard the regularization method as a type of knowledge distillation because the student is trained to mimic the raw output from the teacher. Although the technical details are different, it has been empirically shown that knowledge distillation methods are effective on related tasks such as Single-DGOD [36] and domain adaptive semantic segmentation [38]. We believe that our interpretation revealed one of the reasons knowledge-distillation methods lead to better generalization ability.

We used the artistic style image dataset [13], which has four domains: natural image, clipart, comic, and watercolor. The natural image domain has 16,551 images from PASCAL VOC07&12, and the other domains have 1,000, 2,000, and 2,000 images, respectively. There are six object classes (bike, bird, car, cat, dog, and person), and we removed the images that do not contain these classes.

We conducted the experiments on two patterns of domains. In the first pattern, we set the natural image domain as the labeled domain $s_{1}$ and set clipart and comic as the unlabeled domains $s_{2},s_{3}$. We set watercolor as the target domain $t$. Concretely, we used the trainval set of PASCAL VOC 2007&2012, the train set of clipart, and the train set of comic for training. We then used the test sets of clipart and comic for validation. For evaluation (testing), we used the test set of watercolor.

In the second pattern, we set $(s_{1},s_{2},s_{3},t)=(\mathrm{natural,watercolor,comic,clipart})$. The results on another dataset are shown in the supplementary material.

We used soft pseudo labeling proposed in [5] for the Mean Teacher learning. We used Gaussian FasterRCNN [5] as the object detector, in which the regression output is modified to use the soft labels. We used ResNet101 [11] as the backbone. We applied the same hyperparameters as in a previous study [5] except for the number of iterations. All training (including baseline models) was done with four A100 GPUs. The parameters of the backbone network were initialized with the ResNet101 pre-trained on ImageNet. The hyperparameter $\beta$ in Eq. (18) was set to 0.5 throughout the experiments. During the inference (testing) phase, we used the teacher network. Other details are given in the supplementary material.

As the baseline, we trained the detector Gaussian FasterRCNN on Single-DGOD setting (i.e., supervised learning on $s_{1}$ in Eqs. (1-4)). To show the effectiveness of the EMA update, we trained Gaussian FasterRCNN + EMA with Eqs. (15-16). Gaussian FasterRCNN + EMA + PL is a detector trained with the Mean Teacher learning framework in Eqs. (8-13). Gaussian FasterRCNN + EMA + PL + Regul. is a detector with the Mean Teacher learning framework and the regualization in Eqs. (17-21).

To confirm the upper-bound performance, we also trained Gaussian FasterRCNN on DGOD and Oracle settings. On DGOD, the detector was trained with supervised learning using the ground-truth labels on the domains $s_{1},s_{2}$, and $s_{3}$. On Oracle, the detector was trained with supervised learning on $s_{1},s_{2},s_{3}$, and the target domain $t$.

Because there is only one existing method on SS-DGOD (i.e., CDDMSL [21]), we compared the above detectors with state-of-the-art methods on related task settings such as Single-DGOD and UDA-OD.

Table 2 shows the results on the artistic image style dataset when $(s_{1},s_{2},s_{3},t)=$(natural, clipart, comic, watercolor). We evaluated with the mean average precision (mAP50) when the IoU threshold was 0.5. EMA increased the mAP from 50.5 to 55.5, and this was further boosted to 56.6 with pseudo labeling (PL). We observed additional improvement to 58.2 with the regularization. The regularization improved the performance not only on SS-DGOD but also on WS-DGOD. Those results are comparable to those of the detectors trained on DGOD and Oracle. The detectors trained on SS-DGOD and WS-DGOD also performed comparably to or better than those on UDA-OD, although we did not use the target domain data during the training.

For fair comparisons, we trained CDDMSL with Res101 backbone pre-trained on ImageNet. However, its performance significantly degraded, as reported in a previous study [21], because it requires language-guided training, and initializing the model with RegionCLIP is crucial to achieve good performance.