PEER pressure: Model-to-Model Regularization for Single Source Domain Generalization

TL;DR

The paper addresses mid-train OOD fluctuations in augmentation-based single-source domain generalization and introduces PEER, a model-to-model regularization framework with a frozen task model and a trainable proxy model. PEER uses a shared projection head and a mutual-information-based loss to align proxy features with the task model, while periodically averaging proxy snapshots to update the task model, effectively accumulating diverse augmentation knowledge. Across PACS, Digits, Office-Home, and VLCS, PEER reduces OOD fluctuation and achieves state-of-the-art mean target-domain accuracy, even with simple RandAugment. This approach offers a practical, adaptable strategy to stabilize training and improve generalization without relying on external pre-trained teachers or complex augmentation schemes.

Abstract

Data augmentation is a popular tool for single source domain generalization, which expands the source domain by generating simulated ones, improving generalization on unseen target domains. In this work, we show that the performance of such augmentation-based methods in the target domains universally fluctuates during training, posing challenges in model selection under realistic scenarios. We argue that the fluctuation stems from the inability of the model to accumulate the knowledge learned from diverse augmentations, exacerbating feature distortion during training. Based on this observation, we propose a novel generalization method, coined Parameter-Space Ensemble with Entropy Regularization (PEER), that uses a proxy model to learn the augmented data on behalf of the main model. The main model is updated by averaging its parameters with the proxy model, progressively accumulating knowledge over the training steps. Maximizing the mutual information between the output representations of the two models guides the learning process of the proxy model, mitigating feature distortion during training. Experimental results demonstrate the effectiveness of PEER in reducing the OOD performance fluctuation and enhancing generalization across various datasets, including PACS, Digits, Office-Home, and VLCS. Notably, our method with simple random augmentation achieves state-of-the-art performance, surpassing prior approaches on sDG that utilize complex data augmentation strategies.

Figures (10)

• Figure 1: Despite its generalization effect, data augmentation induces fluctuations in target domain accuracy during the training. This phenomenon becomes more pronounced as the complexity of the augmentation increases, complicating model selection. We address this issue of fluctuation with a simple model-to-model regularization method that cushions the effect of data augmentation.
• Figure 1: Empirical study of (a) target domain accuracy, (b) mid-train OOD fluctuation, and (c) source-target dataset distance. We use MNIST as a source. Large source-target distance (red) coincided with low target accuracy and high OOD fluctuation during training, and vice versa (blue).
• Figure 2: Illustration of pitfalls of augmentation in generalizing to unseen target domains. (a) Augmentation-based methods expand the source domain by providing diverse augmented samples (i.e., Source+). This enhances the model's generalization capability towards the unseen target domain (i.e., Target A). (b) Throughout the course of training, it iteratively simulates diverse unseen domains. However, at the same time, diverse augmentations lead to the distortion of the learned representations, thereby triggering OOD fluctuation.
• Figure 3: OTDD distance [alvarez2020geometric] between the original data (MNIST) and its augmented view.
• Figure 4: Mode connectivity in the proxy model's trajectory. peer benefits parameter-averaging between snapshots of $P$ through its regularization effects.