Domain Generalization with Correlated Style Uncertainty

	$$\displaystyle\Sigma_{\mu}=\frac{1}{B}(\mu-E(\mu))^{T}(\mu-E(\mu)),$$		(3)
	$$\displaystyle\Sigma_{\sigma}=\frac{1}{B}(\sigma-E(\sigma))^{T}(\sigma-E(\sigma)),$$		(4)

where $E(\mu),E(\sigma)$ represents the mean value of $\mu,\sigma$ over batch dimension. It is worth noting that the rank of $\Sigma_{\mu},\Sigma_{\sigma}$ is strictly limited by $min(B,C)\leqslant C$. Previous research indicated that the feature maps are unlikely to be linearly independent over the channel dimension [41, 8]. Without any form of regularization, it becomes difficult to assume a diagonal covariance matrix and a zero correlation across each channel, as suggested in the top row of Figure 2.

We observed that the correlation matrix of style statistics (regardless of the mean or the standard deviation values) is not diagonal. Indeed, there exists a strong correlation among the channels. We applied eigenvalue decomposition over the calculated correlation matrix and find that very few eigenvectors dominate most variance, as shown in the bottom row of the Figure 2. This inspired us to rethink the augmentation of style statistics. The correlation matrix indicates that the combinations of feature statistics are not arbitrary but limited by task objectives and training procedures. Most variances happen within the specific principal directions. Arbitrary augmentation over the style statistics might damage the training itself. Previous research about why InstanceNorm can not outperform the BatchNorm in the discriminative tasks also proves this finding [20].

In reevaluating feature statistics augmentation, we observe that both MixStyle [50] and pAdaIN [22] limit their augmentations within the original feature space, preserving either channel information or channel combinations respectively. DSU [18] attempts to exceed this limitation via uncertainty quantification, but hinges on strict orthogonality between channel feature statistics. Our proposed model, CSU, breaks these boundaries by generating feature statistics beyond the training domain while preserving channel correlations, offering a more robust feature augmentation. CSU builds upon the strengths of these strategies and addresses their limitations, significantly advancing the field of feature statistics augmentation. There are some other style augmentation methods like AdverStyle [42] or Style Neophile [11]. However, direct mathematical correlation relationship analysis for these methods is not feasible due to complex adversarial training or distribution selection.

Given that the correlation matrix is real, symmetric, and positive semi-definite, then we can apply eigenvalue decomposition on $\Sigma_{\mu},\Sigma_{\sigma}$ to analyze its subspaces as:

	$$\displaystyle\Sigma_{\mu}=Q_{\mu}diag(\Lambda_{\mu})Q_{\mu}^{T},$$		(5)
	$$\displaystyle\Sigma_{\sigma}=Q_{\sigma}diag(\Lambda_{\sigma})Q_{\sigma}^{T},$$		(6)
	$$\displaystyle Q_{\mu}Q_{\mu}^{T}=Q_{\sigma}Q_{\sigma}^{T}=I,$$		(7)
	$$\displaystyle\Lambda_{\mu},\Lambda_{\sigma}\in\mathbb{R}^{C},Q_{\mu},Q_{\sigma}\in\mathbb{R}^{C\times C},$$		(8)

where $\Lambda_{\mu,i}\geqslant\Lambda_{\mu,j}\geqslant 0$ , $\Lambda_{\sigma,i}\geqslant\Lambda_{\sigma,j}\geqslant 0$ ($i>j$) represent the sorted eigenvalues, and $Q_{\mu},Q_{\sigma}$ are the corresponding eigenvectors. The eigenvector corresponding to the largest eigenvalue represents the direction that we could apply dense augmentation. Eigenvectors corresponding to the eigenvalues of 0 or close to 0 are not considered in the data augmentation process due to low variance across such directions within the dataset.

Assuming that the $\mu,\sigma$ still follows the multi-variable Gaussian distribution, and $k_{\mu},k_{\sigma}$ represent the independent variable number (or the rank of the corresponding covariance matrix), then we can represent the probability distribution function as:

	$$\displaystyle f_{\mu}=\frac{1}{(2\pi)^{k_{\mu}}\det^{*}(\Sigma_{\mu})}\exp^{-(\mu-E(\mu))^{T}\Sigma_{\mu}^{+}(\mu-E(\mu))},$$		(9)
	$$\displaystyle f_{\sigma}=\frac{1}{(2\pi)^{k_{\sigma}}\det^{*}(\Sigma_{\sigma})}\exp^{-(\sigma-E(\sigma))^{T}\Sigma_{\sigma}^{+}(\sigma-E(\sigma))},$$		(10)

