2025-11-11T23:28:21.956833

Bridging Memory Gaps: Scaling Federated Learning for Heterogeneous Clients

Wu, Li, Tian et al.
Federated Learning (FL) enables multiple clients to collaboratively train a shared model while preserving data privacy. However, the high memory demand during model training severely limits the deployment of FL on resource-constrained clients. To this end, we propose \our, a scalable and inclusive FL framework designed to overcome memory limitations through sequential block-wise training. The core idea of \our is to partition the global model into blocks and train them sequentially, thereby reducing training memory requirements. To mitigate information loss during block-wise training, \our introduces a Curriculum Mentor that crafts curriculum-aware training objectives for each block to steer their learning process. Moreover, \our incorporates a Training Harmonizer that designs a parameter co-adaptation training scheme to coordinate block updates, effectively breaking inter-block information isolation. Extensive experiments on both simulation and hardware testbeds demonstrate that \our significantly improves model performance by up to 84.2\%, reduces peak memory usage by up to 50.4\%, and accelerates training by up to 1.9$\times$.
academic

메모리 격차 해소: 이질적 클라이언트를 위한 연합학습 확장

기본 정보

  • 논문 ID: 2408.10826
  • 제목: Bridging Memory Gaps: Scaling Federated Learning for Heterogeneous Clients
  • 저자: Yebo Wu, Jingguang Li, Chunlin Tian, KaHou Tam, Li Li, Chengzhong Xu (마카오 대학교)
  • 분류: cs.DC (분산 컴퓨팅)
  • 발표 시간: 2024년 8월 (arXiv v2: 2025년 10월)
  • 논문 링크: https://arxiv.org/abs/2408.10826v2

초록

연합학습(FL)은 여러 클라이언트가 데이터 프라이버시를 보호하면서 공유 모델을 협력하여 훈련할 수 있게 한다. 그러나 모델 훈련 과정의 높은 메모리 요구사항은 자원이 제한된 클라이언트에서의 FL 배포를 심각하게 제한한다. 이를 해결하기 위해 본 논문은 순차적 블록 훈련을 통해 메모리 제약을 극복하는 확장 가능하고 포용적인 FL 프레임워크인 SCALEFL을 제안한다. SCALEFL의 핵심 아이디어는 전역 모델을 블록으로 분할하고 순차적으로 훈련하여 훈련 메모리 요구사항을 줄이는 것이다. 블록 훈련에서의 정보 손실을 완화하기 위해 SCALEFL은 각 블록에 대한 커리큘럼 인식 훈련 목표를 수립하는 커리큘럼 멘토(Curriculum Mentor)를 도입한다. 또한 SCALEFL은 훈련 조화기(Training Harmonizer)를 통합하여 매개변수 협동 적응 훈련 방식을 설계하고, 블록 간 정보 격리를 효과적으로 제거한다.

연구 배경 및 동기

핵심 문제

  1. 메모리 벽 문제: 연합학습 훈련 과정에서 모든 중간 활성화, 모델 가중치 및 옵티마이저 상태를 메모리에 유지해야 하므로 높은 메모리 점유율이 발생한다. 예를 들어 ImageNet에서 ResNet34를 훈련하는 데 12GB 이상의 메모리가 소비되지만, 일반적인 모바일 디바이스의 RAM은 보통 4-12GB에 불과하다.
  2. 디바이스 이질성: 자원이 제한된 엣지 디바이스는 로컬 훈련에 참여할 수 없어 전역 모델에 대한 귀중한 데이터 기여가 차단된다.
  3. 기존 방법의 한계:
    • 모델 이질적 훈련: 지식 증류를 위해 고품질 공개 데이터셋이 필요하지만 FL에서는 획득이 어렵다
    • 부분 훈련: 너비 축소는 모델 아키텍처를 손상시키고, 깊이 축소는 최대 메모리 용량 클라이언트에 의해 제한된다

연구 동기

