[논문 리뷰] Towards Understanding and Mitigating Dimensional Collapse in Heterogeneous Federated Learning

Yujun Shi, Jian Liang외 3명|arXiv (Cornell University)|2022. 10. 01.

Privacy-Preserving Technologies in Data인용 수 12

한 줄 요약

이 논문은 데이터 이질성이 표현을 낮은 차원의 다양체로 악화시키는 이질적 피어드 페더레이티드 러닝에서의 차원 축소 현상을 규명하고 해결한다. 저자들은 지역 학습 중 표현 차원 간 상관관계를 최소화하는 경량이며 플러그-앤플레이(Plug-and-Play) 정규화 방법인 FedDecorr를 제안하여, 강한 데이터 이질성 하에서 CIFAR10, CIFAR100 및 TinyImageNet에서 모델 성능을 크게 향상시킨다.

ABSTRACT

Federated learning aims to train models collaboratively across different clients without the sharing of data for privacy considerations. However, one major challenge for this learning paradigm is the {\em data heterogeneity} problem, which refers to the discrepancies between the local data distributions among various clients. To tackle this problem, we first study how data heterogeneity affects the representations of the globally aggregated models. Interestingly, we find that heterogeneous data results in the global model suffering from severe {\em dimensional collapse}, in which representations tend to reside in a lower-dimensional space instead of the ambient space. Moreover, we observe a similar phenomenon on models locally trained on each client and deduce that the dimensional collapse on the global model is inherited from local models. In addition, we theoretically analyze the gradient flow dynamics to shed light on how data heterogeneity result in dimensional collapse for local models. To remedy this problem caused by the data heterogeneity, we propose {\sc FedDecorr}, a novel method that can effectively mitigate dimensional collapse in federated learning. Specifically, {\sc FedDecorr} applies a regularization term during local training that encourages different dimensions of representations to be uncorrelated. {\sc FedDecorr}, which is implementation-friendly and computationally-efficient, yields consistent improvements over baselines on standard benchmark datasets. Code: https://github.com/bytedance/FedDecorr.

연구 동기 및 목표

• 데이터 이질성이 피어드 페더레이티드 러닝에서 표현 학습에 미치는 영향을 조사하는 것.
• 데이터 이질성 하에서 글로벌 및 로컬 모델의 주요 실패 모드로 차원 축소를 규명하는 것.
• 데이터 이질성과 낮은 랭크의 가중치 행렬 및 표현 축소를 연결하는 경량 유동 역학의 이론적 분석을 수행하는 것.
• 피어드 페더레이티드 러닝에서 차원 축소를 완화하기 위한 실용적이고 계산 효율적인 방법을 제안하는 것.

제안 방법

• 지역 학습 중 표현의 상관계수 행렬의 프로베니우스 노름을 최소화하는 정규화 항인 FedDecorr를 제안한다.
• 상관계수 행렬을 계산하기 전에 표현에 z-점수 정규화를 적용하여 학습 안정성을 높인다.
• 기존 피어드 페더레이티드 러닝 기반 모델에 최소한의 수정으로 통합할 수 있는 플러그인 손실 구성 요소로 FedDecorr를 통합한다.
• 정규화된 상관계수 행렬의 제곱 프로베니우스 노름을 정규화 손실로 사용하여 상관관계가 없는 표현을 장려한다.
• 주 작업 손실과 FedDecorr 정규화 손실의 조합으로 로컬 모델을 학습한 후, 표준 FedAvg를 통해 앙상블한다.
• FedDecorr가 계산적으로 효율적이며 기존 피어드 페더레이티드 러닝 프레임워크와 호환됨을 입증한다.

실험 결과

연구 질문

• RQ1 데이터 이질성은 피어드 페더레이티드 러닝 모델의 표현의 내재 차원에 어떤 영향을 미치는가?
• RQ2 글로벌 모델에서의 차원 축소는 이질적 데이터로 훈련된 로컬 모델들로부터 얼마나 상속되는가?
• RQ3 데이터 이질성 하에서 낮은 랭크의 가중치 행렬과 표현 축소를 유도하는 기저의 기울기 역학 메커니즘은 무엇인가?
• RQ4 간단하고 경량인 정규화 항이 계산 비용이나 통신 비용을 증가시키지 않으면서도 피어드 페더레이티드 러닝에서 차원 축소를 효과적으로 완화할 수 있는가?

주요 결과

• 더 강한 데이터 이질성은 표현 공분산 행렬의 특이값 중 0에 가까운 값이 많아지는 것으로 나타나 더 심각한 차원 축소를 초래한다.
• 차원 축소 현상은 글로벌 모델 뿐 아니라 지역적으로 훈련된 모델에서도 관찰되며, 이는 지역 데이터 분포에 뿌리내린 체계적인 문제임을 시사한다.
• 이론적 분석 결과, 데이터 이질성이 지역 모델 가중치를 낮은 랭크의 구조로 편향시키며, 이는 직접적으로 표현 축소를 유도한다.
• FedDecorr는 표현 상관관계를 효과적으로 감소시키고 차원 축소를 완화하여 다양한 벤치마크 데이터셋에서 일관된 성능 향상을 이룬다.
• 높은 데이터 이질성 또는 많은 수의 클라이언트가 있는 도전적인 환경에서 메모리나 계산 비용을 크게 증가시키지 않으면서도 더 뚜렷한 성능 향상이 이루어진다.
• FedDecorr는 극히 미미한 계산 오버헤드를 유발하며, 플러그-앤플레이 구성 요소로 어떤 피어드 페더레이티드 러닝 파이프라인에도 원활하게 통합될 수 있다.