디리클레 프로세스 기반 클러스터형 연합 학습

초록

본 논문은 클러스터형 연합 학습(CFL)에서 사전에 클러스터 수 K를 지정할 필요 없이, 디리클레 프로세스(DP) 사전분포를 이용해 클러스터 수와 클라이언트 할당을 동시에 추정하는 DPMM‑CFL 알고리즘을 제안한다. 매 라운드마다 로컬 SGD 업데이트와 DP 기반 비모수 베이지안 클러스터링을 결합해 클러스터 모델을 갱신하고, 실험을 통해 기존 고정 K 방식보다 정확도와 F1 점수에서 우수함을 확인하였다.

상세 분석

DPMM‑CFL은 기존 CFL 접근법이 갖는 “클러스터 수 K를 사전에 지정해야 한다”는 근본적인 한계를 베이지안 비모수 방법으로 해결한다. 핵심 아이디어는 클러스터 파라미터에 디리클레 프로세스(DP) 사전분포를 부여함으로써, 클러스터 수가 무한히 가능하도록 모델링하고, 실제 클라이언트 수 M에 의해 활성 클러스터 수가 자동으로 제한되도록 하는 것이다. 이를 위해 저자는 클라이언트별 로컬 파라미터 ω_i 를 관측값으로 보고, ω_i 가 클러스터 평균 μ_k 를 중심으로 하는 가우시안 분포를 따른다고 가정한다. DP의 베이스 분포 G₀는 평균이 0이고 공분산이 단위인 구형 가우시안으로 설정되어, 사전 정보가 최소화된 상태에서 클러스터링이 진행된다.

알고리즘 흐름은 크게 세 단계로 나뉜다. 첫째, 각 라운드 시작 시 서버는 현재 클러스터 할당 c_{t‑1}에 따라 클러스터 모델 Ω_{t‑1}를 클라이언트에게 전송한다. 둘째, 클라이언트는 자신이 속한 클러스터 모델을 초기값으로 삼아 Q번의 로컬 SGD를 수행해 업데이트된 파라미터 ω_{t,i} 를 얻는다. 셋째, 서버는 모든 ω_{t,i} 를 모아 DP 기반 클러스터링을 수행한다. 여기서는 단순 Gibbs 샘플링 대신, 지역적인 제안과 메트로폴리스‑헤이스팅 수용률을 갖는 split‑merge MCMC를 사용해 클러스터 할당 c_t 를 샘플링한다. 이 과정에서 새로운 클러스터가 생성되거나 기존 클러스터가 병합될 수 있으며, 이는 DP의 농도 파라미터 α에 의해 조절된다(α가 클수록 새로운 클러스터 생성 확률이 높아진다).

클러스터링이 완료되면, 각 클러스터 k에 속한 클라이언트들의 파라미터를 가중 평균(N_k = Σ_{i∈k} n_i)하여 새로운 클러스터 모델 Ω_{t,k} 를 계산하고, 이를 해당 클라이언트에게 다시 전송한다. 이렇게 클러스터링과 연합 최적화가 라운드마다 교차하면서, 클러스터 구조가 데이터에 맞게 동적으로 적응한다.

이론적 측면에서 저자는 DP의 폴리아우르 모델과 CRP(Chinese Restaurant Process)를 이용해 클러스터 할당의 사전 확률 p(c|α)를 명시하고, 정규–정규(Normal‑Normal) 공액성을 활용해 클러스터별 주변우도 p(W_k|G₀)를 닫힌 형태로 계산한다. 이를 통해 MCMC 샘플링 시 필요한 비율 계산이 효율적으로 수행된다.

실험에서는 Fashion‑MNIST와 CIFAR‑10 두 데이터셋을 사용해 두 가지 비IID 시나리오(Dirichlet 파티셔닝, 클래스‑스플릿 파티셔닝)를 구성하였다. 200개의 클라이언트를 시뮬레이션하고, 각 클라이언트는 CNN의 마지막 완전연결층을 클러스터링 피처로 사용한다. DP의 농도 파라미터는 α=1.0으로 설정했으며, 로컬 SGD는 학습률 0.005, 모멘텀 0.9, 배치 크기 32, 로컬 스텝 Q=10으로 진행하였다.

결과는 고정 K를 여러 값으로 설정한 FeSEM(K‑means 기반 CFL)과 비교했을 때, DPMM‑CFL이 전반적으로 높은 정확도와 macro‑F1 점수를 달성함을 보여준다. 특히 Dirichlet 파티셔닝에서는 DPMM‑CFL이 평균 약 12개의 클러스터를 자동으로 추정했으며, 이는 실제 생성된 K=10에 근접한다. 클래스‑스플릿 파티셔닝에서는 진정한 클러스터 수가 알려지지 않았음에도 불구하고, DPMM‑CFL이 FeSEM이 최고 성능을 보이는 K≈25와 유사한 클러스터 수를 찾아냈다. 또한 클러스터 수 K_t는 초기 몇 라운드 동안 변동하다가 곧 안정화되며, 이후 클러스터 할당이 고정돼 모델 학습이 수렴한다는 점이 실험을 통해 확인되었다.

이 논문은 연합 학습 환경에서 클러스터 수를 사전에 알 수 없거나 동적으로 변할 때, 베이지안 비모수 접근법을 통해 자동으로 구조를 추정하고, 이를 연합 최적화와 효율적으로 결합할 수 있음을 실증적으로 증명한다. 향후 연구에서는 DP의 하이퍼파라미터 자동 튜닝, 더 복잡한 모델 구조에 대한 확장, 그리고 프라이버시 보장을 위한 사후 샘플링 가공 등으로 적용 범위를 넓힐 여지가 있다.