기술 요약: Mixture-of-Experts를 이용한 부분 집합 공유 불변성 학습 기반 도메인 일반화

1. 문제 정의

도메인 일반화(Domain Generalization, DG)는 학습 시 타겟 데이터에 접근하지 않고, 하나 이상의 소스 도메인에서 학습된 모델이 보지 못한 타겟 도메인으로 일반화하는 것을 목표로 한다. DG의 지배적인 접근 방식은 모든 소스 도메인에 대해 조건부 분포 $P(Z|Y)$가 일관되게 유지되도록 강제함으로써 **도메인 불변 표현(domain-invariant representations)**을 학습하는 데 의존한다. 특징 정렬(feature alignment), 불변 위험 최소화(Invariant Risk Minimization, IRM), 그래디언트 매칭(gradient matching)과 같은 방법들은 단일하고 전역적으로 공유되는 예측 구조가 존재한다는 가정하에 작동한다.

그러나 저자들은 이러한 가정이 종종 너무 제한적이라고 주장한다. 실제 시나리오에서 도메인들은 예측 구조를 부분적으로만 공유할 수 있다. 즉, 특정 특징은 일부 도메인에서는 안정적이지만 다른 도m인에서는 변할 수 있다. 이질적인 도메인의 수가 증가함에 따라 전역적 불변성을 강제하는 것은 가능한 표현 공간을 점진적으로 제한한다. 이론적 분석과 경험적 관찰(그림 1)은 소스 도메인의 수가 증가함에 따라 학습된 표현과 레이블 사이의 상호 정보량($I(Z; Y)$)이 단조 감소함을 보여준다. 결과적으로, 특정 도메인 집합 내에서만 안정적인 가치 있는 예측 요인들이 버려지게 되며, 이는 도메인 이질성이 높을 때 일반화 성능 저하를 초낙한다.

2. 방법론: MESSI

전역적 불변성의 한계를 해결하기 위해, 본 논문은 MESSI(Mixture-of-Experts for Subset-Shared Invariances)를 제안한다. 핵심 가설은 불변 구조가 전역적으로 공유되는 것이 아니라 부분 집합 공유(subset-shared) 형태로 모델링되어야 한다는 것이다. 이 방법은 도메인 관계의 구조적 분해를 학습하기 위해 Mixture-of-Experts (MoE) 아키텍처를 채택한다.

2.1 아키텍처

MESSI는 공유 백본 인코더(예: DeiT)와 그 뒤를 잇는 라우터 및 다수의 전문가 네트워크를 활용한다:

1. 공유 인코더(Shared Encoder): 특징 표현 $u = b_\phi(x)$를 생성한다.
2. 라우터(Router): 전문가들에 대한 소프트 분포 $\pi(x) = \text{softmax}(g_\theta(u))$를 계산하며, 이는 잠재적인 부분 집합 변수 $S$를 근사한다.
3. 전문가(Experts): 각 전문가 $m$은 특정 표현 $z^{(m)} = h_m(u)$를 생성한다.
4. 집계(Aggregation): 최종 표현은 가중 합 $z(x) = \sum \pi_m(x) z^{(m)}$이며, 이를 통해 부분 집합 불변 성분들의 적응적 구성을 허용한다.

2.2 학습 목적 함수

모델은 합성 손실 함수(식 5)를 최적화하도록 학습된다:
$$\mathcal{L} = \mathcal{L}_{cls} + \lambda_{ssi}\mathcal{L}_{OT}^{ssi} + \lambda_{sp}\mathcal{L}_{sp} + \lambda_{bal}\mathcal{L}_{bal} + \lambda_{div}\mathcal{L}_{div}$$

• 라우팅 조건부 부분 집합 정렬 ($\mathcal{L}_{OT}^{ssi}$): 모든 도메인 쌍을 정렬하는 대신, 불변성을 선택적으로 강제한다. 각 전문가 $m$에 대해, 해당 전문가가 특정 클래스 $c$에 대해 높은 라우팅 책임을 갖는 도메인 쌍 $(i, j)$ 사이에서만 클래스 조건부 특징 분포를 정렬한다. 정렬 가중치 $a_{ijc}^{(m)}$는 라우팅 확률로부터 유도된다. 불일치는 라우팅 질량에 의해 가중치가 부여된 최적 운송(Optimal Transport, OT) (또는 MMD 변형)을 사용하여 측정된다. 이는 모델이 어디에서 불변성이 유지되어야 하는지를 학습하게 하는 "라우팅 유도 불변성(routing-induced invariance)"을 생성한다.
• 전문가 특화 ($\mathcal{L}_{div}$): 전문가들이 중복된 표현으로 붕괴되는 것을 방방하기 위해, 서로 다른 전문가 간의 특징 방향 사이의 코사인 유사도를 최소화하는 다양성 정규화 항을 사용한다.
• 신뢰성 있고 균형 잡힌 라우팅 ($\mathcal{L}_{sp}, \mathcal{L}_{bal}$):
  • $\mathcal{L}_{sp}$는 인스턴스 수준에서 희소 라우팅(낮은 엔트로피)을 장려하여, 각 샘플이 특정 전문가 집합만을 활성화하도록 보장한다.
  • $\mathcal{L}_{bal}$은 데이터셋 수준에서 균형 잡힌 전문가 활용을 보장하여, 소수의 전문가가 트래픽을 독점하는 것을 방지한다.

