기술 요약: COSMOS

문제 제기

연동 학습 (Federated Learning, FL) 은 클라이언트 간 데이터 분포 (통계적 이질성) 뿐만 아니라 모델 아키텍처 (아키텍처적 이질성) 도 상이한 이질적인 환경에서 상당한 도전에 직면해 있습니다. 개인화 연동 학습 (Personalized Federated Learning, PFL) 방법론은 개별 클라이언트에 맞춰 모델을 조정하기 위해 개발되었으나, 대부분의 기존 접근법은 구조적 호환성을 전제하거나 클라이언트가 독점적인 모델 아키텍처를 공개해야 한다고 가정합니다. 이로 인해 클라이언트가 파라미터나 그래디언트를 공개하지 않고도 임의의 모델이나 독점 모델을 활용하려는 경우가 많아 실질적인 장벽이 발생합니다.

현재의 모델-중립적 (model-agnostic) FL 접근법인 FedMD 와 COMET 은 공유된 레이블이 없는 데이터셋을 사용하여 출력 수준 (예: 소프트 또는 하드 레이블) 통신에 의존합니다. 그러나 이러한 방법론은 다음과 같은 특정 한계를 겪습니다:

1. 수동적 서버 조정: 서버를 수동 조정자로만 취급하여 서버의 계산 자원을 능동적인 개인화에 활용하지 못합니다.
2. 휴리스틱 클러스터링: COMET 과 같은 방법론은 K-평균 (K-means) 클러스터링에 의존하는데, 이는 사전에 클러스터 수 ($K$) 를 알아야 하며 클라이언트 다양성이 높을 때 개인화 보장을 위한 공식적인 메커니즘이 부족합니다.
3. 이론적 보장 부재: 기존 모델-중립적 PFL 문헌은 위험 수축 (risk contraction) 에 관한 엄격한 이론적 보장이 일반적으로 결여되어 있으며, 단순화된 가정 (예: 선형 모델, 가우시안 분포) 에 의존하거나 정상 상태로의 표준 수렴 결과만 제공합니다.

방법론: COSMOS 프레임워크

저자들은 클라이언트 측의 레이블이 있는 데이터와 공유된 레이블이 없는 공개 데이터셋 $U$(예: 웹 스크래핑 또는 합성 데이터) 가 존재한다는 가정 하에, 의사 레이블 (pseudo-label) 통신만을 사용하여 서버 측 개인화를 가능하게 하는 모델-중립적 PFL 프레임워크인 COSMOS(Clustered Output-based Server Models) 를 제안합니다.

COSMOS 는 두 가지 명확한 단계로 진행됩니다:

1. 사전 학습 및 클러스터링 (PTC) 단계

이 단계는 초기 클라이언트 - 서버 매핑 및 클러스터별 모델을 확립합니다:

• 로컬 사전 학습: 클라이언트는 자신의 레이블이 있는 데이터 $D_i$로 로컬 모델 $f_i$를 학습하고 공개 데이터셋 $U$에서 의사 레이블을 생성합니다.
• 탐욕적 클러스터링: 서버는 $U$에서의 의사 레이블 유사성을 기반으로 클라이언트를 클러스터링합니다. K-평균과 달리 COSMOS 는 사전에 $K$를 지정할 필요가 없는 탐욕적 클러스터링 알고리즘을 사용합니다. 이 알고리즘은 거리 임계값 $B_0$ 내의 가장 "가까운" 이웃을 가진 클라이언트를 반복적으로 선택하여 클러스터를 형성하며 모든 클라이언트가 할당될 때까지 진행됩니다. 이 과정은 학습 절차 내재적입니다.
• 서버 측 증류: 각 클러스터 $G_k$에 대해 서버는 구성원 클라이언트들의 의사 레이블을 집계하여 평균 의사 레이블 분포를 생성한 후, 해당 클러스터 전용의 "교사" 모델 $h_k$를 이러한 집계된 레이블과 특정 정규화 항을 사용하여 학습시킴으로써 로컬 부드러움과 견고성을 보장합니다.
• 클라이언트 측 증류: 클라이언트는 할당된 클러스터의 교사 모델로부터 의사 레이블을 수신하고, 이러한 레이블을 사용하여 로컬 모델을 미세 조정합니다.

2. 반복적 연동 미세 조정 (IFFT) 단계

PTC 단계 이후, 시스템은 반복 루프에 진입합니다:

• 로컬 미세 조정: 클라이언트는 레이블이 있는 개인 데이터로 모델을 정제합니다.
• 서버 증류: 클라이언트는 업데이트된 의사 레이블을 서버로 전송합니다. 서버는 새로운 집계된 의사 레이블을 기반으로 클러스터별 교사 모델을 업데이트합니다.
• 클라이언트 증류: 클라이언트는 업데이트된 교사 의사 레이블을 사용하여 로컬 모델을 더 정제합니다.

