MoE Enhanced Federated Learning for Spatiotemporal Prediction

기술 요약: MoE-FedTP를 이용한 시공간 예측

문제 정의
교통 예측은 지능형 교통 시스템에 있어 매우 중요하지만, 많은 도시들이 제한된 센서 배치와 불균형한 도시 개발로 인해 데이터 부족 문제를 겪고 있습니다. 도시 간 지식 전이(cross-city knowledge transfer)가 해결책을 제시할 수 있지만, 기존 방식들은 두 가지 주요 장애물에 직면해 있습니다:

1. 개인정보 보호 문제: 중앙 집중식 방법은 로우 데이터(raw data)의 공유를 요구하며, 이는 심각한 개인정보 보호 위험을 초end합니다.
2. 시공간적 이질성(Spatiotemporal Heterogeneity): 각 도시는 고유한 교통 진화 패턴, 기능적 구역 분포, 이동 행태를 보입니다. 기존의 연합 학습(FL) 및 개인화 연합 학습(PFL) 방식은 대개 통일된 예측 전략에 의존하며, 이는 내재된 차이점을 수용하는 데 어려움을 겪어 모델 드리프트(model drift)와 데이터가 부족한 대상 도시에서의 하위 최적 성능을 초래합니다.

방법론: MoE-FedTP
저자들은 경량화된 Mixture-of-Experts (MoE) 네트워크를 기반으로 한 개인화된 연합 교차 도시 시공간 예측 프레임워크인 MoE-FedTP를 제안합니다. 이 프레임워크는 두 단계로 작동합니다:

1. 소스 도시 학습 (로컬 트레이닝):

   • 각 소스 도시는 독립적인 클라이언트로 동작하며 로컬 시공간 신경망을 학습합니다.
   • 네트워크 아키텍처는 공간적 의존성과 시간적 역학을 모델링하기 위해 그래프 합성곱 네트워크(GCN), 그래프 어텐션 네트워크(GAT) 및 어텐션 모듈을 통합합니다.
   • 특징(features)은 GRU(Gated Recurrent Units)를 통해 시간적 시퀀스로 추출 및 풀링되며, GCN/GAT 레이어의 공간 임베딩과 융합됩니다.
   • 다층 퍼셉트론(MLP)이 최종 예측을 생성하며, 평균 제곱 오차(MSE)를 사용하여 최적화됩니다.

2. 대상 도시 적응 (연합 개인화):

   • 전문가 구성 (Expert Construction): 사전 학습된 여러 소스 도시의 예측 헤드(prediction heads)로부터 핵심 파라미터들을 공유하고 "전문가 네트워크(expert networks)"로 유지합니다. 각 전문가는 해당 소스 도시의 고유한 시공간 패턴을 인코딩합니다.
   • 게이팅 메커니즘 (Gating Mechanism): 게이팅 네트워크는 대상 도시의 입력 특징을 바탕으로 각 전문가에 대한 중요도 가중치를 동적으로 계산합니다.
   • Top-k 선택: 계산 효율성을 보장하기 위해, 시스템은 게이팅 점수가 가장 높은 상위 $k$개의 전문가만을 활성화합니다.
   • 예측 융합 (Prediction Fusion): 최종 예측은 선택된 전문가들의 출력값에 대한 가중 합계입니다. 대상 도시는 자신의 제한된 로컬 데이터를 사용하여 자체 예측 손실을 최소화함으로써, 외부 지식을 활용하는 동시에 개인정보를 보호합니다.

주요 기여

• 프레임워크 설계: 교차 도시 시공간 이질성을 구체적으로 해결하기 위해 경량 MoE 구조를 활용하는 개인화된 연합 프레임워크인 MoE-FedTP를 도입했습니다.
• 파라미터 공유 메커니즘: 소스 도시의 시공간 패턴을 전문가 네트워크로 요약하여 대상 도시로 동적으로 전달함으로써, 데이터가 부족한 시나리오에서의 성능을 향하는 새로운 교차 도시 파라미터 공유 전략을 제시했습니다.
• 실증적 검증: 제안된 방법이 정확도와 일반화 성능 모두에서 최신 베이스라인 모델들을 능가함을 광범에한 실험을 통해 입증했습니다.

실험 결과
본 방법은 네 가지 실제 데이터셋인 PEMS-BAY, METR-LA, DiDi-Chengdu, DiDi-Shenzhen에서 평가되었습니다. 실험 설정에서는 세 개의 다른 데이터셋을 데이터가 풍부한 소스 도시로 사용하여 사전 학습을 진행하는 동안, 단 3일간의 데이터만을 사용하는 데이터 부족 대상 도시 상황을 시뮬레이션했습니다.

• 성능: MoE-FedTP는 STGCN(시공간 GNN), ST-GFSL(Few-shot learning), pFedCTP(적응형 보간 PFL)를 포함한 베이스라인 모델들을 일관되게 능가했습니다.
  • METR-LA에서, 이 모델은 베이스라인 대비 MAE 6.23%, RMSE 5.20%의 개선을 달-성했습니다.
  • DiDi-Shenzhen에서는 평균 **4.68%**의 성능 향상이 관찰되었습니다.
• 절제 연구 (Ablation Studies):
  • MoE 모듈을 제거했을 때 상당한 성능 저하(예: DiDi-Chengdu에서 MAE 7.4% 증가)가 발생하여, 전문가 네트워크 구조의 필요성을 확인했습니다.
  • Top-k 민감도: 활성화되는 최적의 전문가 수($k$)는 데이터셋마다 달랐습니다. 복잡하고 다양한 패턴을 가진 도시(PEMS-BAY, METR-LA)는 $k=2$일 때 가장 좋은 성능을 보였으나, 독특한 패턴을 가진 도시(DiDi-Chengdu, DiDi-Shenzhen)는 $k=1$일 때 최적의 결과를 얻었습니다. 이는 너무 많은 전문가를 활성화하는 것이 특정 유형의 도시에는 노이즈를 유발할 수 있음을 시사합니다.

의의 및 주장
본 논문은 MoE-FedTP가 교차 도시 교통 예측을 위한 효율적이고 개인정보를 보호하는 솔루션을 제공한다고 주장합니다. 연합 학습의 개인정보 보호 이점과 MoE의 세밀한 모델링 능력을 결합함으로써, 이 프레임워크는 시공간적 이질성 문제를 성공적으로 완화합니다. 이를 통해 데이터가 부족한 도시들은 데이터 보안을 위협하거나 로컬 교통 역학을 무시하지 않으면서도 데이터가 풍부한 대응 도시로부터 지식을 활용할 수 있습니다. 저자들은 이 연구가 더 넓은 도시 컴퓨팅 및 시공간 예측 작업 분야에서 연합 전이 학습의 유망한 방향이라고 제안합니다.

한계 및 향후 과제
저자들은 현재의 부분적 파라미터 공유 메커니즘이 극단적인 이질성을 가진 도시나 소스 도시의 수가 제한적인 경우에 불충분할 수 있음을 인정합니다. 또한, 고정된 top-$k$ 선택 전략은 동적인 교통 시나리오에 최적으로 적응하지 못할 수 있습니다. 향후 연구는 더욱 강력한 견고성과 확장성을 확보하기 위해 적응형 전문가 선택 전략(예: 강화 학습 또는 불확실성 인지 게이팅)을 탐색하고 멀티모달 도시 데이터를 통합하는 것을 목표로 합니다.