3. 주요 기여

본 논문은 세 가지 주요 기여를 명시한다:

1. 근본적 한계 식별: 저자들은 불변성을 강제하는 것이 예측 정보를 감소시킨다는 점을 입증하였다. 이들은 **부분 집합 공유 불변성(subset-shared invariance)**을 도입하여, DG를 단일한 전역 불변 표현을 학습하는 문제가 아닌 구조화되고 부분 집합에 의존적인 불변성을 학습하는 문제로 재정의하였다.
2. 라우팅 기반 MoE 프레임워크: 도메인 관계의 데이터 기반 분해를 유도하기 위해 라우팅을 사용하는 방법을 제안한다. 전문가가 담당하는 도메인 부분 집합과만 정렬함으로써, 명시적인 부분 집합 레이블 없이도 불변 성분을 적응적으로 구성할 수 있게 한다.
3. 경험적 검증: 이 접근 방식은 표준 벤치마크, 특히 이질적인 도메인 변화가 있고 소스 도메인의 수가 증가하는 상황에서 개선된 도메인 외 일반화 성능을 보여준다.

4. 실험 결과

실험은 표준 DG 벤치마크(DomainBed: PACS, OfficeHome, TerraIncognita, DomainNet)와 도메인 증가를 테스트하기 위해 설계된 통제된 Rotated-Colored MNIST 데이터셋에서 수행되었다.

• 표준 벤치마크: MESSI(특히 OT를 사용한 DeiT-S/16 변형)는 DomainBed에서 평균 정확도 66.2%를 달성하며, ERM, CORAL, Fishr, SAGM 및 GMoE, OMoE와 같은 강력한 MoE 기반 방법들을 포함한 기존 베이스라인들을 능가했다.
• 도메인 증가에 대한 강건성: 고정된 예산의 도메인 증가 프로토콜(Rotated-Colored MNIST)에서 MESSI는 우수한 강건성을 보였다. 전역 정렬 방법(CORal 등)은 소스 도메인이 증가함에 따라 (과잉 제약으로 인해) 성능 저하를 겪었으나, MESSI는 높은 정확도를 유지하며 최소한의 하락(최고 성능의 CORAL 변형이 2.58% 하락할 때 MESSI는 1.86% 하락)을 기록했다.
• 절제 연구(Ablation Studies):
  • 부분 집합 조건부 정렬($\mathcal{L}_{ssi}$)을 제거했을 때 가장 큰 성능 저하가 발생하여, 이것이 핵심 구성 요소임을 확인했다.
  • 다양성 또는 라우팅 정규화 항을 제거하면 전문가의 중복성과 확산된 라우팅이 발생하여 성능이 감소했다.
  • 대조군(Global-MoE, Random-Subset-MoE)과의 비교를 통해, MESSI가 단순히 전역 목적 함수를 희소화하는 것이 아니라 독특한 정렬 구조를 학습함을 보여주었다.
• 효율성: OT 기반 변형은 학습 시간을 약 12.39배 증가시키지만, MMD 변형을 사용하면 이를 약 4.24배로 줄일 수 있다. 추론 비용은 다른 MoE 방법들과 유사한 수준을 유지한다.

5. 의의 및 주장

본 논문은 도메인 일반화가 단일한 전역 불변성을 강제하는 패러다임을 넘어서야 한다고 주장한다. 대신, 예측 구조는 본질적으로 **부분 집합 의존적(subset-dependent)**이라고 강조한다.

• 이론적 통찰: 이 연구는 상호 정보량 분석을 통해 왜 전역 불변성이 도메인 이질성이 증가할 때 실패하는지에 대한 이론적 근거를 제공한다: 즉, 전역 불변성은 보편적으로 공유되지 않는 예측 요인들을 버리도록 강제하기 때문이다.
• 실무적 시사점: 부분 집합을 통해 불변 구조를 모델링함으로써, MESSI는 더 큰 강건성을 달ace한다. 라우팅 메커니즘을 통해 명시적인 부분 집합 레이블이나 사전 정의된 도메인 그룹화 없이도 데이터로부터 이러한 부분 집합을 직접 발견할 수 있다.
• 한계점: 저자들은 이 방법이 라우터의 전문화 능력에 의존한다는 점을 인정한다. 즉, 소스 도메인이 적거나 클래스 분포가 불균형할 경우 라우팅 신호에 노이의가 생길 수 있다. 또한, 특히 최적 운송(Optimal Transport)을 사용할 경우 쌍별 정렬(pairwise alignment)에 따른 계산 비용이 표준 전역 정렬 방법보다 높다.

결론적으로, MESSI는 부분 집합 수준의 불변성을 보존하고 구성하는 메커니즘을 제공하며, 이는 도메인 일반화의 미래가 균일한 전역 정렬이 아닌 구조화되고 적응적인 분해에 있음을 시사한다.