IBAFormer: Intra-batch Attention Transformer for Domain Generalized Semantic Segmentation

TL;DR

IBAFormer introduces intra-batch attention mechanisms to Transformer models, leveraging inter-sample correlations to significantly improve domain generalization in semantic segmentation tasks.

Contribution

The paper proposes novel intra-batch attention modules, MIBA and EIBA, and integrates them into Transformer networks, advancing DGSS performance beyond existing methods.

Findings

Achieves state-of-the-art results in DGSS

Demonstrates effectiveness of intra-batch attention modules

Shows improved generalization over CNN-based models

Abstract

Domain generalized semantic segmentation (DGSS) is a critical yet challenging task, where the model is trained only on source data without access to any target data. Despite the proposal of numerous DGSS strategies, the generalization capability remains limited in CNN architectures. Though some Transformer-based segmentation models show promising performance, they primarily focus on capturing intra-sample attentive relationships, disregarding inter-sample correlations which can potentially benefit DGSS. To this end, we enhance the attention modules in Transformer networks for improving DGSS by incorporating information from other independent samples in the same batch, enriching contextual information, and diversifying the training data for each attention block. Specifically, we propose two alternative intra-batch attention mechanisms, namely mean-based intra-batch attention (MIBA) and element-wise intra-batch attention (EIBA), to capture correlations between different samples, enhancing feature representation and generalization capabilities. Building upon intra-batch attention, we introduce IBAFormer, which integrates self-attention modules with the proposed intra-batch attention for DGSS. Extensive experiments demonstrate that IBAFormer achieves SOTA performance in DGSS, and ablation studies further confirm the effectiveness of each introduced component.