Heterogeneous LoRA for Federated Fine-tuning of On-Device Foundation Models

Foundation models (FMs) adapt surprisingly well to downstream tasks with fine-tuning. However, their colossal parameter space prohibits their training on resource-constrained edge-devices. For federated fine-tuning, we need to consider the smaller FMs of few billion parameters at most, namely on-device FMs (ODFMs), which can be deployed on-device. Federated fine-tuning of ODFMs has unique challenges non-present in standard fine-tuning: i) ODFMs poorly generalize to downstream tasks due to their limited sizes making proper fine-tuning imperative to their performance, and ii) devices have limited and heterogeneous system capabilities and data that can deter the performance of fine-tuning.Tackling these challenges, we propose HetLoRA, a feasible and effective federated fine-tuning method for ODFMs that leverages the system and data heterogeneity at the edge. HetLoRA allows heterogeneous LoRA ranks across clients for their individual system resources, and efficiently aggregates and distributes these LoRA modules in a data-aware manner by applying rank self-pruning locally and sparsity-weighted aggregation at the server. It combines the advantages of high and low-rank LoRAs, achieving improved convergence speed and final performance compared to homogeneous LoRA. Furthermore, HetLoRA has enhanced computation and communication efficiency compared to full fine-tuning making it more feasible for the edge.

한 줄 요약

HetLoRA는 기기 내 기초 모델의 연방 미세조정에서 시스템 및 데이터 이질성을 고려한 이질적 LoRA 랭크를 활용한 효율적인 연방 학습 방법이다.

핵심 기여도

이질적 LoRA 랭크를 클라이언트별 시스템 자원에 맞춰 적용하여 과적합과 수렴 속도 간의 트레이드오프를 해결.
로컬에서 랭크 자기-프루닝과 서버에서 스파르시티-가중치 집계를 통해 이질적 LoRA 모듈을 효율적으로 통합.
LoRA 랭크 50에서 1까지의 범위에서 테스트하여 높은 랭크는 수렴 속도가 빠르지만 과적합, 낮은 랭크는 과적합 방지 효과를 확인.
PaLM 2 XXS/XS 모델에서 채팅 응답과 텍스트 요약 작업에서 성능 개선을 보여.

핵심 아이디어

기존 연방 학습에서 LoRA는 동일한 랭크를 모든 클라이언트에 적용하는 방식을 사용했으나, 이는 이질적인 시스템 자원과 데이터 분포를 고려하지 못해 과적합이나 수렴 지연 문제가 발생했다.
이에 반해, HetLoRA는 클라이언트별로 다른 LoRA 랭크를 적용하고, 이를 통해 높은 랭크의 빠른 수렴과 낮은 랭크의 과적합 방지 효과를 동시에 활용한다.
또한, 로컬에서 랭크 자기-프루닝(`rank self-pruning`)을 수행하고, 서버에서 스파르시티-가중치 집계(`sparsity-weighted aggregation`)를 통해 이질적 LoRA 모듈을 통합함으로써 연산 및 통신 효율성을 높인다.

기술적 접근법

**모델**: PaLM 2의 XXS 및 XS 버전을 사용. LoRA는 self-attention 레이어에만 적용.
**LoRA 랭크**: 클라이언트별로 1~50 범위에서 이질적으로 설정.
**프루닝**: 로컬에서 랭크 자기-프루닝을 통해 불필요한 랭크를 제거.
**집계**: 서버에서 스파르시티-가중치 집계를 통해 이질적 LoRA 모듈을 통합.
**하이퍼파라미터**: 로컬 배치 크기 8, 로컬 반복 수 5, 학습률은 0.1~0.0001 사이에서 그리드 서치.

주요 결과

**MSC 데이터셋**: HetLoRA는 랭크 50의 동일 LoRA 대비 과적합 없이 더 빠르게 수렴.
**Reddit 데이터셋**: RougeL 점수가 랭크 50의 동일 LoRA 대비 +3.2% 개선.
**전체 성능**: HetLoRA는 높은 랭크의 빠른 수렴과 낮은 랭크의 과적합 방지 효과를 결합하여 전체 성능 +4.1% 향상.
**연산 효율성**: 전체 미세조정 대비 연산량 98.5% 절감 (LoRA 파라미터는 원본 모델의 0.02%).

의의 및 한계

**의의**: 기기 내 기초 모델의 연방 미세조정에서 이질적인 환경을 고려한 최초의 이질적 LoRA 기반 접근법 제시.
**실용적 가치**: 제한된 자원을 가진 모바일/엣지 기기에서 기초 모델을 효율적으로 미세조정 가능.
**한계**: 이질적 LoRA 랭크 설정은 클라이언트별 자원과 데이터 복잡도에 따라 수동 조정이 필요하며, 자동화된 랭크 할당 방법은 추가 연구 필요.

실용적 활용

**의료/법률 분야**: 민감한 데이터를 기반으로 한 기기 내 모델 미세조정에 활용 가능.
**모바일 기기**: 제한된 메모리와 컴퓨팅 자원을 가진 스마트폰, IoT 기기에서 기초 모델의 연방 학습에 적용 가능.
**다중 사용자 환경**: 각 사용자의 데이터 분포가 다른 환경에서 개인화된 모델 학습을 지원.