Hugging Face와 Flower를 조합해 distilBERT 모델을 IMDB 감정 분류 작업에서 다중 클라이언트 간 연합학습으로 훈련

Context

전통적인 머신러닝 훈련은 중앙 서버에서 모든 데이터를 수집해야 하므로 데이터 프라이버시 문제가 발생한다. 개별 클라이언트가 로컬 데이터에서 독립적으로 모델을 훈련하고 파라미터만 서버로 전송하는 방식이 필요했다.

Technical Solution

• Hugging Face datasets 라이브러리로 IMDB 데이터셋 로드 및 100개 샘플로 다운샘플링: random.sample()로 train/test 데이터 선택
• AutoTokenizer를 사용해 텍스트를 토큰화하고 DataCollatorWithPadding으로 배치 처리: 배치 크기 32로 설정
• distilbert-base-uncased 사전학습 모델을 AutoModelForSequenceClassification으로 로드해 2개 클래스 분류 헤드 추가
• Flower의 NumPyClient 클래스 상속으로 IMDBClient 구현: get_parameters()로 모델 파라미터 추출, set_parameters()로 서버 파라미터 수신, fit()으로 로컬 훈련 1 에포크 수행, evaluate()로 정확도와 손실 계산
• Flower 서버에서 FedAvg 전략 적용: fraction_fit=1.0, fraction_evaluate=1.0으로 설정하고 weighted_average() 함수로 클라이언트 간 메트릭 가중평균 계산, 3 라운드 훈련 진행
• AdamW 옵티마이저를 학습률 5e-5로 설정해 로컬 훈련 수행

Key Takeaway

Hugging Face의 표준 PyTorch 훈련 루프를 Flower의 NumPyClient 인터페이스로 래핑하면 기존 딥러닝 워크플로우를 최소한의 수정으로 연합학습으로 전환할 수 있다. 이 패턴은 프라이버시이 필수적인 산업(금융, 헬스케어)에서 모델 성능과 데이터 보호를 동시에 달성하는 기반이 된다.