[논문 리뷰] Compare Where It Matters: Using Layer-Wise Regularization To Improve Federated Learning on Heterogeneous Data

Ha Min Son, Moon Hyun Kim외 1명|arXiv (Cornell University)|2021. 12. 01.

Privacy-Preserving Technologies in Data참고 문헌 20인용 수 3

한 줄 요약

FedCKA는 비독립 동일 분포(non-IID) 데이터에서 성능을 향상시키기 위해 중심 커널 일치(CKA)를 사용하여 자연스럽게 유사한 레이어—특히 첫 두 레이어—만을 선택적으로 정규화하는 계층별 정규화 프레임워크를 제안한다. 이 방법은 ResNet-50와 같은 깊은 모델에서도 FedAvg와 유사한 학습 효율성을 유지하면서 CIFAR-10에서 최대 97.58%의 SOTA 정확도를 달성한다.

ABSTRACT

Federated Learning is a widely adopted method to train neural networks over distributed data. One main limitation is the performance degradation that occurs when data is heterogeneously distributed. While many works have attempted to address this problem, these methods under-perform because they are founded on a limited understanding of neural networks. In this work, we verify that only certain important layers in a neural network require regularization for effective training. We additionally verify that Centered Kernel Alignment (CKA) most accurately calculates similarity between layers of neural networks trained on different data. By applying CKA-based regularization to important layers during training, we significantly improve performance in heterogeneous settings. We present FedCKA: a simple framework that out-performs previous state-of-the-art methods on various deep learning tasks while also improving efficiency and scalability.

연구 동기 및 목표

• 비독립 동일 분포 데이터 분포 하에서 연합 학습의 성능 저하 문제를 해결하기 위해.
• 이종 환경에서 정규화가 가장 중요한 신경망의 레이어를 특정하기 위해.
• 모든 레이어에 대한 정규화를 피하여 학습 효율성과 확장성을 향상시키기 위해.
• 비독립 동일 분포 데이터에서 학습된 모델 간의 표현 유사도를 측정하는 데 CKA가 다른 유사도 측정 지표(예: l2, 코사인)와 비교해 어떻게 성능을 내는지 평가하기 위해.
• 기존 SOTA 방법들을 능가하는 확장성 있고 높은 성능을 보이는 정규화 프레임워크를 개발하기 위해.

제안 방법

• 성능에 가장 큰 영향을 미치며 자연스럽게 유사한 첫 두 레이어에만 계층별 정규화를 적용한다.
• 로컬 모델 간 활성화 표현 간의 유사도를 측정하기 위해 중심 커널 일치(CKA)를 사용한다.
• 로컬 학습 중 첫 두 레이어의 표현을 정렬하기 위해 CKA 기반 정규화 항 ℓcka를 도입한다.
• 레이어 수를 늘리지 않고도 CKA 유사도를 효율적으로 계산하기 위해 클라이언트, 글로벌, 타겟 모델의 세 단계 전방향 전파를 활용한다.
• 깊은 아키텍처 전반에 걸쳐 각 레이어의 연산을 피하기 위해 핵심 레이어에만 정규화를 집중함으로써 계산 비용을 최적화한다.
• 기존의 FedAvg 파이프라인에 최소한의 수정으로 통합 가능한 플러그인 프레임워크로 FedCKA를 설계한다.

실험 결과

연구 질문

• RQ1 신경망의 어느 레이어가 재초기화에 가장 민감하며, 따라서 연합 학습에서 정규화가 가장 중요한가?
• RQ2 비독립 동일 분포 데이터에서 학습된 모델 간의 표현 유사도를 측정할 때 CKA는 l2, 코사인 등 다른 유사도 측정 지표와 비교해 어떻게 성능을 내는가?
• RQ3 자연스럽게 유사한 레이어만 선택적으로 정규화하는 것이 학습 효율성을 떨어뜨리지 않고 성능 향상에 기여할 수 있는가?
• RQ4 특히 모든 레이어를 정규화하는 방법과 비교해, 레이어 깊이가 증가함에 따라 FedCKA는 어떻게 확장되는가?
• RQ5 CKA 기반 정규화를 사용할 경우, CIFAR-10 및 Tiny ImageNet과 같은 표준 벤치마크에서 더 나은 수렴성과 정확도를 달성할 수 있는가?

주요 결과

• α = 5.0일 때 FedCKA는 CIFAR-10에서 97.58%의 정확도를 달성하여 FedAvg(54.82%)와 다른 SOTA 방법들을 크게 앞서간다.
• CKA를 정규화에 사용할 경우 가장 높은 성능를 기록했으며, 커널 CKA는 선형 CKA를 略로 뛰어넘지만 계산 비용이 더 높다.
• Tiny ImageNet에서 ResNet-50를 사용할 경우, FedCKA는 FedAvg와 유사한 학습 시간(750.97s)을 유지하지만, FedProx, SCAFFOLD, MOON는 지수적으로 증가하는 경향을 보였다.
• 첫 두 레이어만 정규화함으로써 학습 오버헤드를 줄이고 깊은 모델에서의 확장성 향상이 두드러진다.
• 표현 정렬의 정확도를 높이기 위해 CKA를 유사도 측정 지표로 사용함으로써, 진짜로 유사하지 않은 업데이트에만 페널티를 주는 효과를 얻었다.
• 자연스럽게 유사한 레이어에만 정규화를 집중하는 것이 모든 레이어에 균일하게 정규화하는 것보다 더 효과적임을 FedCKA가 입증했다.