where the $det^{*}$ is the pseudo-determinant and $\Sigma^{+}$ is the generalized inverse. Based on this distribution function, we further derive the correlated uncertainty augmentation after we sample $\epsilon_{\mu},\epsilon_{\sigma}\in\mathbb{R}^{N\times C}$ from the standard Gaussian distribution $Y\sim\mathcal{N}(0,I)$ as follow:

	$$\displaystyle P_{\mu}=Q_{\mu}diag(\Lambda_{\mu})^{\frac{1}{2}}Q_{\mu}^{T},$$		(11)
	$$\displaystyle P_{\sigma}=Q_{\sigma}diag(\Lambda_{\sigma})^{\frac{1}{2}}Q_{\sigma}^{T},$$		(12)
	$$\displaystyle\hat{\epsilon}_{\mu}=\epsilon_{\mu}P_{\mu},\quad\hat{\epsilon}_{\mu}=\epsilon_{\sigma}P_{\sigma}.$$		(13)

Essentially, we determine the transform matrix $P$ such that the covariance matrix $\Sigma=PP^{T}$ as shown in Eq.11-12, and these transformation matrices allow us to maintain the correlation. We sample the correlated perturbations $\hat{\epsilon}\in\mathbb{R}^{N\times C}$ from independent Gaussian noise $\epsilon\in\mathbb{R}^{N\times C}$ as shown in Eq.13. Note the covariance matrix:

	$$\displaystyle\Sigma=PP^{T}=Qdiag(\Lambda_{\mu})^{\frac{1}{2}}Q^{T}(Qdiag(\Lambda_{\mu})^{\frac{1}{2}}Q^{T})^{T}$$
	$$\displaystyle\quad=Qdiag(\Lambda_{\mu})^{\frac{1}{2}}(Qdiag(\Lambda_{\mu})^{\frac{1}{2}})^{T}.$$

These two transformations $Qdiag(\Lambda_{\mu})^{\frac{1}{2}}Q^{T}$, $Qdiag(\Lambda_{\mu})^{\frac{1}{2}}$ both work in principle. However, in practice, if we use the standard format of eigen-decomposition, stochastic axis swapping in the eigenvector will not influence this decomposition, but it can lead to unstable training. Thus, we apply the decomposition trick in the format of $Qdiag(\Lambda_{\mu})^{\frac{1}{2}}Q^{T}$ rather than using only one set of eigenvectors $Qdiag(\Lambda_{\mu})^{\frac{1}{2}}$ to avoid the random flip issue in traditional eigenvalue decomposition  [8]. While we derive $\epsilon_{\mu},\epsilon_{\sigma}$ from C independent normal variables, the resulting $\hat{\epsilon}{\mu},\hat{\epsilon}{\mu}$ contain only $k_{\mu},k_{\sigma}$ independent components, corresponding to the non-zero components of eigenvalues. It is also worth noticing that the computation of eigenvalue decomposition is conducted by the torch operator with a complexity of $\sim O(N^{3})$ in practice [30]. Given the computation is conducted in polynomial time complexity and $N<512$ for most cases, the the extra time required is negligible.

Based on the previous two sections, we now present the style augmentation with correlated style uncertainty as follows:

	$$\displaystyle\beta(x)=\mu(x)+\lambda*\hat{\epsilon}_{\mu},$$		(14)
	$$\displaystyle\gamma(x)=\sigma(x)+\lambda*\hat{\epsilon}_{\sigma},$$		(15)

where $\lambda\sim Beta(\alpha,\alpha)$ represents the augmentation intensity generated from the Beta distribution. Hyperparameter $\alpha$ controls the shape of the distribution. In the ablation experiments, we further show the influence of hyperparameter selections on the final performance. We can understand the equation in one more intuitive way, the first part is to provide in-domain samples to cover the whole training domain, and the second is to provide the extrapolation while maintaining the same data distribution.

$\displaystyle\beta(x)=\underbrace{\mu(x)}_{\text{In Domain Sample}}+\underbrace{\lambda*\hat{\epsilon}_{\mu}}_{\text{Out Domain extrapolation}}.$

The final augmented instance feature can be defined as:

$$CSU(x)=\gamma(x)(\frac{x-\mu(x)}{\sigma(x)})+\beta(x).$$

(16)

This plug-and-play module can be easily inserted into any current framework. We provide pseudo-code (PyTroch) in the supplementary materials.

We validate our model’s performance on various multidomain classification tasks, including PACS, Office-Home, and Camelyon17. Figure 4 shows some examples with observable domain shifts within the same class. In all experiments, the domain tags are agnostic. Following the MixStyle, we adopt the ResNet-18 [7] with ImageNet [5] pre-training as the backbone for classification. We follow the Leave-One-Domain-Out strategy, which leaves one domain out for evaluation and the rest of the domains participating in the training. Adhering to a fair evaluation framework, we implement the multi-domain classification DASSL setup widely embraced in the works of Zhou et al. [48] for comparison. The batch size is set as 64. We conduct all the experiments on 2 NVIDIA A6000 GPU based on PyTorch [23] framework.

