[논문 리뷰] Federated Full-Parameter Tuning of Billion-Sized Language Models with Communication Cost under 18 Kilobytes

Zhen Qin, Daoyuan Chen외 4명|arXiv (Cornell University)|2023. 12. 11.

DNA and Biological Computing인용 수 8

한 줄 요약

논문은 FedKSeed를 제로차(zeroth-order) 연합 미세조정 방법으로, 고정된 소(seed) 풀과 스칼라 그래디언트를 사용하고 비대역적(전체 파라미터 교환) 없이도 빌리언 규모 LLM의 전체 파라미터 튜닝을 가능하게 한다. 라운드당 통신량이 극도로 낮아 (18 KB 미만)이며, 비균일 시드 샘플링 강화 FedKSeed-Pro 존재.

ABSTRACT

Pre-trained large language models (LLMs) need fine-tuning to improve their responsiveness to natural language instructions. Federated learning offers a way to fine-tune LLMs using the abundant data on end devices without compromising data privacy. Most existing federated fine-tuning methods for LLMs rely on parameter-efficient fine-tuning techniques, which may not reach the performance height possible with full-parameter tuning. However, federated full-parameter tuning of LLMs is a non-trivial problem due to the immense communication cost. This work introduces FedKSeed that employs zeroth-order optimization with a finite set of random seeds. It significantly reduces transmission requirements between the server and clients to just a few random seeds and scalar gradients, amounting to only a few thousand bytes, making federated full-parameter tuning of billion-sized LLMs possible on devices. Building on it, we develop a strategy enabling probability-differentiated seed sampling, prioritizing perturbations with greater impact on model accuracy. Experiments across six scenarios with various LLMs, datasets and data partitions demonstrate that our approach outperforms existing federated LLM fine-tuning methods in both communication efficiency and new task generalization.

연구 동기 및 목표

• 연합 환경에서 최소한의 통신으로 빌리언 규모의 LLM의 전체 파라미터 미세 조정 가능성 시연.
• 전체 파라미터 교환과 역전파 기반 메모리 사용을 피함으로써 서버-클라이언트 간 전송량을 줄임.
• 연합 환경의 비 IID 데이터에서도 정확도를 유지하거나 향상시키는 시드 기반 섭동 전략을 개발.
• 제로차(zeroth-order) 최적화 프레임워크에서 시드 재사용에 대한 이론적 및 실험적 분석 제공.

제안 방법

• 유한 개의 무작위 시드를 사용하여 perturbation을 생성하고 전체 파라미터 튜닝에 활용하는 zeroth-order 최적화(ZOO) 사용.
• 전체 파라미터를 전송하지 않고 최신 글로벌 모델을 재구성하기 위해 서버 측 K개의 시드 풀과 스칼라 그래디언트 누적기를 유지.
• 클라이언트는 선택된 시드에 연결된 두 점 기울기 추정치를 사용하여 로컬 업데이트를 수행하고 전체 모델 업데이트 대신 시드-그래디언트 이력치를 보고한다.
• 서버에서 시드당 스칼라 그래디언트를 집계하여 글로벌 섭동 투영을 업데이트하고 시드 기반 업데이트 규칙으로 로컬에서 최신 모델을 재구성한다.
• 추정된 기울기 크기에 따라 시드에 가중치를 부여하는 비균일 시드 샘플링(FedKSeed-Pro)을 도입하여 효율성과 정확도를 향상시킨다.

실험 결과

연구 질문

• RQ1 10억 규모의 LLM의 전체 파라미터 튜닝이 실용적인 연합 배포를 위해 충분히 통신 효율적으로 될 수 있는가?
• RQ2 시드 풀을 K 개의 섭동으로 한정하는 것이 무제한 시드 방식과 비교해 수렴성 및 정확도에 어떤 영향을 미치는가?
• RQ3 FedKSeed-Pro에서 비균일 시드 샘플링이 정확도를 향상시키고 필요한 시드 기본 수를 감소시키는가?
• RQ4 ZOO 기반 연합학습에서 시드 재사용과 수렴 보장 사이의 이론적 관계는 무엇인가?
• RQ5 다양한 LLM, 데이터 세트 및 데이터 분할 설정에서 FedKSeed와 FedKSeed-Pro가 PEFT 기반 연합 학습 방법과 비교하여 얼마나 잘 작동하는가?

주요 결과

• FedKSeed는 perturbation을 시드로 인코딩하고 스칼라 그래디언트만 전송함으로써 라운드당 통신이 18 킬로바이트 미만인 전체 파라미터 튜닝을 달성한다.
• 수렴 분석에 따르면 제안된 프레임워크에서 시드 다양성은 수렴에 영향을 미치지 않으며 이론적 보장을 해치지 않으면서 시드 재사용을 가능하게 한다.
• 비균일 시드 샘플링(FedKSeed-Pro)은 Rouge-L 점수를 일관되게 향상시키고 필요한 시드 기본 수를 감소시키며, 시나리오 전반에서 기준선 대비 평균적으로 최대 7.26%의 상대 개선을 달성한다.
• 1B–3B 모델을 대상으로 한 여섯 가지 시나리오의 실험 결과, FedKSeed와 FedKSeed-Pro가 PEFT 기반 연합 미세조정 기준선보다 정확도 및 통신/메모리 효율면에서 우수하다.
• FedKSeed-Pro는 더 작은 K(예: 1024–4096)에서도 더 높은 정확도를 달성할 수 있으며, 여러 설정에서 동기화 효율성 면에서 FedKSeed를 능가한다.