[논문 리뷰] Towards Federated Foundation Models: Scalable Dataset Pipelines for Group-Structured Learning

Zachary Charles, Nicole Mitchell외 3명|arXiv (Cornell University)|2023. 07. 18.

Privacy-Preserving Technologies in DataComputer Science인용 수 3

한 줄 요약

이 논문은 기존 데이터셋의 사용자 정의 분할을 가능하게 함으로써 대규모 그룹 구조 데이터셋—특히 피어드 학습을 위한 것—을 생성할 수 있는 확장성 있고 프레임워크에 종속되지 않는 라이브러리인 Dataset Grouper를 소개한다. 이는 수십억 개의 그룹을 포함하는 데이터셋을 지원하며, 수백만에서 수십억 개의 파라미터를 가진 기초 모델의 첫 번째 피어드 학습을 가능하게 하여, FedAvg가 스케일이 클수록 경험적 리스크 최소화보다는 메타학습 방법에 더 가까운 행동을 한다는 것을 드러낸다.

ABSTRACT

We introduce Dataset Grouper, a library to create large-scale group-structured (e.g., federated) datasets, enabling federated learning simulation at the scale of foundation models. This library facilitates the creation of group-structured versions of existing datasets based on user-specified partitions and directly leads to a variety of useful heterogeneous datasets that can be plugged into existing software frameworks. Dataset Grouper offers three key advantages. First, it scales to settings where even a single group's dataset is too large to fit in memory. Second, it provides flexibility, both in choosing the base (non-partitioned) dataset and in defining partitions. Finally, it is framework-agnostic. We empirically demonstrate that Dataset Grouper enables large-scale federated language modeling simulations on datasets that are orders of magnitude larger than in previous work, allowing for federated training of language models with hundreds of millions, and even billions, of parameters. Our experimental results show that algorithms like FedAvg operate more as meta-learning methods than as empirical risk minimization methods at this scale, suggesting their utility in downstream personalization and task-specific adaptation. Dataset Grouper is available at https://github.com/google-research/dataset_grouper.

연구 동기 및 목표

• 피어드 학습 및 기초 모델 연구를 위한 대규모 그룹 구조 데이터셋의 부족을 해결하기 위해.
• TensorFlow 및 HuggingFace Datasets와 같은 기존 데이터셋으로부터 이질적이고 피어드 가능한 데이터셋을 효율적이고 메모리 최적화된 방식으로 생성할 수 있도록 하기 위해.
• 현대 기초 모델의 규모에 맞는 피어드 학습 시뮬레이션을 지원하기 위해, 수백만에서 수십억 개의 파라미터를 가진 모델을 포함한 스케일을 고려하기 위해.
• 메모리에 올라가지 않는 데이터셋까지도 확장 가능한 스트리밍 및 계층적 데이터 포맷을 사용하여, 프레임워크에 종속되지 않는 솔루션을 제공하기 위해.
• FedAvg 및 FedSGD와 같은 피어드 최적화 알고리즘의 행동을 초대규모에서 경험적으로 조사하여, 특히 메타학습 및 개인화와의 관계를 탐색하기 위해.

제안 방법

• Dataset Grouper는 사용자 정의 분할 함수를 사용하여 기본 데이터셋의 예제를 클라이언트별 그룹으로 묶어, 영리하고 커스터마이징 가능한 피어드 데이터 생성을 가능하게 한다.
• 이 라이브러리는 메모리 내, 계층적, 스트리밍의 세 가지 데이터 포맷을 지원하여, 전체 데이터셋을 RAM에 로드하지 않음으로써 메모리 효율성을 높인다.
• 스트리밍 반복을 사용하여 피크 메모리 사용량을 최소화함으로써, 수백만 또는 수십억 개의 그룹을 포함하는 데이터셋에 확장 가능하도록 설계되어 있다.
• TensorFlow 및 HuggingFace와 같은 기존 머신러닝 프레임워크와 통합되어, 훈련 파이프라인에서 직접 사용할 수 있다.
• 긴 꼬리 분포와 높은 시퀀스 길이를 가진 대규모 텍스트 데이터셋인 C4의 그룹 구조 버전을 생성함으로써, 대규모 피어드 언어 모델링을 가능하게 한다.
• 프리트레인 및 파인튜닝 시나리오를 모두 지원하며, 개인화 및 차별적 프라이버시 인식 훈련도 지원한다.

Figure 1: Per-group statistics of the new group-structured (i.e. federated) language modeling datasets.

실험 결과

연구 질문

• RQ1 표준 시스템의 메모리 용량을 초월하는 대규모 그룹 구조 데이터셋을 효율적으로 생성할 수 있는 방법은 무엇인가?
• RQ2 피어드 설정에서 대규모 기초 모델을 훈련시킬 때, FedAvg가 경험적 리스크 최소화 방법이 아니라 메타학습 방법으로서 어떻게 행동하는가?
• RQ3 클라이언트당 로컬 배치 수와 같은 하이퍼파라미터는 대규모 피어드 학습에서 프리- 및 포스트-개인화 성능에 어떤 영향을 미치는가?
• RQ4 기존 최적화 알고리즘과 확장 가능한 데이터 파이프라인을 사용하여 빌리언 파라미터 모델의 피어드 학습을 실제로 달성할 수 있는가?
• RQ5 데이터 분할 전략과 통신 효율성은 대규모 피어드 학습에서 모델 수렴 및 개인화 성능에 어떤 영향을 미치는가?

주요 결과

• Dataset Grouper는 C4 데이터셋의 그룹 구조 버전을 기반으로, 1억에서 10억 파라미터를 가진 디코더 전용 트랜스포머 모델의 첫 번째 피어드 학습을 가능하게 하여, 스케일에서의 실현 가능성을 입증한다.
• 64개의 배치를 클라이언트당 사용할 경우, FedAvg는 포스트-개인화 손실 0.008을 달성하지만, 동일 조건에서 FedSGD는 포스트-개인화 손실 3.3을 유지하여, 이 영역에서 FedAvg가 더 뛰어난 일반화 성능을 보임을 시사한다.
• 처리된 토큰 수를 동일하게 조정했을 때, τ=4인 FedAvg는 전처리 손실(3.8)과 포스트-개인화 손실(0.006)이 모두 최저 수준을 기록하여, 로컬 업데이트 빈도에서 최적의 트레이드오프를 제공함을 나타낸다.
• 통신이 병목이 되지 않는 조건에서, 다양한 τ 값에 대해 포스트-개인화 손실이 안정적으로 유지되며, 충분한 훈련 토큰이 사용될 경우 로컬 업데이트 스케줄에 대한 강건성을 보임을 시사한다.
• 스트리밍 및 계층적 형식의 피크 메모리 사용량은 심지어 수십억 개의 그룹을 포함하는 데이터셋에서도 2MB 미만이었으며, 반면 메모리 내 형식은 수백 GB가 필요했기에, 시스템의 확장성을 입증한다.
• 결과는 FedAvg가 스케일이 클수록 경험적 리스크 최소화가 아니라 메타학습 방법으로서 작동하며, 후행 개인화 작업에서 뛰어난 성능을 보임을 시사한다.

(a) In-Memory.