Variational Consensus Monte Carlo for Bayesian Mixture

기술 요약: 베이지안 혼합 모델을 위한 변분 합의 몬테카를로(Variational Consensus Monte Carlo)

문제 정의
본 논문은 데이터가 자연스럽게 여러 노드(예: 의료 센터)로 분할되어 있으며, 개인정보 보호 또는 규제 문제로 인해 데이터를 통합할 수 없는 "크로스-실로(cross-silo)" 시나리오의 연합 학습(Federated Learning, FL) 환경에서 베이지안 추론을 수행하는 과제를 다룹니다. 합의 몬테카를로(Consensus Monte Carlo, CMC)는 로컬 사후 샘플을 집계함으로써 분산 베이지안 추론을 위한 프레임워크를 제공하지만, 기존 방법들은 비지도 클러스터링에 적용될 때 상당한 한계에 직면합니다. 구체적으로, Rabinovich, Angelino, 그리고 Jordan(2015)이 제안한 변분 합의 몬테카를로(VCMC) 방식은 클러스터의 수와 주요 혼합 매개변수를 알고 있으며, 모든 데이터 샤드에 모든 클러스터가 존재한다고 가정합니다. 이러한 가정은 실제 상황, 특히 클러스터의 출현 빈도가 사이트마다 다르고 진정한 잠재 하위 그룹의 수를 알 수 없는 헬스케어 데이터 응용 분야에서는 거의 성립하지 않습니다. 또한, 표준 CMC 집계(예: 가중 평균)는 전역 사후 분포의 최적 근사를 보장하지 않으며, 기존의 연합 클러스터링 방법들은 공액 사전 분포(conjugate priors)를 요구하는 평균장 근사(mean-field approximations)에 의존하거나 전체 데이터에 접근하지 않고는 비공액(non-conjugate) 설정을 처리하지 못하는 경우가 많습니다.

방법론
저자들은 VCMC를 **범주형 분포의 과적합된 베이지안 혼합 모델(over-fitted Bayesian mixture models of categorical distributions)**로 확장하는 포괄적인 파이프라인을 제안합니다. 이 프레임워크는 두 가지 주요 단계로 작동합니다:

1. 적용 단계 (로컬 추론):

   • 각 데이터 샤드 내에서 독립적으로 MCMC 알고리즘(구체적으로 깁스 샘플러)을 실행합니다.
   • 클러스터 수를 모르는 문제를 해결하기 위해, 모델을 예상되는 실제 클러스터 수($C$)보다 더 큰 수의 혼합 성분($C'$)을 갖도록 "과적합(over-fitted)"합니다.
   • 분할된 사전 분포(fractionated prior)를 사용하여, 전역 사전 분포를 $S$개의 샤드에 걸쳐 $p(\theta)^{1/S}$로 나눕니다. 이는 로컬 사후 분포의 곱이 전역 사후 분포를 근사하도록 보장합니다.
   • 저자들은 로컬 추론을 가속화하기 위해 병렬화된 C++17 깁스 샘플러를 구현했습니다.

2. 집계 단계 (전역 추론):

   • 클러스터 매칭: 혼합 모델은 레이블 순열에 불변(invariant)하므로, 클러스터를 샤드 간에 매칭하는 것이 중요한 단계입니다. 본 논문은 세 가지 알고리즘을 도입하고 평가합니다:
     • 헝가리안 매칭(Hungarian Matching): 일대일 대응을 강제하지만, 모든 샤드에 모든 클러스터가 나타나지 않는 경우 실패합니다.
     • 최소 발산 매칭(Minimum Divergence Matching): 로컬 주변 사후 분포 간의 쿨백-라이블러(KL) 발산을 최소화합니다.
     • 볼 매칭(Ball Matching): 과적합된 혼합 모델의 수렴 속도에 기반하여, 파라미터 공간 내에서 "가까운"(반경 $r$ 이내) 클러스터들을 그룹화하는 새로운 알고리즘입니다. 이 방법은 구별할 수 없는 클러스터들을 샤드 내에서 병합할 수 있게 하여, 클러스터 수의 로컬 과대 추정을 완화합니다.
   • 변분 집계: 집계를 근사된 사후 분포와 실제 전역 사후 분포 사이의 KL 발산을 최소화하는 변분 추론 문제로 프레임링합니다. 집계 함수는 로컬 매개변수의 가중 합($\sum \upsilon_s \theta_s$)으로 제한됩니다.
   • 최적화: 최적의 집계 가중치는 투영 확률적 경사 하강법(projected stochastic gradient descent)을 사용하여 찾습니다. 본 논문은 다양한 FL 제약 조건을 수용하기 위해 네 가지 최적화 전략을 상세히 설명합니다:
     1. 글로벌 노드로 로우 데이터를 공유함 (가장 제한이 적음).
     2. 연합 경사 하강법 (그래디언트만 공유).
     3. 제한된 통신 (여러 시작 지점을 가진 적은 횟수의 반복 실행).
     4. 컨팅전시 행렬(요약 통계량)을 글로벌 노드로 공유.

주요 기여
본 논문은 네 가지 주요 방법론적 기여를 합니다:

1. 과적합 모델로의 확장: 클러스터의 수를 아는 것이 필요하지 않도록, 클러스터 수와 모든 모델 매개변수를 동시에 추론하도록 VCMC를 확장했습니다.
2. 새로운 클러스터 매칭 알고리즘: 표준 헝가리안 매칭이 실패하는 시나리오(일부 샤드에 클러스터가 존재하지 않는 경우)에 적합한 매칭 전략(특히 최소 발산 및 볼 매칭)을 도입했습니다.
3. 유연한 추론 전략: 다양한 연합 제약 조건(예: 제한된 통신 대역폭 또는 엄격한 개인정보 보호 요구 사항)에 맞춤화된 여러 집계 전략을 제안했습니다.
4. 실무 가이드라인: 데이터 특성(예: 클러스터 분리도 및 샤드 이질성)에 따른 매칭 알고리즘 및 추론 전략 선택을 위한 경험적 가이드라인을 제공합니다.

결과
프레임워크는 종합적인 시뮬레이션 연구와 실제 전자 건강 기록(EHR) 데이터에 대한 적용을 통해 평가되었습니다.

• 시뮬레이션 연구:

  • 매칭 성능: 모든 클러스터가 모든 샤드에 나타나지 않는 이질적인 설정에서, 헝가리안 알고리즘은 올바른 매칭을 찾는 데 실패했습니다(중첩되지 않은 시나리오에서 성공률 0%). 볼 매칭과 최소 발산 매칭은 훨씬 더 나은 성능을 보였습니다. 특히 볼 매칭은 로컬 MCMC 노이즈로 인한 클러스터 수 과대 추정을 완화하는 데 효과적이었습니다.
  • 클러스터링 품질: VCMC는 SOTA(state-of-the-art) 대안들(FedMerDel 및 Federated K-modes) 및 "골드 스탠다드"인 풀 데이터(full-data) MCMC와 비슷하거나 약간 더 나은 조정 랜드 지수(Adjusted Rand Index, ARI) 점수를 달성했습니다. 주목할 점은, VCMC가 샤드 간 데이터 분포가 이질적인 경우 작은 클러스터를 복구하는 데 있어 풀 데이터 MCMC보다 뛰어난 성능을 보였다는 것입니다.
  • 매개변수 추정: VCMC는 클러스터 가중치에 대해 낮은 상대적 편향을 제공했습니다. 클러스터 분리도가 낮고 샤드 분포가 이질적인 시나리오에서, VCMC는 풀 데이터 MCMC보다 작은 클러스터의 가중치를 더 정확하게 추정했는데, 이는 해당 작은 클러스터들이 특정 샤드 내에서 충분히 커서 낮은 불확실성으로 식별되었기 때문으로 보입니다.
  • 대안과의 비교: FedMerDel은 더 빠르지만, 평균장 근사에 의존하여 사후 분산을 과소 추정했습니다. Federated K-modes는 범주형 데이터에 대한 클러스터 구조 추정에서 높은 편향과 낮은 성능을 보였습니다.

• 실제 적용 (THIN 코호트):

  • 이 방법은 다중 질환 패턴을 식별하기 위해 영국 THIN 데이터베이스의 289,812명의 노인 환자 데이터에 적용되었습니다.
  • 알고리즘은 27개의 클러스터를 식별했습니다. 뇌졸중, 치매 또는 특정 암 유형이 지배적인 클러스터와 같이 임상적으로 의미 있는 그룹들을 성공적으로 특징지었습니다.
  • 결정적으로, 이 방법은 (췌장염, 류마티스 관절염, 성 기능 장애로 특징지어지는 단 31명만을 포함하는 클러스터와 같이) 문서화된 의학적 상관관계를 따르는 작지만 뚜렷한 클러스터를 식별해 냈으며, 이는 풀링된 분석에서 희석될 수 있는 희귀 하위 그룹을 탐지하는 방법의 능력을 입증했습니다.

의의 및 주장
저자들은 자신들의 프레임워크가 개인정보 보호가 민감한 환경에서 비지도 클러스터링을 위한 강력한 솔루션을 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

• 데이터 풀링 없이 클러스터 수와 모델 매개변수를 추론할 수 있습니다.
• 로컬 데이터 샤드가 특정 하위 그룹에 대해 충분한 신호를 포함하고 있는 경우, 표준 풀 데이터 MCMC보다 더 정확하게 작고 희귀한 클러스터를 복구할 수 있습니다.
• 공액 사전 분포나 전체 데이터에 대한 접근 없이도 엄격한 개인정보 보호 제약(크로스-실로 FL) 하에서 작동할 수 있습니다.

저자들은 본 프레임워크가 이론적으로 다른 혼합 분포로 확장 가능하지만, 현재 검증은 범주형 분포에 국한되어 있다고 겸허히 언급했습니다. 또한 클러스터 매칭 알고리즘의 선택이 특정 목표에 따라 달라짐을 강조했습니다. 즉, 클러스터 수를 추정하는 것이 목적이라면 볼 매칭이 선호되며, 별개의 클러스터가 병합되는 것을 피하는 것이 우선순위라면 최소 발산 매칭이 더 나을 수 있습니다. 저자들은 VCMC가 FedMer-Del과 같은 변분 방법들에 비해 계산 집약적일 수 있지만, 복잡하고 이질적인 연합 환경에서 매개변수 추정 및 클러스터 복구에 있어 우수한 정확도를 제공한다고 결론지었습니다.