모델 아키텍처가 더 깊고 넓어져 더 높은 분석 능력을 얻으면서 메모리 문제가 더욱 악화된다. 본 논문은 훈련 메모리 요구사항을 크게 줄이면서 모델 성능을 유지할 수 있는 FL 프레임워크를 설계하는 것을 목표로 한다.

핵심 기여

  1. SCALEFL 프레임워크 제안: 순차적 블록 훈련을 통해 훈련 메모리 요구사항을 크게 줄여 자원이 제한된 디바이스가 효과적으로 참여할 수 있게 함
  2. 두 가지 핵심 구성 요소 설계: 커리큘럼 멘토와 훈련 조화기가 각 블록의 학습 행동을 협력하여 형성하여 일관된 구조화된 특성 학습을 촉진
  3. 포괄적 실험 검증: 여러 벤치마크 데이터셋에서 SCALEFL의 효과성과 견고성을 입증
  4. 이론적 분석: 수렴성 분석을 제공하여 방법의 이론적 신뢰성을 입증

방법 상세 설명

작업 정의

N개의 클라이언트를 포함하는 FL 시스템에서 각 클라이언트 n은 로컬 데이터셋 Dn을 소유한다. 목표는 각 클라이언트의 메모리 제약을 만족하면서 전역 모델 Θ를 훈련하는 것이다.

순차적 블록 훈련 패러다임

기본 흐름:

  1. 모델 구성: 서버가 현재 훈련 단계 t에 대한 부분 모델 Θg,t = θ1,F, θ2,F, ..., θt, θOp을 구성
  2. 로컬 훈련: 블록 θt와 출력 모듈 θOp만 업데이트
  3. 모델 집계: 가중 평균을 사용하여 매개변수 업데이트 집계
  4. 진행 상황 평가: 블록 θt의 훈련 진행 상황을 모니터링하고 수렴 판단
  5. 모델 성장: 수렴된 블록을 고정하고 새 블록 도입

핵심 기술 구성 요소

1. 커리큘럼 멘토 (Curriculum Mentor)

문제 분석: 정보 병목 이론을 기반으로 순차적 블록 훈련이 심각한 정보 손실을 초래함을 발견했다. nHSIC 평면 동적 분석을 통해 SBT가 첫 번째 블록 훈련 후 대량의 입력 정보를 손실하여 후속 블록이 핵심 특성을 추출할 수 없음을 보여준다.

해결 방안: 커리큘럼 인식 훈련 목표 설계

L_θt = L_CE - λt · nHSIC(X;Zt) - γt · nHSIC(Y;Zt)

여기서:

  • L_CE는 교차 엔트로피 손실
  • nHSIC(X;Zt)는 입력 정보 보존을 측정
  • nHSIC(Y;Zt)는 작업 관련성을 측정
  • λt와 γt는 훈련 단계에 따라 동적으로 조정

전략: 초기 단계에서는 높은 λt와 낮은 γt로 입력 정보 보존을 강조하고, 후기 단계에서는 λt를 점진적으로 낮추고 γt를 증가시켜 작업 특정 특성 추출로 전환한다.

2. 훈련 조화기 (Training Harmonizer)

문제 식별:

  • 전방향 정보 흐름 제한: 다운스트림 블록은 선행 블록이 수렴한 후에만 훈련 시작
  • 역방향 정보 흐름 제한: 그래디언트가 블록 내로 제한되어 그래디언트 격리 초래

매개변수 협동 적응 방안:

  1. 동적 모델 성장: 매 라운드마다 각 블록의 학습 과정을 동적으로 편성하여 다운스트림 블록이 선행 블록 업데이트에 실시간으로 적응
  2. 동시 훈련 전략: 현재 블록과 선행 블록의 후속 레이어를 동시에 훈련하여 그래디언트 흐름 촉진

업데이트 공식:

θ^(k+1)_(n,t) + L^(k+1)_(n,t-1) ← (θ^k_(n,t) + L^k_(n,t-1)) - η · ∂L^k_(n,t)/∂(θ^k_(n,t) + L^k_(n,t-1))

