Gradient-Guided Annealing for Domain Generalization

Overview

• Domain Generalization은 train과 다른 unseen domain에서도 성능을 유지하는 것이 목표
• 기존 학습은 i.i.d 가정 기반으로 gradient descent를 수행하지만 실제 환경에서는 domain shift 존재
• 서로 다른 domain 간 gradient 방향 불일치가 domain-invariant feature 학습을 방해
• GGA는 gradient alignment를 유도하는 annealing 기반 optimization 전략 제안

Key Takeaways

Problem Setting

• 여러 source domain에서 학습하고 unseen target domain에서 평가하는 DG 문제
• domain마다 gradient 방향이 달라지는 gradient disagreement 발생
• gradient 충돌로 인해 domain-specific feature에 과적합되고 일반화 성능 저하
• 기존 방법은 gradient를 수정하거나 일부 제거하지만 근본적인 해결은 어려움

Main Idea

• gradient alignment를 직접 유도하는 annealing 기반 optimization 제안
• Gradient Disagreement
  • domain별 gradient 간 cosine similarity가 음수이면 충돌로 정의
  • domain-invariant 모델에서는 gradient expectation이 동일해야 함
• Gradient-Guided Annealing (GGA)
  • 학습 초기 단계에서 파라미터에 작은 random perturbation 적용
  • 여러 domain에서 gradient alignment가 증가하고 loss도 감소하는 방향 탐색
  • 가장 높은 gradient alignment와 낮은 loss를 동시에 만족하는 파라미터 선택
• Optimization 과정
  • warm-up 이후 gradient similarity 계산
  • parameter neighborhood에서 조건을 만족하는 지점 탐색
  • 일정 반복 후 최적 파라미터 선택
  • 이후에는 일반 optimizer로 학습 진행

Result

• ERM 대비 모든 benchmark에서 일관된 성능 향상
• 기존 DG 방법과 결합 시 대부분 추가적인 성능 개선 발생
• gradient 기반 방법들과 비교해도 경쟁력 있는 성능 달성
• 학습 중 gradient cosine similarity가 증가하며 alignment가 실제로 개선됨
• 초기 학습 단계에서 적용하는 것이 가장 효과적이며 hyperparameter에 민감함