FedProIn: Mitigating Client Drift for Learnable Prototypes in Federated Medical Imaging

기술 요약: FedProIn – 학습 가능한 프로토타입을 통한 연합 의료 영상의 클라이언트 드리프트 완화

문제 정의

연합 학습(Federated Learning, FL)은 스캐너 프로토콜, 획득 방법 및 다양한 의료 센터 간의 다양한 환자군으로 인한 통계적 이질성으로 인해 의료 영상 분야에서 상당한 어려움을 겪고 있습니다. 이러한 비-IID(non-Independent and Identically Distributed) 데이터 분포는 로컬 모델이 최적화 과정 중 글로벌 모델로부터 멀어지는 **클라이언트 드리프트(client drift)**를 유발합니다. 구체적으로, 동일한 인스턴스가 서로 다른 클라이언트에 의해 특징 공간(feature space) 내의 서로 다른 영역으로 매핑되며, 로컬 표현(representation)이 클라이언트 특유의 특성에 편향되게 됩니다.

기존 방식들은 다음과 같은 한계에 직면해 있습니다:

• 제약 기반 정규화(예: FedProx, MOON)는 고차원 공간에서 의미론적 표현을 정렬하는 데 어려움을 겪습니다.
• 표준 프로토 타입 기반 방식은 종종 클래스당 단일 글로벌 프로토타입(평균 특징 벡터를 통해)을 가정하는데, 이는 이질적인 의료 데이터에는 부적절합니다.
• 다중 프로토타입 방식은 높은 계산 복잡성을 도입하거나, 결정 경계를 효과적으로 최적화하지 못하는 단일 모드 표현에 의존하는 경우가 많습니다.

방법론: FedProIn

저자들은 **학습 가능한 프로토타입(Learnable Prototypes, LPs)**을 정적인 평균 벡터가 아닌 최적화 가능한 모델 파라미터로 활용하는 프레임워크인 FedProIn(Influence를 통한 프로토타입 집합의 연합 학습)을 제안합니다. 글로벌 모델은 특징 인코더 $f_\theta$와 $M$을 클래스당 프로토타입 수로 하는 프로토타입 집합 $P \in \mathbb{R}^{C \times M \times d}$로 구성됩니다.

1. 클라이언트 드리프트의 분해

FedProIn은 클라이언트 드리프트를 두 가지 구성 요소로 명시적으로 분해하고 특정 정규화 손실을 적용하여 이를 해결합니다.

• 특징 드리프트(Feature Drift): **특징 발산 손실(Feature Divergence Loss, FDL)**을 통해 해결합니다. 클라이언트는 동결된 글로벌 특징 추출기($\theta^{t-1}$)의 복사본을 유지합니다. 로컬 인코더는 로컬 특징 임베딩과 참조 글로벌 특징 사이의 평균 제곱 오차를 최소화하도록 정규화되어 의미론적 정렬을 보장합니다.
• 프로토타입 드리프트(Prototype Drift): **프로토타입 대조 손실(Prototype Contrastive Loss, PCL)**을 통해 해결합니다. 이 손실은 업데이트된 로컬 프로토타입을 현재의 글로벌 프로토타입에 가깝게 끌어당기는 동시에, 이전의 로컬 상태로부터는 멀어지게 밀어냅니다. 이는 프로토타입이 로컬 데이터 변화에 과도하게 편향되는 것을 방지합니다.

전체 로컬 목적 함수는 기본 분류 손실($\ell_{Task}$)과 이러한 정규화 항들을 결합합니다:
$$\ell_{Total} = \ell_{Task} + \lambda_{FDL}\ell_{FDL} + \lambda_{PCL}\ell_{PCL}$$

롱테일(long-tailed) 의료 데이터셋의 클래스 불균형을 처리하기 위해, 이 방법은 클래스 균형 샘플링과 적응형 증강(adaptive augmentation)을 사용하여 희귀 질환 하위 유형이 프로토타입 업데이트 중에 충분히 표현되도록 합니다.

2. 정규화된 영향력 집합(Normalized Influence Aggregation, NIA)

표준 집합 방식(예: FedAvg)은 프로토타입 활용도를 무시하고 데이터 크기에 의해서만 가중치를 부여하기 때문에 다중 프로토타입을 위한 집합에는 불충분합니다. 이는 미사용 프로토타입으로부터 발생하는 노이즈를 증폭시키거나 특수 사례를 위한 전문화된 표현을 억제할 수 있습니다.