The undertaking of person re-identification, which involves the cross-camera recognition of individuals, poses a substantial domain generalization challenge. Considering each distinct camera output as its own unique domain amplifies the intricacy inherent to the person re-identification task. Following previous research, we conducted this experiment on the commonly used Duke [25] and Market1501 [45] datasets. To evaluate the model’s generalizability, we take one dataset as training and test the performance on the other domain. The camera data from the test domain does not participate in any training process. We adopt the exact framework implementation of MixStyle and test the CSU influence on the ResNet50 [7] and OSNet [49]. Similarly, ranking accuracy and mean average precision (mAP) are performance measures. For a fair comparison, we repeat the pAdaIN and DSU experiments on the same framework with the MixStyle and use the best configuration reported in the original paper.

Table 4 shows the experiment results using two models in the two domains. We could observe that CSU outperforms other methods by a large margin. CSU achieves around $26.9\%,19.6\%$ advancement in mAP using the ResNet-50 model in the Market1501 to Duke and the Duke to Market1501 experiment, correspondingly. Similarly, CSU achieves around $20.1\%,24.2\%$ improvement for the OSNet model experiment. We could also observe similar advancements in ranking accuracy, and CSU achieves impressive improvement over other methods. Nevertheless, to show the effectiveness of CSU rather than position fine-tuning, we insert the permutation in all positions as described in Sec 5. The supplementary materials show that changing the inserting position can achieve even more significant performance advancement.

Insert Position Selection: To answer the question of where we should insert the CSU, we conduct comprehensive experiments on the PACS and Office-Home datasets using the ResNet18 structure. We investigate all possible positions of ResNet18, including the first Convolution, first Max Pooling, and 1, 2, 3, and 4 Res-block, which are named 0-5, respectively. We divide the experiment into several groups according to the inserted CSU number. Within each group, we shift the start position one by one from 0 to end. For example, for the group containing 2 CSU blocks, we will have 01, 12, 23, 34, and 45 potential combinations and five comparison experiments in total. To avoid the influence of hyperparameters, we set $\alpha=0.3$ for all experiments. Thus, we can reasonably and adequately compare the inserting position’s influence on final performance.

Figure 5 shows the ablation experiment results. Inserting 6 blocks of CSU in all potential positions achieves the best results for the PACS dataset while inserting 4 blocks of CSU in the last 4 positions achieves the best results for the Office-Home dataset. Within each group of a fixed number of CSU blocks, the performance tends to increase when we start the inserting position at the medium blocks. This trend is different with the MixStyle which the model prefers the first several blocks [50]. We explain this phenomenon as CSU can provide more reasonable feature (statistics) augmentation due to correlation preservation. This preservation will avoid information loss in the medium or last blocks. We can also notice that compared with inserting 4, 5, and 6 blocks of CSU, inserting 2, or 3 blocks of CSU can not achieve comparable performance. This indicates that a more significant number of CSU blocks can be helpful to increase the model’s generalization ability due to accumulated correlations over the blocks. It is also worth noting that no matter how we choose the inserting position, the performance of the CSU model always shows superiority over the baseline by a large margin. This firmly proves the effectiveness of the proposed model.

Hyper-parameter Selection: As described in the previous section, the hyper-parameter alpha determines the intensity of augmentation during training by manipulating the shape of the Beta distribution. Here, we show the influence of alpha on the PACS, Office-Home using the ResNet18 structure. Similarly, to avoid the influence of different inserting positions, we insert CSU block in all positions for every experiment. We select $\alpha$ from 0.1, 0.2, 0.3, 0.4, 0.5, 0.7, and 0.9 for one comprehensive experiment. As shown in Figure 6, we can find that a smaller number of $\alpha<0.5$ always performs better than the relatively larger number ($>0.5$). Based on these experiment results, we recommend selecting the alpha from 0.1, 0.2, 0.3, and 0.4, and the best configuration may vary according to the tasks.

Effect of Batch Size: The batch size could have a notable impact on the accuracy of correlation data, making it critical to scrutinize its effect on the final model’s generalizability. Accordingly, we’ve introduced the CSU at every stage, as per the earlier section, and set a fixed value of $\alpha=0.3$ in all the experiments. We compare the model’s performance with a batch size of 16, 32, 64, 128, 256, and 512. Figure 7 shows the experimental result. We found that it might be hard to estimate an accurate correlation when the batch size is too small. Similarly, when the batch size is too large, the network tends to converge to sharp minimizers of the training and testing functions, leading to poorer generalizations, as shown in previous research [12].