지식 전달을 이용한 연합학습 환경의 엣지에서 깊은 모델 학습시키기
우리는 CNN 모델을 스케일-업 하면 모델의 정확도를 향상시킬 수 있다는 사실을 알고 있다. 하지만, 연합학습 상황에서는 모델을 학습하는 엣지 클라이언트의 자원 부족으로 CNN 모델을 스케일-업 하는게 어렵다.
연합학습에서 엣지 클라이언트를 지칭하는 대상이 여러개이지만, 기본적으로 스마트폰을 떠올리면 된다. 현대 스마트폰은 GPU Acceleration 기능과 높은 CPU 연산력을 가지지만 배터리 문제로 장기간 연산을 수행하기 힘들고, 기본 성능 또한 서버급 컴퓨터가 가지는 성능보다 한참 뒤떨어지기 때문에 복잡한 CNN모델을 그대로 스마트폰에서 학습하기는 힘들 것이다.
이 깊은 CNN 모델(Large CNN Model)을 학습하기 힘든 문제를 해결하기 위한 여러 시도 중에서 지식 증류(Knowledge Distillation)를 이용해서 해결한 FedGKT 기법 있다. 이번 포스팅에서는 이 FedGKT를 간단하게 다뤄보고자 한다.
FedGKT는 “Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge”[1] 논문에서 제안된 기법으로 NIPS 2020에서 소개되었다.
구성
- FedGKT 개요
- 지식 증류
- FedGKT 알고리즘
- 실험
- 결론
FedGKT 개요
FedGKT는 연합 그룹 지식 전송(Federated Group Knowledge Transfer)의 약자로 연합 학습 환경에서 엣지 클라이언트와 중앙 서버가 지식 증류 알고리즘을 이용해서 학습한다. FedGKT는 이 방법으로 이용해서 엣지에서의 연산 비용은 줄이고 중앙 서버의 연산 비용은 높인다. 중앙 서버는 무거운 CNN 모델을 학습해서 모델 A를 만든다. 엣지에서는 모델 A의 지식을 증류해서 로컬에서 상대적으로 가벼운 모델을 학습한다. 이 순환 구조는 각 학습 라운드마다 실행된다. 이를 구조적으로 보여주면 아래 [그림 1]과 같다.
[그림 1]은 논문에 삽입된 그림으로, 그림에서 Edge영역을 보면 Compact CNN 모델이 f_extractor와 f_classifier(이하 f_c)로 나뉜 것을 볼 수 있는데, f_extractor(이하 f_e)는 신경망에서 마지막 레이어인 Fully Connected 레이어를 제외한 나머지 히든 레이어를 의미한다. f_e의 결과값은 중앙 서버인 f_server(이하 f_s)의 입력값으로 제공된다. f_c의 결과값은 신경망이 이미지를 보고 분류한 소프트 레이블을 의미하는데 이 값은 f_s의 지식 증류(KD) 알고리즘에 이용된다. 마지막으로 f_s의 소프트 레이블은 다시 f_c의 지식 증류 알고리즘에서 사용한다. 이렇게 각 라운드마다 서버와 엣지는 서로의 모델을 학습하면서 정확도를 높여나간다.
그림에 각 모델과 순서를 구분할 수 있도록 색칠이 잘 되어 있기 때문에 각 알고리즘의 실행 순서를 쉽게 볼 수 있을 것이다. (b) 영역을 보면 엣지에서 학습하는 모델보다 서버의 모델이 무거운 것을 알 수 있다. FedGKT 논문은 제목에서 Edge에서 무거운 CNN 모델을 학습하기라고 했지만, 실제로는 서버의 무거운 모델의 지식을 가벼운 모델을 학습하는 엣지로 전송하는 방법으로 구현되었다.
지식 증류(Knowledge Distillation)
지식 증류란 2014년 NIPS에서 Hinton 교수가 제출한 논문[2]에서 소개된 내용으로 앙상블 모델을 단일 모델로 압축하는 방법을 말한다. 딥러닝을 서비스로 도입할 때 정확성을 높이기 위해 가장 쉽게 할 수 있는 방법은 n개의 딥러닝 모델을 학습하고 n개의 답안 중 가장 많이 제출된 답변을 정답으로 쓰는 방법이다. 이렇게 다수의 딥러닝을 운용하는 건 학습 뿐 아니라 서비스를 serving 하는 단계에서도 연산 비용이 많이 소모되기 때문에 이를 하나의 단일 모델로 압축할 필요가 있는데 이를 지식 증류를 이용해서 간단하게 할 수 있다.
[그림 2]는 다수의 모델을 학습한 앙상블을 하나의 단일 모델로 지식 증류를 이용해서 변환하는 모습을 보여준다. 지식증류 알고리즘의 핵심적인 원리는 모델에서 나온 결과값을 소프트레이블로 변환하고 이를 하나의 단일 모델이 학습하는 방식으로 동작한다.
요즘에는 조금 다르지만, 일반적인 딥러닝 분류 모델에서는 마지막 레이어에 교차 엔트로피(Cross Entropy)함수를 취하는게 일반적이다. 이 교차 엔트로피는 정답에 가까운 확률값을 극단적으로 높이기 때문에 일반화가 잘 된것이라고 볼 수 없다는 견해가 있는데 이를 위 [그림 3]에서 Hard-Label과 Soft-Label의 관계로 잘 보여주고 있다.
Hard-Label은 [그림 3]의 숫자 2 그림을 보고 2로 분류할 확률을 극단적으로 99%를 취하는 반면, 지식 증류 알고리즘에서 사용하는 Soft-Label은 2일 확률은 80%, 3일 확률은 19%로 평가함으로써 숫자 2의 모양은 숫자 3의 모양도 일부분 가지고 있다라는 것을 암시한다. 그래서 지식 증류 논문을 보면 앙상블에서 단일 모델로 지식을 전달할 때 숫자 3의 데이터를 학습하지 않아도 단일 모델이 숫자 3을 학습하는 결과를 보여준다. 즉 직관적으로 해석해 보면 Soft-Label을 이용해서 일반화된 지식을 학습할 수 있다고 볼 수 있다.
FedGKT에서는 이 지식증류 알고리즘을 앙상블-단일 모델의 관계로 사용하는 것이 아닌 큰 모델(Large Model)과 작은 모델(Small Model)의 관계로 사용한다. 연산력이 높은 중앙 서버에서는 크고 복잡한 CNN 모델을 학습해서 Soft-Label을 만들고 연산력이 부족한 엣지 기기에서 작은 모델을 Soft-Label을 보고 학습시킨다. 여기서 한가지 재밌는 건 지식 증류의 대상이 되는 단일 모델은 앙상블 모델이 학습한 데이터를 모른다는 사실이다. 자연스럽게 연합학습 상황과 연결지어 생각되지 않는가?
FedGKT 알고리즘
이미 지식 증류(KD)를 알고 있거나, 위에서 이야기한 내용을 이해했다면 FedGKT 방법을 95% 정도 이해한 것과 같다. FedGKT는 지식 증류를 그대로 이용해서 학습에 적용하기 때문에 수식으로 보아도 크게 어려운 점은 없다. 지금부터 나머지 5% 소개해보겠다.
위 [그림 4]에서 a. 영역은 기존 연합 학습의 방식을, c. 영역은 FedGK를 나타낸다. a. 영역에 해당하는 연구로는 구글에서 연구한 FedAvg가 있는데 이는 중앙서버에서 뉴런의 가중치를 평균내서 반영하는 것으로 서버는 따로 딥러닝 모델을 학습하지 않는다. FedGKT에서는 중앙 서버에서도 모델을 학습해야 하기 때문에 정의해야 하는 손실 함수는 서버, 엣지 두 가지이다. 논문을 보면 수식의 notation이 보기 힘들게 작성되어 있기 때문에 여기서는 핵심이 되는 부분만 소개하고 마무리 하겠다.
[그림 5에서 l_s와 l_c는 각각 서버와 엣지 클라이언트의 모델 손실 함수로 정의한다. D_KL은 쿨백-라이블러 발산 함수로 인자로 주어진 두 확률분포가 서로 동일할수록 1에 가까워지고, 서로 다를 수록 0으로 멀어지는 특징을 가지고 있다. 그리고 l_CE는 일반적으로 사용하는 교차-엔트로피 함수이다. 여기서 특징은 D_KL의 인자가 되는 확률 분포는 모델의 마지막 레이어에서 출력한 확률값을 사용하는데 이 때 계산 수식은 지식 증류 알고리즘과 동일하다.
의사코드를 보면 수식의 notation 정의가 직관적이지 않기 때문에 우리는 중요한 특징만 보도록 하자. 먼저 ClientTrain() 함수는 엣지 클라이언트 실행하는 것으로 이미지 데이터 X를 입력값으로 받아 분류 모델을 학습한다. 학습이 끝나면 Hidden Layer의 뉴런값 H, Final Layer의 출력값인 Z 그리고 분류 정답인 Y를 서버로 전달한다. 연합학습이기 때문에 서버는 X에 대해서는 알 수 없다. 그리고 서버는 정보를 가지고 내부적으로 학습을 진행 한다. 의사 코드에서는 많은 부분이 대략적인 실행 흐름만을 보여주지만 l_s, l_c 함수 안에 지식 증류 알고리즘의 수식이 정의되어있다.
결론
위 [그림 7]은 FedGKT와 대표적인 알고리즘인 FedAvg와 학습 수렴 속도를 비교한 것이다. 사실 학습 속도는 크게 차이 나지 않는다. [그림 8]은 FedGKT와 FedAvg의 성능일 비교한 것으로 중앙 서버에서 큰 서버를 사용해도 정확도는 크게 증가하지 않았다.
사실 연합학습은 기본적으로 non-IID 상황을 가정해서 데이터 자체가 많이 편향되어 있기 때문에 모델의 성능과는 크게 관계 없는 것 일수도 있다. 다만, 표를 자세히 보면 FedGKT에서는 ResNet-8로 작은 모델을, FedAvg에서는 ResNet-56과 ResetNet-110의 무거운 모델을 이용한 것을 볼 수 있다. FedGKT에서는 이렇게 엣지 클라이언트에서 연산력의 부담을 줄일 수 있었다. 엣지에서의 상황만 보자면 FedGKT는 FedAvg대비 평균 9–17배 적은 연산력(FLOPS)를 소모하였고 모델의 파라미터는 평균 54–105배 적었다고 한다.
이도현
🎓서강대학교 BASE 연구실
📓 CURG 에디터
🌱 https://github.com/scalalang2
레퍼런스
[1] Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge
[2] Distilling the Knowledge in a Neural Network
