DSDRNet: Disentangling Representation and Reconstruct Network for Domain Generalization

DG requires us to learn a model from a given number of source domains that can make good predictions about the unseen data distribution. MixStyle [18] introduced the mixing of source domains with a certain probability to generate new samples, aiming to enhance the model’s generalization capability. UFDN [19] proposed a unified feature disentangle for multi-domain image translation. MTAE [20] provided a general framework for domain generalization based on multi-task learning. MMD-AAE [21] utilized maximum mean difference and adversarial auto-encoder for domain generalization. DIVA [22] divided the data into domain information, category information, and other information for the structure, and used VAE for data reconstruction to improve the generalization of the system. CCSA [23] proposed a unified architecture of siamese network to solve the generalization problem in the vision domain. CrossGrad [24] utilizes domain information to assist samples, using domain-guided data augmentation to improve the model’s generalization. DDAIG [25] proposed to divide the model into label classifier, DoTNet, and domain classifier, using DoTNet to map source domain data to the unseen domain. L2A-OT [3] proposed to model the distribution between the synthesized pseudo-novel domain and the source domain, and maximize their divergence using the optimal transport. D-SAM [26] proposed a generic architecture that could improve the generalization of the system while maintaining the separation of existing source domain information and at the same time utilizing common information. DRANet [15] proposed a content-adaptive domain transfer network that helped to preserve the scene structure and transferring style. JiGen [27] proposed an algorithm similar to solving a jigsaw puzzle to improve the generalization of the system. DAEL [4] utilized a deep adversarial auto-encoder to extract domain-specific features from class labels. Epi-FCR [28] decomposed the deep network into feature extractor and classifier components and trained each component using the current domain partners.

Disentangling representation was the analysis of potential generative factors inside the data from the perspective of data generation by studying the physical structure within the data. Bengio et al. [29] introduced a variance factor disentanglement problem. VAE [30] proposed modeling real data from a maximum likelihood perspective by calculating the KL divergence between the true posterior and the variational posterior to achieve data disentanglement. ${\beta}$-VAE [31] improved disentanglement performance by increasing the value of ${\beta}$, making the learned posterior distribution statistically similar to the prior distribution. AAE [32] suggested using adversarial networks to measure the similarity between the posterior and prior distributions to enhance the representational capacity of latent variables. DRAW [33] introduced a deep recurrent network architecture to learn historical states for simple detection and object tracking. CDD [34] leveraged a pair of generative adversarial networks to achieve a cross-domain bidirectional image translation model. MLVAE [35] separated latent representations of grouped data-related factors by exchanging and sharing the bottom-level representations of samples. MixNMatch [36] achieved a disentangled representation of pose, texture, and background for single-object scenes through data factorization and layered generation.

Assuming there are M source domains ${Source}=\{{{Source}_{1}},\dots,{{Source}_{M}}\}$, where each source domain ${Source}_{m}=\{{x_{i},y_{i}}\}_{i=1}^{N}$ contains N labeled samples, the goal of the model is to learn a powerful and generalizable predictive function $h$ from the data of M source domains. The objective is to minimize the error on the unseen target domain ${T}$: ${min}\mathbb{E}{(x,y)\mathcal{\in}T}[l(h(x),y)]$.

The main process of DSDRNet is depicted in Figure 2, which includes an encoder $E$, a disentangle $S$, used in this paper as AdaIN [17], a generator $G$, a discriminator $D$, and a classifier $C$. DSDRNet aims to disentangle the semantics and attributes of samples into different latent spaces through a dual-stream network, training a highly robust and generalizable classifier. It utilizes semantic distance supervision to constrain the similarity of images and reconstructs the semantics and attributes of samples. Through adversarial training and consistency principles, it enables flexible manipulation and augmentation of the training data. Ultimately, this approach equips the model with excellent generalization capabilities for unseen samples.

