Distribution Alignment for One-Shot Federated Learning via Optimal Transport

기술 요약: 최적 운송을 통한 원샷 연합 학습의 분포 정렬

1. 문제 정의

본 논문은 클라이언트가 중앙 서버와 단 한 번만 상호작용하는 통신 체계인 **원샷 연방 학습(One-Shot Federated Learning, OSFL)**의 고유한 과제를 다룹니다. 표준 연합 학습(FL)은 반복적인 클라이언트-서버 최적화를 통해 데이터 이질성을 완화하지만, OSFL는 이러한 정교화 능력이 결여되어 있습니다.

식별된 핵심 문제는 두 가지 특정 유형의 통계적 이질성이 결합되어 발생하는 효과입니다:

1. 도메인 시프트(Domain Shift): 클라이언트는 다양한 획득 조건이나 환경으로 인해 서로 다른 입력 분포 $P_k(x)$를 관찰합니다.
2. 레이블 시프트(Label Shift): 클라이언트는 서로 다른 레이블 마진 $P_k(y)$를 나타내며, 이는 클래스 비율의 차이로 이어집니다.

반복적인 교정 기능이 없는 상태에서, 이러한 시프트는 클라이언트 간의 **특징 표현 불일치(misaligned feature representations)**와 일관되지 않은 클래스 사후 결정 함수를 유발합니다. 증류(distillation), 데이터 합성 또는 앙상블에 의존하는 기존 OSFL 방법들은 정렬된 표현을 가정하거나 이러한 시프트를 별도로 처리하는 경우가 많습니다. 결과적으로, 특징 기하학(feature geometry)과 클래스 조건부 편향(class-conditional biases)의 불일치가 교정되지 않아 집계된 글로벌 모델의 성능을 저하시킵니다.

2. 방법론: SLOT-Align

저자들은 클라이언트의 특징 분포를 다운스트림 OSFL 학습 전에 조화시키기 위해 설계된 기하학 인식 프레임워크인 SLOT-Align(Single-round, Learning-free Optimal Transport Alignment)을 제안합니다. 이 방법은 전처리 단계로 작동하며, 모든 로컬 입력을 공통 잠재 공간으로 매핑하기 위해 (외부 데이터셋으로 사전 학습된) 공유된 동결 인코더(shared frozen encoder) $f_\theta$에 의존합니다.

워크플로우는 네 가지 주요 단계로 구성됩니다:

A. 로컬 특징 통계 추출

각 클라이언트 $k$는 로컬 데이터로부터 추출된 특징 $z = f_\theta(x)$의 압축된 1차 및 2차 통계량을 계산합니다:

• 평균(Mean): $\mu_k = \mathbb{E}[f_\theta(x)]$
• 공분산(Covariance): 양의 정치성(positive definiteness)을 보장하기 위해 Ledoit–Wolf 수축법(shrinkage)을 통해 정규화된 $\Sigma_k$.
  이러한 통계량은 도메인 및 레이블 시프트로 인해 발생하는 위치와 확산의 주요 불일치를 포착하는 가우시안 근사치 $N(\mu_k, \Sigma_k)$ 역할을 합니다.

B. 글로벌 참조 구축

서버는 이러한 로컬 통계량을 집계하여 글로벌 참조 분포를 구축합니다:

• 글로벌 평균: 클라이언트 평균들의 가중 평균.
• 글로벌 공분산: 클라이언트 공분산들의 Bures–Wasserstein barycenter로 계산됩니다. 이는 Bures 거리의 제곱 합의 최소화 값을 찾기 위해 고정점 반복(fixed-point iteration)을 통해 해결되며, 클라이언트별 변동성과 글로벌 일관성 사이의 균형을 효과적으로 맞춥니다.

C. 지오데식 최적 운송 매핑 (Geodesic Optimal Transport Mapping)

각 클라이언트는 자신의 로컬 가우시안 근사를 글로벌 참조에 정렬하기 위한 운송 맵을 구축합니다.

• 최적 맵(Optimal Map): 이차 비용(quadratic cost) 하에서, 두 가우시안 사이의 최적 운송 맵은 아핀(affine) 형태이며 공분산 행렬의 제곱근을 포함하는 폐쇄형 해(closed-form solution)를 가집니다.
• 보간(Interpolation): 실제 특징 분포가 엄격하게 가우시안이 아님을 고려하여, 과도한 교정을 방지하기 위해 보간 파라미터 $\tau \in [0, 1]$를 도입합니다. 최종 변환 $T^{(\tau)}_k$는 항등 맵(identity map)과 전체 최적 운송 맵의 볼록 조합(convex combination)입니다. 이는 클라이언트의 분포를 글로벌 참조를 향한 등속 지오데식(constant-speed geodesic)을 따라 이동시키는 Wasserstein displacement interpolation을 생성합니다.