완전한 훈련 목표

L2 정규화를 결합하여 데이터 이질성 처리:

L^r_t = L_θt + (μ/2)||θ^r_t - θ^(r-1)_t||^2_2

실험 설정

데이터셋

  • CIFAR10/CIFAR100: 고전적 이미지 분류 데이터셋
  • CINIC10: CIFAR10 확장 버전
  • Mini-ImageNet: 소규모 ImageNet
  • FEMNIST: 대규모 FL 데이터셋(805,263개 이미지)

모델 아키텍처

  • ResNet18/ResNet34: 깊이 있는 잔차 네트워크
  • VGG11 BN: 고전적 합성곱 네트워크
  • SqueezeNet: 경량 네트워크
  • Vision Transformer (ViT): Transformer 아키텍처

실험 환경

  • 혼합 설정: 시뮬레이션 및 실제 디바이스 테스트베드
  • 디바이스 구성: 100개의 이질적 모바일 디바이스, 매 라운드마다 10% 무작위 선택
  • 메모리 예산: 100-1000MB 무작위 할당
  • 옵티마이저: SGD, 가중치 감소 5e-4, 로컬 에포크=5

비교 방법

  1. AllSmall: 가장 약한 디바이스의 메모리를 기반으로 전역 모델 축소
  2. ExclusiveFL: 메모리가 충분한 디바이스만 참여 허용
  3. DepthFL: 깊이 축소로 이질적 디바이스 적응
  4. HeteroFL: 정적 채널 축소
  5. FedRolex: 동적 너비 축소
  6. SmartFreeze: 간단한 순차적 블록 훈련
  7. ProFL: 분해식 순차 훈련

실험 결과

주요 결과

비IID 시나리오에서의 성능:

방법CIFAR10 (ResNet18/VGG11/SqueezeNet)참여율
AllSmall69.5%/75.1%/49.6%100%/100%/100%
ExclusiveFL76.8%/79.3%/40.6%18%/22%/11%
SCALEFL80.4%/87.6%/58.0%100%/100%/100%

주요 발견:

  1. 현저한 성능 향상: AllSmall 대비 10.9%, 12.5%, 8.4% 향상
  2. 전체 디바이스 참여: 100% 디바이스 참여율 달성, ExclusiveFL은 18-22%만 가능
  3. 메모리 효율성: 피크 메모리 사용량 최대 50.4% 감소
  4. 훈련 가속화: 수렴 속도 1.9배 향상

확장성 분석

다양한 메모리 제약 조건에서의 견고성:

  • ResNet34 시나리오에서 ExclusiveFL은 완전히 불가능(0% 참여율)
  • SCALEFL은 다른 방법 대비 최대 27.4% 향상

대규모 데이터셋:

  • FEMNIST 데이터셋에서 FedAvg 대비 3% 정확도 향상
  • 120-500개 디바이스 규모 지원

Transformer 호환성:

  • ViT 모델에서 이론적 기준선보다 2% 낮지만, 이론적 기준선은 실제로는 불가능
  • 실용적 성능은 우수함

하드웨어 평가

메모리 효율성:

  • Jetson TX2에서의 테스트 결과 피크 메모리 사용량 50.4% 감소
  • 단일 라운드 훈련 시간 1.84-2.31배 감소

훈련 효율성:

  • 엔드투엔드 훈련 대비 단일 라운드 훈련 시간 현저히 감소
  • ViT에서 1.9배 가속화 달성

절제 실험

구성 요소 기여도 분석:

  • 커리큘럼 멘토 제거: CIFAR100 IID 시나리오에서 정확도 1.2% 하락
  • 훈련 조화기 제거: 정확도 9.0% 현저히 하락
  • 두 구성 요소의 협력이 성능에 매우 중요함

관련 연구

자원 제한 FL

  1. 모델 이질적 훈련: FedMD 등의 방법은 지식 증류를 위해 공개 데이터셋 필요
  2. 부분 훈련: HeteroFL, FedRolex는 너비 축소, DepthFL, InclusiveFL은 깊이 축소 사용

