A Simple Recipe for Language-guided Domain Generalized Segmentation

Adaptive Instance Normalization (AdaIN).  For a feature map $\mathbf{f}\in\mathbb{R}^{h\times w\times c}$, AdaIN [22] shows that the channel-wise mean $\boldsymbol{\mu}\in\mathbb{R}^{c}$ and standard deviation $\boldsymbol{\sigma}\in\mathbb{R}^{c}$ capture information about the style of the input image, allowing style transfer between images. Hence, stylizing a source feature ${\mathbf{f}}_{\text{s}}$ with an arbitrary target style $(\mu({\mathbf{f}}_{\text{t}}),\sigma({\mathbf{f}}_{\text{t}}))$ reads:

with $\mu(\cdot)$ and $\sigma(\cdot)$ the mean and standard deviation of input feature; multiplications and additions being element-wise.

Prompt-driven Instance Normalization (PIN).  PIN was introduced for prompt-driven zero-shot domain adaptation in PØDA [11]. It replaces the target style $(\mu({\mathbf{f}}_{\text{t}}),\sigma({\mathbf{f}}_{\text{t}}))$ in AdaIN (1) with two optimizable variables $(\boldsymbol{\mu},\boldsymbol{\sigma})$ guided by a single prompt in natural language. The rationale is to leverage a frozen CLIP [43] to mine visual styles from the prompt representation in the shared space. Given a prompt $P$ and a feature map ${\mathbf{f}}_{\text{s}}$, PIN reads as:

where $\boldsymbol{\mu}$ and $\boldsymbol{\sigma}$ are optimized using gradient descent, such that the cosine distance between the visual feature representation and the prompt representation is minimized.

Style randomization.  During training, randomly cropped images ${\mathcal{I}}_{\text{s}}$ are encoded into ${\mathbf{f}}_{\text{s}}$ using CLIP-I1. Each batch of feature maps ${\mathbf{f}}_{\text{s}}$ is viewed as a grid of $m$ patches, without cropping them. For each patch ${\mathbf{f}}_{\text{s}}{\!}^{(ij)}$ within the grid, the dominant class $c_{p}{\!}^{(ij)}$ is queried using the corresponding ground truth patch ${\mathbf{y}}_{\text{s}}{\!}^{(ij)}$, and a style is randomly sampled from the corresponding mined bank $\mathcal{T}{}^{(c_{p}{\!}^{(ij)})}$. We then apply patch-wise convex combination (i.e., style mixing) of the original style of the patch and the mined style. Specifically, for an arbitrary patch ${\mathbf{f}}_{\text{s}}{\!}^{(ij)}$, our local style mixing reads:

with $(\boldsymbol{\mu}^{(ij)},\boldsymbol{\sigma}^{(ij)})\in\mathcal{T}^{(c_{p}^{(ij)})}$ and $\alpha\in[0,1]^{c}$.

Minimal fine-tuning.  During training, we fine-tune only the last few layers of the backbone. Subsequently, we examine various alternatives and show that the minimal extent of fine-tuning is the crucial factor in witnessing the effectiveness of our local style mixing strategy.

Synthetic datasets.  GTAV [45] and SYNTHIA [46] are used as synthetic datasets. GTAV consists of 24 966 images split into 12 403 images for training, 6 382 for validation and 6 181 for testing. SYNTHIA consists of 9 400 images: 6 580 for training and 2 820 for validation. GTAV and SYNTHIA are denoted by G and S, respectively.

Real datasets.  Cityscapes [8], BDD-100K [58], and Mapillary [39] contain 2 975, 7 000, and 18 000 images for training and 500, 1 000, and 2 000 images for validation, respectively. ACDC [48] is a dataset of driving scenes in adverse conditions: night, snow, rain and fog with respectively 106, 100, 100 and 100 images in the validation sets. C, B, and M denote Cityscapes, BDD-100K and Mapillary, respectively; AN, AS, AR and AF denote night, snow, rain and fog subsets of ACDC, respectively.

Implementation details.  Following previous works [7, 41, 25, 62, 55, 57, 21, 31, 26], we adopt DeepLabv3+ [5] as segmentation model. ResNet-50 and ResNet-101 [19], initialized with CLIP pretrained weights, are used in our experiments as backbones. Specifically, we remove the attention pooling layer and add a randomly initialized decoder head. The output stride is $16$. Single-source and multi-source models are trained respectively for $40K$ and $60K$ iterations with a batch size of $8$. The training images are cropped to $768\times 768$. Stochastic Gradient Descent (SGD) with a momentum of $0.9$ and weight decay of $10^{-4}$ is used as optimizer. Polynomial decay with a power of $0.9$ is used, with an initial learning rate of $10^{-1}$ for the classifier and $10^{-2}$ for the backbone. We use color jittering and horizontal flip as data augmentation. Label smoothing regularization [50] is adopted. For style mining, Layer1 features are divided into $9$ patches. Each patch is resized to $56\times 56$, corresponding to the dimensions of Layer1 features for an input image of size $224\times 224$ (i.e. the input dimension of CLIP). We use ImageNet templates¹¹1https://github.com/openai/CLIP/ for each prompt.

Evaluation metric.  We evaluate our models on the validation sets of the unseen target domains with mean Intersection over Union (mIoU%) of the $19$ shared semantic classes. For each experiment, we report the average of three runs.

Training on GTAV (G) as source, Tab. 2 reports models trained with either ResNet-50 or ResNet-101 backbones. The unseen target datasets are C, B, M, S, and the four subsets of A. Tab. 2 shows that our method significantly outperforms all the baselines on all the datasets for both backbones. We note that WildNet [31] and TLDR [26] use extra-data, while SHADE [62] uses the full G dataset (24,966 images) for training with ResNet-101. Class-wise performances are reported in Appendix B.