Disentangling Representation.    The model’s objective is to reduce semantic differences among generated samples belonging to the same class. Unlike traditional approaches that use a unified encoder and decoder for feature extraction, this paper employs two separate networks: the disentangle $S$ and the encoder $E$. They independently perform the disentangling operation on samples $A$ and $B$ at the pixel level, resulting in $A^{v}$, $A^{s}$, $A^{f}$, $B^{v}$, $B^{s}$, and $B^{f}$, as shown below:

where $A^{v}$ represents the attributes value of sample $A$ extracted by the disentangle $S$, while $B^{v}$ represents the attributes of sample $B$ extracted by the disentangle $S$, $A^{s}$ and $A^{f}$ denote the semantic and feature information of sample $A$ obtained through the encoder $E$, respectively. Similarly, $B^{s}$ and $B^{f}$ refer to the semantic and feature information of sample $B$ acquired through the encoder $E$.

Reconstruction.    After the disentanglement operations in the previous stage, we obtained $A^{v}$, $A^{s}$, $A^{f}$, $B^{v}$, $B^{s}$, and $B^{f}$. We utilize the generator $G$ to generate reconstructed images $\bar{A}$ with the semantics $A^{s}$ and attributes $A^{v}$ of sample $A$, as well as $\bar{B}$ with the semantics $B^{s}$ and attributes $B^{v}$ of sample $B$. The results are as follows:

At the same time, we utilize the generator $G$ to generate reconstruction images $U$ with the semantic $A^{s}$ of sample $A$ and the attribute $B^{v}$ of sample $B$, as well as reconstruction images $Q$ with the semantic $B^{s}$ of sample $B$ and the attribute $A^{v}$ of sample $A$. The results are as follows:

where $U$ and $Q$ refer to the newly generated images.

We generate images $U$ and $Q$ through the first reconstruct operation. Then, we operate on $U$ using the disentangle $S$ to obtain attribute values $U^{v}$. Similarly, by using the encoder $E$, we obtain the semantic value $U^{s}$ and feature value $U^{f}$ of $U$. Likewise, we can obtain the attribute values $Q^{v}$, semantic value $Q^{s}$, and feature value $Q^{f}$ of $Q$. Using these semantics and attributes, we perform another reconstruct operation to obtain images $A^{\prime}$ and $B^{\prime}$. The specific process is as follows:

where $U^{s}$ and $U^{v}$ represent the semantics and attributes of $U$, respectively. Similarly, $Q^{s}$ and $Q^{v}$ represent the semantics and attributes of $Q$. $A^{\prime}$ and $B^{\prime}$ denote the newly generated images after the disentangling and reconstruction operations.

The overall loss of our framework consists of the intra-instance reconstruction loss $\mathcal{L}_{\text{intra}}$, inter-instance reconstruction loss $\mathcal{L}_{\text{inter}}$, self-reconstruct loss $\mathcal{L}_{\text{recon}}$, cross-cycle consistency loss $\mathcal{L}_{\text{cycle}}$, semantics adversarial loss $\mathcal{L}_{\text{adv}}$, classification loss $\mathcal{L}_{\text{kl}}$ and $\mathcal{L}_{\text{ce}}$, and with the balancing term $\alpha$:

Intra-instance Reconstruction Loss.    To ensure consistency between the features of the original image $A^{f}$ and the recombined features $U^{f}$ in the feature space, we need to calculate the angle between the two vectors, $A^{f}$ and $U^{f}$, to measure their similarity. Cosine similarity is well-suited for such calculations, as its values range from -1 to 1, with values closer to 1 indicating higher similarity. Therefore, the final intra-instance reconstruction loss can be expressed as:

This loss function quantifies the similarity between the original and recombined features by measuring the cosine of the angle between the two feature vectors, intending to make them as similar as possible.

Similarly, we can obtain the cosine similarity loss between $B^{f}$ and $Q^{f}$ as $\mathcal{L}_{\text{intra}}^{B}$. The overall intra-instance reconstruction loss can be represented as:

The cosine similarity loss allows us to assess the similarity between the two feature vectors by quantifying their angular difference. By minimizing this loss, we aim to ensure that the reconstructed images remain similar to the original images in terms of feature representations.

