기술 요약: ClientMorpher – 이질적 연합 지시 미세 조정을 위한 MoE 라우터 가이드 클러스터링

1. 문제 정의

연합 지시 미세 조정(Federated instruction fine-tuning)은 분산되어 있고 프라이버시가 민감한 데이터에 중앙 집중화 없이 대규모 언어 모델(LLM)을 적응시킬 수 있게 합니다. 혼합 전문가(Mixture-of-Experts, MoE) 아키텍처는 희소 활성화(sparse activation)를 통해 계산 및 통신량을 줄일 수 있어 연합 학습에서 매력적이지만, 기존의 연합 MoE 방식은 이질적인 클라이언트 분포 하에서 상당한 문제에 직면합니다.

표준 연합 평균(Federated Averaging, FedAvg)과 같은 기존 방식은 모든 클라이언트의 업데이트를 단일 글로벌 모델으로 무차별적으로 집계합니다. 이는 데이터 분포가 균일하다는 것을 가정하지만, 실제로는 그렇지 않은 경우가 많습니다. 클라이언트가 특화된 데이터(예: 요약 vs 분류와 같은 서로 다른 도메인)를 보유하고 있을 때, 무차별적인 집계는 **부정적 전이(negative transfer)**를 유발하여 개인화 성능과 전반적인 모델 품질을 모두 저하시킵니다. 더욱이, 기존의 개인화된 연합 학습 방법들은 최적화 이전에 클라이언트 간 유사성을 식별하기 위한 구조적 신호로서 MoE 모델의 **라우팅 동작(routing behavior)**을 간과하는 경향이 있습니다.

2. 방법론: ClientMorpher

저자들은 라우팅 인지형, 개인화 연합 지시 미세 조정 프레임워크인 ClientMorpher를 제안합니다. 핵심 통찰은 사전 학습된 MoE 모델의 라우터가 클라이언트의 데이터 분포를 나타내는 압축된 시그니처 역할을 한다는 것입니다. 유사한 지시 분포를 가진 클라이언트는 유사한 전문가를 활성화하며, 따라서 라우팅 동작은 모델 집계가 일어나기 전에 효과적인 협업 그룹을 가이드할 수 있습니다.

이 프레임워크는 두 가지 주요 단계로 작동합니다.

단계 0: 공유 클라이언트 전문가 사용 프로파일링 (Shared Client Expert-Usage Profiling)

학습이 시작되기 전, 서버는 동결된 사전 학습된 MoE 백본을 활용합니다. 각 클라이언트 $i$는 파라미터를 업데이트하지 않고 동결된 라우터를 통해 소량의 로컬 데이터($B$ 배치)를 처리합니다.

• 지표: 시스템은 각 클라이언트 $i$와 전문가 $e$에 대한 전문가 선택 빈도($c_{i,e}$)를 계산합니다.
• 프로필: 이 카운트들을 정규화하여 각 클라이언트 $i$에 대한 전문가 사용 프로필 벡터 $u_i \in [0, 1]^E$를 생성합니다.
• 행렬: 이 벡터들은 클라이언트-전문가 사용 행렬 $U \in [0, 1]^{N \times E}$를 형성합니다.

단계 1: 클러스터링 전략

ClientMorpher는 행렬 $U$를 사용하여 협업 그룹을 형성하기 위한 두 가지 상호 보완적인 전략을 조사합니다.

1. ClientMorpher-C (클라이언트 중심 - Client-Centric):

   • 전문가 활성화 프로필($u_i$)을 기반으로 클라이언트를 직접 클러스터링합니다.
   • 클라이언트 프로필에 K-Means를 사용하여 클라이언트를 클러스터 $C_k$에 할당합니다.
   • 유사한 라우팅 빈도를 가진 클라이언트들이 그룹화됩니다.

2. ClientMorpher-E (전문가 중심 - Expert-Centric):

   • 먼저 클라이언트 간 사용 시그니처(행렬 $U$의 열)를 기반으로 전문가를 클러스터링합니다.
   • 전문가들은 그들을 사용하는 클라이언트를 기준으로 $K$개의 클러스터로 그룹화됩니다.
   • 그 후 클라이언트는 자신의 로컬 라우팅 프로필에서 가장 큰 "질량(mass)"을 받는 전문가 클러스터에 할당됩니다.
   • 이 접근 방식은 클라이언트 그룹을 도출하기 전에 클라이언트 전반에 걸쳐 공유되는 전문가 그룹을 발견합니다.

단계 2: 클러스터별 개인화 연합 학습 (Cluster-Personalized Federated Learning)

클러스터가 형성되면, 프레임워크는 각 클러스터 내에서 연합 최적화를 수행합니다.

