Distillation-based Single Source Domain Generalization via Mutual Information Regularization

Abstract

Machine learning frequently suffers from the discrepancy in data distribution, commonly known as domain shift. Single-source Domain Generalization (sDG) is a task designed to simulate domain shift artificially, in order to train a model that can generalize well to multiple unseen target domains from a single source domain. A popular approach is to learn robustness via the alignment of samples generated by data augmentation. However, prior works frequently overlooked what can be learned through such alignment. In this paper, we study the effectiveness of augmentation-based sDG methods by connecting recent identifiability results by Von Kugelgen et al. (2021). We highlight the overlooked issues in using augmentation for OOD generalization and search ways to alleviate them. We introduce a novel sDG method that leverages pretrained models to guide the learning process via a feature-level regularization of mutual information, which we name PROF (Progressive mutual information Regularization for Online distillation of Frozen oracles). PROF can be added to conventional augmentation-based methods to dampen the fluctuation of the OOD performance. We further introduce a data-effective alignment objective as well as a novel augmentation method for fine-grained simulation of domain shift.

일반적으로, 머신러닝 모델은 데이터 분포 상의 변화에 취약한 경향을 보인다. 단일 소스 도메인 일반화(sDG)는 이러한 문제를 해결하기 위해 고안된 연구 태스크로, 인공적으로 설정된 분포 변화에도 강건한 모델을 만드는 것을 목표로 한다. 기존의 sDG 연구는 다양한 데이터 증강 기법을 통해 모델의 일반화 성능을 향상시키는 데 집중하였으나, 이와 같은 증강 기반 접근법의 유효성은 깊이 논의되지 않았다. 본 논문은 최근 Von Kugelgen et al. (2021)의 연구결과를 이용하여 기존에 간과된 증강 기반 접근법의 문제들을 인과적 관점에서 조명하고, 그에 대한 해결책을 탐구한다. 본 연구진은 증강 기반 sDG 방식의 불안정성을 해소하기 위한 "PROF: 상호 정보 정규화를 통한 지식 증류 기반 단일 소스 도메인 일반화" 기법을 제시한다. PROF는 선학습된 모델의 지식을 이용한 증류 기반의 정규화 기법을 통해 모델의 훈련 과정을 지도한다. PROF는 증강 기반 sDG 방식에 추가되어, 모델의 일반화 성능이 안정적으로 증가할 수 있도록 한다. 나아가, 본 논문은 기존 방식에 비해 경제적인 정렬 함수와 개선된 데이터 증강 방식을 제안하였다.