엔비디아 클라라가 구현하는 연합학습

엔비디아의 AI 연학학습 '클라라' 알고리즘

AI는 방대한 양의 데이터를 필요로 한다. 특히 의료분야에선 더욱 그렇다. 예를 들어 자동 종양 진단 시스템을 학습하려면 발생 가능한 해부학과 병리학적 패턴을 전체적으로 파악해야 하기 때문에 그와 관련된 방대한 데이터베이스가 요구된다. 또한 병원과 의료기관이 협력해 현장의 지식을 공유하고 결합해야 하는 경우도 적지 않다.

그러나 환자의 의료 데이터는 개인정보이기 때문에 이를 반출하거나 활용하는 것은 쉬운 일이 아니다. 따라서 덩달아 개인정보를 침해하지 않으면서 알고리즘을 학습하는 방법에 대한 연구도 점차 중요해지고 있다.

엔비디아가 얼마 전 출시한 ‘클라라 트레인 SDK(Clara Train SDK)’는 연합학습(Federated Learning, FL) 기능을 탑재해 엣지 AI 컴퓨팅 플랫폼인 엔비디아 EGX에서 이를 구현한다. 클라라 연합학습 패러다임은 서로 다른 디바이스가 안전하게 협력, 훈련해 글로벌 AI 모델 개선에 기여할 수 있도록 지원한다. 연합학습은 각 디바이스에서 부분적인 모델 가중치만 글로벌에 공유하기 때문에 개인정보를 최대한 보호할 수 있는 새로운 AI 학습 방식으로 각광받고 있다.

연합학습은 P2P(Peer-to-Peer), 순환형(Cyclic), 서버-클라이언트 아키텍처 등 다양한 분산형 아키텍처에서 실행된다. 엔비디아 클라라는 서버-클라이언트 기반으로 실행되는데, 이는 중앙 서버가 다양한 클라이언트가 참여하는 연합학습의 촉매 역할을 한다는 것을 의미한다. 개발자들은 클라라 트레인 SDK의 구성 가능한 MMAR(Medical Model ARchive, 의료모델아카이브) 기능을 이용해 자체 모델과 구성요소를 활용한 연합학습을 진행하고, 로컬(Local) 학습이 단일 GPU 혹은 복수의 GPU에서 실행되는지 여부와 상관없이 제어할 수 있다.

[그림1] 연합학습의 개념을 나타낸 예시, 각각의 정보가 중앙 시스템에 모여 가공된다.

서버는 전체 모델 훈련의 진행 상황을 관리하고 원본 모델을 모든 참여 클라이언트들에게 배포(Broadcast)한다. 모델에 대한 훈련은 각 클라이언트에서 지역적으로 진행되므로, 서버가 훈련 데이터에 직접 접근할 필요가 없다. 클라이언트는 데이터 대신 모델 가중치 업데이트만 부분적으로 공유한다. 또한 각 클라이언트는 서버로 전송되는 모델 가중치 비율에 대한 자체적인 개인정보 통제권을 갖고 있다.

서버가 클라이언트로부터 최신 모델을 수령하면, 모델 집계 방식에 대한 자체 알고리즘을 구축한다. 모든 클라이언트들의 단순 평균값이 될 수도 있고, 클라이언트가 그간 기여해온 것들의 가중치에 기반할 수도 있다.

서버는 연합학습 횟수에 대한 전반적인 통제권을 갖고 수 차례의 학습이 이뤄지는 동안 참여하는 클라이언트들을 새롭게 추가, 제거할 수 있다. 연합학습은 모든 참가자에게 도움이 되는 구조다. 중앙집중화 모델은 견고하고 더욱 일반화되며, 로컬 모델의 정확성은 향상된다.

연합학습 시스템 아키텍처

연합학습 모델 훈련을 실행할 때, 클라이언트는 먼저 서버 훈련 서비스를 시작해야 한다. 서버 세션은 학습 1회를 시작하는데 필요한 최소 클라이언트 수와 학습 1회에 참여할 수 있는 최대 클라이언트 수를 통제한다. 클라이언트가 연합학습 세션에 참여하려는 경우, 먼저 로그인 요청을 제출해 학습 신청을 해야 한다.