FedProIn은 **정규화된 영향력 집합(NIA)**을 도입합니다:

• 영향력 행렬($I_i$): 로컬 학습 중에 각 클라이언트는 특정 프로토타입이 로컬 샘플의 최근접 이웃(nearest neighbor) 역할을 수행한 횟수를 추적하여 모든 프로토타입에 대한 영향력 점수를 기록합니다.
• 집합(Aggregation): 글로벌 프로토타입은 단순히 클라이언트의 데이터 크기가 아니라, 사용 점수($I_{k,c}^i$)에 비례하여 가중치를 두어 집합됩니다. 이는 사용되지 않는 프로토타입으로부터 오는 노이즈 섞인 업데이트를 경감시키고, 특정 클라이언트 서브셋과 관련된 전문화된 프로토타입을 보존합니다.

주요 기여

1. 학습 가능한 프로토타입 공식화: 클래스 분리 마진을 최대화하기 위해 프로토타입을 정적인 평균 벡터가 아닌 (특징 추출기와 함께) 최적화 가능한 파라미터로 채택하였습니다.
2. 이중 정규화 전략: FDL과 PCL을 통해 각각 특징 드리프트와 프로토타입 드리프트를 해결함으로써 클라이언트 드리프트를 명시적으로 완화했습니다.
3. 정규화된 영향력 집합(NIA): 이질적인 활용도와 다중 프로토타입 환경에서의 희소한 활성 문제를 해결하기 위해, 로컬 유용성에 따라 클라이언트 프로토타입에 적응적으로 가중치를 부여하는 새로운 집합 메커니즘을 제안했습니다.
4. 경험적 검증: 다양한 의료 영상 데이터셋에 대한 포괄적인 평가를 통해 통계적 이질성에 대한 강건성을 입증했습니다.

실험 결과

본 프레임워크는 두 가지 공개 롱테일 의료 영상 데이터셋인 HAM10000(더마토스코피 이미지)과 Matek-19(말초 혈액 세포 이미지)에서 평가되었습니다. 실험은 $K=10$명의 클라이언트를 대상으로 IID 및 Non-IID (Dirichlet $\alpha=0.5$) 설정 모두에서 수행되었습니다.

성능 하이라이트:

• HAM10000: FedProIn은 83.5% 정확도(IID) 및 **81.1% 정확도(Non-IID)**를 달성하여 FedAvg, FedProx, MOON 및 기타 프로토타입 기반 방식(FedProto, FedTGP, FedPLVM)을 포함한 모든 베이스라인을 능가했습니다.
• Matek-19: FedProIn은 96.2% 정확도(IID) 및 **95.8% 정확도(Non-IID)**를 달성했습니다.
• 강건성: Non-IID 조건에서 기존 프로토타입 방식들은 상당한 성능 저하를 보였으나(예: HAM10000에서 6~7% 정확도 손실), FedProIn은 HAM10000과 Matek-19에서 각각 2.4%와 0.4%의 상대적 손실로 제한했습니다.
• 절제 연구(Ablation Study): FDL과 PCL 모두 최적의 성능을 위해 필수적임을 확인했습니다. FDL을 제거했을 때 성능이 가장 급격히 떨어졌으며, 이는 특징 드리프트가 주요 실패 모드임을 나타냅니다. 전체 모델은 단일 구성 변형 및 퇴보된 케이스(FedAvg와 동일한 경우)보다 현저히 우수한 성능을 보였습니다.

의의 및 주장

본 논문은 학습 가능한 프로토타입과 영향력 인지 집합을 활용하여 의료 영상의 클라이언트 드리프트를 효과적으로 완화할 수 있다고 주장합니다. 저자들은 자사의 접근 방식이 다음과 같음을 단언합니다:

• 매우 이질적인(Non-IID) 환경에서 기존 FL 베이스라인 및 개인화 방식에 비해 우수한 일반화 성능을 달una합니다.
• 단일 프로토타입 가정과 다중 프로토타입 환경에서의 비가중 집합의 한계를 성공적으로 해결합니다.
• 데이터 이질성과 클래스 불균형이 만연한 의료 영상 분야에 대한 강력한 솔루션을 제공합니다.

저자들은 향후 연구로서 FedProIn을 개인화된 FL(글로벌 모델 의존성 제거)에 적응시키고, 다양한 모델 아키텍처를 가진 이질적 FL 환경으로 프레임워크를 확장하는 것을 탐구할 것이라고 언급했습니다. 재현성을 지원하기 위해 코드는 공개되어 있습니다.