[논문 리뷰] Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning

Liangqiong Qu, Yuyin Zhou외 6명|arXiv (Cornell University)|2021. 06. 10.

Privacy-Preserving Technologies in Data참고 문헌 61인용 수 12

한 줄 요약

이 논문은 비대칭 데이터에서 피어드 학습(federated learning)의 강력한 아키텍처 대안으로 비전 트랜스포머(Vision Transformers, ViTs)를 제안하며, ViTs가 치명적인 잊음(catastrophic forgetting)을 크게 감소시키고 수렴 속도를 가속화함을 입증한다. 주요 기여는 ViT-FL이 다양한 벤치마크와 데이터 분할에서 CNN 기반의 FL을 능가함을 실증한 것이다. 특히 데이터 비균형도가 높은 상황에서 두드러진 성능을 보인다.

ABSTRACT

Federated learning is an emerging research paradigm enabling collaborative training of machine learning models among different organizations while keeping data private at each institution. Despite recent progress, there remain fundamental challenges such as the lack of convergence and the potential for catastrophic forgetting across real-world heterogeneous devices. In this paper, we demonstrate that self-attention-based architectures (e.g., Transformers) are more robust to distribution shifts and hence improve federated learning over heterogeneous data. Concretely, we conduct the first rigorous empirical investigation of different neural architectures across a range of federated algorithms, real-world benchmarks, and heterogeneous data splits. Our experiments show that simply replacing convolutional networks with Transformers can greatly reduce catastrophic forgetting of previous devices, accelerate convergence, and reach a better global model, especially when dealing with heterogeneous data. We release our code and pretrained models at https://github.com/Liangqiong/ViT-FL-main to encourage future exploration in robust architectures as an alternative to current research efforts on the optimization front.

연구 동기 및 목표

• 기기 간 데이터 비균형으로 인한 비수렴 및 치명적인 잊음 문제를 해결하기 위해.
• 특히 비전 트랜스포머 같은 아키텍처 선택이 피어드 학습에서 분포 이탈에 대한 내성을 향상시킬 수 있는지 조사하기 위해.
• 실제 데이터셋과 비균형 데이터 분할을 포함한 다양한 피어드 학습 알고리즘에서 ViT와 CNN을 체계적으로 비교하는 실증 벤치마크를 제공하기 위해.
• 추가 최적화 히우리스틱이나 하이퍼파라미터 튜닝 없이도 아키텍처 개선이 즉각적인 성능 향상을 가져올 수 있음을 입증하기 위해.
• 기존 최적화 기반 방법들과 수직적(orthogonal)이면서도 미래의 피어드 학습 연구에 자연스러운 출발점이 되는 비전 트랜스포머를 정립하기 위해.

제안 방법

• 피어드 학습 알고리즘( FedAVG(병렬) 및 CWT(순차))을 기반으로 비전 트랜스포머(ViTs)와 CNNs(ResNets, EfficientNets)의 대규모 실증 벤치마크를 수행한다.
• 이미지 분류 작업에서 표준 ViT 구현을 사용하며, CIFAR-10 및 Retina 데이터셋에 대해 일관된 데이터 증강 및 정규화 전략을 적용한다.
• 실제 비균형도를 시뮬레이션하기 위해 레이블 분포의 비대칭 정도가 증가하는 세 가지 데이터 분할 전략(Split-1에서 Split-3까지)을 평가한다.
• UMAP 시각화를 통해 ViT와 CNN 모델이 학습한 특징 임베딩을 비교하고, 군집 간 분리도와 표현의 안정성 분석을 수행한다.
• 기존 최적화 기반 피어드 학습 방법(FedProx, FedAVG-Share)과 ViT를 결합하여 호환성과 성능 향상 여부를 테스트한다.
• 사전학습 및 학습 하이퍼파라미터(학습률, 기울기 클리핑, 로컬 에포크 수 등)가 ViT-FL 성능에 미치는 영향을 조사한다.

실험 결과

연구 질문

• RQ1 비대칭 데이터에서 비전 트랜스포머가 CNN보다 치명적인 잊음을 줄일 수 있는가?
• RQ2 다양한 피어드 학습 알고리즘(FedAVG 대비 CWT)과 데이터 분할 전략에서 ViT-FL의 성능은 CNN 기반 FL과 비교해 어떻게 다른가?
• RQ3 비균형 데이터 설정에서 아키텍처 선택(ViT 대비 CNN)이 수렴 속도와 최종 모델 정확도에 얼마나 큰 영향을 미치는가?
• RQ4 대규모 데이터셋이 가용하지 않을 경우 사전학습이 ViT-FL 성능에 상당한 기여를 하는가?
• RQ5 ViT-FL은 기존 최적화 기반 피어드 학습 방법과 효과적으로 통합되어 비균형 데이터에서 성능을 추가로 향상시킬 수 있는가?

주요 결과

• 비대칭 데이터 분할(Split-3)에서 비전 트랜스포머(ViT)는 ResNet(50)-FedAVG 대비 훨씬 높은 테스트 정확도를 기록하며, 사전학습을 적용한 경우 96.40%의 정확도를 달성한 반면, ResNet(50)-FedAVG는 59.68%에 머물렀다.
• ViT-FL은 치명적인 잊음을 감소시킨다: UMAP 시각화 결과, ViT(S)-FedAVG가 학습한 특징는 ResNet(50)-FedAVG보다 더 명확한 클래스 간 분리도를 보였다.
• ViT-FL은 수렴 속도를 가속화하고, 특히 심각한 데이터 비균형 상황에서 더 나은 글로벌 모델에 도달한다. 데이터 비대칭도가 증가할수록 성능 향상 폭이 커진다.
• 사전학습은 ViT-FL에 매우 중요하다: 대규모 사전학습 없이 초기화된 상태에서 Swin(T)-FedAVG는 Split-3에서 64.50%의 정확도를 기록했으며, 이는 ResNet(50)-FedAVG(59.68%)를 초월하는 성능이었다.
• ViT-FL은 최적화 기반 방법과 수직적이다: ViT를 FedProx나 FedAVG-Share와 결합하면 비균형 클라이언트에서 성능 향상이 더욱 뚜렷하게 나타났다.
• 최적의 성능를 얻기 위해서는 고도로 비균형 설정에서는 작은 로컬 에포크 수(E ≤ 5 for ViT-FedAVG, E = 1 for ViT-CWT)를 사용하는 것이 바람직하며, 균형 잡힌 상황에서는 더 큰 E를 사용해 통신 빈도를 줄일 수 있다.