Inter-instance Reconstruction Loss.    Inter-instance reconstruction refers to the measurement of the dissimilarity between the reconstructed images of different instances. It captures the discrepancy between the features and attributes of different instances in the reconstructed space. The inter-instance reconstruction loss can be calculated using various metrics, such as Euclidean distance, L1 distance, or Structural Similarity Index Measure (SSIM). The purpose of the inter-instance reconstruction loss is to encourage the generated images to have distinct attributes and avoid blending or mixing the attributes of different instances in the reconstruction process.

To further constrain the similarity of image reconstruction, we utilize the L1 norm to enforce cosine similarity loss between them. Given the feature information of two images $A$ and $B$ as $A^{f}$ and $B^{f}$ respectively, and the feature information of two distinct images $U$ and $Q$ as $U^{f}$ and $Q^{f}$, the inter-instance reconstruction loss can be expressed as:

where $\text{cos}(A^{f},B^{f})$ presents the cosine similarity loss between $A^{f}$ and $B^{f}$, $\text{cos}(U^{f},Q^{f})$ presents the cosine similarity loss between $U^{f}$ and $Q^{f}$.

Self-reconstruction Loss.    We apply an L1 loss to train $E$, $S$, and $G$, reducing the disparity between the input images $A$, $B$ and their corresponding reconstructed images $\bar{A}$, $\bar{B}$, the reconstructed loss is defined as:

where $\bar{A}=G(A^{s},A^{v})$ and $\bar{B}=G(B^{s},B^{v})$, respectively.

Cross-cycle Consistency Loss.    Consistency refers to maintaining invariant between perturbed images and features to learn robust representations. In this paper, consistency loss attempts to maintain composition invariant of semantics and attributes when the domain-transferred image is re-projected into the representation space. After two translation stages, the reconstructed image should exhibit consistency in semantics and attributes with the original image. To enforce this constraint, we define the cross-cycle consistency loss as follows:

where $A^{\prime}=G(U^{v},Q^{s})$ and $B^{\prime}=G(U^{s},Q^{v})$. This loss explicitly encourages consistency in semantics and the appearance of attributes, which is more conducive to improving the generalization of the model.

Semantics Adversarial Loss.    To ensure semantic consistency between the original image $A$ and the synthesized image $U$, we introduce a content discriminator $D$ to distinguish the differences between real images and synthesized images. Additionally, we employ the generator $G$ to attempt image generation. The same approach is applied to supervise semantic consistency between the original image $B$ and the synthesized image $Q$, the semantics-based adversarial loss can be expressed as:

Kullback-Leibler Divergence Loss.    To further ensure that the model maintains semantic invariant during reconstruction, we also apply a model distribution fitting loss to the classification predictions of the original image $A$, denoted as $A^{p}$, and the classification predictions of the reconstructed image $U$, denoted as $U^{p}$, the results are as follows:

where $P(A_{k}^{p})$ represents the probability distribution of $A_{k}^{p}$, $P(U_{k}^{p})$ represents the probability distribution of $U_{k}^{p}$. Similarly, the Kullback-Leibler Divergence Loss $\mathcal{L}_{\text{kl}}^{B}$ can be obtained for $B_{k}^{p}$ and $Q_{k}^{p}$ over the probability distribution, and the final result can be expressed as:

Classification Loss.    To improve the classification accuracy of the model, we employ multiple class identifiers, denoted as $C$ to enhance classification accuracy. The procedure is as follows, given samples $A$ and $B$, encoded by the encoder $E$ to obtain $A^{f}$ and $B^{f}$, they pass through the classifier $C$ to produce $A^{p}$ and $B^{p}$. The K-way class identifier $C$ is utilized for accurate label predictions, supervised by the cross-entropy loss:

Similarly, the same procedure can be applied to obtain $\mathcal{L}_{\text{ce}}^{B}$, $\mathcal{L}_{\text{ce}}^{U}$, $\mathcal{L}_{\text{ce}}^{Q}$, the final overall cross-entropy loss can be expressed as:

As shown in Table I, the DSDRNet achieved an accuracy of 98.4% on the MNIST dataset, surpassing the ERM baseline by +2.6% and outperforming the best DANN [13] and RSC [40] by +0.6%. On the MNISTM dataset, the model achieved an accuracy of 65.2%, surpassing ERM by +6.4% and outperforming the best DDAIG [25] by +1.1%. On the SVHN dataset, the model’s performance exceeded ERM by +8.4% and outperformed the current best DDAIG [25] and L2A-OT [3] by +1.5%. On the SYN dataset, the model outperformed ERM by +11.2%, but it scored lower than the best GroupDRO [39] by 2.4%. This difference can be attributed to the fact that the SYN dataset contains randomly inserted backgrounds, leading to more pronounced model oscillations. Overall, the DSDRNet achieved an average accuracy of 80.9% across the four datasets, surpassing ERM by +7.2% and outperforming the best L2A-OT [3] by +2.8%.

Based on Table II, in the Art domain, DSDRNet achieves an accuracy of 85.2%, surpassing ERM by 8.2% and outperforming the best-performing method, DAEL [4], by +0.6%. In the Cartoon domain, the model achieves an accuracy of 78.4%, outperforming ERM by +2.5% and slightly below the best method, MixStyle [4], by +0.4%. In the Photo domain, DSDRNet outperforms ERM by +0.8% and is +0.6% better than the current best, L2A-OT [3]. In the Sketch domain, the model achieves an accuracy of 81.2%, surpassing ERM by +12.0% and outperforming the current SOTA CORAL [37] by +0.6%. On the PACS dataset as a whole, DSDRNet reaches an accuracy of 85.4%, significantly outperforming similar comparison methods, exceeding ERM by +5.9%, and outperforming MixStyle by +1.7%.

As shown in Table III, DSDRNet achieves an accuracy of 67.3% on the OfficeHome dataset, surpassing ERM by +2.6%, and outperforming the current best method, DAEL [4], by +1.2%. Specifically, in the Art domain, the model achieves an accuracy of 61.7%, surpassing ERM by +2.8%, and outperforming L2A-OT [3] by +1.1%. In the Clipart domain, DSDRNet reaches 54.6%, surpassing ERM by +5.2%, and slightly below DAEL [4] by 0.5%. In the Product domain, the model outperforms ERM by +0.9% and surpasses L2A-OT [3] by +0.4%. In the RealWorld domain, the model outperforms ERM by +1.6% and surpasses L2A-OT [3] by +0.8%. These results showcase the efficacy of DSDRNet in enhancing domain generalization performance on the OfficeHome dataset.

As shown in Table IV, on the DomainNet dataset, ERM achieves an average accuracy of 41.6% across the 6 domains. In contrast, DSDRNet achieves an accuracy of 43.3%, surpassing ERM by +1.7% and outperforming MLDG [45] by +0.8%. Specifically, in the Clipart domain, the model achieves an accuracy of 60.7%, surpassing ERM by +2.3% and outperforming MLDG [45] by +1.4%. In the Infograph domain, DSDRNet reaches an accuracy of 21.5%, surpassing ERM by +1.7% and outperforming CORAL by +0.7%. The model achieved an accuracy of 47.9% in the Painting domain, which is slightly higher than ERM by +0.6% and slightly lower than MLDG [45] by 0.9%. The model is +2.4% higher than ERM and +1.8% higher than MLDG [45] on the Quickdraw domain. In the Real domain, the model performed +1.5% better than ERM and +1.0% better than MLDG [45]. The model is +2.0% higher than ERM and +0.7% higher than MLDG [45] on the Sketch domain. These results indicate that DSDRNet maintains a relatively high level of generalization when encountering datasets with a large number of samples and complex variations.

