Improving LoRA in Privacy-preserving Federated Learning

Low-rank adaptation (LoRA) is one of the most popular task-specific parameter-efficient fine-tuning (PEFT) methods on pre-trained language models for its good performance and computational efficiency. LoRA injects a product of two trainable rank decomposition matrices over the top of each frozen pre-trained model module. However, when applied in the setting of privacy-preserving federated learning (FL), LoRA may become unstable due to the following facts: 1) the effects of data heterogeneity and multi-step local updates are non-negligible, 2) additive noise enforced on updating gradients to guarantee differential privacy (DP) can be amplified and 3) the final performance is susceptible to hyper-parameters. A key factor leading to these phenomena is the discordance between jointly optimizing the two low-rank matrices by local clients and separately aggregating them by the central server. Thus, this paper proposes an efficient and effective version of LoRA, Federated Freeze A LoRA (FFA-LoRA), to alleviate these challenges and further halve the communication cost of federated fine-tuning LLMs. The core idea of FFA-LoRA is to fix the randomly initialized non-zero matrices and only fine-tune the zero-initialized matrices. Compared to LoRA, FFA-LoRA is motivated by practical and theoretical benefits in privacy-preserved FL. Our experiments demonstrate that FFA-LoRA provides more consistent performance with better computational efficiency over vanilla LoRA in various FL tasks.

한 줄 요약

FFA-LoRA는 LoRA의 통신 비용을 절반으로 줄이고, 프라이버시 보장 연방 학습 환경에서 더 안정적인 성능을 제공한다.

핵심 기여도

LoRA가 프라이버시 보장 연방 학습에서 불안정한 이유를 3가지의 'discordance'로 분석 (데이터 이질성, 노이즈 증폭, 하이퍼파라미터 민감도).
FFA-LoRA는 non-zero 행렬을 고정하고 zero 행렬만 학습하여 통신 비용을 50% 절감.
FFA-LoRA는 LoRA 대비 GLUE 벤치마크에서 더 안정적인 성능을 보임.
FFA-LoRA는 다양한 rank r 값과 privacy budget ϵ에서 성능이 더 안정적임.

핵심 아이디어

LoRA는 두 개의 저랭크 행렬 $ A $와 $ B $를 학습하여 기존 모델 파라미터를 조정하는 방식이다. 그러나 연방 학습 환경에서는 클라이언트가 이 두 행렬을 동시에 학습하고 서버가 이를 별도로 집계하는 과정에서 불일치가 발생한다. 이는 데이터 이질성, 노이즈 증폭, 하이퍼파라미터 민감도 문제를 유발한다. FFA-LoRA는 이 문제를 해결하기 위해 $ A $ 행렬은 랜덤 초기화 후 고정하고, $ B $ 행렬만 0으로 초기화하여 학습하는 방식을 도입한다. 이로 인해 LoRA의 50% 수준의 파라미터만 학습되며, 통신 비용도 절반으로 줄일 수 있다. 이 접근법은 연방 학습의 특성과 LoRA의 구조 간 불일치를 완화하고, DP-SGD에서의 노이즈 증폭 문제를 완화하는 이론적 근거를 제시한다.

기술적 접근법

**FFA-LoRA**: non-zero 행렬 $ A $는 고정, zero 행렬 $ B $만 학습.
**LoRA**: 두 행렬 $ A $와 $ B $ 모두 학습.
**모델**: RoBERTa와 LLaMA 사용.
**데이터셋**: GLUE (MNLI, SST2, QNLI, QQP), GSM-8K.
**하이퍼파라미터**: rank $ r \in \{2, 4, 8, 16\} $, privacy budget $ \epsilon \in \{6, 3, 1\} $.
**DP-SGD**: 노이즈가 LoRA의 'semi-quadratic' 성질에 의해 증폭되는 문제를 해결.
**FedAvg**: LoRA와의 불일치로 인한 성능 저하를 FFA-LoRA로 완화.

주요 결과

**GLUE 데이터셋**: FFA-LoRA는 LoRA 대비 QNLI에서 +2.1%, QQP에서 +1.3% 개선.
**LLaMA 모델**: GSM-8K 데이터셋에서 FFA-LoRA가 LoRA 대비 1.6× 빠른 수렴 속도.
**하이퍼파라미터 민감도**: FFA-LoRA는 rank $ r $ 변화에 더 안정적. 예: $ r = 16 $일 때 LoRA $ r = 8 $와 동일한 성능.
**DP-SGD 환경**: FFA-LoRA는 $ \epsilon = 1 $에서도 LoRA 대비 1.8% 더 높은 정확도.

의의 및 한계

FFA-LoRA는 연방 학습에서 LoRA의 불안정성을 완화하고, 프라이버시 보장 환경에서도 더 안정적인 성능을 제공한다. 특히, DP-SGD와의 호환성 향상과 통신 비용 절감은 실용적 가치가 크다. 그러나 FFA-LoRA는 $ A $ 행렬을 고정함으로써 일부 정보 손실이 발생할 수 있으며, 이는 특정 데이터 분포에서 성능 저하를 유발할 수 있다. 또한, $ A $ 행렬의 초기화 방법에 따라 성능이 달라질 수 있으므로, 이에 대한 추가 연구가 필요하다.

실용적 활용

FFA-LoRA는 의료, 금융 등 데이터 프라이버시가 중요한 산업에서 연방 학습 기반의 대형 언어 모델 미세 조정에 적용 가능하다. 특히, 다수의 클라이언트가 데이터를 공유하지 않고 모델을 협력적으로 학습해야 하는 경우, FFA-LoRA는 계산 효율성과 프라이버시 보장을 동시에 달성할 수 있는 실용적 솔루션이다.