서버는 클라이언트의 자격을 확인하고 클라이언트를 검증하기 위한 인증 절차를 진행한다. 인증에 성공하면, 서버는 다음 단계인 클라이언트-서버 통신에 사용될 연합학습 토큰을 클라이언트에게 전송한다. 클라이언트를 인증할 수 없을 경우엔 인증을 거부한다.

클라이언트들은 연학학습의 각 회차에서 학습반복 횟수(Epoch)를 제어할 수 있다. 클라이언트가 진행 중이던 학습 회차를 마치면 기존 연합학습 토큰을 사용해 업데이트된 모델을 서버로 전송할 수 있다. 서버가 모든 참여 클라이언트로부터 업데이트된 모델을 수신한 뒤에는 가중치 알고리즘에 기반해 모델을 집계하고 전반적인 업데이트를 마친 학습 모델을 받는다. 이로써 연합학습 훈련 1회차가 종료된다. 이어 다음 회차가 시작되고, 서버에 설정된 최대 회차가 모두 종료될 때까지 이 과정이 반복된다.

서버-클라이언트 통신 프로토콜

엔비디아 클라라의 연합학습은 모델 학습 기간 동안 서버와 클라이언트 간에 gPRC 프로토콜을 사용하며, 모델 학습 세션 기간 동안 네 개의 기본 명령어가 사용된다.

서버 사이드 워크플로우

서버 사이드와 클라이언트 사이드에서의 워크플로우를 좀 더 자세히 알아보도록 하자. 연합학습 서버 사이드 서비스를 시작할 때, 연합학습 서비스명, gPRC 통신 포트, SSL 인증키, 최소, 최대 클라이언트 수 등 서버 구성 파일을 사용해 초기 모델을 초기화하고 복원한 뒤 연합학습 서비스를 시작한다.

초기화 후, 서버는 루프(loop)를 시작하고 클라이언트의 참여 요청을 기다리다가 클라이언트에게 모델을 제공한 뒤, 클라이언트가 업데이트된 모델을 반환하기를 기다린다. 서버가 클라이언트로부터 업데이트된 모델을 모두 수신하면 가중집계 알고리즘(Weighted aggregation algorithm)에 기반해 집계를 진행하고, 기존 모델을 전체적으로 업데이트한다. 그리고 업데이트된 모델은 모델 학습의 다음 회차에 사용되며, 서버가 최대 연합학습 훈련 회차를 마칠 때까지 이 과정이 반복된다.

클라이언트 사이드 워크플로우

연합학습 훈련에 참여하는 클라이언트는 우선 클라이언트 구성을 사용해 모델을 초기화한다. 다음으로, 클라이언트 자격(Credential)을 사용해 서버에 로그인 요청을 보내고 연합학습 훈련 토큰을 받는다.

토큰을 수령하면 클라이언트는 서버에 현재 모델을 요청한다. 클라이언트는 다시 글로벌 모델을 사용해 텐서플로우(TensorFlow) 세션을 구축·재건하고, 현재 모델에 맞는 로컬 데이터를 사용한 로컬 학습을 시작한다. 학습 기간 동안, 클라이언트는 연합학습 훈련의 개별 회차에 대한 반복 횟수를 제어할 수 있다. 또, 로컬 학습이 단일 GPU 혹은 복수의 GPU에서 실행될지를 제어할 수 있다.

클라이언트가 로컬 모델 학습의 현재 회차를 마치면, 업데이트된 로컬 모델을 서버로 전송할 수 있다. 클라이언트는 서버 집계를 위해 반환할 가중치 양에 대해 자체적인 개인정보 정책을 구성할 수 있다. 이후 클라이언트는 연합학습 훈련의 다음 회차를 시작하기 위해 서버에 새로운 글로벌 모델을 요청할 수 있다.

연합학습으로 생성된 모델의 품질

BRATS2018 데이터 세트를 사용해 엔비디아 MMAR에서 멀티 모드 멀티 클래스 뇌종양 분할작업(Brain tumour segmentation task)을 훈련할 때, 연합학습은 데이터 중앙 집중식 훈련과 유사한 분할 성능을 보여준다. 아래 그림은 다이스(Dice) 점수로 측정한 헬드아웃(Held-out) 검증 데이트 세트에 대한 분할모델 품질을 학습 반복 횟수와 비교한 그래프다[그림6].