The model involves seven different loss functions. To evaluate the performance of these loss functions on the Digits-DG dataset, we experimented with various combinations of losses to assess their impact on the model’s performance. The strategy for selecting these combinations is as follows: First, the classification loss $\mathcal{L}_{\text{ce}}$, which is essential for the task, is used throughout the entire experiment. Next, we evaluate the impact of the reconstruction loss $\mathcal{L}_{\text{recon}}$. Subsequently, we test the effects of the cycle loss $\mathcal{L}_{\text{cycle}}$ and the adversarial loss $\mathcal{L}_{\text{adv}}$ on the model’s performance. We then separately test the intra-instance loss $\mathcal{L}_{\text{intra}}$ and the inter-instance loss $\mathcal{L}_{\text{inter}}$. Finally, we combine $\mathcal{L}_{\text{recon}}$ and $\mathcal{L}_{\text{kl}}$ to assess the model’s performance. The final results are presented in Table V as shown in the document.

When the model only involves the $\mathcal{L}_{\text{ce}}$ loss, its performance is consistent with ERM and achieves an accuracy of 73.7%. By adding the $\mathcal{L}_{\text{recon}}$ constraint, the performance improves by +0.5%. This effect is similar to the impact of adding only $\mathcal{L}_{\text{cycle}}$ on the model’s performance. When combining $\mathcal{L}_{\text{cycle}}$ and $\mathcal{L}_{\text{adv}}$, the performance reaches 75.4%, which is a +1.7% improvement over ERM. When $\mathcal{L}_{\text{intra}}$ is added, the model’s performance significantly improves from 76.2% to 77.8%. The improvement effect of individual $\mathcal{L}_{\text{inter}}$ is similar to that of $\mathcal{L}_{\text{intra}}$, but when the intra-class $\mathcal{L}_{\text{intra}}$ and inter-class $\mathcal{L}_{\text{inter}}$ constraints are combined, the effect becomes more pronounced, directly boosting the performance to 78.8%. By combining $\mathcal{L}_{\text{cycle}}$ and $\mathcal{L}_{\text{adv}}$, the model’s final performance reaches 80.9%. This indicates that the losses in the model are playing roles to varying degrees, with the constraints of $\mathcal{L}_{\text{intra}}$, $\mathcal{L}_{\text{inter}}$, and $\mathcal{L}_{\text{adv}}$ exerting stronger effects on the model.

The main loss of our DSDRNet is shown in Figure 3, the details are as follows. $\mathcal{L}_{\text{adv}}$ started to oscillate decreasingly from 2.42, after 200 iterations around 0.7 and then kept oscillating between 0.5 and 2. The loss oscillated with a relatively large amplitude, it oscillated violently in the first 100 iterations and tended to stabilize at about 5 after 400 iterations. $\mathcal{L}_{\text{adv}}$ oscillated in the early stage due to relatively few training samples which led to the poor quality of the generated images. $\mathcal{L}_{\text{recon}}$ fell from the peak value of 11.12 and basically stabilized after 600 iterations and finally at around 1.5. $\mathcal{L}_{\text{cycle}}$ declined relatively fast at the beginning of training due to the relatively small amount of data and started to stabilize after 200 iterations, finally oscillating around 0.5. $\mathcal{L}_{\text{intra}}$ fell at the beginning of training from the initial 26, after 400 iterations down to 5.5, finally oscillating around 5.0. This is because, in the early oscillation, the model has not yet learned some common features because of the relatively few samples. In contrast, the later oscillation was caused by the random reading of data from the two domains where the difference in data semantics may be relatively large. As the number of training samples increases, $\mathcal{L}_{\text{ce}}$ decreases because the network learns more useful features with the result of a gradual improvement of the recognition rate of the samples.

To further validate the performance of DSDRNet, we utilize t-SNE to visualize the OfficeHome dataset before and after processing with the model. The resulting visualization is shown in Figure 4. In the left image, the distribution of data from the RealWorld domain and the Product domain is mixed together, making it difficult to distinguish between different categories. However, in the right image, after processing with DSDRNet, the data from different domains are aggregated together based on their respective categories.