중요하게도, 서버 모델은 데이터 변환 (예: 회전, 이동) 하의 예측 불일치를 패널티로 부과하여 견고성을 보장하는 준지도 학습 (SSL) 이론에서 유래된 정규화 항을 사용하여 학습됩니다.

이론적 기여

본 논문은 모델-중립적 PFL 에 대한 개인화 위험 수축의 종단간 (end-to-end) 이론적 분석을 최초로 제공합니다. 준지도 학습 이론의 도구 (특히 확장에 기반한 연결성과 유계된 의사 레이블 오차) 를 활용하여 저자들은 지수적 위험 수축을 위한 충분 조건을 유도합니다.

주요 이론적 구성 요소는 다음과 같습니다:

• 가정: 분석은 공개 분포와 개인 분포 간의 유계된 중요도 가중치, 국소적으로 잘 조건화된 목적 함수, 그리고 클라이언트가 공개 데이터에 대한 예측에서 확신을 갖는 "신뢰 마진"을 가정합니다. 또한, 의사 레이블이 실제 정답에 대해 비자명한 정보를 담고 있음을 보장하는 데이터 분포의 확장 속성에 의존합니다.
• 위험 수축: 저자들은 이러한 조건 하에서 서버 모델과 클라이언트 모델 모두의 개인화 위험이 반복 횟수에 따라 지수적으로 수축함을 증명합니다. 구체적으로, 위험 상한은 클러스터 품질과 공개 데이터셋의 커버리지에 의해 결정되는 속도로 감소하여 결국 로컬 학습 단독의 위험 수준 이하로 떨어집니다.

실험 결과

저자들은 EMNIST-balanced, CIFAR-10, CIFAR-100, Tiny ImageNet 의 네 가지 이미지 분류 벤치마크를 통해 COSMOS 를 평가했습니다. 실험에는 AlexNet, ResNet, MobileNet, SqueezeNet 등 다양한 클라이언트 아키텍처와 VGG16 서버 모델이 혼합되어 사용되었습니다.

• 성능: COSMOS 는 모든 벤치마크에서 모든 모델-중립적 베이스라인 (FedMD, FedCT, COMET) 을 일관되게 능가했습니다. 작업 복잡도가 증가함에 따라 (예: EMNIST 에서 Tiny ImageNet 으로) 성능 격차는 더욱 벌어졌습니다.
• 이질성: COSMOS 는 데이터 이질성의 다양한 정도 (디리클레 파라미터 $\alpha$로 제어) 와 다양한 클라이언트 아키텍처 혼합에 대해 견고하게 유지되었습니다. 클라이언트가 더 약한 모델을 사용하더라도 상당한 성능 우위를 유지했습니다.
• 통신 효율성: 모델 파라미터나 그래디언트 대신 의사 레이블 (하드 또는 소프트) 만을 통신함으로써 COSMOS 는 파라미터 공유 방식 대비 통신 대역폭을 1~2 차수 (orders of magnitude) 감소시켰습니다.
• 동질적 환경: 모든 클라이언트가 동일한 아키텍처를 공유하는 동질적 환경에서도 COSMOS 는 최첨단 PFL 방법론과 경쟁력을 유지했으며, 모델-중립적 설계가 동질성이 존재할 때 성능 패널티를 초래하지 않음을 입증했습니다.

의의 및 주장

본 논문은 COSMOS 가 클라이언트 모델의 구조적 지식을 요구하지 않고 고성능 개인화를 가능하게 함으로써 FL 의 중요한 병목 현상을 해결한다고 주장합니다. 그 의의는 다음 세 가지 주요 영역에 있습니다:

1. 실제 배포: 클라이언트가 독점적이거나 임의의 아키텍처를 사용하는 다양하고 현실적인 생태계에서 PFL 을 제한하는 구조적 장벽을 제거합니다.
2. 이론적 발전: 표준 수렴 결과를 넘어 현실적인 가정 하에서 지수적 위험 감소를 증명하는 모델-중립적 PFL 에 대한 최초의 일반적 위험 수축 보장을 제공합니다.
3. 확장성: 의사 레이블을 사용한 개인화된 서버 측 학습이 고도로 이질적인 환경에서 실현 가능하고 확장 가능한 패러다임임을 보여주며, 상당한 통신 효율성 향상을 제공합니다.

저자들은 현재 작업이 클라이언트 - 서버 구조에 초점을 맞추고 있지만, 이 프레임워크가 향후 탈중앙화 학습 시스템의 기초를 마련하여 연동 학습의 실제적 관련성을 더욱 강화한다고 결론지었습니다.