블록 훈련

  1. ProgFed: 새 블록을 점진적으로 도입하지만 여전히 엔드투엔드 훈련 필요
  2. SmartFreeze: 각 블록을 순차적으로 훈련하지만 정보 손실 문제 무시
  3. ProFL: 축소 및 성장 단계로 분해하지만 핵심 과제 미해결

이론적 분석

수렴성 증명

논문은 표준 가정(평활성, 유계 그래디언트) 하에서 SCALEFL의 수렴성 분석을 제공한다:

(1/R) Σ E[||∇L^r_t(Θ^r_(g,t))||^2] ≤ Ψ/√R

즉, 평균 그래디언트 범수가 0으로 수렴하고 모델이 안정점으로 수렴한다.

결론 및 논의

주요 결론

  1. SCALEFL은 FL의 메모리 벽 문제를 성공적으로 해결하여 자원이 제한된 디바이스가 훈련에 참여할 수 있게 함
  2. 커리큘럼 멘토와 훈련 조화기는 순차적 블록 훈련의 핵심 과제를 효과적으로 완화
  3. 여러 데이터셋과 모델에서 현저한 성능 향상과 메모리 절감 달성

한계

  1. 블록 분할 전략: 논문은 최적 블록 분할 방법에 대해 깊이 있게 논의하지 않음
  2. 통신 오버헤드: 메모리 사용을 줄이지만 통신 라운드 수가 증가할 수 있음
  3. 초매개변수 민감성: λt와 γt의 설정은 신중한 조정 필요

향후 방향

  1. 자적응 블록 분할 전략
  2. 다른 FL 최적화 기법과의 결합
  3. 더 큰 규모 실제 배포에서의 검증

심층 평가

장점

  1. 문제의 중요성: FL 실제 배포의 핵심 병목을 해결
  2. 방법의 창의성: 커리큘럼 인식 훈련 목표와 매개변수 협동 적응 방안이 독창적
  3. 이론적 기초: 정보 병목 이론 기반 분석이 견고한 이론적 지원 제공
  4. 실험의 포괄성: 다양한 모델, 데이터셋 및 실제 하드웨어 테스트 포함
  5. 실용적 가치: 현저한 메모리 절감과 성능 향상이 실제 응용 가치 보유

부족한 점

  1. 복잡성: 도입된 두 구성 요소가 시스템 복잡성 증가
  2. 초매개변수 조정: λt, γt 등 매개변수가 다양한 시나리오에 맞게 조정 필요
  3. 통신 분석: 통신 오버헤드에 대한 상세 분석 부족
  4. 수렴 속도: 단일 라운드는 빠르지만 총 수렴 라운드 수가 증가할 수 있음

영향력

  1. 학술 기여: 자원 제한 FL에 새로운 해결 사고 제공
  2. 실용적 가치: 모바일 디바이스 등 자원 제한 환경에 실제 배포 가능
  3. 재현성: 상세한 실험 설정 및 매개변수 구성 제공

적용 시나리오

  1. 모바일 디바이스 FL: 스마트폰, IoT 디바이스 등 메모리 제한 시나리오
  2. 엣지 컴퓨팅: 엣지 서버 자원이 제한된 환경
  3. 대규모 모델 훈련: 대형 모델 훈련이 필요하지만 디바이스 자원 부족 시나리오

참고문헌

논문은 FL 분야의 중요 연구(FedAvg, HeteroFL, FedRolex 등 고전적 방법)와 정보 병목 이론, HSIC 등 이론적 기초를 인용하여 문헌 인용이 비교적 포괄적이고 권위 있다.

종합 평가: 이는 실제 배포의 핵심 문제에 대해 창의적 해결 방안을 제시하는 고품질의 연합학습 논문이다. 방법 설계가 합리적이고 실험 검증이 충분하며 이론 분석이 완전하여 중요한 학술 가치와 실용 가치를 보유하고 있다.