데이터 중심(Data-centralized) 방식과 연합평균화(Federated averaging) 방식 모두 엔비디아의 MMAR 설정에 따라 동일한 네트워크 아키텍처와 프로세싱 단계를 사용하고 있다. 두 방식의 주요 차이점은 연합평균화 방식의 경우, 데이터를 처음 획득한 소스에 따라 데이터 세트가 13개의 클라이언트로 분리되고, 클라이언트가 프로세스 진행 도중에 학습 데이터를 공유하지 않는 것이다.

[그림6]을 보면 두 가지 학습 프로세스 모두 다이스 점수 0.82 측면에서 유사한 성능을 보였다는 결론을 내릴 수 있다. 또, 연합학습 시스템이 다수의 로컬 학습 세트에서 효과적으로 지식을 추출해, 연합 서버를 통해 그 지식을 집계할 수 있다는 것을 시사한다.

연합학습 MMAR(Medica Model Archive) 통합

연합학습 기능들은 클라라 트레인 SDK v1.0, v1.1과 동일한 MMAR 구조로 구성됐다. 모델 학습 구성 측면에서 연합학습 모델 훈련은 동일한 변환(Transform) 파이프라인 솔루션을 사용한다. 또 연합학습 태도(Behavior)를 설명하기 위한 두 가지 서버-클라이언트 구성이 있다.

서버 트레이너는 전체적인 모델 학습 프로세스에 필요한 연합학습 학습 회차 수를 제어하고, 참가 클라이언트들의 모델을 전체적으로 집계하며, 글로벌 모델 훈련 프로세스를 조정한다. 클라이언트 트레이너는 각 연합학습 회차의 반복 횟수와, 집계를 위해 로컬 훈련을 마친 모델을 연합학습 서버에 발행할 때 사용되는 개인정보보호 정책을 제어한다.

연합학습 모델 훈련을 시작할 때, 우선 명령어 ‘server_train.sh’를 사용해 서버에서 연합합습 훈련 서비스를 시작한다. 이 서비스는 훈련 작업 신원(Training task identity)과 gPRC 통신서비스위치 URL을 관리한다. 연합학습 주기가 진행되는 동안, 서비스는 연합학습에 참여하고자 하는 클라이언트의 요청을 듣고, 글로벌 모델을 배포한 다음 클라이언트의 업데이트된 모델을 집계한다.

클라이언트는 명령어 ‘client_train.sh’를 사용해 연합학습 훈련 작업을 시작한다. 그리고 로그인 요청을 통해 토큰을 받아 서버에서 글로벌 모델을 수령한 뒤, 로컬 보안 데이터를 사용해 로컬 모델을 학습시키고 업데이트한다. 그다음 연합학습 훈련 1회가 종료될 때마다 모델을 서버에 제출해 집계한다.

연합학습 데이터 보안

데이터 보안은 연합학습의 주요 동기부여 요소다. 연합학습 모델 훈련은 일반적으로 다수의 조직이 관여해 진행된다. 개별 조직들은 클라라 트레인의 연합학습을 통해 모델을 로컬에서 훈련할 수 있고, 기밀 데이터를 제외한 부분적인 모델 가중치만 공유할 수 있다. 하지만 데이터와 모델 통신이 왜곡되지 않도록 안전하게 보호하기 위해서는 클라이언트-서버 통신도 매우 중요하다.

연합학습의 보안을 성공적으로 유지하기 위해, 엔비디아는 클라이언트와 서버 간 신뢰성 확보를 위한 연합학습 토큰을 사용한다. 연합학습 토큰은 연합훈련 세션 주기 전체에 걸쳐 사용된다. 클라이언트는 서버의 신원을 검증하고, 서버는 클라이언트를 검증, 그리고 인증해야 한다. 클라이언트와 서버 간 이뤄지는 데이터 교환은 안전한 통신을 위해 HTTPS 프로토콜을 기반으로 한다. 클라이언트와 서버 간 신뢰를 보장하기 위해 자체 서명된 SSL 인증서가 사용된다.

이제 아래 이어지는 내용을 통해 자체 서명된 SSL 인증 기관(SSL Certificate Authority)과 서버-클라이언트 인증서 생성 방법을 확인할 수 있다.

연합학습 서버-클라이언트 구성 예시

연합학습 서버 구성 파일: ‘config_fed_server.json.’