• 서버는 각 클러스터 $k$에 대해 별도의 적응 상태(adaptation state) $\theta_a^{(k)}$를 유지합니다.
• 클러스터 $C_k$에 속한 클라이언트는 클러스터 전용 파라미터를 전달받아, 로컬 데이터(예: LoRA 미세 조정)에 대해 로컬 학습을 수행하고 업데이트를 다시 보냅니다.
• 집계(FedAvg)는 글로벌하게가 아니라 클러스터 내에서만 발생합니다.
• 이는 희소 MoE 미세 조정의 통신 효율성을 유지하면서 서로 다른 클라이언트 간의 부정적 전이를 방지합니다.

3. 주요 기여

본 논문은 다음과 같은 구체적인 기여를 합니다.

• 프레임워크 제안: 모델 집계 전 클라이언트 협업을 조직하기 위해 사전 학습된 MoE 모델의 라우팅 시그니처를 활용하는 프레임워크인 ClientMorpher를 도입했습니다.
• 이중 클러스터링 전략 개발: 두 가지 상호 보완적인 방법을 개발했습니다.
  • ClientMorpher-C: 전문가 활성화 프로필을 사용하여 클라이언트를 직접 클러스터링합니다.
  • ClientMorpher-E: 클라이언트 간 사용 시그니처를 기반으로 전문가를 클러스터링하고 이를 통해 클라이언트 그룹을 도출합니다.
• 실증적 검증: 병리적(클라이언트당 단일 작업) 및 디리클레(Dirichlet) 기반(다양한 왜곡) 이질적 분할 조건 하에서 Databricks Dolly-15K 데이터셋을 사용하여 종합적인 연구를 수행했습니다. 이 연구는 제안된 방법들을 Local Training, MoE-FedAvg 및 기타 베이스라인과 비교합니다.

4. 실험 결과

실험은 네 가지 작업(분류(CLS), 폐쇄형 질의응답(QA), 정보 추출(IE), 요약(SUM))이 포함된 Dolly-15K 데이터셋을 사용하여 수행되었습니다.

• 병리적 분할 (높은 이질성):

  • ClientMorpher-E가 가장 높은 평균 ROUGE-L 점수(0.4986)를 달성하여 Local Training (0.4824) 및 MoE-FedAvg (0.4927)를 능가했습니다.
  • 이 모델은 분류, QA, 요약에서 우수한 성능을 보였으며, ClientMorpher-C는 정보 추출에서 뛰어난 성능을 보였습니다.
  • 모든 연합 학습 방법은 동일한 LoRA 파라미터(모델의 0.0878%)를 교환하였으며, 이는 이 성과가 통신량 증가가 아닌 더 나은 협업 덕분임을 확인시켜 줍니다.

• 디리클레 분할 (가변적 이질성):

  • 매우 이질적인 조건($\alpha=0.1$)에서도 두 ClientMorpher 변형 모두 Local Training과 MoE-FedAvg보다 우수한 성능을 보였습니다.
  • 분포가 더 균형 잡힐수록($\alpha=1.0$), 두 변형 간의 성능 격차는 좁아졌습니다.
  • ClientMorpher-E는 일반적으로 가장 높은 전체 평균(0.4936)과 분류 및 IE에서의 더 강력한 성능을 달성한 반면, ClientMorpher-C는 요약 작업에서 약간 더 나은 성능을 유지했습니다.

5. 의의 및 주장

본 논문은 ClientMorpher가 희소 MoE LLM의 개인화된 연합 지시 미세 조정을 위한 효과적이고 확장 가능한 접근 방식을 제공한다고 주장합니다. 본 연구의 의의는 다음과 같습니다.

• 라우팅을 신호로 활용: 라우팅을 단순히 최적화 중의 개인화를 위한 메커니즘으로 취급하는 것이 아니라, 클라이언트 유사성을 식별하고 협업을 가이드하기 위한 구조적 신호로 명시적으로 사용한 첫 번째 사례입니다.
• 상호 보완적 관점: 클라이언트 중심 클러스터링과 전문가 중심 클러스터링이 작업 유사성의 상호 보완적인 형태를 포착함을 입증했습니다. 클라이언트 중심 클러스터링은 극단적인 이질성 하에서 특히 효과적이며, 전문가 중심 클러스터링은 분포가 덜 왜곡될 때 견고함을 제공합니다.
• 효율성: 본 프레임워크는 희소 MoE 아키텍처의 효율성 이점을 유지하면서, 기존의 연합 평균 및 로컬 학습에 비해 개인화된 성능을 개선하면서도 통신 비용을 증가시키지 않습니다.