DeepFusion: Accelerating MoE Training via Federated Knowledge Distillation from Heterogeneous Edge Devices

1. 문제 정의 (Problem Statement)

• MoE 기반 LLM 의 데이터 요구량: Qwen-MoE, DeepSeek-MoE 와 같은 최신 혼합 전문가 (Mixture-of-Experts, MoE) 기반 대규모 언어 모델 (LLM) 은 방대하고 다양한 학습 데이터를 필요로 합니다. 그러나 공개된 고품질 텍스트 데이터는 2026~2032 년 사이에 고갈될 것으로 예상되어 데이터 부족 문제가 심각해지고 있습니다.
• 연방 학습 (FL) 의 한계: 데이터 프라이버시를 보호하면서 엣지 디바이스의 데이터를 활용하기 위해 연방 학습이 제안되었으나, 기존 방식 (예: FedAvg) 은 모든 디바이스가 글로벌 모델의 전체 사본을 호스팅해야 하므로 리소스 제약이 있는 엣지 디바이스에는 비현실적입니다.
• 기존 Federated MoE 연구의 문제점: 최근 연구들은 엣지 디바이스에 경량화된 로컬 전문가 모델을 배포하여 MoE 학습을 시도했으나, 여전히 MoE 의 복잡한 백본 아키텍처와 게이트 메커니즘을 유지해야 하므로 저사양 엣지 디바이스 (예: 저전력 임베디드 장치) 에서는 실행이 어렵거나 메모리 부족으로 인해 참여가 제한됩니다.
• 이질성 (Heterogeneity) 과 뷰 불일치 (View-mismatch): 엣지 디바이스는 서로 다른 하드웨어와 애플리케이션 요구사항에 따라 다양한 아키텍처의 LLM 을 사용합니다. 이러한 이질적인 로컬 모델 (Teacher) 과 글로벌 MoE 모델 (Student) 간의 아키텍처 차이로 인해 예측 관점 (predictive perspective) 이 달라지며, 기존의 직접적인 로그잇 (logit) 또는 특징 (feature) 정렬 방식은 효과적인 지식 증류가 불가능합니다.

2. 제안 방법론: DEEPFUSION

DEEPFUSION 은 이질적인 엣지 디바이스가 데이터 프라이버시를 유지하면서 MoE 학습에 참여할 수 있도록 설계된 확장 가능한 연방 지식 증류 프레임워크입니다.

핵심 구성 요소

1. 독립적인 온디바이스 LLM 구성 및 학습:

   • 각 엣지 디바이스는 자신의 하드웨어 제약과 애플리케이션 요구사항에 맞춰 독립적으로 경량 LLM 을 구성하고 학습합니다.
   • 학습이 완료된 LLM 은 로컬 지식의 저장소 역할을 하며, 중앙 서버로 업로드됩니다.
   • One-shot FL 설계: 여러 번의 통신 라운드가 아닌, 학습 완료 후 모델 하나만 서버로 전송하는 단일 통신 라운드 방식을 사용하여 통신 오버헤드를 극도로 줄입니다.

2. 로컬 지식 클러스터링 (Local Knowledge Clustering):

   • 서버는 수많은 엣지 디바이스로부터 업로드된 LLM 들을 메타 정보 (저차원 특징 임베딩) 를 기반으로 유사한 도메인 지식 (예: 특정 질병, 금융 분야) 을 가진 그룹으로 클러스터링합니다.
   • 각 클러스터 내에서 모델 가중치를 평균화하여 로컬 지식 프록시 모델 (Local-knowledge Proxy Model) 을 생성합니다. 이를 통해 서버의 계산 부하를 줄이고 확장성을 확보합니다.

3. 교차 아키텍처 지식 증류 및 VAA 모듈:

   • 문제: 프록시 모델 (Teacher) 과 MoE 의 베이스 모델 (Student) 은 아키텍처가 달라 특징 공간과 예측 관점이 다릅니다.
   • 해결책 (VAA - View-Aligned Attention):
     • 학생 모델 (MoE 베이스) 이 여러 단계 (representation stages) 의 특징을 통합하여 교사 모델과 유사한 예측 관점을 갖도록 하는 View-Aligned Attention 모듈을 도입했습니다.
     • 이 모듈은 멀티헤드 어텐션 메커니즘을 사용하여 학생 모델의 특징을 변형하고, 이를 교사 모델의 특징과 정렬 (Feature Matching) 합니다.
     • 이를 통해 아키텍처가 다른 모델 간에도 효과적인 지식 전달이 가능해집니다.