D. 통합

클라이언트는 로컬 특징에 변환 $z' = T^{(\tau)}_k(z)$를 적용합니다. 이렇게 정렬된 특징들은 다운스트림 최적화 절차를 수정하지 않고도 기존의 모든 OSFL 알고리즘(예: FedAvg, FedCGS, FedPFT)에 입력됩니다.

3. 주요 기여

1. 정형화: 본 논문은 OSFL을 공동 도메인 및 레이블 시프트 하에서의 분포 정렬 문제로 정형화하여, 이러한 시프트가 원샷 접근 방식의 효능을 어떻게 제한하는지 강조합니다.
2. 알고리즘 제안: 저자들은 압축된 통계량의 단 한 번의 교환을 통해 클라이언트 특징 분포를 조화시키는 최적 운송 기반의 학습이 필요 없는(learning-free) 프레임워크인 SLOT-Align을 소개합니다. 이는 추가 데이터, 학습 가능한 생성기 또는 반복적인 최적화를 요구하지 않습니다.
3. 실증적 검증: 광범적인 실험을 통해 SLOT-Align이 다양한 벤치마크(Office-Home, Digits, DomainNet), 다양한 사전 학습된 백본(ViT, ResNet), 그리고 최첨단 OSFL 베이스라인 전반에서 일관되게 정확도와 강건성을 개선함을 입증했습니다.

4. 실험 결과

저자들은 다양한 레이블 시프트(디리클레 집중 파라미터 $\alpha$에 의해 제어됨) 및 도메인 다양성 하에서 SLOT-Align을 평가했습니다.

• 성능 향상: 공동 도메인 및 레이블 시프트($\alpha=0.1$) 상황에서, SLOT-Align은 테스트된 모든 베이스라인(O-FedAvg, FedCGS, FedPFT)의 매크로 평균 정확도를 일관되게 향상시켰습니다. 예를 들어, Office-Home 데이터셋에서 O-FedAvg는 +7.86%, FedPFT는 **+6.55%**의 성능 향상을 보였습니다.
• 강건성: 방법론은 더 강력한 레이블 왜곡($\alpha=0.05$ 및 $\alpha=0.01$) 하에서도 효과를 유지하며, 심각한 이질성 환경에서도 성능 향상을 보여주었습니다.
• 정렬 강도에 대한 절제 연구(Ablation): 보간 파라미터 $\tau$가 매우 중요하다는 것이 밝혀졌습니다. 중간 값(예: $\tau=0.4$)이 가장 좋은 트레이드오프를 제공했으며, $\tau=1.0$(전체 운송)은 때때로 성능을 저하시켰는데, 이는 부분적인 정렬이 완전한 교정보다 선호될 수 있음을 시사합니다.
• 도메인 전용 시프트: 레이블이 균형 잡힌 설정(도메인 시프트만 존재)에서 SLOT-Align은 미미한 이득을 보이거나 중립적인 성능을 보였으며, 이는 본 방법이 도메인과 레이블 시프트의 복잡한 상호작용을 구체적으로 타겟팅함을 확인시켜 줍니다.
• 효율성: 통신 비용은 매우 낮으며($\text{클라이언트당 } O(m + m^2)$의 평균 및 공분산), 계산 오버헤드는 특징 추출 및 다운스트림 학습에 비해 무시할 수 있는 수준입니다.

5. 의의 및 주장

본 논문은 SLOT-Align이 OSFL의 결정적 병목 현상인 '단 한 번의 통신 라운드 내에서의 분포 불일치 교정 불가능성'에 대한 모듈식이며 계산 효율적인 솔루션을 제공한다고 주장합니다.

• 모듈성: 기존 OSFL 파이프라인을 대체하는 것이 아니라 호환 가능한 전처리 계층으로 작동하여, 동결된 인코더에 의존하는 광범위한 현재 및 미래의 방법론에 적용 가능합니다.
• 기하학적 인식: 2-Wasserstein 공간과 Bures–Wasserstein barycenter를 활용함으로써, 특징 분포의 위치(평균)와 형태(공분산) 모두에서의 불일치를 명시적으로 교정합니다.
• 실용성: 데이터 합성이나 반복적인 증류와 관련된 개인정보 보호 위험 및 계산 오버헤드를 피하고, 오직 압축된 통계적 요약에만 의존합니다.

저자들은 OSFL이 엄격한 통신 예산에 의해 제약을 받지만, SLOT-Align이 도메인 및 레이블 시프트가 상호작용하여 특징 통계량을 왜곡할 때 발생하는 성능 저하를 효과적으로 완화한다고 결론짓습니다.