Training on Cityscapes (C) as source, Tab. 3 reports performance with ResNet-50 backbone. The unseen target datasets are B, M, G, and S. The table shows that our method outperforms the baseline in average, and is competitive to SOTA on G and M.

Multi-source DGSS.  We also show the effectiveness of FAMix in the multi-source setting, training on G+S and evaluating on C, B and M. The results reported in Tab. 4 for ResNet-50 backbone outperform state-of-the-art.

Qualitative results.  We visually compare the segmentation results with Pin the memory [25], SHADE [62] and WildNet [31] in Fig. 3. FAMix clearly outperforms other DGSS methods on “stuff” (e.g., road and sky) and “things” (e.g., bicycle and bus) classes.

Removing ingredients from the recipe.  FAMix is based on minimal fine-tuning of the backbone (i.e., Freeze), style augmentation and mixing. We show in Tab. 6 that the best generalization results are only obtained when combining the three ingredients. Specifically, when the backbone is fine-tuned (i.e., Freeze ✗), the performances are largely harmed. When minimal fine-tuning is performed (i.e., Freeze ✓), we argue that the augmentations are too strong to be applied without style mixing; the latter brings both effects of domain interpolation and use of the original statistics. Subsequently, when style mixing is not applied (i.e. Freeze ✓, Augment ✓, Mix ✗), the use of mined styles brings mostly no improvement on OOD segmentation compared to training without augmentation (i.e. Freeze ✓, Augment ✗, Mix ✗). Note that for Freeze ✓, Augment ✓, Mix ✗, the line 8 in Algorithm 2 becomes:

Our style mixing is different from MixStyle [67] for being applied: (1) patch-wise and (2) between original styles of the source data and augmented versions of them. Note that the case (Freeze ✓, Augment ✗, Mix ✓) could be seen as a variant of MixStyle, yet applied locally and class-wise. Our complete recipe is proved to be significantly more effective with a boost of $\approx+6$ mean mIoU w.r.t. the baseline of training without augmentation and mixing.

Prompt construction.   Tab. 7 reports results when ablating the prompt construction. In FAMix, the final prompt is derived by concatenating $<$random​ style​ prompt$>$ and $<$class​ name$>$; removing either of those leads to inferior results. Interestingly, replacing the style prompt by random characters – e.g. “ioscjspa” – does not significantly degrade the performance. In certain aspects, using random prompts still induces a randomization effect within the FAMix framework. However, meaningful prompts still consistently lead to the best results.

Number of style prompts.  FAMix uses a set $\mathcal{R}$ of random style prompts which are concatenated with the class names; $\mathcal{R}$ is formed by querying ChatGPT²²2https://chat.openai.com/ using $<$give​ me​ 20​ prompts​ of​ 2​ to​ 5​ words​ describing​ random​ image​ styles$>$. The output prompts are provided in Appendix C. Fig. 4(a) shows that the size of $\mathcal{R}$ has a marginal impact on FAMix performance. Yet, the mIoU scores on C, B, M and AR are higher for $|\mathcal{R}|=20$ compared to $|\mathcal{R}|=1$ and almost equal for the other datasets.

Local vs. global style mining.  To highlight the effect of our class-wise local style mining, we perform an ablation replacing it with global style mining. Specifically, the same set of $<$random​ style​ prompt$>$ are used, though being concatenated with $<$driving$>$ as a global description instead of local class name. Intuitively, local style mining and mixing induces richer style variations and more contrast among patches. The results in Tab. 8 show the effectiveness of our local style mining and mixing strategy, bringing about $3$ mIoU improvement on G $\rightarrow$ C.

What to mix?  Let $\mathcal{S}=\bigcup_{k=1}^{K}\mathcal{S}^{(k)}$ and $\mathcal{T}=\bigcup_{k=1}^{K}\mathcal{T}^{(k)}$ the sets of class-wise source and augmented features, respectively. In FAMix training, for an arbitrary patch ${\mathbf{f}}_{\text{s}}^{(ij)}$, style mixing is performed between the original source statistics and statistics sampled from the augmented set (i.e., $(\boldsymbol{\mu}^{(ij)},\boldsymbol{\sigma}^{(ij)})\in\mathcal{T}^{(c_{p}^{(ij)})}$, see (3) and (4)). In class-wise vanilla MixStyle, $(\boldsymbol{\mu}^{(ij)},\boldsymbol{\sigma}^{(ij)})\in\mathcal{S}^{(c_{p}^{(ij)})}$. In Tab. 9, we show that sampling $(\boldsymbol{\mu}^{(ij)},\boldsymbol{\sigma}^{(ij)})$ from $\mathcal{S}^{(c_{p}^{(ij)})}\cup\mathcal{T}^{(c_{p}^{(ij)})}$ does not lead to better generalization, despite sampling from a set with twice the cardinality. This supports our mixing strategy visualized in Fig. 1. Intuitively, sampling from $\mathcal{S}\cup\mathcal{T}$ could be viewed as applying either MixStyle or our mixing with a probability $p=0.5$.

Minimal fine-tuning.  We argue for minimal fine-tuning as a compromise between pretrained feature preservation and adaptation.  Fig. 4(b) shows an increasing OOD generalization trend with more freezing. Interestingly, only fine-tuning the last layers of the last convolutional block (where the dilation is applied) achieves the best results. When training on Cityscapes (Tab. 3), we observed that freezing all the layers except Layer4 achieves the best results.