4. 글로벌 MoE 모델 튜닝:

   • 증류된 K 개의 MoE 베이스 모델을 글로벌 MoE 모델로 병합합니다.
   • 병합 규칙: 각 전문가 (Expert) 는 해당 베이스 모델의 FFN(Feed-Forward Network) 계층 파라미터를 그대로 가져옵니다. 임베딩, 셀프 어텐션, 게이트, 출력 계층은 모든 베이스 모델의 파라미터를 평균화하여 초기화합니다.
   • 파인튜닝: FFN 기반 전문가 계층은 고정 (Frozen) 하고, 나머지 계층만 공개 벤치마크 데이터를 사용하여 파인튜닝하여 글로벌 MoE 모델을 최적화합니다.

3. 주요 기여 (Key Contributions)

• 최초의 이질적 엣지 기반 Federated MoE 학습: 데이터 프라이버시를 유지하면서 리소스 제약이 있는 다양한 엣지 디바이스로부터 MoE 학습에 필요한 광범위한 데이터를 확보하는 최초의 확장 가능한 프레임워크를 제안했습니다.
• 뷰 정렬 어텐션 (VAA) 모듈: 이질적인 아키텍처 간의 예측 관점 불일치 문제를 해결하기 위해 개발된 새로운 모듈로, 교차 아키텍처 지식 증류의 효율성을 획기적으로 높였습니다.
• 실제 산업 수준의 검증: Qwen-MoE 와 DeepSeek-MoE 와 같은 산업 수준의 모델과 의료/금융 분야의 실제 데이터를 사용하여 광범위한 실험을 수행했습니다.

4. 실험 결과 (Results)

• 성능: DEEPFUSION 은 중앙 집중식 MoE 학습 (DeepSpeed 기준) 에 버금가는 성능을 달성했습니다.
  • Perplexity: 복잡한 오픈형 질문 응답 (Open-ended QA) 태스크에서 기존 연방 MoE 베이스라인 (FedJETS, FedKMT 등) 대비 최대 5.28% 의 토큰 퍼플렉시티 (Perplexity) 개선을 보였습니다.
  • 정확도: 의료 다지선다형 QA 및 금융 오픈형 QA 태스크에서 모든 시스템 규모 (16~128 개 디바이스) 에서 최상위 또는 2 위 수준의 정확도를 기록했습니다.
• 통신 효율성: 기존 방법론에 비해 통신 비용을 최대 71% 까지 절감했습니다. 이는 단일 통신 라운드와 경량 모델 업로드 전략 덕분입니다.
• 자원 효율성: FedJETS 와 비교했을 때, 엣지 디바이스의 메모리 사용량을 3.3 배~9.3 배 줄여 저사양 디바이스 (예: NVIDIA Jetson Nano, iPhone 16) 의 참여를 가능하게 했습니다.

5. 의의 및 결론 (Significance)

DEEPFUSION 은 MoE 기반 LLM 의 학습에 필요한 방대한 데이터 부족 문제를 해결하고, 데이터 프라이버시를 보호하면서도 다양한 엣지 디바이스의 지식을 효과적으로 통합하는 새로운 패러다임을 제시합니다. 특히, 이질적인 하드웨어와 모델 아키텍처를 가진 엣지 디바이스가 MoE 학습에 참여할 수 있는 장벽을 낮추었으며, VAA 모듈을 통해 교차 아키텍처 지식 증류의 핵심 난제를 해결했다는 점에서 학술적, 실용적 가치가 매우 높습니다. 이는 의료, 금융 등 민감한 데이터를 다루는 분야에서 대규모 AI 모델의 분산 학습을 실현 가능한 방향으로 이끌 것입니다.