Augmentation-Aware Expert Design for Domain Generalization

Deep neural networks often suffer from significant performance degradation when encountering domain shifts not present in their training data. While various domain generalization (DG) methods have been proposed to address this issue, they frequently introduce substantial computational overhead during both training and inference. To overcome this limitation, we propose a novel and efficient architecture inspired by the mixture of experts (MoE) mechanism. Our approach redesigns the initial stem layer of the network to include multiple specialized Batch Normalization (BN) layers, each acting as an expert trained on a specific data augmentation. A lightweight gating network adaptively selects the most appropriate BN expert for a given input at test time, enabling augmentation-aware inference with minimal additional parameters and latency. We evaluated our method on five digit classification datasets using a leave-one-domain-out protocol. Experimental results demonstrate that our approach serves as a powerful universal baseline, achieving an average accuracy improvement of 10.5% over Empirical Risk Minimization, a standard baseline in the field.