기반 모델(FMs), 즉 대형 머신 러닝 모델은 강력한 머신 러닝 시스템을 구축하는 데 있어 놀라운 성능을 보입니다. 그러나 FMs은 방대한 양의 훈련 데이터와 계산 능력이 필요합니다. 이러한 대규모 훈련 데이터를 확보하는 것은 데이터 프라이버시, 법적 제약 및 계산 부담과 관련된 문제를 야기합니다.
**연합 학습(FL)**은 FMs의 이러한 문제를 해결할 수 있는 유망한 분산 머신 러닝 프레임워크입니다. FL에서 로컬 클라이언트들은 로컬 훈련 데이터를 이용해 글로벌 모델을 공동으로 훈련하며, 개인 로컬 데이터를 공유하지 않고 로컬에서 업데이트된 모델 매개변수를 집계합니다.
FL에는 다양한 계산 및 통신 능력을 가진 클라이언트들이 있을 수 있습니다. 그 중 일부 자원이 제한된 클라이언트들은 글로벌 모델을 훈련하는 데 어려움을 겪을 수 있습니다. 이러한 경우, 글로벌 모델의 크기를 줄이거나 충분한 자원을 가진 클라이언트만 학습에 참여하도록 하는 방법으로 FL을 수행할 수 있습니다. 그러나 이러한 접근 방식은 글로벌 모델의 성능 저하를 초래할 수 있습니다. 이전 연구에서는 글로벌 모델을 서브 모델로 축소하여 이질적인 시스템 클라이언트를 처리하려 했으나, 모델 크기를 줄이면 모델 용량이 제한되어 성능이 저하됩니다.
이러한 자원 제한 문제를 해결하기 위해, **분할 학습(SL)**은 모델을 축소하지 않고 클라이언트의 원시 데이터에 접근하지 않으면서 서버를 활용해 글로벌 모델을 훈련시키는 방법을 제안합니다. 이는 특정 레이어(분할 지점)에서 모델을 분할하고 클라이언트가 해당 지점까지 부분 모델을 훈련하도록 합니다. 클라이언트는 특정 레이어에서의 출력을 서버로 전송하며, 이를 **'smashed data'**라고 합니다. 그런 다음 서버가 나머지 레이어를 훈련합니다.
**분할 연합 학습(SFL)**은 SL과 FL을 결합하여 훈련 시간을 단축하는 알고리즘을 제안했습니다. 그러나 SL과 SFL은 단일 분할 지점으로 모델을 분할하기 때문에, 분할 지점은 가장 낮은 계산 능력을 가진 클라이언트의 요구를 충족해야 합니다. 그러면 충분한 능력을 가진 클라이언트는 잠재력을 충분히 발휘하지 못하고 서버에 계산 부담이 발생할 수 있습니다. 이는 대형 모델을 훈련할 때 더 치명적일 수 있습니다.
본 논문에서는 제한된 계산 능력을 가진 시스템 이질성을 해결하기 위해 새로운 FL 프레임워크인 FedSplitX를 제안합니다. FedSplitX에서는 대형 모델을 클라이언트 측 모델과 서버 측 모델로 여러 분할 지점에서 분할하여 이질적인 능력에 맞춥니다. 여러 분할 지점을 사용함으로써, 다양한 분할 지점에 따라 여러 쌍의 클라이언트 측 모델과 서버 측 모델이 생성됩니다. 또한, 다양한 분할 지점을 위한 보조 네트워크를 도입합니다. 이 보조 네트워크는 클라이언트 측 모델의 로컬 손실 기반 학습을 가능하게 하여 통신 비용과 지연을 줄입니다. 더욱이, 보조 네트워크로부터의 협력 손실이 성능을 향상시킵니다. 로컬 모델 최적화를 위해, 우리의 프레임워크는 클라이언트 측 모델과 서버 측 모델의 매개변수를 각각 집계 방법으로 별도로 집계합니다.
우리의 실험 결과, FedSplitX는 클라이언트의 능력보다 더 많은 계산 능력이 필요한 대형 모델을 서버의 능력을 활용해 훈련시킬 수 있으며, 글로벌 모델 크기가 클라이언트의 계산 능력에 맞춰 제한되는 기준 모델보다 더 나은 성능을 보여줍니다. 특히, 각 분할 지점에서 보조 네트워크와의 협력 손실이 성능을 효과적으로 향상시킬 수 있음을 확인했습니다.
우리의 기여는 다음과 같이 요약됩니다: • 우리는 서버의 능력을 활용하여 이질적인 클라이언트들이 협력할 수 있는 프레임워크인 FedSplitX를 제안합니다. • 우리는 서버의 지원을 받아 대형 모델을 훈련시키는 FedSplitX가 모든 클라이언트의 요구에 맞춘 기준 모델보다 우수한 성능을 보임을 입증합니다. • 우리는 FedSplitX가 모델을 여러 분할 지점으로 분할하여 클라이언트가 자신의 능력을 최대한 활용하면서 서버의 계산 부담을 줄일 수 있음을 보여줍니다.
그림 1: M = 3의 깊이 수준을 갖는 클라이언트 시스템에서의 FedSplitX 프레임워크. (상단) 훈련될 모델의 전체 구조와 깊이 수준에 따른 분할 지점, (중단) 깊이 수준에 따라 분할 지점에서 분할된 클라이언트 측 모델과 서버 측 모델, (하단) 분할 작업 후의 훈련 과정.