[논문 리뷰] TCT: Convexifying Federated Learning using Bootstrapped Neural Tangent Kernels

Yaodong Yu, Alexander Wei외 3명|arXiv (Cornell University)|2022. 07. 13.

Privacy-Preserving Technologies in Data인용 수 5

한 줄 요약

이 논문은 TCT를 제안한다. TCT는 먼저 FedAvg를 사용해 연합 학습을 통해 특징을 추출한 후, 경험적 신경 장력 커널(eNTK)을 사용해 모델을 볼록화하고 그에 따라 기울기 보정 최적화를 적용하는 훈련-볼록화-다시 훈련 프레임워크이다. 이 방법은 데이터 이질성이 극도로 높은 상황에서도 CIFAR10에서 최대 37%의 정확도 향상을 달성하며, 중심화된 학습과 연합 학습 간의 성능 격차를 완전히 해소한다.

ABSTRACT

State-of-the-art federated learning methods can perform far worse than their centralized counterparts when clients have dissimilar data distributions. For neural networks, even when centralized SGD easily finds a solution that is simultaneously performant for all clients, current federated optimization methods fail to converge to a comparable solution. We show that this performance disparity can largely be attributed to optimization challenges presented by nonconvexity. Specifically, we find that the early layers of the network do learn useful features, but the final layers fail to make use of them. That is, federated optimization applied to this non-convex problem distorts the learning of the final layers. Leveraging this observation, we propose a Train-Convexify-Train (TCT) procedure to sidestep this issue: first, learn features using off-the-shelf methods (e.g., FedAvg); then, optimize a convexified problem obtained from the network's empirical neural tangent kernel approximation. Our technique yields accuracy improvements of up to +36% on FMNIST and +37% on CIFAR10 when clients have dissimilar data.

연구 동기 및 목표

• 데이터 이질성 하에서 비볼록 깊은 신경망에 대해 성능 저하 문제를 해결하기 위해.
• 기존의 연합 최적화 방법이 초기 네트워크 레이어에서 학습된 특징을 유용하게 활용하지 못하는 이유를 조사하기 위해.
• 비볼록성 하에서 최종 레이어에서 안정적이고 높은 정확도로 훈련을 유지하면서 학습된 특징을 보존하는 방법을 개발하기 위해.
• eNTK를 통한 볼록화가 이질적인 연합 환경에서 중심화된 학습 수준에 가까운 성능을 복원할 수 있는지 보여주기 위해.
• 제안된 TCT 프레임워크가 극도의 데이터 이질성 하에서도 하이퍼파ram터와 통신 효율성에 대해 강건한지 보여주기 위해.

제안 방법

• 첫 번째 단계로, 클라이언트 데이터로부터 유용한 특징을 연합적으로 추출하기 위해 FedAvg를 사용해 딥 신경망을 훈련한다.
• 두 번째 단계로, 훈련된 모델의 경험적 신경 장력 커널(eNTK) 근사를 계산하여 최종 레이어를 선형화한다.
• 세 번째 단계로, eNTK 근사된 특징을 기반으로 얻어진 볼록 문제를 해결하기 위해 기울기 보정 최적화(예: SCAFFOLD)를 적용한다.
• 두 번째 단계에서 수렴 속도와 통신 효율성을 향상시키기 위해 eNTK 특징을 정규화한다.
• 첫 번째 단계에서의 모델 체크포인트를 사용해 eNTK 기반 볼록 최적화를 초기화함으로써 특징 일관성을 확보한다.
• 최종 레이어를 고정된 eNTK 특징 위에서 엔드 투 엔드로 훈련함으로써 사전에 학습된 표현 위에 선형 헤드를 효과적으로 학습한다.

실험 결과

연구 질문

• RQ1 왜 현재의 연합 최적화 방법들은 데이터가 i.i.d.가 아니더라도 중심화된 학습 성능에 비해 성능이 떨어지는가?
• RQ2 신경망의 초기 레이어가 연합 환경에서 얼마나 많은 유용한 특징을 학습하는가? 그리고 왜 최종 레이어는 이를 활용하지 못하는가?
• RQ3 경험적 신경 장력 커널(eNTK)을 통한 볼록화가 데이터 이질성 하에서 비볼록 연합 학습의 성능을 복원할 수 있는가?
• RQ4 TCT 프레임워크에서 첫 번째 단계의 사전 훈련 기간 선택이 최종 모델 정확도에 어떤 영향을 미치는가?
• RQ5 eNTK 특징의 정규화가 두 번째 단계 볼록 최적화의 수렴 속도와 통신 효율성을 향상시키는가?

주요 결과

• 클라이언트의 데이터 분포가 상이할 경우 FMNIST에서 최대 36%의 정확도 향상과 CIFAR10에서 최대 37%의 정확도 향상을 달성한다.
• 첫 번째 단계에서 사전 훈련된 특징 추출기를 사용할 경우 중심화된 학습과 연합 학습 간의 성능 격차를 완전히 해소한다.
• FedAvg로 훈련된 모델에서 추출한 eNTK 특징을 사용한 TCT는 표준 FedAvg보다 테스트 정확도 약 20% 향상된다.
• eNTK 특징의 정규화는 수렴 속도를 크게 향상시켜 TCT가 약 40회의 통신 라운드 내에 거의 100%의 훈련 정확도에 도달할 수 있도록 한다.
• TCT의 성능는 첫 번째 단계에서의 통신 라운드 수가 다양하더라도 안정적이며, T₁ ≥ 60일 경우 하이퍼파ram터 선택에 대해 강건함을 보인다.
• FedAvg는 국소 스텝 수가 증가함에 따라 성능이 저하되는 반면, TCT에서 사용된 SCAFFOLD는 수렴 성능을 유지하거나 향상시키며, 기울기 보정의 중요성을 입증한다.